diff --git a/RUNBOOK.md b/RUNBOOK.md index d837f9423b..fadded123d 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -123,6 +123,8 @@ The archives require glibc 2.28 or newer and an NVIDIA driver providing `libcuda CUDA serves Nemotron 3.5 Lightning, qualified on GB10 (`linux-aarch64`, greedy and sampled) and on Ampere cards with compute capability 8.6 (`linux-x86_64`, greedy: the RTX 30 series, RTX A6000, A10 and A40), and checked on an RTX PRO 6000 Blackwell (compute capability 12.0, `linux-x86_64`). +Flash Next also serves on CUDA from its MLX 4-bit checkpoint, on the built-in kernels, run on a GB10: see +[Flash Next on CUDA](docs/recipes/flash-next-cuda.md). On a GB10 arm64 host (an x86_64 host uses `linux-x86_64` in the same commands): diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index e625bcb2fd..0b80446a42 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -108,3 +108,7 @@ CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel`), and the reduce-scatter's `tf_rank_sum`. + +## CUDA compiler resource baseline + +The registered usage arrays were re-derived with NVCC 13.3.73 (`cuda_13.3.r13.3/compiler.38244171_0`) for SM 80, 86, 89, 120 and 121. `kernel-contract-limits` reads each compiled image and reports its register count, static shared memory and thread bound. Source changes or a different compiler may change these values; re-derive the arrays and run the live audit rather than changing the audit tolerance. Launch dimensions, dynamic shared memory, device limits and planted controls remain independent checks. Only SM 120 was executed on hardware in this update; the other four records come from their compiled images, not hardware execution. diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..7130603ec8 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -9,7 +9,7 @@ the native engine, or the Python engine frozen at 0.6.6. | Nemotron 3.5 Super VL | [Native CUDA on one or two GB10s, images](nemotron-3.5-super-vl.md) | | Qwen3.8-27B | [MLX, quantization and CUDA](qwen3.8-27b.md) | | Qwen3.5-2B | [Native Metal preview, tied affine 4-bit](qwen3.5-2b-native.md) | -| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [CUDA images](flash-next-vision.md) | +| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [native CUDA](flash-next-cuda.md), [CUDA images](flash-next-vision.md) | | Ternary Bonsai 2 27B | [MLX](ternary-bonsai-2.md) | | GLM-5.3-Flash | [MLX on a 256 GB Mac, two-rank CUDA](glm-5.3-flash.md) | | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) | diff --git a/docs/recipes/adding-a-zig-family.md b/docs/recipes/adding-a-zig-family.md index 7252bd9cf7..8ce8c1609d 100644 --- a/docs/recipes/adding-a-zig-family.md +++ b/docs/recipes/adding-a-zig-family.md @@ -38,7 +38,7 @@ Some of the host side already lives in core, but so far each piece serves one fa - staggered segments (`zig/src/core/segments.zig`, and `zig/src/cuda/segments.zig` on CUDA) serve Flash Next on Metal and Nemotron on CUDA; - every host returns the HTTP contract in `zig/src/core/engine_api.zig`, and `zig/src/core/lane_host.zig` serves Nemotron on Metal - and CUDA. Flash Next has its own host and serves one reply at a time. + and CUDA, and Flash Next on CUDA. Flash Next on Metal has its own host and serves one reply at a time. The GPU side is barely shared yet, and that costs time. In the GLM-5.3 port each extra row in a window cost about 4.4 ms of GPU time, against about 2.3 ms for the row's own experts. Row kernels read every replicated weight again for @@ -71,11 +71,11 @@ lane kernels go in, check against the engine's own plain output (step 4). | Piece | Shared now | Family copies or gaps | | --- | --- | --- | -| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels | +| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels. On CUDA, `zig/kernels/cuda/experts_rows.cu` runs affine 4-bit group-32 experts over a window's rows with the decode kernel's sums, and `route_rows.cu` the router logits and softmax top-k (Flash Next) | | GPU-side rounds | the host loop, depth rule and copy proposer in `zig/src/core/lanes`, used by Nemotron | Flash Next: its own host loop, plus `fz_accept`, the ring and event chaining (`zig/src/families/flashnext/replay.zig`). Nemotron on Metal: `zig/src/families/nemotron/gpu_round.zig` for a lone stream. Kimi K3: none | | Several Macs | the MCDMA fabric (`zig/src/fabric`) | Flash Next: `zig/src/families/flashnext/tp.zig` splits rows and holds the whole model on each Mac. Kimi K3: `zig/src/families/kimi_k3/parallel.zig` is a split plan with no link yet | | Prompt path | segments; the prompt cache, used by Flash Next, GLM, Nemotron and the Qwen3.8-27B | per-family prompt kernels and snapshots | -| Serving | `engine_api` for every host; `lane_host` for Nemotron | Flash Next: its own host, one reply at a time | +| Serving | `engine_api` for every host; `lane_host` for Nemotron, and for Flash Next on CUDA | Flash Next on Metal: its own host, one reply at a time | **The order.** 1. GLM-5.3 builds the core pieces first: @@ -194,6 +194,12 @@ On M5 tensor units. Chips before M5 need their own path. 250-500 ms, on its own queue, inside the engine process. Without it macOS unwires the set after 1-2 s idle, and the next request pays a full re-wire: 2.7 s on K3. +Native CUDA families may export `healthy(ctx: *anyopaque) bool`. The native host forwards this callback to +lane health and admission. It must read thread-safe host state without binding a GPU context, waiting on device +work or taking the inference thread's lock. Families without it retain the lane host's normal closing/stale checks. +Unavailable or busy admission returns HTTP 503 before a streaming response opens. Context-window refusals retain +their existing streamed error policy; failures after admission remain stream errors. + ## Across several Macs From Kimi K3 on four Macs: diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index c777a1cfb7..aeb28ec541 100644 --- a/docs/recipes/cuda.md +++ b/docs/recipes/cuda.md @@ -8,7 +8,7 @@ They are not the native Zig release's support matrix. The kernel paths, the prom (`--prefill-fp8`, `--no-prefill-fp8`, `--precision`) and the request and memory notes below are that engine's too, and the native binary refuses those flags. -The native binary registers only Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40). +The native binary registers Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), and Flash Next from its MLX 4-bit checkpoint, run on a GB10 ([Flash Next on CUDA](flash-next-cuda.md)). A user's receipt for 1.0.4 on a bare-metal RTX 3090, at 225 W and 350 W, is in [#592](https://github.com/ashhart/TensorFold/issues/592). GLM-5.3-Flash on CUDA, EXL3 loading and DFlash2 drafting remain Python-engine features, not native upgrade paths. Inspect the installed binary with `tensorfold-native --version` and `tensorfold-native capabilities --json`. diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md new file mode 100644 index 0000000000..fb0661be9b --- /dev/null +++ b/docs/recipes/flash-next-cuda.md @@ -0,0 +1,159 @@ +# Qwen3.8 Flash Next on CUDA + +The native CUDA build serves Flash Next (`qwen4_exp`) on one NVIDIA GPU from an MLX affine 4-bit, group-32 checkpoint +with its MTP head, on kernels built in this repository. A set of the Python 0.6.1 engine's Triton kernels, captured +and packed, can stand in for its glue to compare against that engine's teacher path. @AdrianBinDC started the port in +[#476](https://github.com/ashhart/TensorFold/pull/476); it is tracked in +[#472](https://github.com/ashhart/TensorFold/issues/472). It has been run on a GB10 (sm_121) only. + +## Experimental ModelOpt and two-rank path + +The CUDA family also reads ModelOpt NVFP4 routed experts, including CSF-compressed scale planes, with separate +MXFP8 shared experts and dense projections and the checkpoint's BF16 activations. Two local CUDA ranks partition +expert intermediate channels and gather disjoint vocabulary ranges without reducing their BF16 bits. Each rank +owns its recurrent state, and peer failure makes the serving rank unavailable rather than returning partial results. + +This path remains experimental pending independent full-model numerical qualification and complete GPU kernel +contract receipts. Compile success and host tests do not establish checkpoint fidelity. Performance and context +limits must be measured with the exact checkpoint, binary, device topology and request settings. Prefix retention +with MTP is outside this path's qualified scope; evaluate with `--prompt-cache-gib 0`. + +## Run + +```bash +tensorfold-native serve /path/to/checkpoint --temperature 0 +``` + +- `--context N` takes up to the model's 262,144 tokens; the default is the whole window. Past 2,048 keys each attention + layer reads 512 blocks of 4 keys a row, chosen by the model's block indexer, as the Python engine does. +- `--parallel N` admits N streams. Each stream's caches are its own: about 110 MiB of DeltaNet state and the n-gram + tail, plus 28 KiB a token of keys, values and indexer keys, and with drafts its MTP head's cache. +- `--no-drafts` serves without the MTP head; `"draft": false` on a request does the same for that request. +- `TF_FLASHNEXT_DRAFTS` (1 to 15, default 6) and `TF_FLASHNEXT_CONFIDENCE` (default 0.60) set the head's chain. +- The engine uploads 73.3 GiB of the checkpoint: the n-gram tables stay mapped from the shards, and the vision tower is + not loaded. +- It refuses another quantization, other shapes, and semantic settings its kernels fix (RMS epsilon, RoPE base, rotated + share, type and scaling, the gate and activation types, attention bias, the layer roles, the n-gram layer, kernel + and dilation): config.json is checked at load and the log names the field. It also refuses `--segments`, `--slide` + and a folder whose shared experts carry Living Weights neurons. +- `TENSORFOLD_CUDA_KERNELS=DIR` swaps the glue for a captured set (see [A captured set](#a-captured-set)) when the set + holds every entry point the engine launches; a set without them (Nemotron's, say) is passed over, and the startup + line says which glue runs. With a set loaded, `TF_FLASHNEXT_OWNED` (`all`, or Triton names such as + `_chunks,_merge`) runs those entry points on our kernels and the rest on the set's, to put a difference down to one + kernel. + +## What the core runs and what the family runs + +The family plugs into the core's pieces and adds no scheduler: + +- The core's round loop and lane host drive every token. Drafted windows, shared rounds of streams without drafts, and + prompts a chunk at a time between other streams' rounds (`prefill_step`) all go through them. +- The core's prompt cache keeps and resumes a stream's state at the marks it asks for. The family copies the state + (`cuda_snap.zig`), with the MTP head's cache rows for a drafting stream; the head reads one token past a row, so + states keep with a lookahead of 1, as Nemotron's do. +- Shared row kernels: the token rows (`embed_rows.cu`, groups of 32 and a copy a hyper-connection stream), the MTP + head's RMS rows with their group sums (`norms.cu`), the router logits and softmax top-k (`route_rows.cu`), and the + routed experts over a window's or a chunk's rows on the core's plan (`cuda/grouped.zig`, `experts_rows.cu`). + +The family supplies its config check, weight loading and packing, the layer graph on one batched pass +(`cuda_batch.zig`), and kernels for what is its own: the hyper-connections and the n-gram gate and conv +(`flashnext_hc.cu`), attention with its gate and the sparse indexer's pooling, scores and block selection +(`flashnext_attn.cu`, `flashnext_batch.cu`), the Gated DeltaNet chain, the n-gram gather and the MTP head. + +Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any width: drafted windows equal +serial decoding, rounds equal lone streams, and resumed prompts equal fresh ones. The glue kernels round to bf16 where +the Python engine's do, with IEEE division and square root and a within-one-ulp exp; the router and the 4-bit +hyper-connection dots run on mma tiles with K in ascending order. + +## Checks + +- `tf-cuda-test flashnext-hc`, `flashnext-attn` and `route-rows`, and `norm-rows` for the shared rows: each kernel + against a float64 chain that rounds where it rounds, and every row count's bits equal (tolerances in the test files). +- `zig build flashnext-prompt -Dflashnext_model=DIR -Dflashnext_kernels=native -Dflashnext_ids=FILE + -Dflashnext_count=32`: 32 greedy tokens through the batched pass, against the Python engine's teacher path + (`zig/tests/cuda/flashnext/lr7/results.json`, `results-sparse.json`); a captured set's folder in place of `native` + runs its glue. `TF_FLASHNEXT_DRAFTS=6` decodes each prompt serially and drafted and compares them token for token + (`TF_FLASHNEXT_SEED` for keyed sampled draws). +- `zig build flashnext-compare -Dflashnext_model=DIR -Dflashnext_kernels=SET -Dflashnext_ids=FILE`: our glue against a + captured set on one loaded engine: every prompt row's scores (top-1 agreement, KL, each path's perplexity on the + prompt's own next tokens, dense rows and rows past 2,048 apart), then each path's greedy reply and where they part. + With `TF_FLASHNEXT_OWNED` it compares the set against the set with those entries swapped. +- `zig build flashnext-select -Dflashnext_fixtures=DIR`: the native block selection against Triton's lists. +- `tf-cuda-test experts-rows`: the rows kernel against the decode kernel's one-row launches, byte for byte. +- Served: `tools/bench_concurrent.py --alone --serial` (concurrent equal to solo, drafted equal to `"draft": false`), + and a prompt sent twice (the second resumed from a kept state, equal to the first). + +## A captured set + +A set is built from the Python 0.6.1 engine in a container (`IMAGE`, default tag `tensorfold-qwen38:v0.6.1`) with the +scripts in `zig/tests/cuda/flashnext`, then packed with `tools/zig/triton_aot_pack.py`: + +```bash +zig/tests/cuda/flashnext/capture.sh WORK/cap /path/to/hf/snapshot # the model's own run, about 78 GiB of GPU +zig/tests/cuda/flashnext/capture_rows.sh WORK/rows # the batched pass's forms, no checkpoint +zig/tests/cuda/flashnext/capture_mtp.sh WORK/mtp # the MTP head's own kernels +zig/tests/cuda/flashnext/capture_sparse.sh WORK/sparse # the sparse attention's kernels +zig/tests/cuda/flashnext/pack.sh KERNELS WORK/cap/out/cap WORK/rows/out WORK/mtp/out WORK/sparse/out +``` + +`capture.sh` also writes the prompts and the Python engine's references the checks read. The set has 76 kernels; +`capture_sparse.sh` writes the selection fixtures `flashnext-select` reads. + +## Measurements + +On a GB10 (DGX Spark, driver 580.178, CUDA 13.0), an MLX 4-bit group-32 Flash Next checkpoint with its MTP head, +greedy unless said, `tools/bench_openai.py` (256 tokens, median of 3) and `tools/bench_concurrent.py` (256 tokens, +aggregate tok/s), on our kernels and, for comparison, on the 76-kernel captured set: + +| | ours | captured | +| --- | --- | --- | +| One stream, drafted, code / chat | 84.2 / 62.6 | 82.5 / 62.1 | +| One stream, `--no-drafts`, code / chat | 35.1 / 35.1 | 35.2 / 35.2 | +| `--no-drafts`, 1 / 2 / 4 streams, sampled code | 35.0 / 53.1 / 77.9 | 35.2 / 53.0 / 77.1 | +| drafted, 1 / 2 / 4 streams, greedy code | 83.8 / 82.7 / 82.2 | 82.3 / 81.1 / 80.7 | +| Cold first token at 3,001 / 12,001 tokens (s) | 2.67 / 10.40 | 2.67 / 10.30 | +| 45,030-token prompt, no cache (s) | 36.5 | 35.8 | + +Every concurrent reply equalled its solo run, and every drafted one its `"draft": false` run. Drafted streams verify +their windows a pass each, so drafted throughput does not grow with streams yet. While a 12,001-token prompt runs, +another stream's decode waits at most 1.3 s between tokens (one prompt chunk). Each stream at a 16,384-token window +takes 0.74 GiB with drafts, 0.54 GiB without. + +Our kernels' sums are their own, so tokens part from the captured set's (the Python engine's teacher path) where two +candidates are close: 32 greedy tokens equal it on four of `zig/tests/cuda/flashnext/lr7`'s five prompts and part at +token 15 on the 3,001-token one. Over 4,096 tokens of this repository's documentation, the next-token perplexity is +6.168 on the captured set and 6.178 on ours. + +### The ModelOpt checkpoint on one GB10 + +`nvidia/Qwen3.8-Flash-Next-FP4-Local` (revision `b9ce18a6`) on a GB10 (128 GB, driver 580.178, CUDA 13.3), greedy, +`--prompt-cache-gib 0`, the tools above, three runs of the defaults at a 32,768-token window alternated with other +settings in one session: + +| | defaults | +| --- | --- | +| One stream, drafted, code / chat | 64.1-64.6 / 40.7-41.1 | +| drafted, 1 / 4 / 8 streams, greedy code | 64.1-64.4 / 44.9-45.2 / 66.2-66.8 | +| drafted, 1 / 4 / 8 streams, greedy chat | 40.6-41.1 / 35.0-35.2 / 52.4-52.8 | +| Cold first token at 2,048 / 8,192 / 32,768 / 65,536 / 131,072 tokens (s) | 1.83 / 7.11 / 29.2 / 60.1 / 131.7 | + +Every concurrent reply equalled its solo run and every drafted one its `"draft": false` run. A reply cancelled +mid-stream, or sent again after a restart, equalled a fresh one, and a second conversation turn resumed from the prompt +cache (`--no-drafts`) equalled the same turn computed without it. Prompt rows equal one-token rows bit for bit at seven +fixed histories of 20 to 6,496 tokens. Three streams draft at a time; the rest share plain rounds. + +The engine holds 73.2 GiB and keeps a 12.2 GiB reserve. A 32,768-token window then admits 7 or 8 streams of 1.21 GiB +and a 131,136-token window (the cold prompts above) 2 of 4.03 GiB; a 262,144-token stream takes 7.79 GiB, which fitted +in one start of four on this machine. After the 131,072-token prompt the process had peaked at 6.5 GB of host and +83.0 GB of device memory. Absolute speeds on this machine moved by about a tenth between sessions. + +## Not yet + +- The dense projections run a copy of the Python engine's group-32 lane matmul (`qmm.cu`); the core's `qmm_group` and + `lane_gemv` are built for groups of 64. +- No CUDA graphs or GPU-side rounds; drafted streams verify their windows one pass each, and the MTP head drafts one + stream at a time. +- Prompts run slower than the Python engine's: a cold 12,001-token prompt takes about 10 s against its 6.7 s. +- Draws are made on the host from each row's scores. +- Our glue's sums are not the captured kernels' bit for bit, so replies part from the Python engine's teacher path + after a while, as that engine's own served path does (10 to 36 tokens, from its routed-experts kernel's rounding). diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index af65c4fccd..8da3709b72 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -2,7 +2,8 @@ The `qwen4_exp` family has Gated DeltaNet, sparse attention, MoE, hyper-connections and hashed n-gram embeddings. The native engine's supported Metal checkpoint uses MLX affine 6-bit weights in groups of 32 -and includes an MTP head. The native engine refuses the 4-bit/group-32 checkpoint at load. +and includes an MTP head. On a Mac the native engine refuses the 4-bit/group-32 checkpoint at load; on CUDA it serves +that one ([Flash Next on CUDA](flash-next-cuda.md)). ```bash tensorfold pull TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @@ -10,8 +11,8 @@ tensorfold serve TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP --name bench --tempe ``` For native serving, see the [Metal runbook](../../RUNBOOK.md#flash-next-and-paired-metal-serving). -Native Flash Next has no CUDA backend yet; its port is tracked in -[#472](https://github.com/ashhart/TensorFold/issues/472). +Native Flash Next on CUDA serves the MLX 4-bit group-32 checkpoint on the built-in kernels: see +[Flash Next on CUDA](flash-next-cuda.md) ([#472](https://github.com/ashhart/TensorFold/issues/472)). The sections below describe the Python 0.6.6 engine, available at the [`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6). Their flags, variables and paths are that engine's, diff --git a/tools/zig/triton_aot_pack.py b/tools/zig/triton_aot_pack.py index 61d617a1bb..c9510e9f4c 100644 --- a/tools/zig/triton_aot_pack.py +++ b/tools/zig/triton_aot_pack.py @@ -60,7 +60,8 @@ def entry(h: str, k: dict, nospec: set[str], cache: Path | None, out: Path) -> d "global_scratch": md.get("global_scratch_size", 0), "global_align": md.get("global_scratch_align", 1), "profile_scratch": md.get("profile_scratch_size", 0), "pdl": bool(md.get("launch_pdl", False)), "params": [{"name": n, "type": sig[n], "div16": n in div16, "nospec": n in nospec} for n in runtime], - "consts": {n: const(v) for n, v in k["constexprs"].items() if n not in runtime}, + # a None argument Triton folded into the kernel (an unused pointer) leaves nothing to match + "consts": {n: const(v) for n, v in k["constexprs"].items() if n not in runtime and v is not None}, } @@ -69,14 +70,15 @@ def main() -> int: ap.add_argument("--manifest", action="append", required=True, help="a manifest.json or a Recorder's launches.json (repeatable)") ap.add_argument("--cache", action="append", help="each capture's Triton cache as seen here, in --manifest order") - ap.add_argument("--jit", help="the capture's jit.json; without it no parameter is marked do_not_specialize") + ap.add_argument("--jit", action="append", + help="a capture's jit.json (repeatable, merged); without one no parameter is do_not_specialize") ap.add_argument("--out", required=True) a = ap.parse_args() docs = [json.loads(Path(m).read_text()) for m in a.manifest] caches = a.cache or [d.get("cache") for d in docs] if len(caches) != len(docs): raise SystemExit("--cache is given once per --manifest, or not at all") - jit = json.loads(Path(a.jit).read_text()) if a.jit else None + jit = {f: v for j in a.jit for f, v in json.loads(Path(j).read_text()).items()} if a.jit else None out = Path(a.out) (out / "cubins").mkdir(parents=True, exist_ok=True) kernels, seen = [], set() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..88330fad81 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -14,6 +14,14 @@ const torch_ops = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false" }; const glue = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false", "--prec-div=true", "--prec-sqrt=true" }; const kernels = [_]Kernel{ + .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row + .{ .name = "flashnext_batch", .flags = &.{ "-O3", "--fmad=false" } }, // ours: Flash Next's batched prompt pass glue + .{ .name = "flashnext_mtp", .flags = &.{"-O3"} }, // ours: Flash Next's MTP drafts, a row's best score and log-sum-exp + .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 + .{ .name = "experts_rows", .flags = &.{"-O3"} }, // experts.cu's decode sums over a window or prompt chunk's rows + .{ .name = "route_rows", .flags = glue }, // ours: router logits and softmax top-k rows (cuda/route_rows.zig) + .{ .name = "flashnext_hc", .flags = glue }, // ours: Flash Next's hyper-connection and n-gram glue + .{ .name = "flashnext_attn", .flags = glue }, // ours: Flash Next's attention glue and its indexer's pooling and scores .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 @@ -143,6 +151,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("core", mods.core); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); + runner.addImport("flashnext", flashnextModule(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("lfm2", lfm2Module(b, target, optimize, cuda, mods.core)); runner.addImport("radio", mods.radio); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); @@ -156,6 +165,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("glm_cuda", glm_cuda); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); + flashnextChecks(b, target, optimize, cuda, mods.core, mods.lanes); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @import("kernel_contract.zig").cuda(b, runner_exe, controls); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, build_options, true).root_module.strip = strip; @@ -183,6 +193,15 @@ 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 } } }); } +/// Flash Next's CUDA family, sharing the maintained runtime and lane core. +fn flashnextModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_launch.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("core", core); + m.addImport("lanes", lanes); + 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 }} }); @@ -198,6 +217,7 @@ 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 = "flashnext", .module = flashnextModule(b, target, optimize, cuda, core, lanes) }, .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, }, }); @@ -240,6 +260,65 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); + const writeback_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_flashnext_writeback.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&writeback_contract.step); + const utility_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_utility_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&utility_contract.step); + const flash_contract = b.step("test-flashnext-contract-host", "Flash Next contract stress and rounding references without a GPU"); + flash_contract.dependOn(&writeback_contract.step); + flash_contract.dependOn(&utility_contract.step); + const row_best_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_row_best_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&row_best_contract.step); + flash_contract.dependOn(&row_best_contract.step); + const qsa_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_qsa_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&qsa_contract.step); + flash_contract.dependOn(&qsa_contract.step); + const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)"); + const projection_step = b.step("test-cuda-projection-fixture", "Projection fixture dispatch controls with CPU-only CUDA calls"); + for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| { + const reference = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/quant_ref_host_tests.zig"), + .target = host, + .optimize = mode, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&reference.step); + reference_step.dependOn(&reference.step); + const projection = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/projection_f32_host_tests.zig"), + .target = host, + .optimize = mode, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&projection.step); + projection_step.dependOn(&projection.step); + } const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too @@ -268,7 +347,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); // 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); @@ -331,3 +410,41 @@ pub fn checkDistFatbins(b: *std.Build, dir: []const u8) *std.Build.Step { for (kernels) |k| check.addFileArg(b.graph.cwdRelativePath(b.pathJoin(&.{ dir, b.fmt("{s}.fatbin", .{k.name}) }))); return &check.step; } + +/// Explicit model diagnostics; these tools do not qualify a checkpoint merely by compiling. +fn flashnextChecks(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) void { + const flashnext = flashnextModule(b, target, optimize, cuda, core, lanes); + // not installed: one embedding row through the captured _embed cubin + const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + embed_check.addImport("cuda", cuda); + embed_check.addImport("core", core); + embed_check.addImport("flashnext", flashnext); + const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check })); + const embed_model = b.option([]const u8, "flashnext_model", "Flash Next checkpoint directory for the embedding check"); + const embed_kernels = b.option([]const u8, "flashnext_kernels", "captured Triton pack for the embedding check"); + const embed_token = b.option([]const u8, "flashnext_token", "token id for the embedding check") orelse "7"; + if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); + b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); + const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + prompt_check.addImport("cuda", cuda); + prompt_check.addImport("core", core); + prompt_check.addImport("flashnext", flashnext); + const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check })); + const prompt_ids = b.option([]const u8, "flashnext_ids", "file of prompt token ids for the Flash Next prompt check"); + const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; + if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); + b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); + // not installed: our own glue against a captured set on one loaded engine (scores of every prompt row, replies) + const compare_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_compare_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + compare_check.addImport("cuda", cuda); + compare_check.addImport("flashnext", flashnext); + const compare_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-compare", .root_module = compare_check })); + if (embed_model != null and embed_kernels != null and prompt_ids != null) compare_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.? }); + b.step("flashnext-compare", "Our Flash Next glue against a captured set: every prompt row's scores, perplexity, greedy replies").dependOn(&compare_run.step); + // not installed: the native block selection against Triton's lists (capture_sparse.sh's fixtures) + const select_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_select_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + select_check.addImport("cuda", cuda); + const select_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-select", .root_module = select_check })); + if (b.option([]const u8, "flashnext_fixtures", "capture_sparse.sh's selection fixtures folder")) |dir| select_run.addArgs(&.{dir}); + b.step("flashnext-select", "fn_qsa_select against Triton's _select / _select_tiles lists on captured score rows").dependOn(&select_run.step); +} diff --git a/zig/kernels/cuda/embed_rows.cu b/zig/kernels/cuda/embed_rows.cu index 2ed32c9665..664e6a41ec 100644 --- a/zig/kernels/cuda/embed_rows.cu +++ b/zig/kernels/cuda/embed_rows.cu @@ -11,18 +11,19 @@ constexpr int NT = 256; using namespace tf_embed; -// out[r, i] = bf16(q * s + b) of token ids[r]: q 4 bits (8 a word, low first), s and b bf16 a group of 64; d % 64 == 0. +// out[r, k, i] = bf16(q * s + b) of token ids[r], k < copies: q 4 bits (low first), s and b a group's (32 or 64). extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int* __restrict__ ids, const uint32_t* __restrict__ w, const __nv_bfloat16* __restrict__ s, const __nv_bfloat16* __restrict__ b, - __nv_bfloat16* __restrict__ out, int d) { + __nv_bfloat16* __restrict__ out, int d, int group, + int copies) { const int64_t r = blockIdx.x; const int c = blockIdx.y * NT + threadIdx.x; // 8 values, one word if (8 * c >= d) return; const int64_t tok = ids[r]; const uint32_t word = w[tok * (d / 8) + c]; - const int64_t g = tok * (d / 64) + c / 8; + const int64_t g = tok * (d / group) + 8 * c / group; const float sc = __bfloat162float(s[g]), bi = __bfloat162float(b[g]); uint4 o; __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&o); @@ -30,7 +31,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(__fmaf_rn(static_cast((word >> (8 * i)) & 0xFu), sc, bi), __fmaf_rn(static_cast((word >> (8 * i + 4)) & 0xFu), sc, bi)); - *reinterpret_cast(out + r * d + 8 * c) = o; + for (int k = 0; k < copies; ++k) *reinterpret_cast(out + (r * copies + k) * d + 8 * c) = o; } // out[r] = w[ids[r]] for a bf16 table, 16 bytes a thread (d % 8 == 0). @@ -42,3 +43,16 @@ extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16(const int* _ if (8 * c >= d) return; *reinterpret_cast(out + r * d + 8 * c) = *reinterpret_cast(w + static_cast(ids[r]) * d + 8 * c); } + +// Exact BF16 rows repeated across residual streams; no floating-point conversion. +extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16_copies(const int* __restrict__ ids, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ out, int d, + int copies) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + if (8 * c >= d) return; + const uint4 value = *reinterpret_cast(w + static_cast(ids[r]) * d + 8 * c); + for (int k = 0; k < copies; ++k) + *reinterpret_cast(out + (r * copies + k) * d + 8 * c) = value; +} diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..9e46b57a19 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -280,3 +280,8 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); TF_EXPERT(1) TF_EXPERT(0) +// Groups of 32: SwiGLU gate and up together (epilogue 2), and the down projection to fp32. +template __global__ void tf_experts::expert_kernel<32, 2, 2, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_experts::expert_kernel<32, 1, 0, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..327517f5b4 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -37,6 +37,8 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack -// Groups of 64 inputs. +// Groups of 64 inputs, then groups of 32. template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<1>(const uint32_t*, const uint16_t*, const uint16_t*, + uint32_t*, int, int, int, int); diff --git a/zig/kernels/cuda/experts_rows.cu b/zig/kernels/cuda/experts_rows.cu new file mode 100644 index 0000000000..8ec6d0d162 --- /dev/null +++ b/zig/kernels/cuda/experts_rows.cu @@ -0,0 +1,144 @@ +// experts.cu's decode expert kernel over up to 16 * RT pair rows a warp; a row's bits do not depend on its window. + +#include +#include + +#include "experts.cuh" + +namespace tf_experts_rows { + +__device__ __forceinline__ uint4 ld_w(const uint4* p) { + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +__device__ __forceinline__ uint4 ld_x(const __nv_bfloat16* p) { return __ldg(reinterpret_cast(p)); } + +// experts.cu's group_sum: a lane's 8 inputs a row in order, then the quad's partners. +template +__device__ __forceinline__ float group_sum(const uint4 (&v)[XV]) { + float s = 0.f; +#pragma unroll + for (int c = 0; c < XV; ++c) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t u = comp(v[c], i); + s += lo_f(u); + s += hi_f(u); + } + s += __shfl_xor_sync(0xffffffffu, s, 1); + s += __shfl_xor_sync(0xffffffffu, s, 2); + return s; +} + +template +__global__ void __launch_bounds__(WARPS * 32) + expert_rows(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, + int NB, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, float limit) { + using G = Geo; + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x * WARPS + (threadIdx.x >> 5); unit < units; unit += gridDim.x * WARPS) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + bool v0[RT], v1[RT]; + int pr0[RT], pr1[RT]; + const __nv_bfloat16* x0[RT]; + const __nv_bfloat16* x1[RT]; + int nt = 0; // row tiles of 16 this item fills +#pragma unroll + for (int r = 0; r < RT; ++r) { + v0[r] = 16 * r + gq < cnt; + v1[r] = 16 * r + gq + 8 < cnt; + pr0[r] = v0[r] ? __ldg(members + first + 16 * r + gq) : 0; + pr1[r] = v1[r] ? __ldg(members + first + 16 * r + gq + 8) : 0; + const int r0 = slots ? pr0[r] / slots : pr0[r], r1 = slots ? pr1[r] / slots : pr1[r]; + x0[r] = X + (size_t)r0 * x_stride + t * (GS / 4); + x1[r] = X + (size_t)r1 * x_stride + t * (GS / 4); + if (16 * r < cnt) nt = r + 1; + } + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (M * G::BLOCK); + float acc[M][RT][NTW][4]; +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[m][r][j][0] = acc[m][r][j][1] = acc[m][r][j][2] = acc[m][r][j][3] = 0.f; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + for (int g = 0; g < KG; ++g) { + const uint4* b = blk + (size_t)g * (M * G::BLOCK); + uint4 w[M][G::WV], sb[M][G::SBV]; +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int c = 0; c < G::WV; ++c) w[m][c] = ld_w(b + m * G::BLOCK + c * 32 + lane); +#pragma unroll + for (int c = 0; c < G::SBV; ++c) sb[m][c] = ld_w(b + m * G::BLOCK + 32 * G::WV + t * G::SBV + c); + } + const int k0 = g * GS; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + uint4 xa[G::XV], xb[G::XV]; +#pragma unroll + for (int c = 0; c < G::XV; ++c) { + xa[c] = v0[r] ? ld_x(x0[r] + k0 + 8 * c) : zero; + xb[c] = v1[r] ? ld_x(x1[r] + k0 + 8 * c) : zero; + } + const float xsa = group_sum(xa); + const float xsb = group_sum(xb); +#pragma unroll + for (int m = 0; m < M; ++m) { + float p[NTW][4]; +#pragma unroll + for (int j = 0; j < NTW; ++j) p[j][0] = p[j][1] = p[j][2] = p[j][3] = 0.f; +#pragma unroll + for (int ks = 0; ks < G::KS; ++ks) { + const uint32_t a0 = comp(xa[ks >> 1], 2 * (ks & 1)); + const uint32_t a2 = comp(xa[ks >> 1], 2 * (ks & 1) + 1); + const uint32_t a1 = comp(xb[ks >> 1], 2 * (ks & 1)); + const uint32_t a3 = comp(xb[ks >> 1], 2 * (ks & 1) + 1); + const int sh = (ks & 1) ? 8 : 0; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int wi = j * (GS / 32) + (ks >> 1); + const uint32_t word = comp(w[m][wi >> 2], wi & 3); + mma(p[j], a0, a1, a2, a3, nib2(word, sh), nib2(word, sh + 4)); + } + } +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const uint32_t sp = comp(sb[m][j >> 2], j & 3); + const uint32_t bp = comp(sb[m][(NTW + j) >> 2], (NTW + j) & 3); + const float s0 = lo_f(sp), s1 = hi_f(sp), b0 = lo_f(bp), b1 = hi_f(bp); + float(&a)[4] = acc[m][r][j]; + a[0] = fmaf(xsa, b0, fmaf(p[j][0], s0, a[0])); + a[1] = fmaf(xsa, b1, fmaf(p[j][1], s1, a[1])); + if (v1[r]) { + a[2] = fmaf(xsb, b0, fmaf(p[j][2], s0, a[2])); + a[3] = fmaf(xsb, b1, fmaf(p[j][3], s1, a[3])); + } + } + } + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr0[r], pr1[r], v0[r], v1[r], limit); + } + } +} + +} // namespace tf_experts_rows + +// Groups of 32, two row tiles (32 rows) a warp: SwiGLU gate and up together (epilogue 2), down to fp32 (epilogue 0). +template __global__ void tf_experts_rows::expert_rows<32, 2, 2, 2, 4>(const __nv_bfloat16*, int, int, const uint4*, + int, int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_experts_rows::expert_rows<32, 1, 0, 2, 4>(const __nv_bfloat16*, int, int, const uint4*, + int, int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu new file mode 100644 index 0000000000..71aaad7026 --- /dev/null +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -0,0 +1,370 @@ +// Flash Next's attention glue and its sparse indexer's pooling and scores: each row's own fixed-order sums. + +// glue.py and attention.py's arithmetic and bf16 rounding points; the attention dots on mma tiles, K ascending. + +#include +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace fn_attn { + +constexpr int TK = 32; // keys an attention tile holds +constexpr int HD = 256; // head width the attention tiles are built for +constexpr int QR = 16; // query heads a KV head's tile holds (G <= 16, the rest zero) + +// A block's fixed-order sum: each warp's butterfly, then the warps' sums in warp order. +__device__ __forceinline__ float block_sum(float v, float* scratch) { + v = warp_sum(v); + const int warps = blockDim.x >> 5; + if ((threadIdx.x & 31) == 0) scratch[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < warps; ++w) total = __fadd_rn(total, scratch[w]); + __syncthreads(); + return total; +} + +__device__ __forceinline__ float inv_rms(float ss, int width, float eps) { + return __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ss, static_cast(width)), eps))); +} + +// Rotate-half RoPE over the first 2 * half features at `pos`: cos and sin of the fp32 angle (full range reduction). +__device__ __forceinline__ float rope(float xn, float xpn, int d, int half, int pos, const float* __restrict__ inv) { + if (d >= 2 * half) return xn; + const int i = d < half ? d : d - half; + float s, c; + sincosf(__fmul_rn(static_cast(pos), inv[i]), &s, &c); + return d < half ? __fsub_rn(__fmul_rn(xn, c), __fmul_rn(xpn, s)) : __fadd_rn(__fmul_rn(xpn, s), __fmul_rn(xn, c)); +} + +// ldmatrix .x4 .trans: four 8x8 bf16 tiles of shared memory, each thread's pair down a column (the B layout of rows). +__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(qmm_frag::smem(p))); +} + +__device__ __forceinline__ int partner(int d, int half) { + return d < half ? d + half : (d < 2 * half ? d - half : d); +} + +} // namespace fn_attn + +using namespace fn_attn; + +// glue._attn_prep (text positions, bf16 cache), program (row, head): heads normed, rounded, roped, rounded, stored. +extern "C" __global__ void __launch_bounds__(256) fn_attn_prep(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, + const float* __restrict__ QW, const float* __restrict__ KW, + const float* __restrict__ IW, const float* __restrict__ INV, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, + __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, + __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, + int NI, int IHD, int HALF) { + __shared__ float scratch[32]; + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + const int pos = POS0[0] + r; + const long long row = static_cast(r) * PW; + if (head >= NQ + NKV + NI) { + if (d < IHD) IKC[static_cast(pos) * IHD + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; + return; + } + const bool index = head >= NQ + NKV; + const int width = index ? IHD : HD; + const bool live = d < width; + const long long src = row + (head < NQ ? head * 2 * HD : (head < NQ + NKV ? NQ * 2 * HD + (head - NQ) * HD : NQ * 2 * HD + 2 * NKV * HD + (head - NQ - NKV) * IHD)); + const float* w = head < NQ ? QW : (head < NQ + NKV ? KW : IW); + const float x = live ? bf(P[src + d]) : 0.0f; + const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), width, eps); + if (!live) return; + const int p = partner(d, HALF); + const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), w[d])); + const float xpn = rbf(__fmul_rn(__fmul_rn(bf(P[src + p]), rinv), w[p])); + const __nv_bfloat16 out = __float2bfloat16_rn(rope(xn, xpn, d, HALF, pos, INV)); + if (head < NQ) { + Q[(static_cast(r) * NQ + head) * HD + d] = out; + } else if (!index) { + const long long slot = static_cast(pos) * NKV + head - NQ; + KC[slot * HD + d] = out; + VC[slot * HD + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; + } else { + IQ[(static_cast(r) * NI + head - NQ - NKV) * IHD + d] = out; + } +} + +// glue._attn_gate, program (row, head) of HD threads: bf16(o / (1 + e^-gate)) and its 32-group sums. +extern "C" __global__ void __launch_bounds__(256) fn_attn_gate(const __nv_bfloat16* __restrict__ O, const __nv_bfloat16* __restrict__ P, + __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, int PW, int NQ) { + const int r = blockIdx.x, h = blockIdx.y, d = threadIdx.x; + const float o = bf(O[(static_cast(r) * NQ + h) * HD + d]); + const float g = bf(P[static_cast(r) * PW + h * 2 * HD + HD + d]); + const float out = rbf(__fdiv_rn(o, __fadd_rn(1.0f, tf_exp(-g)))); + OUT[static_cast(r) * NQ * HD + h * HD + d] = __float2bfloat16_rn(out); + const float sum = warp_sum(out); + if ((d & 31) == 0) XS[static_cast(r) * (NQ * HD / 32) + h * (HD / 32) + d / 32] = sum; +} + +// attention._pool, a block a completed block of RATIO indexer keys: their mean, rounded, normed and roped at its start. +extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, const float* __restrict__ W, + const float* __restrict__ INV, float eps, int R, int DI, int HALF, + int RATIO) { + __shared__ float scratch[32]; + __shared__ float xs[256]; + const int p0 = POS0[0]; + const int b = p0 / RATIO + blockIdx.x, d = threadIdx.x; + if (RATIO * b + RATIO > p0 + R) return; + float acc = bf(IKC[static_cast(RATIO) * b * DI + d]); + for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[(static_cast(RATIO) * b + k) * DI + d])); + const float x = rbf(__fdiv_rn(acc, static_cast(RATIO))); + xs[d] = x; + const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), DI, eps); + const int p = partner(d, HALF); + const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), W[d])); + const float xpn = rbf(__fmul_rn(__fmul_rn(xs[p], rinv), W[p])); + POOLED[static_cast(b) * DI + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, RATIO * b, INV)); +} + +// attention._scores, program (row, 64 blocks): the heads' sum of relu(q . k) / sqrt(DI), rows past TOP (DI 128). +extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, float* __restrict__ SC, int NB, int HI, + int RATIO, int TOP) { + constexpr int DI = 128; + __shared__ float qs[8][DI]; + const int r = blockIdx.x, j = blockIdx.y; + const int complete = (POS0[0] + r + 1) / RATIO; + if (complete <= TOP || j * 64 >= complete) return; + for (int e = threadIdx.x; e < HI * DI; e += blockDim.x) qs[e / DI][e % DI] = bf(IQ[static_cast(r) * HI * DI + e]); + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const float scale = __fsqrt_rn(static_cast(DI)); + uint2 raw[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int b = j * 64 + warp * 8 + i; + raw[i] = b < complete ? *reinterpret_cast(POOLED + static_cast(b) * DI + 4 * lane) : make_uint2(0, 0); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int b = j * 64 + warp * 8 + i; + if (b >= complete) break; + const __nv_bfloat16* kb = reinterpret_cast(&raw[i]); + float total = 0.0f; + for (int h = 0; h < HI; ++h) { + float part = 0.0f; +#pragma unroll + for (int e = 0; e < 4; ++e) part = __fmaf_rn(bf(kb[e]), qs[h][4 * lane + e], part); + total = __fadd_rn(total, fmaxf(warp_sum(part), 0.0f)); + } + if (lane == 0) SC[static_cast(r) * NB + b] = __fdiv_rn(total, scale); + } +} + +// attention._chunks, program (row, KV head, chunk): the chunk's keys TK at a time, online softmax, partial o, m, l. +template +__device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, + int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + typedef __nv_bfloat16 Row[HD + 8]; + Row* qs = reinterpret_cast(smem_raw); // [QR] the row's query heads + Row* tiles_kv = qs + QR; // [NB][2][TK] each buffer's keys, then values + __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + NB * 2 * TK); // [QR] bf16 p + __shared__ float red_m[4][QR], red_l[4][QR]; + __shared__ int key_ids[NB][TK]; + const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + int n = POS0[0] + r + 1; + bool sparse = false; + if (QSA) { + sparse = SPR[r] != 0; + if (sparse) n = NKR[r]; + } + const int start = c * CH; + if (start >= n) return; + for (int e = threadIdx.x; e < QR * HD / 8; e += blockDim.x) { + const int i = e / (HD / 8), v = e % (HD / 8); + qmm_frag::cp16z(&qs[i][8 * v], Q + (static_cast(r) * H + hk * G + min(i, G - 1)) * HD + 8 * v, i < G); + } + // a tile's key ids, then its keys and values staged as two cp.async groups + auto load = [&](int t0, int buf) { + if (threadIdx.x < TK) { + const int ki = start + t0 + threadIdx.x; + key_ids[buf][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + } + __syncthreads(); + for (int pass = 0; pass < 2; ++pass) { + const __nv_bfloat16* src = pass == 0 ? KC : VC; + Row* dst = tiles_kv + (buf * 2 + pass) * TK; + for (int e = threadIdx.x; e < TK * HD / 8; e += blockDim.x) { + const int k = e / (HD / 8), v = e % (HD / 8), key = key_ids[buf][k]; + qmm_frag::cp16z(&dst[k][8 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 8 * v, key >= 0); + } + qmm_frag::commit(); + } + }; + float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.0f, 0.0f}; + float o[8][4] = {}; + const int tiles = min(n - start, CH); + if (NB == 2) load(0, 0); + for (int t0 = 0, it = 0; t0 < tiles; t0 += TK, ++it) { + const int buf = NB == 2 ? it & 1 : 0; + const bool ahead = NB == 2 && t0 + TK < tiles; // the next tile's loads go out before this tile's math + if (NB == 1) load(t0, 0); + if (ahead) load(t0 + TK, buf ^ 1); + Row* kt = tiles_kv + buf * 2 * TK; + Row* vt = kt + TK; + const int* keys = key_ids[buf]; + if (ahead) qmm_frag::wait<3>(); else qmm_frag::wait<1>(); + __syncthreads(); + constexpr int JT = TK / 32; // a warp's 8-key tiles + float s[JT][4] = {}; +#pragma unroll 4 + for (int ks = 0; ks < HD / 16; ks += 2) { + uint32_t a0[4], a1[4]; + qmm_frag::ldmatrix4(a0, &qs[lane % 16][ks * 16 + (lane / 16) * 8]); + qmm_frag::ldmatrix4(a1, &qs[lane % 16][(ks + 1) * 16 + (lane / 16) * 8]); +#pragma unroll + for (int jt = 0; jt < JT; ++jt) { + uint32_t b[4]; // two k16 steps of the tile's 8 keys + qmm_frag::ldmatrix4(b, &kt[warp * 8 * JT + jt * 8 + lane % 8][ks * 16 + (lane / 8) * 8]); + qmm_frag::mma(s[jt], a0, b[0], b[1]); + qmm_frag::mma(s[jt], a1, b[2], b[3]); + } + } + float tile_m[2] = {-INFINITY, -INFINITY}; +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const bool valid = keys[warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] >= 0; + s[jt][e] = valid ? __fmul_rn(s[jt][e], SCALE) : -INFINITY; + tile_m[e >> 1] = fmaxf(tile_m[e >> 1], s[jt][e]); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 1)); + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 2)); + if (t == 0) red_m[warp][g + 8 * h] = tile_m[h]; + } + __syncthreads(); + float next_m[2], alpha[2], rows[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float tm = fmaxf(fmaxf(red_m[0][row], red_m[1][row]), fmaxf(red_m[2][row], red_m[3][row])); + const bool active = tm != -INFINITY; + next_m[h] = active ? fmaxf(m[h], tm) : m[h]; + alpha[h] = active ? (m[h] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m[h], next_m[h]))) : 1.0f; + rows[h] = 0.0f; + } +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int h = e >> 1; + const float p = s[jt][e] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(s[jt][e], next_m[h])); + rows[h] = __fadd_rn(rows[h], p); + ps[g + 8 * h][warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] = __float2bfloat16_rn(p); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 1)); + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 2)); + if (t == 0) red_l[warp][g + 8 * h] = rows[h]; + } + if (ahead) qmm_frag::wait<2>(); else qmm_frag::wait<0>(); + __syncthreads(); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float sum = __fadd_rn(__fadd_rn(__fadd_rn(red_l[0][row], red_l[1][row]), red_l[2][row]), red_l[3][row]); + l[h] = __fmaf_rn(l[h], alpha[h], sum); + m[h] = next_m[h]; + } +#pragma unroll + for (int q = 0; q < 8; ++q) +#pragma unroll + for (int e = 0; e < 4; ++e) o[q][e] = __fmul_rn(o[q][e], alpha[e >> 1]); + for (int ks = 0; ks < TK / 16; ++ks) { + uint32_t a[4]; + qmm_frag::ldmatrix4(a, &ps[lane % 16][ks * 16 + (lane / 16) * 8]); +#pragma unroll + for (int q = 0; q < 8; q += 2) { + uint32_t b[4]; // two 8-column tiles of the values, transposed into the B layout + ldmatrix4t(b, &vt[ks * 16 + lane % 16][warp * 64 + q * 8 + (lane / 16) * 8]); + qmm_frag::mma(o[q], a, b[0], b[1]); + qmm_frag::mma(o[q + 1], a, b[2], b[3]); + } + } + __syncthreads(); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = g + 8 * h; + if (i >= G) continue; + const long long base = (static_cast(r) * NCH + c) * H + hk * G + i; +#pragma unroll + for (int q = 0; q < 8; ++q) { + float2 v = make_float2(o[q][2 * h], o[q][2 * h + 1]); + *reinterpret_cast(PO + base * HD + warp * 64 + q * 8 + 2 * t) = v; + } + if (warp == 0 && t == 0) { + PM[base] = m[h]; + PL[base] = l[h]; + } + } +} + +// attn_chunks one tile at a time (prompt chunks: two blocks a multiprocessor). +extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +// attn_chunks with the next tile's loads in flight (a decode round's few rows): the same sums. +extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_ahead(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +// attention._merge, program (row, head) of HD threads: the head's chunk partials merged in chunk order, o / l. +extern "C" __global__ void __launch_bounds__(256) fn_attn_merge(const float* __restrict__ PO, const float* __restrict__ PM, + const float* __restrict__ PL, const int* __restrict__ POS0, + __nv_bfloat16* __restrict__ OUT, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + int QSA) { + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + int n = POS0[0] + r + 1; + if (QSA && SPR[r] != 0) n = NKR[r]; + const int chunks = (n + CH - 1) / CH; + { + float m = -INFINITY, l = 0.0f, o = 0.0f; + for (int c = 0; c < chunks; ++c) { + const long long base = (static_cast(r) * NCH + c) * H + head; + const float cm = PM[base], cl = PL[base], co = PO[base * HD + d]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; + o = __fmaf_rn(co, b, __fmul_rn(o, a)); + l = __fmaf_rn(cl, b, __fmul_rn(l, a)); + m = next; + } + OUT[(static_cast(r) * H + head) * HD + d] = __float2bfloat16_rn(__fdiv_rn(o, l)); + } +} diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu new file mode 100644 index 0000000000..c4ad0e2691 --- /dev/null +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -0,0 +1,316 @@ +// The small device steps of Flash Next's batched prompt pass (cuda_batch.zig), built without FMA contraction. + +#include +#include + +// out[i] = the sum of the sk slices' part[s][i] in order, rounded once (qmm.cu's reduce_kernel): bf16 if out_bf16. +extern "C" __global__ void fn_sum_slices(const float* __restrict__ part, void* __restrict__ out, int out_bf16, + long long n, int sk) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + float acc = part[i]; + for (int s = 1; s < sk; ++s) acc = acc + part[static_cast(s) * n + i]; + if (out_bf16) reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); + else reinterpret_cast(out)[i] = acc; +} + +// Shifts the last-`taps`-rows window [taps, c] in place past `rows`, n rows `stride` apart (conv history, n-gram tail). +extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, const __nv_bfloat16* __restrict__ rows, + long long stride, int n, int taps, int c) { + const int ch = blockIdx.x * blockDim.x + threadIdx.x; + if (ch >= c) return; + __nv_bfloat16 next[16]; + for (int j = 0; j < taps; ++j) { + const int src = j + n; // row src of [window (taps rows); new rows] + next[j] = src < taps ? win[static_cast(src) * c + ch] + : rows[static_cast(src - taps) * stride + ch]; + } + for (int j = 0; j < taps; ++j) win[static_cast(j) * c + ch] = next[j]; +} + +// ---- the DeltaNet chain over a chunk: flashnext_gdn.cu's chain cut into front, rows and norm, with its bits ---- + +namespace fn_gdn { +constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; +constexpr int C = 2 * NK * DK + NV * DV; +constexpr int PW = C + NV * DV + 2 * NV; + +static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +static __device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } + +static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, int r, int c) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]); + acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x; + } + return bf(acc / (1.0f + expf(-acc))); +} +} // namespace fn_gdn + +// Program (row, unit), 256 threads: key head u's normed q and k (u < NK), else a value head's v, decay g and beta. +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn, + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + using namespace fn_gdn; + const int r = blockIdx.x, u = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; + if (u < NK) { + __shared__ float x[2][DK]; + const int hk = u; + x[t / DK][t % DK] = conv_act(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK)); + __syncthreads(); + if (warp < 2) { + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[warp][lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); + float* dst = (warp == 0 ? qn : kn) + ((size_t)r * NK + hk) * DK; +#pragma unroll + for (int i = 0; i < 4; ++i) dst[lane * 4 + i] = v4[i] * inv; + } + } else { + const int hv = u - NK; + if (t < DV) v[(size_t)r * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); + else if (t == DV) { + const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]); + const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]); + gb[((size_t)r * NV + hv) * 2] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gb[((size_t)r * NV + hv) * 2 + 1] = bf(sigmoidf_(b)); + } + } +} + +// The chain through the chunk from state_in to state_out (in place, or apart for a rollback): a warp 4 state rows. +static __device__ __forceinline__ void gdn_rows( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + using namespace fn_gdn; + const int hv = blockIdx.x, hk = hv / (NV / NK), warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int row0 = blockIdx.y * 16 + warp * 4; + const size_t sbase = (size_t)hv * DV * DK; + float s[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(row0 + j) * DK + lane * 4 + i]; + float kk[4], qq[4], vv[4], g = 0.0f, beta = 0.0f; + auto load = [&](int r, float (&k_)[4], float (&q_)[4], float (&v_)[4], float& g_, float& b_) { + const float4 k4 = *reinterpret_cast(kn + ((size_t)r * NK + hk) * DK + lane * 4); + const float4 q4 = *reinterpret_cast(qn + ((size_t)r * NK + hk) * DK + lane * 4); + const float4 v4 = *reinterpret_cast(v + (size_t)r * NV * DV + hv * DV + row0); + k_[0] = k4.x; k_[1] = k4.y; k_[2] = k4.z; k_[3] = k4.w; + q_[0] = q4.x; q_[1] = q4.y; q_[2] = q4.z; q_[3] = q4.w; + v_[0] = v4.x; v_[1] = v4.y; v_[2] = v4.z; v_[3] = v4.w; + g_ = gb[((size_t)r * NV + hv) * 2]; + b_ = gb[((size_t)r * NV + hv) * 2 + 1]; + }; + if (rows > 0) load(0, kk, qq, vv, g, beta); + for (int r = 0; r < rows; ++r) { + float nk[4], nq[4], nv[4], ng = 0.0f, nb = 0.0f; + if (r + 1 < rows) load(r + 1, nk, nq, nv, ng, nb); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vv[j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } + float o4[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o4[j] = bf(warp_sum(o)); + } + if (lane == 0) *reinterpret_cast(ys_out + ((size_t)r * NV + hv) * DV + row0) = make_float4(o4[0], o4[1], o4[2], o4[3]); + if (r + 1 < rows) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + kk[i] = nk[i]; + qq[i] = nq[i]; + vv[i] = nv[i]; + } + g = ng; + beta = nb; + } + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(row0 + j) * DK + lane * 4 + i] = s[j][i]; +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) { + gdn_rows(qn, kn, v, gb, state, state, rows, ys_out); +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + gdn_rows(qn, kn, v, gb, state_in, state_out, rows, ys_out); +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( + const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, + float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { + using namespace fn_gdn; + const int r = blockIdx.x, hv = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float ys[DV]; + __shared__ float rinv; + ys[t] = ys_in[((size_t)r * NV + hv) * DV + t]; + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t])); + const float z = __bfloat162float(P[(size_t)r * PW + C + hv * DV + t]); + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; +} + +// ---- sparse attention (QSA): each row's blocks ---------------------------------------------------------------------- + +// attention._select's order-preserving key of a block score: negative floats flipped, positive ones topped. +__device__ __forceinline__ uint32_t fn_qsa_key(float v) { + const uint32_t bits = __float_as_uint(v); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// An inclusive block-wide sum of one int a thread (blockDim.x a multiple of 32, at most 1024); *total gets the sum. +__device__ __forceinline__ int fn_block_scan(int v, int* warp_sums, int* total) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps = blockDim.x >> 5; + for (int o = 1; o < 32; o <<= 1) { + const int n = __shfl_up_sync(0xffffffffu, v, o); + if (lane >= o) v += n; + } + if (lane == 31) warp_sums[warp] = v; + __syncthreads(); + if (warp == 0) { + int w = lane < warps ? warp_sums[lane] : 0; + for (int o = 1; o < 32; o <<= 1) { + const int n = __shfl_up_sync(0xffffffffu, w, o); + if (lane >= o) w += n; + } + if (lane < warps) warp_sums[lane] = w; + } + __syncthreads(); + if (warp > 0) v += warp_sums[warp - 1]; + *total = warp_sums[warps - 1]; + __syncthreads(); + return v; +} + +// attention._select and _select_tiles, a row a block: a row past `top` complete blocks lists the best, then its tail. +extern "C" __global__ void __launch_bounds__(1024) fn_qsa_select(const float* __restrict__ sc, + const int* __restrict__ pos0, int* __restrict__ ids, + int* __restrict__ nkr, int* __restrict__ spr, int nb, + int idw, int top, int ratio) { + __shared__ int hist[256]; + __shared__ int warp_sums[32]; + __shared__ uint32_t cut_s; + __shared__ int need_s; + const int r = blockIdx.x; + const int end = *pos0 + r + 1; + const int complete = end / ratio; + if (complete <= top) { + if (threadIdx.x == 0) { + nkr[r] = end; + spr[r] = 0; + } + return; + } + const float* row = sc + static_cast(r) * nb; + int* out = ids + static_cast(r) * idw; + uint32_t cut = 0; + int need = top; + for (int p = 0; p < 4; ++p) { + const int shift = 24 - 8 * p; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int b = threadIdx.x; b < complete; b += blockDim.x) { + const uint32_t k = fn_qsa_key(row[b]); + if (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); + } + __syncthreads(); + if (threadIdx.x == 0) { + int above = 0; // keys under the prefix with a larger byte + int d = 255; + for (; d > 0; --d) { + if (above + hist[d] >= need) break; + above += hist[d]; + } + need_s = need - above; + cut_s = cut | (static_cast(d) << shift); + } + __syncthreads(); + cut = cut_s; + need = need_s; + __syncthreads(); + } + // every key above the cut, then the first `need` keys equal to it, placed in block order + int seen = 0; + int placed = 0; + for (int t0 = 0; t0 < complete && placed < top; t0 += blockDim.x) { + const int b = t0 + threadIdx.x; + const bool ok = b < complete; + const uint32_t k = ok ? fn_qsa_key(row[b]) : 0u; + const bool above = ok && k > cut; + const bool equal = ok && k == cut; + int eq_total = 0; + const int rank = seen + fn_block_scan(equal ? 1 : 0, warp_sums, &eq_total); + const bool chosen = above || (equal && rank <= need); + int ch_total = 0; + const int place = placed + fn_block_scan(chosen ? 1 : 0, warp_sums, &ch_total) - 1; + if (chosen) + for (int j = 0; j < ratio; ++j) out[place * ratio + j] = b * ratio + j; + seen += eq_total; + placed += ch_total; + } + if (threadIdx.x < ratio) { + const int tail = ratio * complete + threadIdx.x; + if (tail < end) out[top * ratio + threadIdx.x] = tail; + } + if (threadIdx.x == 0) { + nkr[r] = top * ratio + end - ratio * complete; + spr[r] = 1; + } +} diff --git a/zig/kernels/cuda/flashnext_gdn.cu b/zig/kernels/cuda/flashnext_gdn.cu new file mode 100644 index 0000000000..681528a89a --- /dev/null +++ b/zig/kernels/cuda/flashnext_gdn.cu @@ -0,0 +1,142 @@ +// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm, without FMA contraction. + +#include +#include + +constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; + +static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +static __device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } + +static __device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float* vrow, int warp, float g, + float beta) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vrow[warp * 4 + j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } +} + +extern "C" __global__ void __launch_bounds__(1024) flashnext_gdn_chain( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ state_in, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ norm_w, float eps, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ state_out) { + constexpr int C = 2 * NK * DK + NV * DV; + constexpr int PW = C + NV * DV + 2 * NV; + const int hv = blockIdx.x, hk = hv / (NV / NK); + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK], vs[DV], ys[DV]; + __shared__ float gates[2], rinv; + int c = -1; + if (t < DK) c = hk * DK + t; + else if (t < 2 * DK) c = NK * DK + hk * DK + (t - DK); + else if (t < 2 * DK + DV) c = 2 * NK * DK + hv * DV + (t - 2 * DK); + float s[4][4]; + const size_t sbase = (size_t)hv * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + const __nv_bfloat16* pz = P + C + hv * DV + (t < DV ? t : 0); + const __nv_bfloat16* pb = P + C + NV * DV + hv; + const __nv_bfloat16* pa = pb + NV; + for (int r = 0; r < rows; ++r) { + if (c >= 0) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]); + acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x; + } + const float act = bf(acc / (1.0f + expf(-acc))); + if (t < DK) qs[t] = act; + else if (t < 2 * DK) ks[t - DK] = act; + else vs[t - 2 * DK] = act; + } + __syncthreads(); + if (warp < 2) { + float* x = warp == 0 ? qs : ks; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); + __syncwarp(); +#pragma unroll + for (int i = 0; i < 4; ++i) x[lane * 4 + i] = v4[i] * inv; + } else if (warp == 2 && lane == 0) { + const float b = __bfloat162float(pb[(size_t)r * PW]); + const float a = __bfloat162float(pa[(size_t)r * PW]); + gates[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gates[1] = bf(sigmoidf_(b)); + } + __syncthreads(); + const float g = gates[0], beta = gates[1]; + float kk[4], qq[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + kk[i] = ks[lane * 4 + i]; + qq[i] = qs[lane * 4 + i]; + } + update(s, kk, vs, warp, g, beta); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o = warp_sum(o); + if (lane == 0) ys[warp * 4 + j] = bf(o); + } + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + if (t < DV) { + const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t])); + const float z = __bfloat162float(pz[(size_t)r * PW]); + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; + } + __syncthreads(); + } + if (state_out != nullptr) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; + } +} diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu new file mode 100644 index 0000000000..2865adde56 --- /dev/null +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -0,0 +1,449 @@ +// Flash Next's hyper-connection and n-gram glue: each row's own fixed-order fp32 sums, IEEE division and square root. + +// glue.py and qmm.py's arithmetic and bf16 rounding points, the 4-bit dots on mma tiles with K in ascending order. + +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace fn_hc { + +constexpr int TILE_N = 64; // columns a stored weight tile holds ([N/64][K/32][64][4] words) + +__device__ __forceinline__ float bsig(float x) { return rbf(tf_sigmoid(x)); } +__device__ __forceinline__ float bsilu(float x) { return rbf(tf_silu(x)); } + +// bf16x2 of nibbles i (low half) and i + 1 (high half) of an MLX word: the integers 0..15, exact. +__device__ __forceinline__ uint32_t nibbles(uint32_t w, int i) { + const float lo = static_cast((w >> (4 * i)) & 0xFu), hi = static_cast((w >> (4 * i + 4)) & 0xFu); + return (__float_as_uint(lo) >> 16) | (__float_as_uint(hi) & 0xFFFF0000u); +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + return static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(lo))) | + (static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(hi))) << 16); +} + +// A [16, 32] bf16 tile at `x` (row stride `ld` elements) as the two k16 steps' A fragments. +__device__ __forceinline__ void a_frags(const __nv_bfloat16* x, int ld, uint32_t (&a)[2][4]) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int s = 0; s < 2; ++s) { + const __nv_bfloat16* p = x + 16 * s + 2 * t; + a[s][0] = *reinterpret_cast(p + g * ld); + a[s][1] = *reinterpret_cast(p + (g + 8) * ld); + a[s][2] = *reinterpret_cast(p + g * ld + 8); + a[s][3] = *reinterpret_cast(p + (g + 8) * ld + 8); + } +} + +// p = x[16, 32] . q[32, 8]: one column's four words (32 nibbles), two chained k16 steps from zero. +__device__ __forceinline__ void group_dot(float (&p)[4], const uint32_t (&a)[2][4], uint4 words) { + const int t = (threadIdx.x & 31) & 3; + qmm_frag::mma0(p, a[0], nibbles(words.x, 2 * t), nibbles(words.y, 2 * t)); + qmm_frag::mma(p, a[1], nibbles(words.z, 2 * t), nibbles(words.w, 2 * t)); +} + +// A block's fixed-order sum: each warp's butterfly, then the warps' sums in warp order. +__device__ __forceinline__ float block_sum(float v, float* scratch) { + v = warp_sum(v); + const int warps = blockDim.x >> 5; + if ((threadIdx.x & 31) == 0) scratch[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < warps; ++w) total = __fadd_rn(total, scratch[w]); + __syncthreads(); + return total; +} + +} // namespace fn_hc + +using namespace fn_hc; + +// glue._hc_writeback, program (row, 256 columns): MODE 0 sums only, 2 MoE slots (fp32 y and weights), 4 K slices. +template +__device__ __forceinline__ void hc_writeback(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* __restrict__ PSS, + const float* __restrict__ BR, const __nv_bfloat16* __restrict__ INJ, + const float* __restrict__ Y, const float* __restrict__ WTS, long long RS, int D, + int S, int TOPK, int SLOTS, int WORLD) { + constexpr int MAXS = 8; + __shared__ float red[8][MAXS]; + const int r = blockIdx.x, c = blockIdx.y, nc = D / 256; + const int d = c * 256 + threadIdx.x; + float branch = 0.0f; + if constexpr (MODE == 4) { + float acc = BR[static_cast(r) * D + d]; + for (int k = 1; k < WORLD; ++k) acc = __fadd_rn(acc, BR[k * RS + static_cast(r) * D + d]); + branch = rbf(acc); + } else if constexpr (MODE == 2) { + float acc = 0.0f; + for (int k = 0; k <= TOPK; ++k) { + const long long at = static_cast(r) * SLOTS + k; + acc = __fmaf_rn(Y[at * D + d], WTS[at], acc); + } + branch = rbf(acc); + } + const int warp = threadIdx.x >> 5; + for (int s = 0; s < S; ++s) { + const long long at = static_cast(r) * S * D + static_cast(s) * D + d; + float hv = bf(H[at]); + if constexpr (MODE != 0) { + hv = rbf(__fmaf_rn(branch, bf(INJ[r * S + s]), hv)); + HOUT[at] = __float2bfloat16_rn(hv); + } + const float ss = warp_sum(__fmul_rn(hv, hv)); + if ((threadIdx.x & 31) == 0) red[warp][s] = ss; + } + __syncthreads(); + if (threadIdx.x < S) { + float total = 0.0f; + for (int w = 0; w < 8; ++w) total = __fadd_rn(total, red[w][threadIdx.x]); + PSS[(static_cast(r) * nc + c) * S + threadIdx.x] = total; + } +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback0(const __nv_bfloat16* H, float* PSS, int D, int S) { + hc_writeback<0>(H, nullptr, PSS, nullptr, nullptr, nullptr, nullptr, 0, D, S, 0, 1, 1); +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_moe(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, + const __nv_bfloat16* INJ, const float* Y, + const float* WTS, int D, int S, int TOPK, int SLOTS) { + hc_writeback<2>(H, HOUT, PSS, nullptr, INJ, Y, WTS, 0, D, S, TOPK, SLOTS, 1); +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_slices(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, + const float* BR, const __nv_bfloat16* INJ, + long long RS, int D, int S, int WORLD) { + hc_writeback<4>(H, HOUT, PSS, BR, INJ, nullptr, nullptr, RS, D, S, 0, 1, WORLD); +} + +// qmm._qmm_hcdown, block (16 rows, 64 columns, K slice): the slice's normed rows times the down rows, fp32 to PART. +extern "C" __global__ void __launch_bounds__(128) fn_hc_down(const __nv_bfloat16* __restrict__ H, const float* __restrict__ PSS, + const float* __restrict__ SCALE, __nv_bfloat16* __restrict__ NORMED, + const uint4* __restrict__ W, const __nv_bfloat16* __restrict__ S, + const __nv_bfloat16* __restrict__ B, float* __restrict__ PART, + int M, float eps, int N, int K, int D, int NC, int SS, int SK) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + __shared__ float rinv[16]; + const int kg = K / 32, per = kg / SK, width = per * 32, ld = width + 8; + __nv_bfloat16* xt = reinterpret_cast<__nv_bfloat16*>(smem_raw); // [16][width + 8] the slice's normed rows + float* xsum = reinterpret_cast(smem_raw + 16 * ld * 2); // [16][per] their 32-group sums + const int m0 = blockIdx.x * 16, n0 = blockIdx.y * TILE_N, slice = blockIdx.z, k0 = slice * width; + const int st = k0 / D; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + if (threadIdx.x < 16) { + const int r = m0 + threadIdx.x; + float total = 0.0f; + if (r < M) + for (int c = 0; c < NC; ++c) total = __fadd_rn(total, PSS[(static_cast(r) * NC + c) * SS + st]); + rinv[threadIdx.x] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); + } + __syncthreads(); + const int vecs = width / 8; + for (int e = threadIdx.x; e < 16 * vecs; e += blockDim.x) { + const int rr = e / vecs, v = e % vecs, r = m0 + rr; + uint4 out = make_uint4(0, 0, 0, 0); + if (r < M) { + const long long at = static_cast(r) * K + k0 + 8 * v; + const uint4 raw = *reinterpret_cast(H + at); + const __nv_bfloat16* hv = reinterpret_cast(&raw); + float x[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) x[i] = rbf(__fmul_rn(__fmul_rn(bf(hv[i]), rinv[rr]), SCALE[k0 + 8 * v + i])); + out = make_uint4(pack(x[0], x[1]), pack(x[2], x[3]), pack(x[4], x[5]), pack(x[6], x[7])); + if (blockIdx.y == 0) *reinterpret_cast(NORMED + at) = out; + } + *reinterpret_cast(xt + rr * ld + 8 * v) = out; + } + __syncthreads(); + for (int e = threadIdx.x; e < 16 * per; e += blockDim.x) { + const int rr = e / per, i = e % per; + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, bf(xt[rr * ld + 32 * i + k])); + xsum[rr * per + i] = sum; + } + __syncthreads(); + float acc[2][4] = {}; + const uint4* tile = W + static_cast(n0 / TILE_N) * kg * TILE_N; +#pragma unroll 2 + for (int i = 0; i < per; ++i) { + const int grp = slice * per + i; + uint32_t a[2][4]; + a_frags(xt + 32 * i, ld, a); + const float xs0 = xsum[g * per + i], xs1 = xsum[(g + 8) * per + i]; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int local = warp * 16 + j * 8 + g; + float p[4]; + group_dot(p, a, tile[static_cast(grp) * TILE_N + local]); + const int n = n0 + warp * 16 + j * 8 + 2 * t; + const float s0 = n < N ? bf(S[static_cast(grp) * N + n]) : 0.0f; + const float s1 = n + 1 < N ? bf(S[static_cast(grp) * N + n + 1]) : 0.0f; + const float b0 = n < N ? bf(B[static_cast(grp) * N + n]) : 0.0f; + const float b1 = n + 1 < N ? bf(B[static_cast(grp) * N + n + 1]) : 0.0f; + acc[j][0] = __fmaf_rn(xs0, b0, __fmaf_rn(p[0], s0, acc[j][0])); + acc[j][1] = __fmaf_rn(xs0, b1, __fmaf_rn(p[1], s1, acc[j][1])); + acc[j][2] = __fmaf_rn(xs1, b0, __fmaf_rn(p[2], s0, acc[j][2])); + acc[j][3] = __fmaf_rn(xs1, b1, __fmaf_rn(p[3], s1, acc[j][3])); + } + } +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int n = n0 + warp * 16 + j * 8 + 2 * t; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = m0 + g + 8 * h; + if (r >= M) continue; + float* out = PART + (static_cast(slice) * M + r) * N; + if (n < N) out[n] = acc[j][2 * h]; + if (n + 1 < N) out[n + 1] = acc[j][2 * h + 1]; + } + } +} + +// glue._hc_reduce_act, a thread a column: SK (<= 32) slices in order, bf16 SiLU and its sums, the inject gates. +extern "C" __global__ void __launch_bounds__(1024) fn_hc_reduce_act(const float* __restrict__ PART, __nv_bfloat16* __restrict__ ACT, + float* __restrict__ XS, __nv_bfloat16* __restrict__ INJ, + int SK, int M, int S, int LOW, int HAS_INJ, int NDN) { + __shared__ float a_s[1024]; + const int r = blockIdx.x, i = threadIdx.x; + if (i < LOW + (HAS_INJ ? S : 0)) { + float part[32]; +#pragma unroll + for (int s = 0; s < 32; ++s) part[s] = s < SK ? PART[(static_cast(s) * M + r) * NDN + i] : 0.0f; + float v = part[0]; +#pragma unroll + for (int s = 1; s < 32; ++s) + if (s < SK) v = __fadd_rn(v, part[s]); + v = rbf(__fdiv_rn(rbf(v), static_cast(S))); + if (i < LOW) { + const float a = bsilu(v); + ACT[static_cast(r) * LOW + i] = __float2bfloat16_rn(a); + a_s[i] = a; + } else { + INJ[static_cast(r) * S + (i - LOW)] = __float2bfloat16_rn(__fmul_rn(2.0f, bsig(v))); + } + } + __syncthreads(); + if (i < LOW / 32) { + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, a_s[i * 32 + k]); + XS[static_cast(r) * (LOW / 32) + i] = sum; + } +} + +// qmm._qmm_upmix, block (16 rows, 32 dims), two warps a stream (SS 4): the up rows, then the streams mixed in order. +extern "C" __global__ void __launch_bounds__(256) fn_hc_upmix(const __nv_bfloat16* __restrict__ X, const float* __restrict__ XS, + const uint4* __restrict__ W, const __nv_bfloat16* __restrict__ S, + const __nv_bfloat16* __restrict__ B, const __nv_bfloat16* __restrict__ NORMED, + __nv_bfloat16* __restrict__ MIXED, float* __restrict__ XSM, + int M, int N, int K, int D, int SS) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + const int kg = K / 32; + const int ld = K + 8; + __nv_bfloat16* xt = reinterpret_cast<__nv_bfloat16*>(smem_raw); // [16][K + 8] + float* part = reinterpret_cast(smem_raw + 16 * ld * 2); // [SS][16][32] each stream's terms + const int m0 = blockIdx.x * 16, j = blockIdx.y; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + for (int e = threadIdx.x; e < 16 * (K / 8); e += blockDim.x) { + const int rr = e / (K / 8), v = e % (K / 8), r = m0 + rr; + *reinterpret_cast(xt + rr * ld + 8 * v) = r < M ? *reinterpret_cast(X + static_cast(r) * K + 8 * v) : make_uint4(0, 0, 0, 0); + } + __syncthreads(); + const int st = warp >> 1; + const int row0 = st * D + j * 32; + const int q0 = 2 * (warp & 1); // this warp's two of the stream's four 8-column tiles + float acc[4][4] = {}; +#pragma unroll 5 + for (int grp = 0; grp < kg; ++grp) { + uint32_t a[2][4]; + a_frags(xt + grp * 32, ld, a); + const float xs0 = m0 + g < M ? XS[static_cast(m0 + g) * kg + grp] : 0.0f; + const float xs1 = m0 + g + 8 < M ? XS[static_cast(m0 + g + 8) * kg + grp] : 0.0f; +#pragma unroll + for (int q = q0; q < q0 + 2; ++q) { + const int col = row0 + q * 8 + g; + float p[4]; + group_dot(p, a, W[(static_cast(col / TILE_N) * kg + grp) * TILE_N + col % TILE_N]); + const int n = row0 + q * 8 + 2 * t; + const float s0 = bf(S[static_cast(grp) * N + n]), s1 = bf(S[static_cast(grp) * N + n + 1]); + const float b0 = bf(B[static_cast(grp) * N + n]), b1 = bf(B[static_cast(grp) * N + n + 1]); + acc[q][0] = __fmaf_rn(xs0, b0, __fmaf_rn(p[0], s0, acc[q][0])); + acc[q][1] = __fmaf_rn(xs0, b1, __fmaf_rn(p[1], s1, acc[q][1])); + acc[q][2] = __fmaf_rn(xs1, b0, __fmaf_rn(p[2], s0, acc[q][2])); + acc[q][3] = __fmaf_rn(xs1, b1, __fmaf_rn(p[3], s1, acc[q][3])); + } + } +#pragma unroll + for (int q = q0; q < q0 + 2; ++q) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int rr = g + 8 * (e >> 1), dd = q * 8 + 2 * t + (e & 1), r = m0 + rr; + const float nv = r < M ? bf(NORMED[static_cast(r) * N + row0 + q * 8 + 2 * t + (e & 1)]) : 0.0f; + part[(st * 16 + rr) * 32 + dd] = rbf(__fmul_rn(bsig(rbf(acc[q][e])), nv)); + } + __syncthreads(); + __shared__ float mixed_s[16][32]; + for (int e = threadIdx.x; e < 16 * 32; e += blockDim.x) { + const int rr = e >> 5, dd = e & 31; + float total = 0.0f; + for (int s = 0; s < SS; ++s) total = __fadd_rn(total, part[(s * 16 + rr) * 32 + dd]); + const float m = rbf(__fdiv_rn(total, static_cast(SS))); + mixed_s[rr][dd] = m; + if (m0 + rr < M) MIXED[static_cast(m0 + rr) * D + j * 32 + dd] = __float2bfloat16_rn(m); + } + __syncthreads(); + if (threadIdx.x < 16 && m0 + threadIdx.x < M) { + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, mixed_s[threadIdx.x][k]); + XSM[static_cast(m0 + threadIdx.x) * (D / 32) + j] = sum; + } +} + +// glue._ple_gate, a block a row: per stream normed keys and queries, their dot's signed root gating the values. +extern "C" __global__ void __launch_bounds__(512) fn_ple_gate(const __nv_bfloat16* __restrict__ KEYS, const __nv_bfloat16* __restrict__ VALS, + const __nv_bfloat16* __restrict__ H, const float* __restrict__ NK, + const float* __restrict__ NQ, __nv_bfloat16* __restrict__ GATED, + float* __restrict__ PSS, float eps, int D, int S) { + __shared__ float scratch[32]; + const int r = blockIdx.x; + const float fd = static_cast(D); + for (int s = 0; s < S; ++s) { + const long long row = static_cast(r) * S * D + static_cast(s) * D; + float kss = 0.0f, qss = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float kv = bf(KEYS[row + d]), hv = bf(H[row + d]); + kss = __fmaf_rn(kv, kv, kss); + qss = __fmaf_rn(hv, hv, qss); + } + kss = block_sum(kss, scratch); + qss = block_sum(qss, scratch); + const float krinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(kss, fd), eps))); + const float qrinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(qss, fd), eps))); + float dot = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float kn = rbf(__fmul_rn(__fmul_rn(bf(KEYS[row + d]), krinv), NK[s * D + d])); + const float qn = rbf(__fmul_rn(__fmul_rn(bf(H[row + d]), qrinv), NQ[s * D + d])); + dot = __fadd_rn(dot, rbf(__fmul_rn(kn, qn))); + } + dot = block_sum(dot, scratch); + float gate = rbf(__fdiv_rn(dot, __fsqrt_rn(fd))); + const float mag = __fsqrt_rn(fmaxf(fabsf(gate), 1e-6f)); + gate = rbf(gate > 0.0f ? mag : (gate < 0.0f ? -mag : 0.0f)); + const float sg = bsig(gate); + float gss = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float gv = rbf(__fmul_rn(sg, bf(VALS[static_cast(r) * D + d]))); + GATED[row + d] = __float2bfloat16_rn(gv); + gss = __fmaf_rn(gv, gv, gss); + } + gss = block_sum(gss, scratch); + if (threadIdx.x == 0) PSS[static_cast(r) * S + s] = gss; + } +} + +// glue._ple_conv, program (row, 512 columns): the dilated conv over [TAIL; rows], bf16 SiLU, the adds, NROW kept. +extern "C" __global__ void __launch_bounds__(512) fn_ple_conv(const __nv_bfloat16* __restrict__ GATED, const float* __restrict__ PSS, + const float* __restrict__ NC, const __nv_bfloat16* __restrict__ TAIL, + const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* H, + __nv_bfloat16* HOUT, __nv_bfloat16* __restrict__ NROW, float eps, + int D, int S, int TAPS, int DIL) { + const int r = blockIdx.x; + const int w = S * D, nt = (TAPS - 1) * DIL; + const int e = blockIdx.y * blockDim.x + threadIdx.x; + if (e >= w) return; + const int s = e / D; + const float fd = static_cast(D); + float acc = 0.0f; + for (int t = 0; t < TAPS; ++t) { + const int at = r + t * DIL; + float xv; + if (at < nt) { + xv = bf(TAIL[static_cast(at) * w + e]); + } else { + const int rr = at - nt; + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(PSS[rr * S + s], fd), eps))); + xv = rbf(__fmul_rn(__fmul_rn(bf(GATED[static_cast(rr) * w + e]), rinv), NC[e])); + } + acc = __fmaf_rn(bf(CW[static_cast(e) * TAPS + t]), xv, acc); + } + const float act = bsilu(rbf(acc)); + const float g = bf(GATED[static_cast(r) * w + e]); + const float ple = rbf(__fadd_rn(g, act)); + const long long at = static_cast(r) * w + e; + HOUT[at] = __float2bfloat16_rn(__fadd_rn(bf(H[at]), ple)); + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(PSS[r * S + s], fd), eps))); + NROW[at] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(g, rinv), NC[e])); +} + +// glue._add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. +extern "C" __global__ void fn_add_streams(const __nv_bfloat16* __restrict__ E, const __nv_bfloat16* __restrict__ HS, + __nv_bfloat16* __restrict__ OUT, int D, int S) { + const int r = blockIdx.x, d = blockIdx.y * blockDim.x + threadIdx.x; + if (d >= D) return; + const float e = bf(E[static_cast(r) * D + d]); + for (int s = 0; s < S; ++s) { + const long long at = (static_cast(r) * S + s) * D + d; + OUT[at] = __float2bfloat16_rn(__fadd_rn(e, bf(HS[at]))); + } +} + +// Dense HC uses the same stored stream sums and FP32 gamma before one BF16 normalization rounding. +extern "C" __global__ void fn_hc_normalize(const __nv_bfloat16* H, const float* PSS, const float* SCALE, + __nv_bfloat16* OUT, float eps, int D, int SS) { + __shared__ float rinv; + const int r = blockIdx.x, stream = blockIdx.y, nc = D / 256; + if (threadIdx.x == 0) { + float total = 0.0f; + for (int c = 0; c < nc; ++c) total = __fadd_rn(total, PSS[(static_cast(r) * nc + c) * SS + stream]); + rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); + } + __syncthreads(); + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const int col = stream * D + d; + const long long at = static_cast(r) * SS * D + col; + OUT[at] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(bf(H[at]), rinv), SCALE[col])); + } +} + +// Dense up-projection values are already BF16: sigmoid, product and stream mean keep their original rounding seams. +extern "C" __global__ void fn_hc_dense_mix(const __nv_bfloat16* UP, const __nv_bfloat16* NORMED, + __nv_bfloat16* MIXED, float* XSM, int D, int SS) { + __shared__ float row[32]; + const int r = blockIdx.x, group = blockIdx.y, d = group * 32 + threadIdx.x; + float total = 0.0f; + for (int stream = 0; stream < SS; ++stream) { + const long long at = (static_cast(r) * SS + stream) * D + d; + const float value = rbf(__fmul_rn(bsig(bf(UP[at])), bf(NORMED[at]))); + total = __fadd_rn(total, value); + } + const float mixed = rbf(__fdiv_rn(total, static_cast(SS))); + MIXED[static_cast(r) * D + d] = __float2bfloat16_rn(mixed); + row[threadIdx.x] = mixed; + __syncthreads(); + if (threadIdx.x == 0) { + float sum = 0.0f; + for (int i = 0; i < 32; ++i) sum = __fadd_rn(sum, row[i]); + XSM[static_cast(r) * (D / 32) + group] = sum; + } +} + +// Shared expert uses its own BF16 projection outputs and rounding boundaries. +extern "C" __global__ void fn_shared_swiglu(const __nv_bfloat16* GATE, const __nv_bfloat16* UP, + __nv_bfloat16* ACT, int W) { + const int r = blockIdx.x, j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= W) return; + const long long at = static_cast(r) * W + j; + ACT[at] = __float2bfloat16_rn(__fmul_rn(bsilu(bf(GATE[at])), bf(UP[at]))); +} + +// Write only the final reserved slot; routed slots retain their FP32 partial sums. +extern "C" __global__ void fn_shared_scatter(const float* SHARED, float* Y, int D, int SLOTS) { + const int r = blockIdx.x, j = blockIdx.y * blockDim.x + threadIdx.x; + if (j < D) Y[(static_cast(r) * SLOTS + SLOTS - 1) * D + j] = SHARED[static_cast(r) * D + j]; +} diff --git a/zig/kernels/cuda/flashnext_mtp.cu b/zig/kernels/cuda/flashnext_mtp.cu new file mode 100644 index 0000000000..1ac2e6fa4b --- /dev/null +++ b/zig/kernels/cuda/flashnext_mtp.cu @@ -0,0 +1,70 @@ +// Flash Next's MTP drafts (cuda_mtp.zig): a score row's best value, its first id and the row's log-sum-exp. + +#include +#include +#include + +namespace fn_mtp { + +static __device__ __forceinline__ void better(float& v, int& i, float ov, int oi) { + if (ov > v || (ov == v && oi < i)) { + v = ov; + i = oi; + } +} + +} // namespace fn_mtp + +// Program (row), 1024 threads: out [rows, 4] fp32 holds the best of a row's first V values, its log-sum-exp, its id, 0. +extern "C" __global__ void __launch_bounds__(1024) fn_row_best(const __nv_bfloat16* __restrict__ L, int V, long long ld, + float* __restrict__ out) { + using namespace fn_mtp; + const int r = blockIdx.x, t = threadIdx.x, lane = t & 31, warp = t >> 5; + const __nv_bfloat16* row = L + static_cast(r) * ld; + float v = -CUDART_INF_F; + int at = 0x7fffffff; + for (int i = t; i < V; i += blockDim.x) { + const float x = __bfloat162float(row[i]); + if (x > v) { + v = x; + at = i; + } + } + for (int o = 16; o; o >>= 1) better(v, at, __shfl_xor_sync(0xffffffffu, v, o), __shfl_xor_sync(0xffffffffu, at, o)); + __shared__ float sv[32]; + __shared__ int si[32]; + __shared__ float top; + __shared__ int top_at; + if (lane == 0) { + sv[warp] = v; + si[warp] = at; + } + __syncthreads(); + if (warp == 0) { + v = lane < (int)(blockDim.x >> 5) ? sv[lane] : -CUDART_INF_F; + at = lane < (int)(blockDim.x >> 5) ? si[lane] : 0x7fffffff; + for (int o = 16; o; o >>= 1) better(v, at, __shfl_xor_sync(0xffffffffu, v, o), __shfl_xor_sync(0xffffffffu, at, o)); + if (lane == 0) { + top = v; + top_at = at; + } + } + __syncthreads(); + const float m = top; + float sum = 0.0f; + for (int i = t; i < V; i += blockDim.x) sum += expf(__bfloat162float(row[i]) - m); + for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o); + __shared__ float ss[32]; + if (lane == 0) ss[warp] = sum; + __syncthreads(); + if (warp == 0) { + sum = lane < (int)(blockDim.x >> 5) ? ss[lane] : 0.0f; + for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o); + if (lane == 0) { + out[r * 4 + 0] = m; + out[r * 4 + 1] = m + logf(sum); + out[r * 4 + 2] = __int_as_float(top_at); + out[r * 4 + 3] = 0.0f; + } + } +} diff --git a/zig/kernels/cuda/norms.cu b/zig/kernels/cuda/norms.cu index be264f6d17..73902985bd 100644 --- a/zig/kernels/cuda/norms.cu +++ b/zig/kernels/cuda/norms.cu @@ -175,6 +175,27 @@ extern "C" __global__ void __launch_bounds__(NT) tf_rms_rows(const void* x, int #undef TF_RMS } +// bf16 rows `x_ld` apart: y = bf16(x * rsqrt(mean(x^2) + eps) * w), then xs = the stored y's sums over each `group`. +extern "C" __global__ void __launch_bounds__(NT) tf_rms_rows_sums(const __nv_bfloat16* __restrict__ x, long long x_ld, + const float* __restrict__ w, __nv_bfloat16* __restrict__ y, + float* __restrict__ xs, int width, int group, float eps) { + const int64_t at = static_cast(blockIdx.x) * x_ld, out = static_cast(blockIdx.x) * width; + float part = 0.f; + for (int i = threadIdx.x; i < width; i += NT) { + const float a = __bfloat162float(x[at + i]); + part = __fmaf_rn(a, a, part); + } + const float inv = inv_rms(part, width, eps); + for (int i = threadIdx.x; i < width; i += NT) + y[out + i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(__bfloat162float(x[at + i]), inv), w[i])); + __syncthreads(); + for (int j = threadIdx.x; j < width / group; j += NT) { + float sum = 0.f; + for (int k = 0; k < group; ++k) sum = __fadd_rn(sum, __bfloat162float(y[out + j * group + k])); + xs[static_cast(blockIdx.x) * (width / group) + j] = sum; + } +} + // res (fp32, in place) += rms(x) * w1 (a sandwich norm; null: res += x), then y = rms(res) * w2 in y's kind. extern "C" __global__ void __launch_bounds__(NT) tf_add_rms_rows(const void* x, int x_kind, float* res, const float* w1, const float* w2, void* y, int y_kind, int width, diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index be724dead8..afbaac8925 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/zig/kernels/cuda/nvfp4_experts.cu @@ -315,6 +315,12 @@ template __global__ void nvfp4_expert_kernel<2, 2, 4>(const __nv_bfloat16* __res 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); + +// One warp per CTA exposes more independent expert tiles for narrow rank-local projections. +template __global__ void nvfp4_expert_kernel<2, 2, 1>(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, 1>(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, 1>(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, 1>(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. diff --git a/zig/kernels/cuda/qmm.cu b/zig/kernels/cuda/qmm.cu new file mode 100644 index 0000000000..4901978a3a --- /dev/null +++ b/zig/kernels/cuda/qmm.cu @@ -0,0 +1,230 @@ +// Device code of src/tensorfold/cuda/kernels/qmm.cu (lines 1-230, comments and ATen includes dropped). + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_qmm { + +using namespace qmm_frag; + +template +__global__ void __launch_bounds__(WM * WN * 32) qmm_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + using T = LaneTile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, per = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = blockIdx.z, g0 = slice * per; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int g, int m0) { + unsigned char* p = stage(s); + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(p + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS / 2; // one stored 64-column tile's group block + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = pw + T::W; + for (int c = tid; c < 2 * (T::S / 16); c += T::THREADS) { + const int which = c / (T::S / 16), off = c % (T::S / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + cp16(ps + which * T::S + off * 16, src); + } + float* px = reinterpret_cast(ps + 2 * T::S); + for (int r = tid; r < BM; r += T::THREADS) + cp4(px + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + }; + + { + float acc[T::MT][T::NT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + auto epilogue = [&](float (&acc_)[T::MT][T::NT][4], const float (&d_)[T::MT][T::NT][4], + const float (&s_)[T::NT][2], const float (&b_)[T::NT][2], const float (&x_)[T::MT][2]) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc_[i][j][e] = __fmaf_rn(x_[i][e >> 1], b_[j][e & 1], __fmaf_rn(d_[i][j][e], s_[j][e & 1], + acc_[i][j][e])); + }; + float dp[T::MT][T::NT][4], sp[T::NT][2], bp[T::NT][2], xp[T::MT][2]; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s, m0); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next, m0); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[T::NT][GS / 32]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) + words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v]; + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const uint32_t b0 = pair(words[j][kt / 2], (kt & 1) * 8); + const uint32_t b1 = pair(words[j][kt / 2], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if (PIPE && kt == 0 && it > 0) epilogue(acc, dp, sp, bp, xp); // the previous group's, between MMAs + } + float sv[T::NT][2], bv[T::NT][2], xv[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + sv[j][0] = __low2float(s2); sv[j][1] = __high2float(s2); + bv[j][0] = __low2float(b2); bv[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[i][0] = px[row]; + xv[i][1] = px[row + 8]; + } + if constexpr (PIPE) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) dp[i][j][e] = d[i][j][e]; + xp[i][0] = xv[i][0]; xp[i][1] = xv[i][1]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + sp[j][0] = sv[j][0]; sp[j][1] = sv[j][1]; bp[j][0] = bv[j][0]; bp[j][1] = bv[j][1]; + } + } else { + epilogue(acc, d, sv, bv, xv); + } + } + if (PIPE && per > 0) epilogue(acc, dp, sp, bp, xp); + wait<0>(); + __syncthreads(); + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ < 900 + __trap(); // no clusters before sm_90: the host never launches this +#else + constexpr int E = T::MT * T::NT * 4; + auto cluster = cooperative_groups::this_cluster(); + float* mine = reinterpret_cast(buf); + if (slice != 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) mine[((i * T::NT + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e]; + } + cluster.sync(); + if (slice == 0) { + for (int peer = 1; peer < SK; ++peer) { + const float* theirs = cluster.map_shared_rank(mine, peer); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[i][j][e] = acc[i][j][e] + theirs[((i * T::NT + j) * 4 + e) * T::THREADS + tid]; + } + } + cluster.sync(); // peers keep their memory until slice 0 has read it + if (slice != 0) return; +#endif + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = acc[i][j][2 * h], v1 = acc[i][j][2 * h + 1]; + if (SK > 1 && !CLUSTER) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else if (F32) { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + } + } + } + } +} + +template +__global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < SK; ++s) acc = acc + part[s * total + i]; + if (F32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} +} // namespace tf_qmm + +// Decode lane: groups of 32, 16 rows, 64 columns, one by four warps, four stages, bf16 out. +template __global__ void tf_qmm::qmm_kernel<32, 16, 64, 1, 4, 4, false, false, false>( + const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, + void*, float*, int, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 65cf08fd22..6ef58c2406 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -32,6 +32,11 @@ __device__ __forceinline__ uint32_t fp8pair(uint32_t w) { return r; } +// UE8M0 has no zero encoding: byte 0 is 2^-127, byte 255 is NaN. +__device__ __forceinline__ float e8m0f(uint8_t b) { + return __uint_as_float(b == 0 ? 0x00400000u : b == 255 ? 0x7FC00000u : uint32_t(b) << 23); +} + __device__ __forceinline__ uint32_t comp(const uint4& v, int c) { return c == 0 ? v.x : c == 1 ? v.y : c == 2 ? v.z : v.w; } __device__ __forceinline__ float e4m3f(uint8_t b) { @@ -193,8 +198,8 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel( s0 = e4m3f(static_cast(sq[j][0] >> (8 * blk))); s1 = e4m3f(static_cast(sq[j][1] >> (8 * blk))); } else { - s0 = __int_as_float(static_cast((sq[j][0] >> (8 * blk)) & 0xFFu) << 23); - s1 = __int_as_float(static_cast((sq[j][1] >> (8 * blk)) & 0xFFu) << 23); + s0 = e8m0f(static_cast(sq[j][0] >> (8 * blk))); + s1 = e8m0f(static_cast(sq[j][1] >> (8 * blk))); } #pragma unroll for (int i = 0; i < T::MT; ++i) @@ -293,6 +298,14 @@ __global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); } +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); + // The instances the Zig launchers load by name (the Python host instantiates these through its launch switch). template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); @@ -343,6 +356,42 @@ template __global__ void qmmf_kernel( template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); -template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); + + + + + + + +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); + } diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu new file mode 100644 index 0000000000..9407493f75 --- /dev/null +++ b/zig/kernels/cuda/route_rows.cu @@ -0,0 +1,167 @@ +// Softmax top-k routing rows for every CUDA family (float64 reference: core/route_ref.zig), a row's bits its own. + +// Router logits on mma tiles, K ascending from zero (a row's sums never depend on the row count), and top-k rows. + +#include +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace tf_route { + +constexpr int NT = 128; +constexpr int MOST = 8; // logits a thread holds: up to NT * MOST experts + +// The block's largest (value, index), lower index on ties; slots alternate by `parity`: one barrier an iteration. +__device__ __forceinline__ void block_best(float& v, int& i, float (*sv)[NT / 32], int (*si)[NT / 32], int parity) { + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, v, off); + const int oi = __shfl_xor_sync(0xffffffffu, i, off); + if (ov > v || (ov == v && oi < i)) { + v = ov; + i = oi; + } + } + const int warp = threadIdx.x >> 5; + if ((threadIdx.x & 31) == 0) { + sv[parity][warp] = v; + si[parity][warp] = i; + } + __syncthreads(); + v = sv[parity][0]; + i = si[parity][0]; + for (int w = 1; w < NT / 32; ++w) + if (sv[parity][w] > v || (sv[parity][w] == v && si[parity][w] < i)) { + v = sv[parity][w]; + i = si[parity][w]; + } +} + +} // namespace tf_route + +using namespace tf_route; + +// OUT[m, e] = x[m] . w[e], a warp of 8 experts: one mma chain over K from zero, STAGES of KC staged (D % 16 == 0). +template +__device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, + float* __restrict__ OUT, int M, int D, int E) { + constexpr int LD = KC + 8, EB = 8 * WARPS, ROWS = 16 + EB; + extern __shared__ __align__(16) unsigned char smem_raw[]; + auto tile = reinterpret_cast<__nv_bfloat16 (*)[ROWS][LD]>(smem_raw); // [STAGES] 16 input rows, then EB expert rows + const int m0 = blockIdx.x * 16, e0 = blockIdx.y * EB; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + const int stages = (D + KC - 1) / KC; + const int xr = min(16, M - m0); // input rows past M stay zero: written once, never staged + for (int v = threadIdx.x; v < STAGES * (16 - xr) * (KC / 8); v += blockDim.x) { + const int per = (16 - xr) * (KC / 8); + *reinterpret_cast(&tile[v / per][xr + (v % per) / (KC / 8)][8 * (v % (KC / 8))]) = make_uint4(0, 0, 0, 0); + } + auto stage = [&](int st) { + const int k0 = st * KC, buf = st % STAGES; + for (int v = threadIdx.x; v < (xr + EB) * (KC / 8); v += blockDim.x) { + const int i = v / (KC / 8), k = k0 + 8 * (v % (KC / 8)); + const bool x_row = i < xr; + const int row = x_row ? i : 16 + i - xr; + const int r = x_row ? m0 + i : e0 + i - xr; + const bool read = k < D && (x_row || r < E); + const __nv_bfloat16* src = x_row ? X + static_cast(r) * x_ld + (read ? k : 0) + : W + static_cast(read ? r : 0) * D + (read ? k : 0); + qmm_frag::cp16z(&tile[buf][row][8 * (v % (KC / 8))], src, read); + } + }; + float acc[4] = {}; + for (int st = 0; st < STAGES - 1; ++st) { + if (st < stages) stage(st); + qmm_frag::commit(); + } + for (int st = 0; st < stages; ++st) { + if (st + STAGES - 1 < stages) stage(st + STAGES - 1); + qmm_frag::commit(); + qmm_frag::wait(); + __syncthreads(); + const int buf = st % STAGES; + const int steps = min(KC, D - st * KC) / 16; + for (int kk = 0; kk < steps; ++kk) { + uint32_t a[4], b[2]; + qmm_frag::ldmatrix4(a, &tile[buf][lane % 16][kk * 16 + (lane / 16) * 8]); + qmm_frag::ldmatrix2(b, &tile[buf][16 + warp * 8 + lane % 8][kk * 16 + ((lane / 8) % 2) * 8]); + qmm_frag::mma(acc, a, b[0], b[1]); + } + __syncthreads(); + } + const int e = e0 + warp * 8 + 2 * t; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = m0 + g + 8 * h; + if (r >= M) continue; + if (e < E) OUT[static_cast(r) * E + e] = acc[2 * h]; + if (e + 1 < E) OUT[static_cast(r) * E + e + 1] = acc[2 * h + 1]; + } +} + +// route_logits over 32 experts a block (prompt chunks: the input rows are read once per 32 experts). +extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int M, int D, int E) { + route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); +} + +// route_logits with three stages of 256 in flight (up to 16 rows, where the router rows' reads set its time). +extern "C" __global__ void __launch_bounds__(128) tf_route_logits_few(const __nv_bfloat16* __restrict__ X, int x_ld, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int M, int D, int E) { + route_logits<4, 3, 256>(X, x_ld, W, OUT, M, D, E); +} + +// A row's TOPK largest of NE logits (lower id on ties), bf16(e^(l_k - l_0) / sum); `shared`: slot TOPK gated by NE. +extern "C" __global__ void __launch_bounds__(NT) tf_route_topk(const float* __restrict__ L, int NL, int NE, int TOPK, + int* __restrict__ PICK, float* __restrict__ WTS, int SLOTS, + int shared) { + __shared__ float sv[2][NT / 32]; + __shared__ int si[2][NT / 32]; + const int r = blockIdx.x; + const float* row = L + static_cast(r) * NL; + float v[MOST]; +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int i = threadIdx.x + j * NT; + v[j] = i < NE ? row[i] : -INFINITY; + } + float top = 0.0f, total = 0.0f, mine = 0.0f; + int pick = 0; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int i = threadIdx.x + j * NT; + if (v[j] > bv || (v[j] == bv && i < bi)) { + bv = v[j]; + bi = i; + } + } + block_best(bv, bi, sv, si, k & 1); + if (k == 0) top = bv; +#pragma unroll + for (int j = 0; j < MOST; ++j) + if (threadIdx.x + j * NT == bi) v[j] = -INFINITY; + const float ex = tf_exp(__fsub_rn(bv, top)); + total = __fadd_rn(total, ex); + if (threadIdx.x == k) { + mine = ex; + pick = bi; + } + } + if (threadIdx.x < TOPK) { + PICK[static_cast(r) * SLOTS + threadIdx.x] = pick; + WTS[static_cast(r) * SLOTS + threadIdx.x] = rbf(__fdiv_rn(mine, total)); + } + if (shared && threadIdx.x == 0) { + PICK[static_cast(r) * SLOTS + TOPK] = NE; + WTS[static_cast(r) * SLOTS + TOPK] = rbf(tf_sigmoid(rbf(row[NE]))); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..14aab5d740 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -19,24 +19,94 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. pub const cuda = [_]CudaRecord{ + .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to", .symbol = "fn_gdn_rows_to", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = .{ .{ .regs = 22, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_down", .symbol = "fn_hc_down", .image = "flashnext_hc", .threads = 128, .dynamic = 11136 }, .usage = .{ .{ .regs = 56, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize", .symbol = "fn_hc_normalize", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_reduce_act", .symbol = "fn_hc_reduce_act", .image = "flashnext_hc", .threads = 352 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 60, .static = 4096, .bound = 1024 }, .{ .regs = 60, .static = 4096, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_upmix", .symbol = "fn_hc_upmix", .image = "flashnext_hc", .threads = 256, .dynamic = 18688 }, .usage = .{ .{ .regs = 69, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 62, .static = 2048, .bound = 256 }, .{ .regs = 62, .static = 2048, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback0", .symbol = "fn_hc_writeback0", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 37, .static = 256, .bound = 256 }, .{ .regs = 37, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_moe", .symbol = "fn_hc_writeback_moe", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 47, .static = 256, .bound = 256 }, .{ .regs = 47, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 23, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 17, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_shift_window", .symbol = "fn_shift_window", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 31, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032, .opt_in = true }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = .{ .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, - .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "qmm_group cluster of 8", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmm_group one CTA a tile", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 0 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4 fp32 out", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 1 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408 }, .usage = .{ .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 96, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600 }, .usage = .{ .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out fused slices", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 1 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280 }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 50, .static = 0, .bound = 128 }, .{ .regs = 50, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 32768 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 32768, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280 }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25600 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33792 }, .usage = .{ .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33792, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 122, .static = 0, .bound = 128 }, .{ .regs = 122, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true }, .usage = .{ .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 40960 }, .usage = .{ .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 40960, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 55, .static = 0, .bound = 128 }, .{ .regs = 55, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "rank_sum tf_rank_sum", .symbol = "tf_rank_sum", .image = "rank_sum", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit vision_kernel<80, 8> (RADIO)", .symbol = vision80_symbol, .image = "prefill_attention", .threads = 256, .dynamic = 45056 }, .usage = .{ .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 } } }, diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 243b3386a4..d494d29651 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -346,7 +346,7 @@ pub const Decision = struct { }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque, follow: ?*const fn (*anyopaque) anyerror!void = null }; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; @@ -362,27 +362,8 @@ pub const MemorySource = struct { read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, }; -pub const Status = struct { - /// Requests in prefill or decode, and those waiting for a lane. - running: u32 = 0, - waiting: u32 = 0, - decode_tokens_per_second: f64 = 0, - prefill_tokens_per_second: f64 = 0, - /// Requests that gave a lane up to a later one; null when the engine does not count them. - preemptions: ?u64 = null, - warming: bool = false, - /// Live streams written to the caller's buffer: tokens each holds. - streams: usize = 0, - /// Generated tokens held by live streams, excluding their prompts. - generation_tokens: u64 = 0, - /// Filling prompts' total tokens and those still to prefill (null: no prompt is going in). - fill_prompt_tokens: ?u32 = null, - fill_tokens_left: ?u32 = null, - /// A prompt pass is running whose backend reports no `fill_at`: the totals stay null rather than guessed. - fill_unknown: bool = false, - /// Changes whenever the filling set does: a reader pairs its drain samples with one fill. - fill_run: u64 = 0, -}; +pub const HealthSource = @import("engine_status.zig").HealthSource; +pub const Status = @import("engine_status.zig").Status; pub const SubmitError = error{ Closed, Busy, InvalidSpans, Stale }; // Stale: a learner's put-back has not reached the device diff --git a/zig/src/core/engine_status.zig b/zig/src/core/engine_status.zig new file mode 100644 index 0000000000..fdbbbc96ec --- /dev/null +++ b/zig/src/core/engine_status.zig @@ -0,0 +1,30 @@ +//! Engine readiness and status values, read without entering a device context. + +/// Thread-safe backend readiness; no GPU context or blocking device work is allowed on HTTP threads. +pub const HealthSource = struct { + ctx: *anyopaque, + read: *const fn (*anyopaque) bool, +}; + +pub const Status = struct { + healthy: bool = true, + /// Requests in prefill or decode, and those waiting for a lane. + running: u32 = 0, + waiting: u32 = 0, + decode_tokens_per_second: f64 = 0, + prefill_tokens_per_second: f64 = 0, + /// Requests that gave a lane up to a later one; null when the engine does not count them. + preemptions: ?u64 = null, + warming: bool = false, + /// Live streams written to the caller's buffer: tokens each holds. + streams: usize = 0, + /// Generated tokens held by live streams, excluding their prompts. + generation_tokens: u64 = 0, + /// Filling prompts' total tokens and those still to prefill (null: no prompt is going in). + fill_prompt_tokens: ?u32 = null, + fill_tokens_left: ?u32 = null, + /// A prompt pass is running whose backend reports no `fill_at`: the totals stay null rather than guessed. + fill_unknown: bool = false, + /// Changes whenever the filling set does: a reader pairs its drain samples with one fill. + fill_run: u64 = 0, +}; diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4a8f7d3493..ce38de489c 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -48,6 +48,7 @@ pub const LaneHost = struct { cache: ?*pc.Store = null, // kept prompt states (engine thread only); the backend restores and saves them recall: ?*lanes.recall_store.Store = null, // Recall's record: drafts from earlier requests (lane_host_recall.zig) decoded_rows: bool = false, // the backend prefills Request.decode_spans with decoded rows' arithmetic, cache or not + health: ?api.HealthSource = null, memory: ?api.MemorySource = null, // the backend's memory counts; null: Engine.memory reports none explain: ?api.Explain = null, // the backend's words for a request it refuses; null: the error's name learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused @@ -160,9 +161,10 @@ pub const LaneHost = struct { job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); - if (h.closing or h.stale) { + const unavailable = h.closing or (if (h.health) |source| !source.read(source.ctx) else false); + if (unavailable or h.stale) { h.gpa.destroy(job); - return if (h.closing) error.Closed else error.Stale; + return if (unavailable) error.Closed else error.Stale; } // foreground before background, each in arrival order (the Python job queue's priority) var at = h.queued.items.len; diff --git a/zig/src/core/lane_host_health_test.zig b/zig/src/core/lane_host_health_test.zig new file mode 100644 index 0000000000..06b978a9fe --- /dev/null +++ b/zig/src/core/lane_host_health_test.zig @@ -0,0 +1,49 @@ +//! Backend readiness and learner staleness retain their distinct admission failures. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; + +test "an unhealthy backend fails health and refuses a request before queue admission" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{}, 1, 0); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + const Health = struct { + available: bool = true, + fn read(ctx: *anyopaque) bool { + return @as(*@This(), @ptrCast(@alignCast(ctx))).available; + } + fn event(_: *anyopaque, _: Id, _: *const Event) void { + @panic("a refused request must never reach its sink"); + } + }; + var health: Health = .{}; + const engine = host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + host.health = .{ .ctx = &health, .read = Health.read }; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + const request: Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + host.stale = true; + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + try std.testing.expectError(error.Stale, engine.submit(1, &request, .{ .ctx = &health, .event = Health.event })); + host.stale = false; + health.available = false; + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &health, .event = Health.event })); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len); + try std.testing.expectEqual(@as(usize, 0), host.admitted.items.len); +} diff --git a/zig/src/core/lane_host_status.zig b/zig/src/core/lane_host_status.zig index ed9f5baa5b..e989f25ea0 100644 --- a/zig/src/core/lane_host_status.zig +++ b/zig/src/core/lane_host_status.zig @@ -20,6 +20,7 @@ pub fn read(ctx: *anyopaque, out: *api.Status, stream_tokens: []u32) void { const n = @min(stream_tokens.len, h.live_tokens.items.len); @memcpy(stream_tokens[0..n], h.live_tokens.items[0..n]); out.* = .{ + .healthy = !h.closing and !h.stale and (if (h.health) |source| source.read(source.ctx) else true), .running = @intCast(h.admitted.items.len), .waiting = @intCast(h.queued.items.len), .decode_tokens_per_second = @as(f64, @floatFromInt(tokens)) / 2.0, diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 85a34c0e15..a49e0c7e3c 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -564,6 +564,7 @@ test { _ = @import("lane_host_recall_test.zig"); _ = @import("lane_host_fill_test.zig"); _ = @import("lane_host_status_test.zig"); + _ = @import("lane_host_health_test.zig"); _ = @import("lane_host_stale_test.zig"); _ = @import("lane_host_reuse_test.zig"); _ = @import("lane_host_cache_lifetime_test.zig"); diff --git a/zig/src/core/lanes/sampling.zig b/zig/src/core/lanes/sampling.zig index e9d469756c..43bddfa3a1 100644 --- a/zig/src/core/lanes/sampling.zig +++ b/zig/src/core/lanes/sampling.zig @@ -75,7 +75,7 @@ pub fn pairwiseSum(a: []const f64) f64 { return pairwiseSum(a[0..n2]) + pairwiseSum(a[n2..]); } -const Candidate = struct { value: f64, id: u64 }; +pub const Candidate = struct { value: f64, id: u64 }; fn byValueThenId(_: void, a: Candidate, b: Candidate) bool { return a.value > b.value or (a.value == b.value and a.id < b.id); @@ -88,6 +88,12 @@ pub fn choose(gpa: Allocator, values: []const f64, ids: []const u64, position: u defer gpa.free(order); for (order, values, ids) |*o, v, id| o.* = .{ .value = v, .id = id }; std.mem.sort(Candidate, order, {}, byValueThenId); + return chooseOrdered(gpa, order, position, s); +} + +/// ``choose`` over candidates already in its order (by value, then the lower id at a tie). +pub fn chooseOrdered(gpa: Allocator, order: []const Candidate, position: u64, s: Sampling) !u64 { + const width = order.len; if (s.temperature <= 0) return order[0].id; // temperature 0 is greedy: the first maximum, ties to the lower id const k = @max(1, @min(if (s.top_k != 0) @as(usize, s.top_k) else width, width)); const scaled = try gpa.alloc(f64, k); diff --git a/zig/src/core/nvfp4_csf.md b/zig/src/core/nvfp4_csf.md new file mode 100644 index 0000000000..02e2cdf40e --- /dev/null +++ b/zig/src/core/nvfp4_csf.md @@ -0,0 +1,36 @@ +# Lossless NVFP4 scale planes + +`nvfp4_csf.Plane` reads `byte-window4-fixed-stream-u24-exceptions/1` storage. +It preserves the original scale bytes, including signed zeros and NaN encodings; +it performs no floating-point conversion or requantization. + +Each 16-row slab contains 16 byte bases, followed by row-major pairs of four-bit +offsets. The even column occupies the low nibble. A decoded byte is its row base +plus its offset. Bases cannot exceed 240. Exceptions are little-endian uint32 +words: the low 24 bits name a row-major position, and the high byte replaces its +value. Positions must be unique, sorted, and inside the logical plane. + +Rows must be a positive multiple of 16, columns a positive multiple of two, and +the plane can contain at most 2^24 bytes. Stream extents must match exactly. +`Plane.init` validates the entire plane before returning a borrowed view. The +caller keeps the streams alive and immutable while that view is used. + +```zig +const plane = try core.nvfp4_csf.Plane.init(rows, columns, fixed_bytes, exception_bytes); +try plane.decode(rank_scales, row_begin, row_count, column_begin, column_count); +``` + +`decode` writes exactly the selected rectangle in row-major order, with no heap +allocation. Both row and column partitions can end inside a storage slab or a +packed byte. Invalid output extents and overlap with either borrowed stream are +rejected before writing. Sorted exceptions are searched per selected row, so a +rank does not need to materialize another rank's scales. + +This is a host format primitive. It does not change checkpoint admission, +projection arithmetic, CUDA kernels, or the runtime capability advertisement. +Callers must validate safetensors dtypes/shapes and select this exact codec before +constructing the view. + +The public format reference is the [byte-window implementation in b12x](https://github.com/local-inference-lab/b12x/blob/2cc7f66aaef76a9c7cbb660a1efb865cf3f26c88/b12x/_lib/quant/nvfp4_csf.py). +Tests use independently generated synthetic source bytes, handwritten slab +examples, every replacement byte, both partition axes, and malformed streams. diff --git a/zig/src/core/nvfp4_csf.zig b/zig/src/core/nvfp4_csf.zig new file mode 100644 index 0000000000..49d6e24319 --- /dev/null +++ b/zig/src/core/nvfp4_csf.zig @@ -0,0 +1,90 @@ +//! Lossless byte-window4 scale planes, with sorted u24-position/u8-value exceptions. + +const std = @import("std"); + +pub const codec = "byte-window4-fixed-stream-u24-exceptions/1"; +pub const Error = error{ BadGeometry, BadFixedStream, BadExceptions, BadSlice, OverlappingOutput }; + +/// A validated, borrowed checkpoint plane; callers keep both streams immutable until decoding finishes. +pub const Plane = struct { + rows: usize, + columns: usize, + fixed: []const u8, + exceptions: []const u8, + + /// Slabs hold 16 row bases followed by those rows' low-even/high-odd packed offsets. + pub fn init(rows: usize, columns: usize, fixed: []const u8, exceptions: []const u8) Error!Plane { + if (rows == 0 or rows % 16 != 0 or columns == 0 or columns % 2 != 0) return error.BadGeometry; + if (columns > 1 << 24 or rows > (1 << 24) / columns) return error.BadGeometry; + const stride = 16 + 8 * columns; + if (fixed.len != rows / 16 * stride) return error.BadFixedStream; + for (0..rows / 16) |slab| { + for (fixed[slab * stride ..][0..16]) |base| if (base > 240) return error.BadFixedStream; + } + if (exceptions.len % 4 != 0) return error.BadExceptions; + var previous: ?u32 = null; + var offset: usize = 0; + while (offset < exceptions.len) : (offset += 4) { + const position = word(exceptions, offset) & 0xffffff; + if (position >= rows * columns) return error.BadExceptions; + if (previous) |p| if (position <= p) return error.BadExceptions; + previous = position; + } + return .{ .rows = rows, .columns = columns, .fixed = fixed, .exceptions = exceptions }; + } + + /// Decode a rectangular TP slice into exactly sized row-major bytes, without allocating or rounding. + pub fn decode(self: Plane, out: []u8, row: usize, rows: usize, column: usize, columns: usize) Error!void { + if (row > self.rows or rows > self.rows - row or column > self.columns or columns > self.columns - column) + return error.BadSlice; + if (out.len != rows * columns) return error.BadSlice; + if (overlaps(out, self.fixed) or overlaps(out, self.exceptions)) return error.OverlappingOutput; + if (rows == 0 or columns == 0) return; + const stride = 16 + 8 * self.columns; + for (0..rows) |r| { + const source_row = row + r; + const slab = source_row / 16 * stride; + const base = self.fixed[slab + source_row % 16]; + const codes = slab + 16 + source_row % 16 * (self.columns / 2); + for (0..columns) |c| { + const source_column = column + c; + const codes_byte = self.fixed[codes + source_column / 2]; + const shift: u3 = if (source_column % 2 == 0) 0 else 4; + out[r * columns + c] = base + ((codes_byte >> shift) & 15); + } + var i = self.lowerBound(source_row * self.columns + column); + const end = source_row * self.columns + column + columns; + while (i < self.exceptions.len) : (i += 4) { + const entry = word(self.exceptions, i); + const position = entry & 0xffffff; + if (position >= end) break; + out[r * columns + position - (source_row * self.columns + column)] = @intCast(entry >> 24); + } + } + } + + fn lowerBound(self: Plane, position: usize) usize { + var lo: usize = 0; + var hi = self.exceptions.len / 4; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (word(self.exceptions, mid * 4) & 0xffffff < position) lo = mid + 1 else hi = mid; + } + return lo * 4; + } +}; + +fn word(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +fn overlaps(out: []u8, source: []const u8) bool { + if (out.len == 0 or source.len == 0) return false; + const a = @intFromPtr(out.ptr); + const b = @intFromPtr(source.ptr); + return if (a <= b) b - a < out.len else a - b < source.len; +} + +test { + _ = @import("nvfp4_csf_test.zig"); +} diff --git a/zig/src/core/nvfp4_csf_test.zig b/zig/src/core/nvfp4_csf_test.zig new file mode 100644 index 0000000000..13e63c88b7 --- /dev/null +++ b/zig/src/core/nvfp4_csf_test.zig @@ -0,0 +1,145 @@ +//! Independent byte fixtures and malformed-storage boundaries for lossless scale decoding. + +const std = @import("std"); +const Plane = @import("nvfp4_csf.zig").Plane; +const expect = std.testing.expectEqualSlices; + +test "handwritten slab preserves nibble order, bases and exception bytes" { + var fixed: [32]u8 = @splat(0); + fixed[0] = 240; + fixed[1] = 1; + fixed[16] = 0xf0; + fixed[17] = 0x21; + const exceptions = [_]u8{ 2, 0, 0, 255, 31, 0, 0, 128 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 2); + try expect(u8, &.{ 240, 255, 255, 3 }, out[0..4]); + try std.testing.expectEqual(@as(u8, 128), out[31]); + var odd: [2]u8 = undefined; + try plane.decode(&odd, 0, 2, 1, 1); + try expect(u8, &.{ 255, 3 }, &odd); +} + +const Fixture = struct { + fixed: []u8, + exceptions: []u8, + source: []u8, + rows: usize, + columns: usize, + + fn init(rows: usize, columns: usize) !Fixture { + const gpa = std.testing.allocator; + const source = try gpa.alloc(u8, rows * columns); + errdefer gpa.free(source); + const fixed = try gpa.alloc(u8, rows + source.len / 2); + errdefer gpa.free(fixed); + const exceptions = try gpa.alloc(u8, ((source.len + 96) / 97) * 4); + errdefer gpa.free(exceptions); + @memset(fixed, 0); + var cursor: usize = 0; + for (0..rows) |row| { + const base: u8 = @intCast((row * 37 + 19) % 241); + const slab = row / 16 * (16 + columns * 8); + fixed[slab + row % 16] = base; + for (0..columns) |column| { + const i = row * columns + column; + const delta: u8 = @intCast((column * 11 + row * 3) % 16); + source[i] = base + delta; + const offset = slab + 16 + row % 16 * (columns / 2) + column / 2; + fixed[offset] |= delta << @as(u3, if (column % 2 == 0) 0 else 4); + if (i % 97 == 0) { + source[i] = @intCast((i / 97) % 256); + const record = @as(u32, @intCast(i)) | (@as(u32, source[i]) << 24); + std.mem.writeInt(u32, exceptions[cursor..][0..4], record, .little); + cursor += 4; + } + } + } + return .{ .fixed = fixed, .exceptions = exceptions, .source = source, .rows = rows, .columns = columns }; + } + + fn deinit(self: Fixture) void { + const gpa = std.testing.allocator; + gpa.free(self.fixed); + gpa.free(self.exceptions); + gpa.free(self.source); + } + + fn check(self: Fixture, row: usize, rows: usize, column: usize, columns: usize) !void { + const out = try std.testing.allocator.alloc(u8, rows * columns); + defer std.testing.allocator.free(out); + @memset(out, 0xd6); + const plane = try Plane.init(self.rows, self.columns, self.fixed, self.exceptions); + try plane.decode(out, row, rows, column, columns); + for (0..rows) |r| try expect(u8, self.source[(row + r) * self.columns + column ..][0..columns], out[r * columns ..][0..columns]); + } +}; + +test "full planes and row/column TP partitions reproduce independent source bytes" { + for ([_][2]usize{ .{ 16, 2 }, .{ 32, 20 }, .{ 640, 160 }, .{ 2560, 40 } }) |shape| { + const fixture = try Fixture.init(shape[0], shape[1]); + defer fixture.deinit(); + try fixture.check(0, shape[0], 0, shape[1]); + for (0..2) |rank| { + try fixture.check(rank * shape[0] / 2, shape[0] / 2, 0, shape[1]); + try fixture.check(0, shape[0], rank * shape[1] / 2, shape[1] / 2); + } + try fixture.check(15, 1, 1, shape[1] - 1); + } +} + +test "all 256 byte values survive exception decoding unchanged" { + const fixed: [144]u8 = @splat(0); + var exceptions: [1024]u8 = undefined; + for (0..256) |i| std.mem.writeInt(u32, exceptions[i * 4 ..][0..4], @as(u32, @intCast(i)) * 0x1000001, .little); + const plane = try Plane.init(16, 16, &fixed, &exceptions); + var out: [256]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 16); + for (out, 0..) |value, i| try std.testing.expectEqual(@as(u8, @intCast(i)), value); +} + +test "the last u24 position is valid and larger planes are refused" { + const columns = 1 << 20; + const fixed = try std.testing.allocator.alloc(u8, 16 + 8 * columns); + defer std.testing.allocator.free(fixed); + @memset(fixed, 0); + const plane = try Plane.init(16, columns, fixed, &.{ 255, 255, 255, 173 }); + var out: [1]u8 = undefined; + try plane.decode(&out, 15, 1, columns - 1, 1); + try expect(u8, &.{173}, &out); + try std.testing.expectError(error.BadGeometry, Plane.init(32, columns, fixed, &.{})); +} + +test "bad geometry, truncated streams, overflow bases and unordered exceptions are refused" { + var fixed: [32]u8 = @splat(0); + const bad = std.testing.expectError; + for ([_][2]usize{ .{ 0, 2 }, .{ 1, 2 }, .{ 16, 0 }, .{ 16, 3 }, .{ 16, 1 << 24 }, .{ 1 << 24, 2 } }) |shape| + try bad(error.BadGeometry, Plane.init(shape[0], shape[1], &fixed, &.{})); + try bad(error.BadFixedStream, Plane.init(16, 2, fixed[0..31], &.{})); + fixed[15] = 241; + try bad(error.BadFixedStream, Plane.init(16, 2, &fixed, &.{})); + fixed[15] = 0; + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{0})); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 32, 0, 0, 1 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 1, 0, 0, 2, 1, 0, 0, 3 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 2, 0, 0, 2, 1, 0, 0, 3 })); +} + +test "invalid slices and aliased storage fail before output changes" { + var fixed: [32]u8 = @splat(0); + var exceptions = [_]u8{ 0, 0, 0, 1 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = @splat(0xd6); + const bad = std.testing.expectError; + try bad(error.BadSlice, plane.decode(&out, 17, 0, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 17, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 3, 0)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 1, 2)); + try bad(error.BadSlice, plane.decode(out[0..31], 0, 16, 0, 2)); + const poison: [32]u8 = @splat(0xd6); + try expect(u8, &poison, &out); + try bad(error.OverlappingOutput, plane.decode(&fixed, 0, 16, 0, 2)); + try bad(error.OverlappingOutput, plane.decode(&exceptions, 0, 2, 0, 2)); + try plane.decode(out[0..0], 16, 0, 2, 0); +} diff --git a/zig/src/core/nvfp4_rows.zig b/zig/src/core/nvfp4_rows.zig new file mode 100644 index 0000000000..5a03b7b2c8 --- /dev/null +++ b/zig/src/core/nvfp4_rows.zig @@ -0,0 +1,57 @@ +//! Borrowed NVFP4 table rows: decode requested mapped rows to BF16 without expanding the table. +const std = @import("std"); +const Tensor = @import("safetensors.zig").Tensor; +const quant = @import("quant/quant.zig"); + +pub const Error = error{ BadNvfp4Table, BadNvfp4Rows, AliasedNvfp4Rows }; + +pub const Table = struct { + codes: []const u8, + scales: []const u8, + rows: usize, + width: usize, + global: f32, + + /// Borrow exact row-major U8 codes and E4M3 scales per 16 values; the caller retains the maps. + pub fn init(weight: Tensor, scales: Tensor, global: f32) Error!Table { + if (weight.rank != 2 or weight.dtype != .u8 or weight.dim(0) == 0 or weight.dim(1) == 0) + return error.BadNvfp4Table; + const width = std.math.mul(usize, weight.dim(1), 2) catch return error.BadNvfp4Table; + const size = std.math.mul(usize, weight.dim(0), weight.dim(1)) catch return error.BadNvfp4Table; + if (width % 16 != 0 or weight.bytes.len != size or !std.math.isFinite(global)) return error.BadNvfp4Table; + if (!scales.is(.f8_e4m3, &.{ weight.dim(0), width / 16 }) or scales.bytes.len != size / 8) + return error.BadNvfp4Table; + return .{ .codes = weight.bytes, .scales = scales.bytes, .rows = weight.dim(0), .width = width, .global = global }; + } + + /// Token-major requested rows, with FP32 products rounded once to BF16; malformed requests write nothing. + pub fn gather(t: Table, ids: []const usize, out: []u16) Error!void { + const count = std.math.mul(usize, ids.len, t.width) catch return error.BadNvfp4Rows; + if (out.len != count) return error.BadNvfp4Rows; + for (ids) |id| if (id >= t.rows) return error.BadNvfp4Rows; + const bytes = std.mem.sliceAsBytes(out); + if (overlaps(bytes, t.codes) or overlaps(bytes, t.scales) or overlaps(bytes, std.mem.sliceAsBytes(ids))) + return error.AliasedNvfp4Rows; + for (ids, 0..) |id, row| { + for (0..t.width) |col| { + const code = t.codes[id * (t.width / 2) + col / 2]; + const nib = if (col % 2 == 0) code & 15 else code >> 4; + const a: f32 = @floatCast(quant.nvfp4.e2m1(nib)); + const b: f32 = @floatCast(quant.fp8.e4m3(t.scales[id * (t.width / 16) + col / 16])); + const product: f32 = a * b; + out[row * t.width + col] = quant.convert.bf16(product * t.global); + } + } + } +}; + +fn overlaps(a: []const u8, b: []const u8) bool { + if (a.len == 0 or b.len == 0) return false; + const x = @intFromPtr(a.ptr); + const y = @intFromPtr(b.ptr); + return if (x <= y) y - x < a.len else x - y < b.len; +} + +test { + _ = @import("nvfp4_rows_test.zig"); +} diff --git a/zig/src/core/nvfp4_rows_test.zig b/zig/src/core/nvfp4_rows_test.zig new file mode 100644 index 0000000000..f6006c08ee --- /dev/null +++ b/zig/src/core/nvfp4_rows_test.zig @@ -0,0 +1,76 @@ +const std = @import("std"); +const Table = @import("nvfp4_rows.zig").Table; +const Tensor = @import("safetensors.zig").Tensor; + +fn tensor(dtype: @import("safetensors.zig").DType, rows: usize, cols: usize, bytes: []const u8) Tensor { + return .{ .dtype = dtype, .rank = 2, .shape = .{ rows, cols, 1, 1, 1 }, .bytes = bytes }; +} + +fn rounded(value: f64) u16 { + const v: f32 = @floatCast(value); + if (std.math.isNan(v)) return 0x7fc0; + const bits: u32 = @bitCast(v); + const high = bits >> 16; + const low = bits & 65535; + return @intCast(high + @intFromBool(low > 32768 or (low == 32768 and high % 2 == 1))); +} + +fn scale(byte: u8) f64 { + const mag = byte & 127; + if (mag == 127) return std.math.nan(f64); + const e: i32 = @as(i32, mag >> 3) - 7; + const value = if (mag < 8) @as(f64, @floatFromInt(mag)) / 512 else std.math.ldexp(1 + @as(f64, @floatFromInt(mag & 7)) / 8, e); + return if (byte & 128 != 0) -value else value; +} + +test "mapped NVFP4 rows preserve every nibble and scale encoding in reordered repeated gathers" { + var codes: [256 * 8]u8 = undefined; + var scales: [256]u8 = undefined; + for (0..256) |row| { + scales[row] = @intCast(row); + for (0..8) |col| codes[row * 8 + col] = @intCast(col * 2 + ((col * 2 + 1) << 4)); + } + const t = try Table.init(tensor(.u8, 256, 8, &codes), tensor(.f8_e4m3, 256, 1, &scales), 1.375); + var ids: [258]usize = undefined; + for (0..256) |i| ids[i] = 255 - i; + ids[256] = 17; + ids[257] = 17; + var out: [258 * 16]u16 = undefined; + try t.gather(&ids, &out); + const values = [_]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (ids, 0..) |id, row| for (0..16) |col| { + const sign: f64 = if (col >= 8) -1 else 1; + const want = rounded(sign * values[col % 8] * scale(@intCast(id)) * 1.375); + const actual = out[row * 16 + col]; + if (id & 127 == 127) { + try std.testing.expect(actual & 0x7fc0 == 0x7fc0); + } else try std.testing.expectEqual(want, actual); + }; + try t.gather(&.{}, &.{}); +} + +test "a bad row or destination refuses before writing and maps are never aliased" { + var storage: [128]u8 align(8) = @splat(0x22); + const scales = [_]u8{ 0x38, 0x40 }; + const t = try Table.init(tensor(.u8, 2, 8, storage[0..16]), tensor(.f8_e4m3, 2, 1, &scales), 1); + var out: [32]u16 = @splat(0xbeef); + try std.testing.expectError(error.BadNvfp4Rows, t.gather(&.{ 0, 2 }, &out)); + for (out) |v| try std.testing.expectEqual(@as(u16, 0xbeef), v); + try std.testing.expectError(error.BadNvfp4Rows, t.gather(&.{0}, &out)); + try std.testing.expectError(error.AliasedNvfp4Rows, t.gather(&.{0}, std.mem.bytesAsSlice(u16, storage[0..32]))); + var ids: [4]usize = .{ 0, 0, 0, 0 }; + try std.testing.expectError(error.AliasedNvfp4Rows, t.gather(ids[0..1], std.mem.bytesAsSlice(u16, std.mem.sliceAsBytes(&ids)))); +} + +test "invalid table dtype shape extent and scale cannot create a borrowed view" { + const codes: [16]u8 = @splat(0); + const scales = [_]u8{0x38}; + const w = tensor(.u8, 1, 16, &codes); + const s = tensor(.f8_e4m3, 1, 1, &scales); + try std.testing.expectError(error.BadNvfp4Table, Table.init(w, s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 7, codes[0..7]), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.bf16, 1, 8, &codes), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 8, &codes), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 8, codes[0..8]), s, std.math.inf(f32))); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, std.math.maxInt(usize), 8, &codes), s, 1)); +} diff --git a/zig/src/core/nvfp4_slice.md b/zig/src/core/nvfp4_slice.md new file mode 100644 index 0000000000..b3d4355835 --- /dev/null +++ b/zig/src/core/nvfp4_slice.md @@ -0,0 +1,29 @@ +# Rank-local NVFP4 checkpoint reads + +`nvfp4_slice.read(checkpoint, allocator, module, rows, inputs)` reads a logical +rectangle without expanding the whole projection. The result owns its packed +FP4 code bytes and row-major E4M3 scale bytes; call `deinit` with the same +allocator. `global` preserves the stored FP32 tensor scale. There is no numeric +conversion or requantization. + +The source weight is U8 `[N, K/2]`. Plain scales are F8_E4M3 `[N, K/16]`. +Alternatively, the source may contain the fixed U8 and U32 exception tensors +described by [the CSF codec](nvfp4_csf.md). Both representations cannot be present +at once, and a partial compressed pair is refused. Tensor scale and optional +input calibration must be FP32 scalars. Calibration is consumed as checkpoint +metadata; this reader does not implement an activation quantizer. + +Row and input spans are half-open, nonempty and bounded by the original tensor. +Input cuts must align to groups of 16. Row cuts can cross a compressed 16-row +slab. Source geometry, dtype, byte extents and complete compressed streams are +validated before allocation. Allocation failures release every intermediate +buffer. The result contains only the requested code and scale rectangle. + +`dimensions(source)` and `slice(allocator, source, rows, inputs)` also accept +borrowed tensor views for callers with their own checkpoint mapping. The caller +owns those mappings; the returned buffers are independent of them. + +Tests compare plain and compressed row/column partitions with the independent +FP64 decoder, exercise allocator failures and malformed inputs, and read a real +synthetic safetensors file through the checkpoint reader. These are host format +checks, not tensor-parallel execution or GPU qualification. diff --git a/zig/src/core/nvfp4_slice.zig b/zig/src/core/nvfp4_slice.zig new file mode 100644 index 0000000000..c76e3ed7a5 --- /dev/null +++ b/zig/src/core/nvfp4_slice.zig @@ -0,0 +1,105 @@ +//! Plain or CSF-compressed ModelOpt scales and packed codes, materialized only for one tensor-parallel rectangle. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const csf = @import("nvfp4_csf.zig"); +const Tensor = ckpt.Tensor; +const Span = @import("quant/slice.zig").Span; + +pub const Scales = union(enum) { plain: Tensor, csf: struct { fixed: Tensor, exceptions: Tensor } }; +pub const Source = struct { weight: Tensor, scales: Scales, global: Tensor, input: ?Tensor = null }; +pub const Dims = struct { n: usize, k: usize }; + +/// Owned rank-local bytes: low-even FP4 nibbles, row-major E4M3 scales and the unchanged FP32 global scale. +pub const Linear = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + global: f32, + + pub fn deinit(self: *Linear, gpa: std.mem.Allocator) void { + gpa.free(self.codes); + gpa.free(self.scales); + self.* = undefined; + } +}; + +/// Validate a projection's original geometry before allocating rank-local storage. +pub fn dimensions(source: Source) !Dims { + const w = source.weight; + if (w.rank != 2 or w.dtype != .u8 or w.shape[0] == 0 or w.shape[1] == 0) return error.BadNvfp4; + const k = std.math.mul(usize, w.shape[1], 2) catch return error.BadNvfp4; + const count = std.math.mul(usize, w.shape[0], w.shape[1]) catch return error.BadNvfp4; + if (w.bytes.len != count or k % 16 != 0) return error.BadNvfp4; + const d: Dims = .{ .n = w.shape[0], .k = k }; + if (!scalar(source.global)) return error.BadNvfp4; + if (source.input) |a| if (!scalar(a)) return error.BadNvfp4; + switch (source.scales) { + .plain => |s| if (!s.is(.f8_e4m3, &.{ d.n, d.k / 16 }) or s.bytes.len != count / 8) return error.BadNvfp4, + .csf => |s| { + if (s.exceptions.dtype != .u32 or s.exceptions.rank != 1) return error.BadNvfp4; + if (s.exceptions.shape[0] > std.math.maxInt(usize) / 4 or s.exceptions.bytes.len != s.exceptions.shape[0] * 4) + return error.BadNvfp4; + if (!s.fixed.is(.u8, &.{ d.n / 16, 16 + 8 * (d.k / 16) })) return error.BadNvfp4; + _ = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + }, + } + return d; +} + +/// Rows and input channels are logical weight axes; input cuts must preserve complete groups of 16. +pub fn slice(gpa: std.mem.Allocator, source: Source, rows: Span, inputs: Span) !Linear { + const d = try dimensions(source); + if (rows.from >= rows.to or rows.to > d.n or inputs.from >= inputs.to or inputs.to > d.k) + return error.BadNvfp4Slice; + if (inputs.from % 16 != 0 or inputs.to % 16 != 0) return error.BadNvfp4Slice; + const n = rows.to - rows.from; + const k = inputs.to - inputs.from; + const codes = try gpa.alloc(u8, n * (k / 2)); + errdefer gpa.free(codes); + const scales = try gpa.alloc(u8, n * (k / 16)); + errdefer gpa.free(scales); + for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 2) + inputs.from / 2; + @memcpy(codes[row * (k / 2) ..][0 .. k / 2], source.weight.bytes[begin..][0 .. k / 2]); + } + switch (source.scales) { + .plain => |s| for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 16) + inputs.from / 16; + @memcpy(scales[row * (k / 16) ..][0 .. k / 16], s.bytes[begin..][0 .. k / 16]); + }, + .csf => |s| { + const plane = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + try plane.decode(scales, rows.from, n, inputs.from / 16, k / 16); + }, + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales, .global = @bitCast(std.mem.readInt(u32, source.global.bytes[0..4], .little)) }; +} + +/// Read exact stored tensors, consuming optional calibration and refusing ambiguous plain/compressed scales. +pub fn read(ck: anytype, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { + var buf: [512]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const input_name = try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module}); + const a: ?Tensor = if (ck.has(input_name)) try ck.get(input_name) else null; + const fixed_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_fixed", .{module}); + const fixed: ?Tensor = if (ck.has(fixed_name)) try ck.get(fixed_name) else null; + const exception_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_exceptions", .{module}); + const exceptions: ?Tensor = if (ck.has(exception_name)) try ck.get(exception_name) else null; + if ((fixed != null) != (exceptions != null)) return error.BadNvfp4; + const plain_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}); + const plain = ck.has(plain_name); + if (plain == (fixed != null)) return error.BadNvfp4; + const scales: Scales = if (plain) .{ .plain = try ck.get(plain_name) } else .{ .csf = .{ .fixed = fixed.?, .exceptions = exceptions.? } }; + return slice(gpa, .{ .weight = w, .scales = scales, .global = g, .input = a }, rows, inputs); +} + +fn scalar(t: Tensor) bool { + return t.dtype == .f32 and t.bytes.len == 4 and (t.rank == 0 or (t.rank == 1 and t.shape[0] == 1)); +} + +test { + _ = @import("nvfp4_slice_test.zig"); +} diff --git a/zig/src/core/nvfp4_slice_test.zig b/zig/src/core/nvfp4_slice_test.zig new file mode 100644 index 0000000000..1ca0a128df --- /dev/null +++ b/zig/src/core/nvfp4_slice_test.zig @@ -0,0 +1,141 @@ +//! Rank-local NVFP4 bytes against the independent FP64 decoder and real safetensors headers. + +const std = @import("std"); +const reader = @import("nvfp4_slice.zig"); +const st = @import("safetensors.zig"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); + +const Fixture = struct { + codes: [512]u8 = undefined, + scales: [64]u8 = undefined, + fixed: [48]u8 = @splat(0), + exceptions: [4]u8 = .{ 3, 0, 0, 0x40 }, + global: [4]u8 = undefined, + + fn init() Fixture { + var f: Fixture = .{}; + std.mem.writeInt(u32, &f.global, @bitCast(@as(f32, 0.25)), .little); + for (&f.codes, 0..) |*b, i| b.* = @truncate(i * 29 + 3); + for (0..16) |row| { + const base: u8 = @intCast(0x30 + row % 8); + f.fixed[row] = base; + f.fixed[16 + row * 2] = 0x10; + f.fixed[17 + row * 2] = 0x32; + for (0..4) |column| f.scales[row * 4 + column] = base + @as(u8, @intCast(column)); + } + f.scales[3] = 0x40; + return f; + } + + fn source(f: *const Fixture, compressed: bool) reader.Source { + return .{ + .weight = tensor(.u8, &.{ 16, 32 }, &f.codes), + .scales = if (compressed) .{ .csf = .{ + .fixed = tensor(.u8, &.{ 1, 48 }, &f.fixed), + .exceptions = tensor(.u32, &.{1}, &f.exceptions), + } } else .{ .plain = tensor(.f8_e4m3, &.{ 16, 4 }, &f.scales) }, + .global = tensor(.f32, &.{}, &f.global), + }; + } +}; + +fn tensor(dtype: st.DType, shape: []const usize, bytes: []const u8) st.Tensor { + var out: st.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(out.shape[0..shape.len], shape); + return out; +} + +fn reference(linear: reader.Linear, out: []f64) !void { + try quant.nvfp4.reference(.{ + .weight = tensor(.u8, &.{ linear.n, linear.k / 2 }, linear.codes), + .scale = tensor(.f8_e4m3, &.{ linear.n, linear.k / 16 }, linear.scales), + .global = &.{linear.global}, + }, out); +} + +test "plain and compressed rows/columns preserve every FP64 decoded value" { + const f = Fixture.init(); + const a = std.testing.allocator; + var whole = try reader.slice(a, f.source(false), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 64 }); + defer whole.deinit(a); + var expected: [1024]f64 = undefined; + try reference(whole, &expected); + for ([_]bool{ false, true }) |compressed| { + for (0..2) |axis| for (0..2) |rank| { + const r0 = if (axis == 0) rank * 8 else 0; + const rn: usize = if (axis == 0) 8 else 16; + const c0 = if (axis == 1) rank * 32 else 0; + const cn: usize = if (axis == 1) 32 else 64; + var part = try reader.slice(a, f.source(compressed), .{ .from = r0, .to = r0 + rn }, .{ .from = c0, .to = c0 + cn }); + defer part.deinit(a); + var out: [512]f64 = undefined; + try reference(part, &out); + for (0..rn) |row| try std.testing.expectEqualSlices(f64, expected[(r0 + row) * 64 + c0 ..][0..cn], out[row * cn ..][0..cn]); + }; + } +} + +test "bad compressed tensor metadata and cuts cannot become a rank projection" { + const f = Fixture.init(); + var source = f.source(true); + const bad = std.testing.expectError; + source.scales.csf.exceptions.dtype = .i32; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(true); + source.scales.csf.fixed.shape[1] -= 1; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(false); + source.global.dtype = .bf16; + try bad(error.BadNvfp4, reader.dimensions(source)); + const a = std.testing.allocator; + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 17 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 8, .to = 8 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 1, .to = 32 })); +} + +test "allocation failures free every rank-local buffer" { + const f = Fixture.init(); + for (0..3) |fail_at| { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = fail_at }); + const a = failing.allocator(); + if (reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 32 })) |value| { + var owned = value; + owned.deinit(a); + } else |e| try std.testing.expectEqual(error.OutOfMemory, e); + try std.testing.expectEqual(failing.allocated_bytes, failing.freed_bytes); + } +} + +test "checkpoint reading consumes CSF streams and preserves optional calibration" { + const f = Fixture.init(); + const io = std.testing.io; + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"l.weight":{"dtype":"U8","shape":[16,32],"data_offsets":[0,512]}, + \\ "l.weight_scale.nvfp4_csf_fixed":{"dtype":"U8","shape":[1,48],"data_offsets":[512,560]}, + \\ "l.weight_scale.nvfp4_csf_exceptions":{"dtype":"U32","shape":[1],"data_offsets":[560,564]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[564,568]}, + \\ "l.input_scale":{"dtype":"F32","shape":[1],"data_offsets":[568,572]}} + ; + var image: [8 + header.len + 572]u8 = undefined; + std.mem.writeInt(u64, image[0..8], header.len, .little); + @memcpy(image[8..][0..header.len], header); + const data = image[8 + header.len ..]; + @memcpy(data[0..512], &f.codes); + @memcpy(data[512..560], &f.fixed); + @memcpy(data[560..564], &f.exceptions); + @memcpy(data[564..568], &f.global); + @memcpy(data[568..572], &f.global); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &image }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try ckpt.Checkpoint.openModel(a, io, dir); + defer ck.close(); + var part = try reader.read(&ck, a, "l", .{ .from = 8, .to = 16 }, .{ .from = 32, .to = 64 }); + defer part.deinit(a); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + for (0..8) |r| try std.testing.expectEqualSlices(u8, f.scales[(8 + r) * 4 + 2 ..][0..2], part.scales[r * 2 ..][0..2]); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 122ffe9e4c..47f7609449 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -16,6 +16,10 @@ pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const lora_merge = @import("lora_merge_dir.zig"); pub const quant = @import("quant/quant.zig"); +pub const nvfp4_rows = @import("nvfp4_rows.zig"); +pub const route_ref = @import("route_ref.zig"); +pub const nvfp4_csf = @import("nvfp4_csf.zig"); +pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -36,6 +40,10 @@ test { _ = shard_edit; _ = lora_merge; _ = quant; + _ = nvfp4_rows; + _ = route_ref; + _ = nvfp4_csf; + _ = nvfp4_slice; _ = registry; _ = tuning; _ = norms_ref; diff --git a/zig/src/core/route_ref.zig b/zig/src/core/route_ref.zig new file mode 100644 index 0000000000..21200dae9f --- /dev/null +++ b/zig/src/core/route_ref.zig @@ -0,0 +1,42 @@ +//! Softmax top-k routing in float64, the reference every backend's routing kernels are checked against. +const std = @import("std"); + +/// x . w for one router row. +pub fn dot(x: []const f64, w: []const f64) f64 { + var t: f64 = 0; + for (x, w) |a, b| t += a * b; + return t; +} + +/// An fp32 dot's tolerance: every product exact, at most one rounding a term of 2^-24 of the running |sum|. +pub fn dotTolerance(x: []const f64, w: []const f64) f64 { + var t: f64 = 0; + for (x, w) |a, b| t += @abs(a * b); + return t * @as(f64, @floatFromInt(x.len)) * std.math.pow(f64, 2, -24) + 1e-30; +} + +/// The `ids.len` largest logits (the lower id among equal ones), largest first, and weights e^(l_k - l_0) / their sum. +pub fn topk(logits: []const f64, ids: []usize, weights: []f64) void { + for (ids, 0..) |*id, k| { + var best: ?usize = null; + for (logits, 0..) |v, i| { + if (std.mem.indexOfScalar(usize, ids[0..k], i) != null) continue; + if (best == null or v > logits[best.?]) best = i; + } + id.* = best.?; + } + var total: f64 = 0; + for (ids) |id| total += @exp(logits[id] - logits[ids[0]]); + for (ids, weights) |id, *w| w.* = @exp(logits[id] - logits[ids[0]]) / total; +} + +test "the top picks are the largest, ties to the lower id, weights summing to one" { + const l = [_]f64{ 1, 3, 3, -1, 2 }; + var ids: [3]usize = undefined; + var w: [3]f64 = undefined; + topk(&l, &ids, &w); + try std.testing.expectEqualSlices(usize, &.{ 1, 2, 4 }, &ids); + try std.testing.expectApproxEqAbs(@as(f64, 1), w[0] + w[1] + w[2], 1e-12); + try std.testing.expectEqual(w[0], w[1]); + try std.testing.expect(dotTolerance(&.{ 1, 2 }, &.{ 3, -4 }) > 0); +} diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index 64d436f1b0..c957cc9601 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -139,6 +139,12 @@ pub const Set = struct { return error.MissingTritonVariant; } + /// Whether the set holds any variant of `function`. + pub fn has(self: *const Set, function: []const u8) bool { + for (self.variants) |v| if (std.mem.eql(u8, v.spec.@"fn", function)) return true; + return false; + } + /// The smallest value of constexpr `name` at or above `at_least` among `function`'s variants. pub fn smallestConst(self: *const Set, function: []const u8, name: []const u8, at_least: i64) ?i64 { var best: ?i64 = null; diff --git a/zig/src/cuda/embed_rows.zig b/zig/src/cuda/embed_rows.zig index a8d4a67b4c..1fc5cda5b7 100644 --- a/zig/src/cuda/embed_rows.zig +++ b/zig/src/cuda/embed_rows.zig @@ -1,4 +1,4 @@ -//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 64) or bf16 tables. +//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 32 or 64) or bf16 tables. const std = @import("std"); const kernels = @import("kernels.zig"); const module = @import("module.zig"); @@ -16,12 +16,13 @@ pub const EmbedRows = struct { mod: Module, affine4_fn: Function, bf16_fn: Function, + bf16_copies_fn: Function, pub fn load(d: *const Driver) !EmbedRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.embed_rows); errdefer mod.unload(); - return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16") }; + return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16"), .bf16_copies_fn = try mod.function("tf_embed_rows_bf16_copies") }; } pub fn unload(g: *EmbedRows) void { @@ -30,11 +31,16 @@ pub const EmbedRows = struct { /// out [rows, d] bf16 from ids (i32 [rows]): w u32 [vocab, d / 8], scales and biases bf16 [vocab, d / 64]. pub fn affine4(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize) !void { - if (d % 64 != 0 or out % 16 != 0) return error.RowShape; + return g.affine4Groups(s, ids, w, scales, biases, out, rows, d, 64, 1); + } + + /// out [rows, copies, d] bf16, each row `copies` times: scales and biases a group of `group` (32 or 64) values. + pub fn affine4Groups(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize, group: usize, copies: usize) !void { + if ((group != 32 and group != 64) or d % group != 0 or out % 16 != 0 or copies == 0) return error.RowShape; if (rows == 0) return; var a: launch_.Args = .{}; for ([_]u64{ ids, w, scales, biases, out }) |v| a.add(v); - a.add(@as(i32, @intCast(d))); + for ([_]usize{ d, group, copies }) |v| a.add(@as(i32, @intCast(v))); try launch_.launch(g.affine4_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); } @@ -47,4 +53,15 @@ pub const EmbedRows = struct { a.add(@as(i32, @intCast(d))); try launch_.launch(g.bf16_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); } + + /// Copy exact BF16 bits into out [rows, copies, d], including signed zeros and NaN payloads. + pub fn bf16Copies(g: *const EmbedRows, s: Stream, ids: u64, w: u64, out: u64, rows: usize, d: usize, copies: usize) !void { + if (d == 0 or d % 8 != 0 or d > std.math.maxInt(i32) or copies == 0 or copies > std.math.maxInt(i32) or + rows > std.math.maxInt(u32) or out % 16 != 0 or w % 16 != 0) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ ids, w, out }) |v| a.add(v); + for ([_]usize{ d, copies }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(g.bf16_copies_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } }; diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 8a6664cbd6..68dd741e69 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -126,7 +126,7 @@ pub const Affine4Experts = struct { return if (up) .{ .ptr = l.up, .kg = l.dims / 64, .n = l.width } else .{ .ptr = l.down, .kg = l.width / 64, .n = l.dims }; } - fn args(x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize) launch_.Args { + pub fn argsOf(x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize) launch_.Args { var a: launch_.Args = .{}; a.add(x); a.add(int(x_stride)); @@ -144,7 +144,7 @@ pub const Affine4Experts = struct { pub fn run(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { const grid = @min((max_units + 3) / 4, e.resident[if (up) 0 else 1]); if (grid < 1) return; - var a = args(x, x_stride, slots, w, kg, nb, p, out, n); + var a = argsOf(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); } @@ -153,7 +153,7 @@ pub const Affine4Experts = struct { pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { const grid = max_items * ((nb + 3) / 4); if (grid < 1) return; - var a = args(x, x_stride, slots, w, kg, nb, p, out, n); + var a = argsOf(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; try launch_.launch(if (up) e.pre_up_fn else e.pre_down_fn, cfg, s, &a); } @@ -174,6 +174,38 @@ pub const Affine4Experts = struct { } }; +/// Affine-4 SwiGLU experts at groups of 32 over pair rows (experts_rows.cu): each row's bits are those it gets alone. +pub const RowsExperts = struct { + up_fn: Function, + down_fn: Function, + + /// Pairs a plan item holds: the kernel's two row tiles of 16. + pub const tile: usize = 32; + + /// The instances' names (cuobjdump -symbols of the experts_rows fatbin): gate and up with SwiGLU, then down. + pub const symbols = struct { + pub const up = "_ZN15tf_experts_rows11expert_rowsILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const down = "_ZN15tf_experts_rows11expert_rowsILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + /// experts.cu's one-row forms at groups of 32, whose bits these give every row (tf-cuda-test experts-rows). + pub const decode_up = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const decode_down = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + }; + + /// From a loaded experts_rows module. + pub fn resolve(mod: Module) !RowsExperts { + return .{ .up_fn = try mod.function(symbols.up), .down_fn = try mod.function(symbols.down) }; + } + + /// `items` plan items: `up` maps token rows to SwiGLU's bf16, else pair rows to fp32; K in `kg` groups of 32. + pub fn run(e: *const RowsExperts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, items: usize) !void { + const grid = @min((items * nb + 3) / 4, 4096); + if (grid < 1) return; + var a = Affine4Experts.argsOf(x, x_stride, slots, w, kg, nb, p, out, n); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; + try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); + } +}; + /// Block-FP8 experts (fp8/experts.cu): the Python packing, kernels and grid choice, decode and prompt forms. pub const Fp8Experts = struct { mod: Module, @@ -273,6 +305,7 @@ pub const Nvfp4Experts = struct { 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 + warps: usize = 4, 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 @@ -295,14 +328,36 @@ pub const Nvfp4Experts = struct { return stages * stage; } + pub fn decodeWarps(value: ?[]const u8) !usize { + const v = value orelse return 4; + if (std.mem.eql(u8, v, "1")) return 1; + if (std.mem.eql(u8, v, "4")) return 4; + return error.BadExpertWarps; + } + pub fn load(d: *const Driver, sms: usize) !Nvfp4Experts { + const value = if (std.c.getenv("TF_NVFP4_DECODE_WARPS")) |v| std.mem.span(v) else null; + return loadWarps(d, sms, try decodeWarps(value)); + } + + /// Every warp owns independent tiles; changing CTA grouping preserves each tile's arithmetic. + pub fn loadWarps(d: *const Driver, sms: usize, warps: usize) !Nvfp4Experts { + if (warps != 1 and warps != 4) return error.BadExpertWarps; if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.nvfp4_experts); errdefer mod.unload(); - var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined }; + var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined, .warps = warps }; for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2 }, 0..) |name, i| { - e.fns[i] = try mod.function(name); - e.resident[i] = @max(1, try e.fns[i].occupancy(128, 0)) * sms; + var narrow: [256]u8 = undefined; + const selected = if (warps == 1) blk: { + const marker = std.mem.indexOf(u8, name, "ELi4E") orelse return error.MissingSymbol; + @memcpy(narrow[0..name.len], name); + narrow[marker + 3] = '1'; + narrow[name.len] = 0; + break :blk narrow[0..name.len :0]; + } else name; + e.fns[i] = try mod.function(selected); + e.resident[i] = @max(1, try e.fns[i].occupancy(@intCast(warps * 32), 0)) * sms; } for (symbols.prompt, &e.prompt_fns, [3]usize{ 1, 2, 1 }, prompt_wn) |name, *f, parts, wn| { f.* = try mod.function(name); @@ -321,7 +376,7 @@ pub const Nvfp4Experts = struct { 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 grid = @min((items * nb + 3) / 4, e.resident[ei]); + const grid = @min((items * nb + e.warps - 1) / e.warps, e.resident[ei]); if (grid < 1) return; var a: launch_.Args = .{}; a.add(in.x); @@ -335,7 +390,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); - try launch_.launch(e.fns[ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); + try launch_.launch(e.fns[ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = @intCast(e.warps * 32), .y = 1, .z = 1 } }, s, &a); } /// Prompt rows on a plan of at most `items` items of 64 pairs: up (relu^2 or SwiGLU by the layer), bf16 sums down. @@ -492,3 +547,10 @@ test "a format that is not loaded, or an option it was not built with, is refuse try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3)); } + +test "NVFP4 decode CTA geometry is explicit, bounded and defaults to four warps" { + try std.testing.expectEqual(@as(usize, 4), try Nvfp4Experts.decodeWarps(null)); + try std.testing.expectEqual(@as(usize, 1), try Nvfp4Experts.decodeWarps("1")); + try std.testing.expectEqual(@as(usize, 4), try Nvfp4Experts.decodeWarps("4")); + for ([_][]const u8{ "0", "2", "8", "", "garbage" }) |v| try std.testing.expectError(error.BadExpertWarps, Nvfp4Experts.decodeWarps(v)); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..a999eb3ed3 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -13,12 +13,17 @@ fn Blob(comptime import_name: []const u8) type { pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; +pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; +pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; +pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; +pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; +pub const experts_rows: []const u8 = if (available) &Blob("fatbin_experts_rows").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const lfm2: []const u8 = if (available) &Blob("fatbin_lfm2").bytes else &.{}; pub const vit: []const u8 = if (available) &Blob("fatbin_vit").bytes else &.{}; @@ -56,6 +61,9 @@ pub const ck_lane: []const u8 = if (available) &Blob("fatbin_ck_lane").bytes els pub const ck_experts: []const u8 = if (available) &Blob("fatbin_ck_experts").bytes else &.{}; pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; +pub const route_rows: []const u8 = if (available) &Blob("fatbin_route_rows").bytes else &.{}; +pub const flashnext_hc: []const u8 = if (available) &Blob("fatbin_flashnext_hc").bytes else &.{}; +pub const flashnext_attn: []const u8 = if (available) &Blob("fatbin_flashnext_attn").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/cuda/mxfp8.md b/zig/src/cuda/mxfp8.md new file mode 100644 index 0000000000..5878df43eb --- /dev/null +++ b/zig/src/cuda/mxfp8.md @@ -0,0 +1,27 @@ +# Native MXFP8 projections + +`cuda.mxfp8` packs row-major E4M3 code bytes and one UE8M0 scale byte per +32 input channels for `qlinear.Weight.mxfp8`. Both prompt and decode calls use +the lane matmul, with BF16 inputs, FP32 accumulation and BF16 output. Packing +permutes bytes; it does not requantize weights or expand them to BF16. + +The logical shape is `[N, K]`, with positive dimensions and K divisible by 64. +`packCodes` writes `padded(N) * K` bytes; `packScales` writes +`padded(N) * K/32` bytes. Rows are padded to 128; padded codes are zero and +padded scales are one. `weight` returns the borrowed device view. The caller +owns its device buffers and the lane's split-K scratch. + +UE8M0 bytes 0 through 254 mean `2^(byte - 127)`; byte 255 is NaN. In +particular, byte zero is a nonzero subnormal FP32 scale. See the +[CUDA datatype contract](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-math-api/cuda_math_api/struct____nv__fp8__e8m0.html). + +`tf-cuda-test mxfp8` checks all 256 scale encodings through the production +matmul, then compares synthetic projections with an independent FP64 +calculation. It also checks strided activations, partial output tiles, +row counts 1/7/17/33/256, prompt/decode dispatch, and cluster versus explicit +split reduction. Host tests cover scale permutation, padding, shared-memory +sizes, symbols and format refusals. + +This primitive does not admit new checkpoint families or certify model +correctness, tensor parallelism or performance. Rebuild the qmmf fatbin with +the native runtime; previous fatbins lack the MXFP8 instantiations. diff --git a/zig/src/cuda/mxfp8.zig b/zig/src/cuda/mxfp8.zig new file mode 100644 index 0000000000..7c57aa50df --- /dev/null +++ b/zig/src/cuda/mxfp8.zig @@ -0,0 +1,41 @@ +//! E4M3 codes with one UE8M0 byte per 32 inputs, repacked losslessly for the lane matmul. +const std = @import("std"); +const fp8 = @import("fp8.zig"); +const qmmf = @import("qmmf.zig"); + +pub const group = 32; +pub const padded = fp8.padded; +pub const packCodes = fp8.packCodes; + +/// Row-major [n, k/32] bytes to [npad/64][k/64][64][2]; padded rows have scale 1. +pub fn packScales(out: []u8, scales: []const u8, n: usize, k: usize) void { + const npad = padded(n); + const kg = k / qmmf.group; + std.debug.assert(k % qmmf.group == 0 and out.len == npad * (k / group) and scales.len == n * (k / group)); + for (0..npad / 64) |t| for (0..kg) |g| for (0..64) |r| for (0..2) |h| { + const row = t * 64 + r; + out[((t * kg + g) * 64 + r) * 2 + h] = if (row < n) scales[row * (k / group) + g * 2 + h] else 127; + }; +} + +pub fn weight(codes: u64, scales: u64, n: usize, k: usize) qmmf.Weight { + return .{ .mode = .mxfp8, .codes = codes, .scales = scales, .n = @intCast(n), .k = @intCast(k), .npad = @intCast(padded(n)) }; +} + +test "every scale byte lands in its independent 32-input group, including padded rows" { + const a = std.testing.allocator; + const n = 130; + const k = 256; + const scales = try a.alloc(u8, n * (k / group)); + defer a.free(scales); + for (scales, 0..) |*v, i| v.* = @truncate(i * 13 + 7); + const out = try a.alloc(u8, padded(n) * (k / group)); + defer a.free(out); + packScales(out, scales, n, k); + for (0..padded(n)) |row| for (0..k / group) |g| { + const at = (((row / 64) * (k / 64) + g / 2) * 64 + row % 64) * 2 + g % 2; + try std.testing.expectEqual(if (row < n) scales[row * (k / group) + g] else @as(u8, 127), out[at]); + }; + const w = weight(1, 2, n, k); + try std.testing.expect(w.mode == .mxfp8 and w.npad == 256 and w.scale == 1.0); +} diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index b0d527ce75..ef5bf6d177 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -10,6 +10,31 @@ pub const UniqueId = extern struct { internal: [128]u8 }; pub const DataType = enum(c_int) { i8 = 0, u8 = 1, i32 = 2, u32 = 3, i64 = 4, u64 = 5, f16 = 6, f32 = 7, f64 = 8, bf16 = 9 }; pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; +/// NCCL 2.30 public ncclConfig_v23000 ABI. New nonblocking rank owners explicitly require 2.30 or later. +pub const Config = extern struct { + size: usize = @sizeOf(Config), + magic: u32 = 0xcafebeef, + version: u32 = 23000, + blocking: c_int = 0, + cga_cluster_size: c_int = std.math.minInt(c_int), + min_ctas: c_int = std.math.minInt(c_int), + max_ctas: c_int = std.math.minInt(c_int), + net_name: ?[*:0]const u8 = null, + split_share: c_int = std.math.minInt(c_int), + traffic_class: c_int = std.math.minInt(c_int), + comm_name: ?[*:0]const u8 = null, + collnet_enable: c_int = std.math.minInt(c_int), + cta_policy: c_int = std.math.minInt(c_int), + shrink_share: c_int = std.math.minInt(c_int), + nvls_ctas: c_int = std.math.minInt(c_int), + n_channels_per_net_peer: c_int = std.math.minInt(c_int), + nvlink_centric_sched: c_int = std.math.minInt(c_int), + graph_usage_mode: c_int = std.math.minInt(c_int), + num_rma_ctx: c_int = std.math.minInt(c_int), + max_p2p_peers: c_int = std.math.minInt(c_int), + graph_stream_ordering: c_int = std.math.minInt(c_int), +}; + pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; const R = Result; @@ -71,3 +96,11 @@ pub const Library = struct { return v; } }; + +test "nonblocking communicator config preserves the public ABI defaults" { + const c: Config = .{}; + try std.testing.expectEqual(@as(usize, 96), @sizeOf(Config)); + try std.testing.expectEqual(@as(usize, 32), @offsetOf(Config, "net_name")); + try std.testing.expectEqual(@as(usize, 48), @offsetOf(Config, "comm_name")); + try std.testing.expect(c.size == 96 and c.blocking == 0 and c.max_ctas == std.math.minInt(c_int)); +} diff --git a/zig/src/cuda/norms.zig b/zig/src/cuda/norms.zig index fd343b4635..6f774d50d3 100644 --- a/zig/src/cuda/norms.zig +++ b/zig/src/cuda/norms.zig @@ -42,6 +42,7 @@ pub const Norms = struct { add_rms_h_fn: Function, add_rms_h_f32w_fn: Function, gated_fns: [3]Function, // head dim 64, 128, 256 + rms_sums_fn: Function, pub fn load(d: *const Driver) !Norms { if (!kernels.available) return error.BuiltWithoutKernels; @@ -56,6 +57,7 @@ pub const Norms = struct { .add_rms_h_fn = try mod.function("tf_add_rms_rows_bf16h"), .add_rms_h_f32w_fn = try mod.function("tf_add_rms_rows_bf16h_f32w"), .gated_fns = .{ try mod.function("tf_gated_rms_rows_64"), try mod.function("tf_gated_rms_rows_128"), try mod.function("tf_gated_rms_rows_256") }, + .rms_sums_fn = try mod.function("tf_rms_rows_sums"), }; } @@ -87,6 +89,20 @@ pub const Norms = struct { try launch_.launch(n.rms_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); } + /// bf16 rows `x_ld` apart: y = bf16(rms(x) * w) [rows, width], fp32 w, and xs the stored y's sums of each `group`. + pub fn rmsSums(n: *const Norms, s: Stream, x: u64, x_ld: usize, w: u64, y: u64, xs: u64, rows: usize, width: usize, group: usize, eps: f32) !void { + if (group == 0 or width % group != 0 or x_ld < width) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(x_ld))); + for ([_]u64{ w, y, xs }) |v| a.add(v); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(group))); + a.add(eps); + try launch_.launch(n.rms_sums_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + /// The residual step: res (fp32, in place) += rms(x) * w1 (w1 null: res += x), then y = rms(res) * w2. pub fn addRms(n: *const Norms, s: Stream, x: u64, x_kind: Kind, res: u64, w1: ?u64, w2: u64, y: u64, y_kind: Kind, rows: usize, width: usize, eps: f32) !void { if (rows == 0) return; diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index dbb86bb8d0..a5031b6bcd 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -34,6 +34,7 @@ pub const Affine4 = struct { pub const Weight = union(enum) { affine4: Affine4, fp8g: qmmf.Weight, + mxfp8: qmmf.Weight, // e4m3 with UE8M0 scale bytes per 32 inputs nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values @@ -41,14 +42,14 @@ pub const Weight = union(enum) { pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.n, }; } pub fn inputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.k, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.k, }; } }; @@ -81,6 +82,7 @@ pub const Linear = struct { return a.decode(s, in.x, in.sums, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return if (q.act != 0 and l.ck != null) l.checkpoint(s, l.ck.?, q, in, out, rows) else l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -96,6 +98,7 @@ pub const Linear = struct { return a.prompt(s, in.x, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return if (q.act != 0 and l.ck != null) l.checkpoint(s, l.ck.?, q, in, out, rows) else l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -385,6 +388,8 @@ test "a lane weight whose mode is not its format's is refused before any launch" const fp4: qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.prompt(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); const f8: qmmf.Weight = .{ .mode = .fp8g, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; try std.testing.expectError(error.EpilogueNotBuilt, l.decodeF32(s, .{ .fp8g = f8 }, .{ .x = 0 }, 0, 1)); diff --git a/zig/src/cuda/qmmf.md b/zig/src/cuda/qmmf.md new file mode 100644 index 0000000000..82f5c5a1ad --- /dev/null +++ b/zig/src/cuda/qmmf.md @@ -0,0 +1,19 @@ +# FP32 projection outputs + +`cuda.qmmf.Lane.matmulF32` has the same inputs, packed weights, stride, +shape-driven split-K plan and scratch contract as `matmul`. Its output is +`[M, N]` FP32, requiring `M * N * 4` bytes. It preserves FP32 output for a +later residual or tensor-rank sum; the caller chooses when to round to BF16. + +All five lane formats are supported: NVFP4, tensor-scaled FP8, MXFP8, +block-scaled FP8 and BF16. The change exposes the existing kernel template's +FP32 output path. It does not change weight packing or accumulation order. +Rebuild the qmmf fatbin together with the native runtime: older fatbins do +not contain the FP32 instances or reduction symbol. + +`tf-cuda-test projection-f32` compares each format against an independent +FP64 calculation, checks conversion to the existing BF16 result, and varies +row count, activation stride, partial output tiles, fused reduction, +cluster reduction and explicit split reduction. The fixture is synthetic; +model correctness and multi-device collective behavior require separate +qualification. diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index ec7fcd6e04..8fb23389ad 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -1,4 +1,4 @@ -//! The lane matmul (qmmf.cu) for FP4, FP8, FP8G and BF16 weights, with its kernel, tile and slices. +//! The lane matmul (qmmf.cu) for FP4, FP8, MXFP8, FP8G and BF16 weights, with its kernel, tile and slices. const std = @import("std"); const driver = @import("driver.zig"); const module = @import("module.zig"); @@ -17,9 +17,9 @@ const bn = 64; const threads = 128; // WM 1 x WN 4 warps /// The packed formats (qmmf.cu Mode): FP4 with e4m3 scales a 16, e4m3 with one scale or fp32 a 64, and bf16. -pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, fp8g = 3, bf16 = 4 }; +pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, mxfp8 = 2, fp8g = 3, bf16 = 4 }; -const modes = [_]Mode{ .fp4, .fp8, .fp8g, .bf16 }; +const modes = [_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }; fn modeIndex(mode: Mode) usize { return std.mem.indexOfScalar(Mode, &modes, mode).?; @@ -29,10 +29,10 @@ fn modeIndex(mode: Mode) usize { pub fn sharedBytes(mode: Mode, bm: u32) u32 { const w: u32 = switch (mode) { .fp4 => bn * group / 2, - .fp8, .fp8g => bn * group, + .fp8, .mxfp8, .fp8g => bn * group, .bf16 => bn * group * 2, }; - const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else 0; + const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else if (mode == .mxfp8) bn * 2 else 0; const stage = (bm * group * 2 + w + s + 127) / 128 * 128; const partials = (bm / 16) * 2 * 4 * threads * 4; return @max(4 * stage, partials); @@ -85,7 +85,7 @@ pub const reduce_symbols = [2][:0]const u8{ }; const row_tiles = [_]u32{ 16, 32, 64 }; -/// What a call writes: bf16 rows, or the fp32 sums unrounded (FP4 only, a tensor-parallel rank's partial rows). +/// What a call writes: bf16 rows, or the fp32 sums unrounded (a tensor-parallel rank's partial rows). pub const Out = enum { bf16, f32 }; fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function { @@ -103,7 +103,7 @@ fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused - fp4_f32: [7]Function, // FP4's tiles writing fp32 + f32_fns: [modes.len][7]Function, reduce: [2]Function, // by Out clusters: bool, // sm_90 on: slices meet in a cluster's shared memory @@ -114,11 +114,14 @@ pub const Lane = struct { var l: Lane = .{ .mod = mod, .fns = undefined, - .fp4_f32 = try loadTiles(mod, .fp4, true), + .f32_fns = undefined, .reduce = undefined, .clusters = capability_major >= 9, }; - inline for (modes, 0..) |mode, mi| l.fns[mi] = try loadTiles(mod, mode, false); + inline for (modes, 0..) |mode, mi| { + l.fns[mi] = try loadTiles(mod, mode, false); + l.f32_fns[mi] = try loadTiles(mod, mode, true); + } for (&l.reduce, reduce_symbols) |*f, name| f.* = try mod.function(name); return l; } @@ -138,11 +141,15 @@ pub const Lane = struct { return l.matmulOut(s, x, ldx, m, w, y, .bf16, part); } - /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + /// Preserve FP32 outputs for a later residual or tensor-rank sum. + pub fn matmulF32(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulOut(s, x, ldx, m, w, y, .f32, part); + } + + /// matmul() writing y as `out` says: bf16, or fp32. pub fn matmulOut(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, out: Out, part: u64) !void { - if (out == .f32 and w.mode != .fp4) return error.EpilogueNotBuilt; const p = plan(m, w.n, w.k, l.clusters); - const set = if (out == .f32) l.fp4_f32 else l.fns[modeIndex(w.mode)]; + const set = (if (out == .f32) l.f32_fns else l.fns)[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -220,6 +227,9 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp8g, 16)); try std.testing.expectEqual(@as(u32, 33792), sharedBytes(.fp8g, 32)); try std.testing.expectEqual(@as(u32, 50176), sharedBytes(.fp8g, 64)); + try std.testing.expectEqual(@as(u32, 25088), sharedBytes(.mxfp8, 16)); + try std.testing.expectEqual(@as(u32, 33280), sharedBytes(.mxfp8, 32)); + try std.testing.expectEqual(@as(u32, 49664), sharedBytes(.mxfp8, 64)); try std.testing.expectEqual(@as(u32, 17408), sharedBytes(.fp4, 16)); try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp4, 32)); try std.testing.expectEqual(@as(u32, 41984), sharedBytes(.fp4, 64)); @@ -243,16 +253,9 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings( - "_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", - symbol(.fp8g, 16, true, false, false), - ); - try std.testing.expectEqualStrings( - "_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", - symbol(.fp4, 64, false, true, false), - ); - try std.testing.expectEqualStrings( - "_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", - symbol(.fp4, 32, true, false, true), - ); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.bf16, 64, false, true, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 32, true, false, true)); } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6c4e0991b6..aa88e11902 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -26,6 +26,7 @@ pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); +pub const mxfp8 = @import("mxfp8.zig"); pub const exlr8 = @import("exlr8.zig"); pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); @@ -33,6 +34,7 @@ pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); +pub const route_rows = @import("route_rows.zig"); pub const vit = @import("vit.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); @@ -48,6 +50,7 @@ test { _ = @import("qlinear.zig"); _ = @import("norms.zig"); _ = @import("embed_rows.zig"); + _ = @import("route_rows.zig"); _ = @import("vit.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); @@ -58,6 +61,7 @@ test { _ = segments; _ = @import("qmmf.zig"); _ = @import("fp8.zig"); + _ = @import("mxfp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig new file mode 100644 index 0000000000..665463a1df --- /dev/null +++ b/zig/src/cuda/route_rows.zig @@ -0,0 +1,73 @@ +//! route_rows.cu for every CUDA family: router logits on mma tiles and softmax top-k rows, a row's bits its own. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Logits a top-k row reads at most (the kernel's NT * MOST) and picks it keeps (its scratch). +pub const most_experts = 1024; +pub const most_picks = 128; + +/// tf_route_logits' staged rows (16 inputs and 32 experts, padded by 8): two stages of 128, or three of 256. +const many_shared: u32 = 2 * 48 * (128 + 8) * 2; +const few_shared: u32 = 3 * 48 * (256 + 8) * 2; + +pub const RouteRows = struct { + mod: Module, + logits_fn: Function, + few_fn: Function, + topk_fn: Function, + + pub fn load(d: *const Driver) !RouteRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.route_rows); + errdefer mod.unload(); + const r: RouteRows = .{ .mod = mod, .logits_fn = try mod.function("tf_route_logits"), .few_fn = try mod.function("tf_route_logits_few"), .topk_fn = try mod.function("tf_route_topk") }; + try r.few_fn.allowDynamicShared(few_shared); + return r; + } + + pub fn unload(g: *RouteRows) void { + g.mod.unload(); + } + + /// out [rows, e] fp32 = x [rows, d] (bf16, x_ld apart) times w [e, d] bf16 transposed; d % 16 == 0, 16-byte rows. + pub fn logits(g: *const RouteRows, s: Stream, x: u64, x_ld: usize, w: u64, out: u64, rows: usize, d: usize, e: usize) !void { + if (d % 16 != 0 or x_ld % 8 != 0 or x % 16 != 0 or w % 16 != 0) return error.RowShape; + if (rows == 0 or e == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i32, @intCast(x_ld))); + a.add(w); + a.add(out); + for ([_]usize{ rows, d, e }) |v| a.add(@as(i32, @intCast(v))); + // up to 16 rows, deeper staging (the same sums: the router rows' reads set the time) + const few = rows <= 16; + try launch_.launch(if (few) g.few_fn else g.logits_fn, .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = if (few) few_shared else many_shared }, s, &a); + } + + /// Each row's top_k of its first `experts` logits into picks and weights [rows, slots]; `shared`: slot top_k too. + pub fn topk(g: *const RouteRows, s: Stream, l: u64, ld: usize, experts: usize, top_k: usize, picks: u64, wts: u64, slots: usize, rows: usize, shared: bool) !void { + const need = experts + @intFromBool(shared); + if (experts > most_experts or top_k == 0 or top_k > most_picks or top_k > experts or slots < top_k + @intFromBool(shared) or ld < need) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(l); + for ([_]usize{ ld, experts, top_k }) |v| a.add(@as(i32, @intCast(v))); + a.add(picks); + a.add(wts); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intFromBool(shared))); + try launch_.launch(g.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } +}; + +test "the top-k scratch holds a 512-expert router's ten picks" { + try std.testing.expect(most_experts >= 513 and most_picks >= 11); +} diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig new file mode 100644 index 0000000000..e93e1a2e30 --- /dev/null +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -0,0 +1,442 @@ +//! Layer 47's full attention for one decode token at position 0: project, prepare, attend, gate, project out. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const q_heads: usize = 24; +const kv_heads: usize = 2; +const head_dim: usize = 256; +const index_heads: usize = 4; +const index_dim: usize = 128; +const half: usize = 32; +const eps: f32 = 1e-6; +const attn_scale: f32 = 0.0625; +const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; +const out_k: usize = q_heads * head_dim; +const nch: usize = 3; +const prefix = "language_model.model.layers.47.self_attn."; + +const Gap = struct { off: usize, steps: u32, at: usize }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn readBf(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn treeSum(buf: []f32) f32 { + var n = buf.len; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn sumSq(x: []const u16) f32 { + var buf: [head_dim]f32 = undefined; + for (x, 0..) |bits, i| { + const v = promote(bits); + buf[i] = v * v; + } + return treeSum(buf[0..x.len]); +} + +/// glue._attn_prep at position 0: RMSNorm, round, rotate-half. A zero angle stores the normalized row. +fn normRot(x: []const u16, scale_w: []const f32, inv: []const f32, pos: f32, out: []u16) void { + const width = x.len; + const rinv = 1.0 / @sqrt(sumSq(x) / @as(f32, @floatFromInt(width)) + eps); + var xn: [head_dim]f32 = undefined; + for (0..width) |d| xn[d] = promote(toBf16(promote(x[d]) * rinv * scale_w[d])); + for (0..width) |d| { + const partner = if (d < half) d + half else if (d < 2 * half) d - half else d; + const axis_i = if (d < half) d else if (d < 2 * half) d - half else 0; + const ang = pos * inv[axis_i]; + const c = @cos(ang); + const s = @sin(ang); + const a = xn[d]; + const b = xn[partner]; + const rot = if (d < half) a * c - b * s else if (d < 2 * half) b * s + a * c else a; + out[d] = toBf16(rot); + } +} + +fn copyRow(src: []const u16, dst: []u16) void { + @memcpy(dst, src); +} + +/// One key: the softmax weight is 1, so each query head copies its KV head's value. +fn attend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + _ = q; + _ = k; + copyRow(v, out); +} + +fn applyGate(o: []const u16, g: []const u16, out: []u16, xs: []f32) void { + for (o, g, out) |ov, gv, *y| y.* = toBf16(promote(ov) / (1.0 + @exp(-promote(gv)))); + for (xs, 0..) |*s, i| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(out[i * 32 + j]); + s.* = treeSum(buf[0..]); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (got, want, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + return .{ .off = off, .steps = steps, .at = at }; +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "q_proj.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +const Triple = struct { w: core.safetensors.Tensor, s: core.safetensors.Tensor, b: core.safetensors.Tensor }; + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +fn stitch(gpa: std.mem.Allocator, parts: []const Triple) !struct { words: []u8, scales: []u8, biases: []u8 } { + var wb: usize = 0; + var sb: usize = 0; + for (parts) |p| { + wb += p.w.bytes.len; + sb += p.s.bytes.len; + } + const words = try gpa.alloc(u8, wb); + errdefer gpa.free(words); + const scales = try gpa.alloc(u8, sb); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u8, sb); + var wo: usize = 0; + var so: usize = 0; + for (parts) |p| { + @memcpy(words[wo..][0..p.w.bytes.len], p.w.bytes); + wo += p.w.bytes.len; + @memcpy(scales[so..][0..p.s.bytes.len], p.s.bytes); + @memcpy(biases[so..][0..p.b.bytes.len], p.b.bytes); + so += p.s.bytes.len; + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +fn promoteScale(bytes: []const u8, out: []f32) void { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); +} + +fn invFreq(out: []f32) void { + const n: f64 = @floatFromInt(out.len); + for (out, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / n)); +} + +fn prepare(pa: []const u16, q_scale: []const f32, k_scale: []const f32, i_scale: []const f32, inv: []const f32, q: []u16, k: []u16, v: []u16, iq: []u16, ik: []u16) void { + const q_pairs = q_heads * 2 * head_dim; + const kv = kv_heads * head_dim; + for (0..q_heads) |h| normRot(pa[h * 2 * head_dim ..][0..head_dim], q_scale, inv, 0, q[h * head_dim ..][0..head_dim]); + for (0..kv_heads) |h| { + normRot(pa[q_pairs + h * head_dim ..][0..head_dim], k_scale, inv, 0, k[h * head_dim ..][0..head_dim]); + copyRow(pa[q_pairs + kv + h * head_dim ..][0..head_dim], v[h * head_dim ..][0..head_dim]); + } + const iq0 = q_pairs + 2 * kv; + for (0..index_heads) |h| normRot(pa[iq0 + h * index_dim ..][0..index_dim], i_scale, inv, 0, iq[h * index_dim ..][0..index_dim]); + copyRow(pa[iq0 + index_heads * index_dim ..][0..index_dim], ik); +} + +fn hostAttend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + const g = q_heads / kv_heads; + for (0..q_heads) |h| { + const hk = h / g; + attend(q[h * head_dim ..][0..head_dim], k[hk * head_dim ..][0..head_dim], v[hk * head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); + } +} + +fn gatesOf(pa: []const u16, out: []u16) void { + for (0..q_heads) |h| copyRow(pa[h * 2 * head_dim + head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); +} + +/// The projection, prep, one-key attention, gate and output projection, against the host of the same mixed row. +pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, io: std.Io, model_dir: []const u8, tri: Tri, mixed: u64, xs: u64) !u32 { + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const parts = [_]Triple{ + try triple(&file, prefix ++ "q_proj", q_heads * 2 * head_dim, dims / 8), + try triple(&file, prefix ++ "k_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "v_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "indexer.index_qk_proj", (index_heads + 1) * index_dim, dims / 8), + }; + const out_w = try triple(&file, prefix ++ "o_proj", dims, out_k / 8); + const qn = file.get(prefix ++ "q_norm.weight") orelse return error.MissingTensor; + const kn = file.get(prefix ++ "k_norm.weight") orelse return error.MissingTensor; + const iqn = file.get(prefix ++ "indexer.q_layernorm.weight") orelse return error.MissingTensor; + if (!qn.is(.bf16, &.{head_dim}) or !kn.is(.bf16, &.{head_dim}) or !iqn.is(.bf16, &.{index_dim})) return error.UnexpectedTensor; + if (proj_n != 13952) return error.UnexpectedTensor; + + const stacked = try stitch(gpa, &parts); + defer gpa.free(stacked.words); + defer gpa.free(stacked.scales); + defer gpa.free(stacked.biases); + var lane = try qmm.pack(gpa, stacked.words, stacked.scales, stacked.biases, proj_n, dims / 8); + defer lane.deinit(gpa); + var ow = try qmm.pack(gpa, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k / 8); + defer ow.deinit(gpa); + + var q_scale: [head_dim]f32 = undefined; + var k_scale: [head_dim]f32 = undefined; + var i_scale: [index_dim]f32 = undefined; + var inv: [half]f32 = undefined; + promoteScale(qn.bytes, &q_scale); + promoteScale(kn.bytes, &k_scale); + promoteScale(iqn.bytes, &i_scale); + invFreq(&inv); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var pa_b = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + defer pa_b.free(); + try qmm.matmul(driver, stream.*, mixed, xs, w_b.ptr, s_b.ptr, b_b.ptr, pa_b.ptr, 1, proj_n, dims); + + var qw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&q_scale)); + defer qw_b.free(); + var kw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&k_scale)); + defer kw_b.free(); + var iw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&i_scale)); + defer iw_b.free(); + var inv_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + defer inv_b.free(); + var pos_host: i32 = 0; + var pos_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.asBytes(&pos_host)); + defer pos_b.free(); + var q_b = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + defer q_b.free(); + var kc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer kc_b.free(); + var vc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer vc_b.free(); + var iq_b = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + defer iq_b.free(); + var ik_b = try cuda.DeviceBuffer.alloc(driver, index_dim * 2); + defer ik_b.free(); + var dummy = try cuda.DeviceBuffer.alloc(driver, 256); + defer dummy.free(); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + + var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + defer po_b.free(); + var pm_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pm_b.free(); + var pl_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pl_b.free(); + var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer o_b.free(); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false); + try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, false); + + var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer gated_b.free(); + var gxs_b = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + defer gxs_b.free(); + try tri.attnGate(o_b.ptr, pa_b.ptr, gated_b.ptr, gxs_b.ptr, 1, proj_n, q_heads, head_dim); + + var ow_b = try cuda.DeviceBuffer.fromHost(driver, ow.weight); + defer ow_b.free(); + var os_b = try cuda.DeviceBuffer.fromHost(driver, ow.scales); + defer os_b.free(); + var ob_b = try cuda.DeviceBuffer.fromHost(driver, ow.biases); + defer ob_b.free(); + var branch_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer branch_b.free(); + try qmm.matmul(driver, stream.*, gated_b.ptr, gxs_b.ptr, ow_b.ptr, os_b.ptr, ob_b.ptr, branch_b.ptr, 1, dims, out_k); + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + const xs_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xs_raw); + const pa_raw = try gpa.alloc(u8, proj_n * 2); + defer gpa.free(pa_raw); + const q_raw = try gpa.alloc(u8, q_heads * head_dim * 2); + defer gpa.free(q_raw); + const k_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(k_raw); + const v_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(v_raw); + const iq_raw = try gpa.alloc(u8, index_heads * index_dim * 2); + defer gpa.free(iq_raw); + const ik_raw = try gpa.alloc(u8, index_dim * 2); + defer gpa.free(ik_raw); + const o_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(o_raw); + const gated_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(gated_raw); + const gxs_raw = try gpa.alloc(u8, (out_k / 32) * 4); + defer gpa.free(gxs_raw); + const branch_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(branch_raw); + try stream.synchronize(); + var mixed_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + var xs_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = xs, .len = xs_raw.len }; + try mixed_view.download(0, row_raw); + try xs_view.download(0, xs_raw); + try pa_b.download(0, pa_raw); + try q_b.download(0, q_raw); + try kc_b.download(0, k_raw); + try vc_b.download(0, v_raw); + try iq_b.download(0, iq_raw); + try ik_b.download(0, ik_raw); + try o_b.download(0, o_raw); + try gated_b.download(0, gated_raw); + try gxs_b.download(0, gxs_raw); + try branch_b.download(0, branch_raw); + + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + readBf(row, row_raw); + const xs_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const host_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(host_pa); + try qmm.dotRow(row, xs_f, stacked.words, stacked.scales, stacked.biases, proj_n, dims, host_pa); + const gpu_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(gpu_pa); + readBf(gpu_pa, pa_raw); + const proj = gap(gpu_pa, host_pa); + + const host_q = try gpa.alloc(u16, q_heads * head_dim); + defer gpa.free(host_q); + const host_k = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_k); + const host_v = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_v); + const host_iq = try gpa.alloc(u16, index_heads * index_dim); + defer gpa.free(host_iq); + const host_ik = try gpa.alloc(u16, index_dim); + defer gpa.free(host_ik); + prepare(gpu_pa, &q_scale, &k_scale, &i_scale, &inv, host_q, host_k, host_v, host_iq, host_ik); + const gpu_q = try gpa.alloc(u16, host_q.len); + defer gpa.free(gpu_q); + const gpu_k = try gpa.alloc(u16, host_k.len); + defer gpa.free(gpu_k); + const gpu_v = try gpa.alloc(u16, host_v.len); + defer gpa.free(gpu_v); + const gpu_iq = try gpa.alloc(u16, host_iq.len); + defer gpa.free(gpu_iq); + const gpu_ik = try gpa.alloc(u16, host_ik.len); + defer gpa.free(gpu_ik); + readBf(gpu_q, q_raw); + readBf(gpu_k, k_raw); + readBf(gpu_v, v_raw); + readBf(gpu_iq, iq_raw); + readBf(gpu_ik, ik_raw); + const qg = gap(gpu_q, host_q); + const kg = gap(gpu_k, host_k); + const vg = gap(gpu_v, host_v); + const iqg = gap(gpu_iq, host_iq); + const ikg = gap(gpu_ik, host_ik); + + const host_o = try gpa.alloc(u16, out_k); + defer gpa.free(host_o); + hostAttend(gpu_q, gpu_k, gpu_v, host_o); + const gpu_o = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_o); + readBf(gpu_o, o_raw); + const og = gap(gpu_o, host_o); + + const gate_row = try gpa.alloc(u16, out_k); + defer gpa.free(gate_row); + gatesOf(gpu_pa, gate_row); + const host_gated = try gpa.alloc(u16, out_k); + defer gpa.free(host_gated); + const host_gxs = try gpa.alloc(f32, out_k / 32); + defer gpa.free(host_gxs); + applyGate(gpu_o, gate_row, host_gated, host_gxs); + const gpu_gated = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_gated); + readBf(gpu_gated, gated_raw); + const gg = gap(gpu_gated, host_gated); + var xs_ulp: u32 = 0; + for (host_gxs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const gxs_f = try gpa.alloc(f32, out_k / 32); + defer gpa.free(gxs_f); + for (gxs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const host_branch = try gpa.alloc(u16, dims); + defer gpa.free(host_branch); + try qmm.dotRow(gpu_gated, gxs_f, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k, host_branch); + const gpu_branch = try gpa.alloc(u16, dims); + defer gpa.free(gpu_branch); + readBf(gpu_branch, branch_raw); + const bg = gap(gpu_branch, host_branch); + + std.debug.print("attn proj n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ proj_n, proj.off, proj.steps, proj.at, gpu_pa[proj.at], host_pa[proj.at], gpu_pa[0] }); + std.debug.print("attn prep q {d}/{d} k {d}/{d} v {d}/{d} iq {d}/{d} ik {d}/{d} head {x:0>4} host {x:0>4}\n", .{ qg.off, qg.steps, kg.off, kg.steps, vg.off, vg.steps, iqg.off, iqg.steps, ikg.off, ikg.steps, gpu_q[0], host_q[0] }); + std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); + std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); + std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); + if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.off > 1 or kg.steps > 1 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 8192 or bg.steps != 0) return 1; + return 0; +} + +test "position zero stores the normalized query" { + var x: [head_dim]u16 = @splat(0); + x[0] = 0x3f80; + var scale_w: [head_dim]f32 = @splat(1); + var inv: [half]f32 = @splat(1); + var out: [head_dim]u16 = undefined; + normRot(&x, &scale_w, &inv, 0, &out); + try std.testing.expectEqual(@as(u16, 0x4180), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[half]); +} + +test "the gate rounds one times sigmoid of zero to one half" { + var o: [head_dim]u16 = @splat(0); + var g: [head_dim]u16 = @splat(0); + o[0] = 0x3f80; + var out: [head_dim]u16 = undefined; + var xs: [head_dim / 32]f32 = undefined; + applyGate(&o, &g, &out, &xs); + try std.testing.expectEqual(@as(u16, 0x3f00), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig new file mode 100644 index 0000000000..0930545bd8 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -0,0 +1,581 @@ +//! The batched pass: a chunk of rows through each layer at once, each row's bits those of a one-row pass. + +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); +const parts = @import("cuda_batch_state.zig"); + +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const vocab = wt.vocab; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const Layer = load.Layer; + +/// Normed-stream partial sums a row: D / 256 blocks of each stream. +pub const pss_n = (dims / 256) * streams; +/// The widest projection row: the DeltaNet's q|k|v|z|b|a (attention's q|k|v|indexer is 13952). +pub const proj_row = @max(gdn.proj_width, proj_n); +/// Attention rows a _chunks launch takes (its scratch: rows x nch x heads x dim fp32). +pub const att_rows = parts.att_rows; +/// Rows a pass holds at most unless asked for more. +pub const default_chunk = parts.default_chunk; +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +const qmm_smem: u32 = 9472; +/// Pairs an expert item takes: the core rows kernel's tile. +const rows_tile: usize = cuda.experts.RowsExperts.tile; +const plan_small: usize = cuda.grouped.small; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} +pub const Logits = parts.Logits; +pub const LayerKeep = parts.LayerKeep; +pub const Keep = parts.Keep; +pub const Hooks = parts.Hooks; +pub const Kernels = parts.Kernels; +pub const Scratch = parts.Scratch; +pub const SeqState = parts.SeqState; +pub const scoreStride = parts.scoreStride; +pub const scoreBytes = parts.scoreBytes; +pub const Lane = parts.Lane; +pub const Model = parts.Model; + +/// TF_FLASHNEXT_PROFILE=1: each chunk's wall time (synchronized) and its host n-gram gather on stderr. +fn profile() bool { + const v = std.c.getenv("TF_FLASHNEXT_PROFILE") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +pub fn Pass(comptime Tri: type) type { + return struct { + const Self = @This(); + + gpa: std.mem.Allocator, + io: std.Io, + stream: cuda.Stream, + tri: Tri, + k: *const Kernels, + s: *Scratch, + m: Model, + seq: SeqState, + hooks: Hooks, + /// A decode round: row r is lanes[r]'s token at its own position and state; empty, the rows continue `seq`. + lanes: []const Lane = &.{}, + rows: usize = 0, + pos0: usize = 0, + /// The MoE write-back the next layer (or the head) applies first, and the layer it came from. + pending: ?usize = null, + prepared: bool = false, + + // ---- launches ---------------------------------------------------------------------------------------- + + /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. + pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !void { + if (w.native) |native| { + if (sk == 1) return self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); + // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. + try self.k.native_lane.matmulF32(self.stream, x, w.k, m, native, part, self.s.native_part); + const bytes = m * w.n * 4; + if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); + return; + } + const bm: usize = 16; + const bn: usize = 64; + const k = w.k; + const rows_t = cdiv(m, bm); + const cols_t = cdiv(w.n, bn); + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (w.n + 127) / 128 * 128; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(w.w.ptr); + a.add(w.s.ptr); + a.add(w.b.ptr); + a.add(out); + a.add(part); + for ([_]usize{ m, w.n, k, sk, npad, k, sweep }) |v| a.add(cint(v)); + try cuda.launch.launch(self.k.lane, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, .block = .{ .x = 128 }, .shared = qmm_smem }, self.stream, &a); + } + + pub fn sumSlices(self: *Self, part: u64, out: u64, bf16: bool, n: usize, sk: usize) !void { + var a: cuda.Args = .{}; + a.add(part); + a.add(out); + a.add(@as(c_int, @intFromBool(bf16))); + a.add(@as(c_longlong, @intCast(n))); + a.add(cint(sk)); + try cuda.launch.launch(self.k.sum, .{ .grid = .{ .x = @intCast(cdiv(n, 256)) }, .block = .{ .x = 256 } }, self.stream, &a); + } + + pub fn shiftWindow(self: *Self, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { + var a: cuda.Args = .{}; + a.add(win); + a.add(rows_ptr); + a.add(@as(c_longlong, @intCast(stride))); + a.add(cint(n)); + a.add(cint(taps)); + a.add(cint(c)); + try cuda.launch.launch(self.k.shift, .{ .grid = .{ .x = @intCast(cdiv(c, 256)) }, .block = .{ .x = 256 } }, self.stream, &a); + } + + pub fn copy(self: *Self, dst: u64, src: u64, n: usize) !void { + if (dst == 0 or n == 0) return; + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoDAsync_v2(dst, src, n, self.stream.handle), "cuMemcpyDtoDAsync"); + } + + /// A queued copy from pinned host bytes (the scratch's staging), which must stay unchanged until it lands. + pub fn upload(self: *Self, dst: u64, bytes: []const u8) !void { + try self.s.upload(dst, bytes); + } + + // ---- the blocks ---------------------------------------------------------------------------------------- + + /// The pending MoE output (or none: mode 0, the streams' sums alone) written back into rows [r0, r0 + n). + pub fn writeback(self: *Self, r0: usize, n: usize) !void { + const s = self.s; + const dummy = s.dummy; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const inj = s.inj + r0 * streams * 2; + const rs = self.rows * dims; + if (self.pending == null) { + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", rs, n, dims, streams, 0, 1, 1, 1); + return; + } + self.pending = null; + const y = s.y + r0 * slots * dims * 4; + const wts = s.wts + r0 * slots * 4; + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); + } + + /// A read-out over rows [r0, r0 + n) of the streams: mixed and its group sums, and the inject gates. + pub fn readout(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { + const s = self.s; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const inj = s.inj + r0 * streams * 2; + const norm = s.norm + r0 * wide * 2; + const act = s.act + r0 * low * 2; + const xs_act = s.xs_act + r0 * (low / 32) * 4; + if (hc_w.down.native) |down| { + const up = hc_w.up.native orelse return error.UnexpectedTensor; + try self.tri.hcNormalize(h, pss, hc_w.scale.ptr, norm, eps, n, dims, streams); + try self.k.native_lane.matmulF32(self.stream, norm, wide, n, down, s.hc_part, s.native_part); + const inj_ptr = if (hc_w.has_inj == 1) inj else act; + try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); + try self.k.native_lane.matmul(self.stream, act, low, n, up, s.hc_up, s.native_part); + try self.tri.hcDenseMix(s.hc_up, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); + return; + } + try self.tri.hcDown(h, pss, hc_w.scale.ptr, norm, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn, s.hc_part, eps, n, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else act; + try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, hc_sk, n, streams, low, hc_w.n_down, hc_w.has_inj); + try self.tri.hcUpmix(act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, wide, low, dims, streams); + } + + fn pleBlock(self: *Self, lw: *const Layer, ids: []const u32) !void { + const s = self.s; + const n = self.rows; + // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` + _ = ids; + try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); + try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); + try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, 4); + try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); + if (self.lanes.len == 0) { + try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); + try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); + return; + } + // a round: each row's convolution over its own stream's tail, the one-row launch a lone decode makes + for (self.lanes, 0..) |l, r| { + const g = r * wide * 2; + try self.tri.pleConv(s.ple_gated + g, s.ple_pss + r * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, 1, dims, streams, 4, 3); + try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, 1, ple_tail_n, wide); + } + } + + fn gdnBlock(self: *Self, lw: *const Layer, li: usize, keep: LayerKeep) !void { + const s = self.s; + const n = self.rows; + try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1); + try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2); + try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); + // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits + if (self.lanes.len == 0) { + try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); + try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); + } else for (self.lanes, 0..) |l, r| { + // a round: each row's front and recurrence on its own stream's conv window and state + const qk = r * gdn.nk * gdn.dk * 4; + const v = r * gdn.value_dim * 4; + const gb = r * gdn.nv * 2 * 4; + try self.gdnFront(lw, s.proj + r * gdn.proj_width * 2, l.seq.conv[li].ptr, 1, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); + try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, 0, 1, s.ys + v); + } + var b: cuda.Args = .{}; + b.add(s.ys); + b.add(s.proj); + b.add(lw.norm.ptr); + b.add(eps); + b.add(s.gout); + b.add(s.gxs); + try cuda.launch.launch(self.k.gnorm, .{ .grid = .{ .x = @intCast(n), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, self.stream, &b); + if (self.lanes.len == 0) { + try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim); + } else for (self.lanes, 0..) |l, r| { + try self.shiftWindow(l.seq.conv[li].ptr, s.proj + r * gdn.proj_width * 2, gdn.proj_width, 1, gdn.taps - 1, gdn.conv_dim); + } + try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk); + } + + /// fn_gdn_front over `rows` projection rows after the conv window `conv`. + fn gdnFront(self: *Self, lw: *const Layer, proj: u64, conv: u64, rows: usize, gq: u64, gk: u64, gv: u64, ggb: u64) !void { + var f: cuda.Args = .{}; + f.add(proj); + f.add(conv); + f.add(lw.cw.ptr); + f.add(lw.a_log.ptr); + f.add(lw.dt.ptr); + f.add(gq); + f.add(gk); + f.add(gv); + f.add(ggb); + try cuda.launch.launch(self.k.front, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, self.stream, &f); + } + + /// fn_gdn_rows over `rows` rows from `state`, in place or (state_out) to fn_gdn_rows_to's; readouts into ys. + fn gdnChain(self: *Self, gq: u64, gk: u64, gv: u64, ggb: u64, state: u64, state_out: u64, rows: usize, ys: u64) !void { + var a: cuda.Args = .{}; + a.add(gq); + a.add(gk); + a.add(gv); + a.add(ggb); + a.add(state); + if (state_out != 0) a.add(state_out); + a.add(cint(rows)); + a.add(ys); + try cuda.launch.launch(if (state_out != 0) self.k.chain_to else self.k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, self.stream, &a); + } + + pub fn attnBlock(self: *Self, lw: *const Layer, ai: usize) !void { + const s = self.s; + const n = self.rows; + const pa = s.proj; + try self.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); + if (self.lanes.len > 0) { + // a round: each row's rope, cache writes, selection and attention on its own stream's caches + const own_seq = self.seq; + const own_pos0 = self.pos0; + defer { + self.seq = own_seq; + self.pos0 = own_pos0; + } + for (self.lanes, 0..) |l, r| { + self.seq = l.seq; + self.pos0 = l.pos; + try self.attnRows(lw, ai, pa + r * proj_n * 2, 1, .{ .pos = s.pos + r * 16, .q = s.q + r * q_heads * head_dim * 2, .iq = s.iq + r * index_heads * index_dim * 2, .ao = s.ao + r * out_k * 2 }); + } + } else try self.attnRows(lw, ai, pa, n, .{ .pos = s.pos, .q = s.q, .iq = s.iq, .ao = s.ao }); + try self.tri.attnGate(s.ao, pa, s.gated, s.gxs_attn, n, proj_n, q_heads, head_dim); + try self.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, out_sk); + } + + /// Where attnRows reads its rows' positions (16 bytes apart) and writes queries, indexer queries and outputs. + const RowsAt = struct { pos: u64, q: u64, iq: u64, ao: u64 }; + + /// `n` rows of self.seq from self.pos0: rope, cache writes and pooling, then scores, selection and attention. + fn attnRows(self: *Self, lw: *const Layer, ai: usize, pa: u64, n: usize, at: RowsAt) !void { + const s = self.s; + const kc = self.seq.kc[ai].ptr; + const vc = self.seq.vc[ai].ptr; + const ikc = self.seq.ikc[ai].ptr; + const pooled = self.seq.pooled[ai].ptr; + try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0); + // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention + try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n); + const nb = scoreStride(self.seq.cap); + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = self.pos0 + r0 + m; + if (ends > self.seq.cap) return error.UnexpectedTensor; + const pos = at.pos + b * 16; + const iq = at.iq + r0 * index_heads * index_dim * 2; + // rows past `top` complete blocks score theirs; _scores leaves the rest alone, so skip it when none is + if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(self.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + a.add(cint(nb)); + a.add(cint(idw)); + a.add(cint(top)); + a.add(cint(ratio)); + try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true); + try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + } + } + + /// Router, each row's ten experts and the shared one, the grouped expert launches: y [rows, 11, dims] fp32. + pub fn moeBlock(self: *Self, lw: *const Layer) !void { + const s = self.s; + const n = self.rows; + const pairs = n * slots; + try self.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try self.tri.topkMany(s.router_l, s.picks, s.wts, n); + const e = experts_n + 1; + const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; + if (lw.native_experts) |*native| { + const owned = &native.experts; + try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); + try self.k.router.route(self.stream, s.picks, pairs, e, 16, plan); + const items = cuda.grouped.maxItems(pairs, e, 16); + try self.k.nvfp4.run(self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, experts_n); + try self.k.nvfp4.run(self.stream, .down_f32, .{ .x = s.moe_act, .stride = owned.width }, owned.view(), plan, s.y, items, experts_n); + if (native.rank == 0) { + const glue = self.tri.own orelse return error.NativeGlueRequired; + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part); + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part); + try glue.sharedSwiGLU(self.stream, s.shared_gate, s.shared_up, s.shared_act, n, moe_w); + try self.k.native_lane.matmulF32(self.stream, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part); + try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, slots); + } + if (native.world != 1) { + const collective = self.m.collective orelse return error.TensorParallelUnavailable; + try collective.sum(self.stream, s.y, pairs * dims); + } + return; + } + try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); + const items = cdiv(pairs, rows_tile) + e; + try self.k.moe.run(self.stream, true, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items); + try self.k.moe.run(self.stream, false, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items); + } + + // ---- the pass ------------------------------------------------------------------------------------------ + + /// Token ids and positions onto the device from the scratch's pinned copies, which stay put until they land. + pub fn stage(self: *Self, ids: []const u32) !void { + const s = self.s; + const n = ids.len; + if (!self.prepared) try s.ready(); + const id32 = s.hostIds(n); + for (id32, ids) |*o, v| o.* = @intCast(v); + try self.upload(s.ids, std.mem.sliceAsBytes(id32)); + const blocks = if (self.lanes.len > 0) n else cdiv(n, att_rows); + // one int a block (a round: a row), 16 bytes apart: _chunks and _merge take POS0 16-byte aligned + const pos = s.hostPos(blocks * 4); + @memset(pos, 0); + for (0..blocks) |b| pos[4 * b] = @intCast(if (self.lanes.len > 0) self.lanes[b].pos else self.pos0 + b * att_rows); + try self.upload(s.pos, std.mem.sliceAsBytes(pos)); + if (!self.prepared) try s.uploaded(); + } + + /// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. + const Gather = struct { + ngram: *const ple.Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). + lanes: []const Lane = &.{}, + + fn main(g: *Gather) void { + const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); + defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; + if (g.lanes.len == 0) { + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + return; + } + for (g.lanes, 0..) |l, r| { + g.ngram.rows(l.seq.history[0..], g.ids[r..][0..1], g.emb[r * dims ..][0..dims], g.xs[r * (dims / 32) ..][0 .. dims / 32]) catch |err| { + g.failed = err; + return; + }; + } + } + }; + + /// One chunk of rows at positions pos0.. through every layer and the head. + pub fn chunk(self: *Self, ids: []const u32, pos0: usize, logits_host: ?[]u8) !void { + const s = self.s; + const n = ids.len; + const t0 = std.Io.Clock.awake.now(self.io).toNanoseconds(); + defer if (profile()) { + self.complete() catch {}; + std.debug.print("batch: {d} rows at {d} in {d:.1} ms\n", .{ n, pos0, @as(f64, @floatFromInt(std.Io.Clock.awake.now(self.io).toNanoseconds() - t0)) / 1e6 }); + }; + if (n == 0 or n > s.cap) return error.UnexpectedTensor; + if (self.lanes.len > 0) { + // a round reads every row's scores and keeps nothing for a rollback + if (n != self.lanes.len or self.hooks.logits != .all or self.hooks.keep != null or self.hooks.streams != 0) return error.UnexpectedTensor; + for (self.lanes) |l| if (l.pos + 1 > l.seq.cap) return error.UnexpectedTensor; + } else if (pos0 + n > self.seq.cap) return error.UnexpectedTensor; + self.rows = n; + self.pos0 = pos0; + self.pending = null; + // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first + if (!self.prepared) try s.ready(); + const emb = s.hostEmb(n); + const xs = s.hostXs(n); + var gather: Gather = .{ .ngram = self.m.ngram, .history = self.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = self.io, .lanes = self.lanes }; + var thread: ?std.Thread = if (self.prepared) null else try std.Thread.spawn(.{}, Gather.main, .{&gather}); + defer if (thread) |th| th.join(); + try self.stage(ids); + if (self.m.table_dense) { + try self.tri.embedDense(s.ids, self.m.table_w, s.h, dims, streams, n); + } else { + try self.tri.embed(s.ids, self.m.table_w, self.m.table_s, self.m.table_b, s.h, dims, streams, n); + } + for (self.m.layers, 0..) |*lw, layer| { + const keep: LayerKeep = if (self.hooks.keep) |kp| kp.layers[layer] else .{}; + if (layer == ple_layer) { + if (thread) |th| th.join(); + thread = null; + if (gather.failed) |err| return err; + if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); + try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); + try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); + if (!self.prepared) try s.uploaded(); + if (self.pending != null) try self.writeback(0, n); + try self.pleBlock(lw, ids); + } + try self.writeback(0, n); + try self.readout(lw.attn_hc, 0, n); + if (lw.linear) { + try self.gdnBlock(lw, wt.linearIndex(layer), keep); + } else { + try self.attnBlock(lw, wt.attnIndex(layer)); + } + try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, out_sk); + try self.readout(lw.mlp_hc, 0, n); + try self.moeBlock(lw); + self.pending = layer; + } + try self.head(logits_host); + } + + fn complete(self: *Self) !void { + if (self.m.collective) |c| return c.finish(self.stream); + try self.stream.synchronize(); + } + + /// Column tiles keep their own reduction order; each gathered row restores global vocabulary order. + fn projectHead(self: *Self, out: u64, n: usize) !void { + const s = self.s; + if (self.m.head.n == vocab) return self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; + const c = self.m.collective orelse return error.TensorParallelUnavailable; + try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); + for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); + } + + fn head(self: *Self, logits_host: ?[]u8) !void { + const s = self.s; + const n = self.rows; + switch (self.hooks.logits) { + .none => { + if (self.hooks.streams == 0) return; + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + return; + }, + .all => { + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + try self.readout(self.m.mixer, 0, n); + const out = if (self.hooks.logits_out != 0) self.hooks.logits_out else return error.UnexpectedTensor; + try self.projectHead(out, n); + if (logits_host) |lh| { + try self.complete(); + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, out + (n - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + } + }, + .last => { + // the last row's streams, sums, gates and MoE slots moved to row 0, which keeps 16-byte alignment + const r = n - 1; + if (self.hooks.streams != 0) { + // every row written back and kept, then the last row's streams and sums at row 0 + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + if (r > 0) { + try self.copy(s.h, s.h + r * wide * 2, wide * 2); + try self.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + } + } else if (r > 0) { + try self.copy(s.h, s.h + r * wide * 2, wide * 2); + try self.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + try self.copy(s.inj, s.inj + r * streams * 2, streams * 2); + try self.copy(s.y, s.y + r * slots * dims * 4, slots * dims * 4); + try self.copy(s.wts, s.wts + r * slots * 4, slots * 4); + } + if (self.hooks.streams == 0) try self.writeback(0, 1); + try self.readout(self.m.mixer, 0, 1); + try self.projectHead(s.logits, 1); + if (logits_host) |lh| { + try self.complete(); + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, s.logits, vocab * 2), "cuMemcpyDtoH"); + } + }, + } + } + }; +} + +test "a pass's scratch at the dense cap stays near a gigabyte and a half" { + try std.testing.expect(Scratch.bytes(default_chunk) < 2 << 30); + try std.testing.expect(Scratch.bytes(1) < Scratch.bytes(2)); +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig new file mode 100644 index 0000000000..f837db0887 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -0,0 +1,527 @@ +//! The batched pass's kernels, scratch, a sequence's view, and the hooks a verify window or the MTP head passes. + +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); + +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const vocab = wt.vocab; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const Layer = load.Layer; + +/// Normed-stream partial sums a row: D / 256 blocks of each stream. +pub const pss_n = (dims / 256) * streams; +/// The widest projection row: the DeltaNet's q|k|v|z|b|a (attention's q|k|v|indexer is 13952). +pub const proj_row = @max(gdn.proj_width, proj_n); +/// Attention rows a _chunks launch takes (its scratch: rows x nch x heads x dim fp32). +pub const att_rows: usize = 256; +/// Rows a pass holds at most unless asked for more. +pub const default_chunk: usize = 1536; +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +const qmm_smem: u32 = 9472; +/// Pairs an expert item takes: the core rows kernel's tile. +const rows_tile: usize = cuda.experts.RowsExperts.tile; +const plan_small: usize = cuda.grouped.small; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// Which rows get the head: the last (a prompt), every row (a verify window or a round), or none. +pub const Logits = enum { last, all, none }; + +/// A DeltaNet layer's window keep: projection rows, conv history before the chunk, and where its end state goes. +pub const LayerKeep = struct { + proj: u64 = 0, + hist: u64 = 0, + /// The chain's end state goes here, the sequence's own left as before the chunk; 0: in place. + state_out: u64 = 0, +}; + +/// What a drafted window keeps for its rollback (cuda_mtp.zig): each DeltaNet layer's rows and states. +pub const Keep = struct { + layers: [layers_n]LayerKeep = @splat(.{}), +}; + +/// A caller's view into the pass. Without any, the pass is the server's prompt pass. +pub const Hooks = struct { + logits: Logits = .last, + /// Device [rows, vocab] bf16 for `.all` (else the pass's own buffer for the last row). + logits_out: u64 = 0, + keep: ?*Keep = null, + /// Device [rows, wide] bf16: every row's streams after the last MoE write-back (the MTP head's input); 0: none. + streams: u64 = 0, +}; + +/// The kernels the pass launches besides the captured Triton ones, loaded once. +pub const Kernels = struct { + qmm: cuda.Module, + experts: cuda.Module, + rows: cuda.Module, + batch: cuda.Module, + /// The core's expert plan and its rows kernels. + router: cuda.grouped.Router, + moe: cuda.experts.RowsExperts, + lane: cuda.Function, + native_lane: cuda.qmmf.Lane, + nvfp4: cuda.experts.Nvfp4Experts, + front: cuda.Function, + chain: cuda.Function, + chain_to: cuda.Function, + gnorm: cuda.Function, + sum: cuda.Function, + shift: cuda.Function, + /// Each sparse row's top blocks (attention._select / _select_tiles), native. + select: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + k.qmm = try cuda.Module.load(d, cuda.kernels.qmm); + errdefer k.qmm.unload(); + k.experts = try cuda.Module.load(d, cuda.kernels.experts); + errdefer k.experts.unload(); + k.rows = try cuda.Module.load(d, cuda.kernels.experts_rows); + errdefer k.rows.unload(); + k.batch = try cuda.Module.load(d, cuda.kernels.flashnext_batch); + errdefer k.batch.unload(); + k.lane = try k.qmm.function(cubin_qmm); + try k.lane.allowDynamicShared(qmm_smem); + k.router = try cuda.grouped.Router.resolve(k.experts); + k.moe = try cuda.experts.RowsExperts.resolve(k.rows); + k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); + errdefer k.native_lane.unload(); + var device: c_int = 0; + var sms: c_int = 0; + try d.check(d.api.cuDeviceGet(&device, 0), "cuDeviceGet"); + try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, device), "cuDeviceGetAttribute"); + k.nvfp4 = try cuda.experts.Nvfp4Experts.load(d, @intCast(sms)); + errdefer k.nvfp4.unload(); + k.sum = try k.batch.function("fn_sum_slices"); + k.shift = try k.batch.function("fn_shift_window"); + k.select = try k.batch.function("fn_qsa_select"); + k.front = try k.batch.function("fn_gdn_front"); + k.chain = try k.batch.function("fn_gdn_rows"); + k.chain_to = try k.batch.function("fn_gdn_rows_to"); + k.gnorm = try k.batch.function("fn_gdn_norm"); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.nvfp4.unload(); + k.native_lane.unload(); + k.batch.unload(); + k.rows.unload(); + k.experts.unload(); + k.qmm.unload(); + } +}; + +/// The compute stream as a pass's pinned staging uses it: copies from it, an event behind them, and a drain. +pub const Uploads = struct { + stream: cuda.Stream, + event: cuda.Event, + + pub fn copy(u: Uploads, dst: u64, bytes: []const u8) !void { + const d = u.stream.d; + try d.check(d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, u.stream.handle), "cuMemcpyHtoDAsync"); + } + pub fn record(u: Uploads) !void { + try u.event.record(u.stream); + } + pub fn wait(u: Uploads) !void { + try u.event.synchronize(); + } + pub fn drain(u: Uploads) !void { + try u.stream.synchronize(); + } +}; + +/// Pinned bytes' uploads and the fence that says when the bytes may change (`Q`: `Uploads`, or a test's stand-in). +pub fn Staging(comptime Q: type) type { + return struct { + const Self = @This(); + q: Q, + /// An upload was queued after the event's last record (a copy or the record failed): only a drain covers it. + open: bool = false, + + /// Queues a copy from the staged bytes; until `fence` records behind it, `ready` drains the stream. + pub fn upload(g: *Self, dst: u64, bytes: []const u8) !void { + g.open = true; + try g.q.copy(dst, bytes); + } + + /// Records the event behind every upload queued so far. + pub fn fence(g: *Self) !void { + try g.q.record(); + g.open = false; + } + + /// Returns once no queued upload reads the staged bytes: the event's wait, or a drain if one came after it. + pub fn ready(g: *Self) !void { + if (!g.open) return g.q.wait(); + try g.q.drain(); + g.open = false; + } + }; +} + +/// Device scratch for chunks of up to `cap` rows. +pub const Scratch = struct { + cap: usize, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + ids: u64 = 0, + pos: u64 = 0, + h: u64 = 0, + pss: u64 = 0, + norm: u64 = 0, + dn: u64 = 0, + hc_part: u64 = 0, + hc_up: u64 = 0, + native_part: u64 = 0, + act: u64 = 0, + xs_act: u64 = 0, + inj: u64 = 0, + mixed: u64 = 0, + xsm: u64 = 0, + proj: u64 = 0, + part: u64 = 0, + gout: u64 = 0, + gxs: u64 = 0, + ys: u64 = 0, + gq: u64 = 0, + gk: u64 = 0, + gv: u64 = 0, + ggb: u64 = 0, + q: u64 = 0, + iq: u64 = 0, + po: u64 = 0, + pm: u64 = 0, + pl: u64 = 0, + /// An attention block's rows' key lists, key counts and sparse flags (attention.AttnScratch ids/nk/sparse). + qids: u64 = 0, + qnk: u64 = 0, + qspr: u64 = 0, + ao: u64 = 0, + gated: u64 = 0, + gxs_attn: u64 = 0, + router_l: u64 = 0, + picks: u64 = 0, + wts: u64 = 0, + members: u64 = 0, + items: u64 = 0, + counts: u64 = 0, + rank: u64 = 0, + hist: u64 = 0, + moe_act: u64 = 0, + shared_gate: u64 = 0, + shared_up: u64 = 0, + shared_act: u64 = 0, + shared_out: u64 = 0, + y: u64 = 0, + branch: u64 = 0, + ple_emb: u64 = 0, + ple_xs: u64 = 0, + ple_keys: u64 = 0, + ple_vals: u64 = 0, + ple_part: u64 = 0, + ple_gated: u64 = 0, + ple_pss: u64 = 0, + ple_nrow: u64 = 0, + logits: u64 = 0, + head_local: u64 = 0, + inv: u64 = 0, + dummy: u64 = 0, + /// Pinned host copies of what a pass uploads (ids, positions, n-gram rows and sums), owned until the copies land. + host: ?cuda.HostBuffer = null, + /// The uploads from `host` and the event recorded after a pass's last one: the next pass waits before writing it. + staging: ?Staging(Uploads) = null, + + fn add(s: *Scratch, gpa: std.mem.Allocator, d: *const cuda.Driver, len: usize) !u64 { + var b = try cuda.DeviceBuffer.alloc(d, @max(len, 256)); + errdefer b.free(); + try s.bufs.append(gpa, b); + return b.ptr; + } + + /// Byte offsets of the staged ids, positions, n-gram rows and their sums in `host`, and its length. + fn hostAt(cap: usize) [5]usize { + const pos = std.mem.alignForward(usize, cap * 4, 16); + const emb = pos + (cap + 1) * 16; + const xs = emb + cap * dims * 2; + return .{ 0, pos, emb, xs, xs + cap * (dims / 32) * 4 }; + } + + /// Waits until the last pass's uploads have read `host` (they are queued behind its kernels), so it may change. + pub fn ready(s: *Scratch) !void { + if (s.staging) |*g| try g.ready(); + } + + /// A queued copy from `host` to the device, which must stay unchanged until it lands. + pub fn upload(s: *Scratch, dst: u64, src: []const u8) !void { + try s.staging.?.upload(dst, src); + } + + /// Marks `host` read once the stream reaches here: every upload queued before this point has copied its bytes. + pub fn uploaded(s: *Scratch) !void { + if (s.staging) |*g| try g.fence(); + } + + /// The staged ids [n], positions [n] (in 16-byte rows), n-gram rows [n, dims] and their 32-group sums. + pub fn hostIds(s: *const Scratch, n: usize) []i32 { + return @alignCast(std.mem.bytesAsSlice(i32, s.host.?.bytes[hostAt(s.cap)[0]..][0 .. n * 4])); + } + pub fn hostPos(s: *const Scratch, n: usize) []i32 { + return @alignCast(std.mem.bytesAsSlice(i32, s.host.?.bytes[hostAt(s.cap)[1]..][0 .. n * 4])); + } + pub fn hostEmb(s: *const Scratch, n: usize) []u16 { + return @alignCast(std.mem.bytesAsSlice(u16, s.host.?.bytes[hostAt(s.cap)[2]..][0 .. n * dims * 2])); + } + pub fn hostXs(s: *const Scratch, n: usize) []f32 { + return @alignCast(std.mem.bytesAsSlice(f32, s.host.?.bytes[hostAt(s.cap)[3]..][0 .. n * (dims / 32) * 4])); + } + + /// Bytes `alloc(cap)` takes on the device. + pub fn bytes(cap: usize) usize { + const pairs = cap * slots; + const nblk = cdiv(pairs, plan_small); + const items = cdiv(pairs, tile) + experts_n + 1; + var n: usize = 0; + n += cap * 4 + (cap + 1) * 16; + n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; + n += cap * wide * 2 + 8 * cap * dims * 4; + n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; + n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; + n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; + n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; + n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; + n += att_rows * (idw + 2) * 4; + n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; + n += cap * moe_w * 6 + cap * dims * 4; + n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; + n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; + n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; + return n; + } + + pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { + var s: Scratch = .{ .cap = cap }; + errdefer s.free(gpa); + const pairs = cap * slots; + const nblk = cdiv(pairs, plan_small); + const items = cdiv(pairs, tile) + experts_n + 1; + s.ids = try s.add(gpa, d, cap * 4); + // a position every 16 bytes: one a block of att_rows rows, or one a row of a decode round + s.pos = try s.add(gpa, d, (cap + 1) * 16); + s.h = try s.add(gpa, d, cap * wide * 2); + s.pss = try s.add(gpa, d, cap * pss_n * 4); + s.norm = try s.add(gpa, d, cap * wide * 2); + s.dn = try s.add(gpa, d, cap * (low + 4) * 2); + s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); + s.hc_up = try s.add(gpa, d, cap * wide * 2); + s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); + s.act = try s.add(gpa, d, cap * low * 2); + s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); + s.inj = try s.add(gpa, d, cap * streams * 2); + s.mixed = try s.add(gpa, d, cap * dims * 2); + s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); + s.proj = try s.add(gpa, d, cap * proj_row * 2); + s.part = try s.add(gpa, d, out_sk * cap * dims * 4); + s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); + s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); + s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); + s.gq = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); + s.gk = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); + s.gv = try s.add(gpa, d, cap * gdn.value_dim * 4); + s.ggb = try s.add(gpa, d, cap * gdn.nv * 2 * 4); + s.q = try s.add(gpa, d, cap * q_heads * head_dim * 2); + s.iq = try s.add(gpa, d, cap * index_heads * index_dim * 2); + s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); + s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); + s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); + s.qids = try s.add(gpa, d, att_rows * idw * 4); + s.qnk = try s.add(gpa, d, att_rows * 4); + s.qspr = try s.add(gpa, d, att_rows * 4); + s.ao = try s.add(gpa, d, cap * out_k * 2); + s.gated = try s.add(gpa, d, cap * out_k * 2); + s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); + s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); + s.picks = try s.add(gpa, d, pairs * 4); + s.wts = try s.add(gpa, d, pairs * 4); + s.members = try s.add(gpa, d, pairs * 4); + s.items = try s.add(gpa, d, items * 3 * 4); + s.counts = try s.add(gpa, d, 8); + s.rank = try s.add(gpa, d, pairs * 4); + s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); + s.shared_gate = try s.add(gpa, d, cap * moe_w * 2); + s.shared_up = try s.add(gpa, d, cap * moe_w * 2); + s.shared_act = try s.add(gpa, d, cap * moe_w * 2); + s.shared_out = try s.add(gpa, d, cap * dims * 4); + s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); + s.y = try s.add(gpa, d, pairs * dims * 4); + s.branch = try s.add(gpa, d, cap * dims * 2); + s.ple_emb = try s.add(gpa, d, cap * dims * 2); + s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); + s.ple_keys = try s.add(gpa, d, cap * wide * 2); + s.ple_vals = try s.add(gpa, d, cap * dims * 2); + s.ple_part = try s.add(gpa, d, 4 * cap * dims * 4); + s.ple_gated = try s.add(gpa, d, cap * wide * 2); + s.ple_pss = try s.add(gpa, d, cap * streams * 4); + s.ple_nrow = try s.add(gpa, d, cap * wide * 2); + s.logits = try s.add(gpa, d, vocab * 2); + s.head_local = try s.add(gpa, d, cap * vocab); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try s.add(gpa, d, half * 4); + try s.bufs.items[s.bufs.items.len - 1].uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); + s.dummy = try s.add(gpa, d, 4096); + try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); + try stream.synchronize(); + s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); + s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; + return s; + } + + pub fn free(s: *Scratch, gpa: std.mem.Allocator) void { + if (s.staging) |*g| { + // no queued upload may still read `host`; a wait that fails has lost the context, and its copies with it + g.ready() catch {}; + g.q.event.deinit(); + } + if (s.host) |*h| h.free(); + for (s.bufs.items) |*b| b.free(); + s.bufs.deinit(gpa); + s.* = undefined; + } +}; + +/// A sequence's state the pass reads and leaves as the row-by-row path would. +pub const SeqState = struct { + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + /// Each attention layer's pooled indexer keys [cap / 4, 128] bf16. + pooled: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + history: *[2]i64, + /// Tokens the caches hold. + cap: usize, + /// Block scores for an attention block's rows [min(rows, att_rows), scoreStride(cap)] fp32 (the engine's). + scores: u64, +}; + +/// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). +pub fn scoreStride(cap: usize) usize { + return (cdiv(cap, ratio) + 15) / 16 * 16; +} + +/// Bytes of block scores a pass over a cache of `cap` tokens needs for chunks of up to `rows` rows. +pub fn scoreBytes(cap: usize, rows: usize) usize { + return @min(rows, att_rows) * scoreStride(cap) * 4; +} + +/// One row of a decode round: its stream's next token, on that stream's own caches at its own position. +pub const Lane = struct { seq: SeqState, pos: usize }; + +/// The engine's resident weights the pass reads. +pub const Model = struct { + collective: ?*@import("cuda_collective.zig").Collective = null, + table_w: u64, + table_s: u64, + table_b: u64, + table_dense: bool = false, + layers: []const Layer, + mixer: HcW, + head: Q4, + ngram: *const ple.Ngram, +}; + +/// A stand-in for the stream: which copies are still queued, and a failure at a chosen copy or at the record. +const FakeQueue = struct { + queued: usize = 0, + covered: usize = 0, + copies: usize = 0, + fail_copy: ?usize = null, + fail_record: bool = false, + + pub fn copy(q: *FakeQueue, _: u64, _: []const u8) !void { + q.copies += 1; + if (q.fail_copy) |k| if (k == q.copies) return error.CopyFailed; + q.queued += 1; + } + pub fn record(q: *FakeQueue) !void { + if (q.fail_record) return error.RecordFailed; + q.covered = q.queued; + } + pub fn wait(q: *FakeQueue) !void { + q.queued -= q.covered; // the event's copies have landed; any queued after it may still be reading + q.covered = 0; + } + pub fn drain(q: *FakeQueue) !void { + q.queued = 0; + q.covered = 0; + } +}; + +test "a pass whose second upload or whose record fails leaves no copy reading the staging when it is reused" { + const Case = struct { fail_copy: ?usize, fail_record: bool }; + for ([_]Case{ .{ .fail_copy = 2, .fail_record = false }, .{ .fail_copy = null, .fail_record = true }, .{ .fail_copy = null, .fail_record = false } }) |c| { + var fq: FakeQueue = .{}; + var g: Staging(*FakeQueue) = .{ .q = &fq }; + // a pass that stages and fences cleanly, then one that fails part way, as `stage()` uploads ids and positions + try g.ready(); + try g.upload(0, "ids"); + try g.upload(0, "pos"); + try g.fence(); + fq.fail_copy = if (c.fail_copy) |k| fq.copies + k else null; + fq.fail_record = c.fail_record; + try g.ready(); + const failed = blk: { + g.upload(0, "ids") catch break :blk true; + g.upload(0, "pos") catch break :blk true; + g.fence() catch break :blk true; + break :blk false; + }; + try std.testing.expectEqual(c.fail_copy != null or c.fail_record, failed); + // the next pass (or `free`) waits before it writes the bytes: nothing may still be queued + try g.ready(); + try std.testing.expectEqual(@as(usize, 0), fq.queued); + } +} diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig new file mode 100644 index 0000000000..c453cd900e --- /dev/null +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -0,0 +1,103 @@ +//! One two-rank NCCL owner. Nonblocking initialization and collectives observe the control peer and a deadline. +const std = @import("std"); +const cuda = @import("cuda"); +const wire = @import("cuda_rank_wire.zig"); +const nccl = cuda.nccl; + +const Advanced = struct { + ncclCommInitRankConfig: *const fn (*nccl.Comm, c_int, nccl.UniqueId, c_int, *nccl.Config) callconv(.c) nccl.Result, + ncclCommAbort: *const fn (nccl.Comm) callconv(.c) nccl.Result, + ncclCommGetAsyncError: *const fn (nccl.Comm, *nccl.Result) callconv(.c) nccl.Result, +}; + +pub const Collective = struct { + lib: nccl.Library, + api: Advanced, + comm: nccl.Comm = null, + link: wire.Link, + timeout_ms: u64, + failed: bool = false, + + /// `link` has completed the source/settings handshake; it stays owned by the caller. + pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { + if (rank > 1 or timeout_ms == 0) return error.BadRank; + var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms }; + errdefer c.lib.close(); + if (try c.lib.version() < 23000) return error.NcclVersion; + const info = @typeInfo(Advanced).@"struct"; + inline for (info.field_names, info.field_types) |name, T| @field(c.api, name) = c.lib.lib.lookup(T, name) orelse return error.MissingSymbol; + var id: nccl.UniqueId = undefined; + const until = wire.deadline(timeout_ms); + if (rank == 0) { + try c.lib.check(c.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try link.send(std.mem.asBytes(&id), until); + } else try link.receive(std.mem.asBytes(&id), until); + var config: nccl.Config = .{}; + errdefer c.abort(); + try c.submitted(c.api.ncclCommInitRankConfig(&c.comm, 2, id, @intCast(rank), &config)); + try c.wait(null, until); + return c; + } + + pub fn abort(c: *Collective) void { + c.failed = true; + if (c.comm) |comm| { + _ = c.api.ncclCommAbort(comm); + c.comm = null; + } + } + + pub fn deinit(c: *Collective) void { + // Abort is local and also releases an idle communicator if its peer exited first. + c.abort(); + c.lib.close(); + } + + fn submitted(c: *Collective, result: nccl.Result) !void { + if (result == 0 or result == 7) return; + return c.lib.check(result, "NCCL nonblocking operation"); + } + + fn wait(c: *Collective, stream: ?cuda.Stream, until: u64) !void { + errdefer c.abort(); + while (true) { + try c.link.checkPeer(); + if (wire.now() >= until) return error.RankTimeout; + var state: nccl.Result = 0; + try c.lib.check(c.api.ncclCommGetAsyncError(c.comm, &state), "ncclCommGetAsyncError"); + if (state != 0 and state != 7) return c.lib.check(state, "NCCL asynchronous operation"); + if (state == 0) { + if (stream) |s| { + if (try s.done()) return; + } else return; + } + wire.pause(); + } + } + + /// Wait at a host-visible boundary while retaining peer, asynchronous-error and timeout checks. + pub fn finish(c: *Collective, stream: cuda.Stream) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + try c.wait(stream, wire.deadline(c.timeout_ms)); + } + + /// Disjoint BF16 vocabulary columns are concatenated without an arithmetic reduction. + pub fn gather(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.lib.api.ncclAllGather(local, out, count, .bf16, c.comm, stream.handle)); + try c.wait(null, wire.deadline(c.timeout_ms)); + } + + /// FP32 sum is enqueued on the compute stream; later kernels on that stream observe its result. + pub fn sum(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.lib.api.ncclAllReduce(values, values, count, .f32, .sum, c.comm, stream.handle)); + try c.wait(null, wire.deadline(c.timeout_ms)); + } +}; diff --git a/zig/src/families/flashnext/cuda_compare_main.zig b/zig/src/families/flashnext/cuda_compare_main.zig new file mode 100644 index 0000000000..b1f31d54a2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_compare_main.zig @@ -0,0 +1,210 @@ +//! Our own glue against a captured set on one loaded engine: every prompt row's scores, then each path's greedy reply. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const prompt = flash.prompt; +const vocab = prompt.vocabulary; + +/// Prompt rows a teacher-forced pass takes at once (every row's scores come back). +const rows_at_once = 256; + +fn score(raw: u16) f64 { + return @as(f32, @bitCast(@as(u32, raw) << 16)); +} + +/// What two paths' scores over the same rows say: top-1 agreement, the largest score gap, KL and each path's NLL. +const Tally = struct { + rows: usize = 0, + agree: usize = 0, + max_abs: f64 = 0, + kl: f64 = 0, + kl_max: f64 = 0, + nll_a: f64 = 0, + nll_b: f64 = 0, + scored: usize = 0, + /// Among rows whose top-1 differs, the largest gap between the captured path's best two scores. + tie_gap: f64 = 0, + /// Each row's KL, for its median and the share of rows past 0.1, and each scored row's NLL own minus captured. + kls: std.ArrayList(f64) = .empty, + dnll: std.ArrayList(f64) = .empty, + + fn row(t: *Tally, a: []const u16, b: []const u16, next: ?u32) void { + var best_a: usize = 0; + var best_b: usize = 0; + var max_a = -std.math.inf(f64); + var max_b = -std.math.inf(f64); + for (a, b, 0..) |ra, rb, i| { + const va = score(ra); + const vb = score(rb); + t.max_abs = @max(t.max_abs, @abs(va - vb)); + if (va > max_a) { + max_a = va; + best_a = i; + } + if (vb > max_b) { + max_b = vb; + best_b = i; + } + } + var sum_a: f64 = 0; + var sum_b: f64 = 0; + for (a, b) |ra, rb| { + sum_a += @exp(score(ra) - max_a); + sum_b += @exp(score(rb) - max_b); + } + const lse_a = max_a + @log(sum_a); + const lse_b = max_b + @log(sum_b); + var kl: f64 = 0; + for (a, b) |ra, rb| { + const la = score(ra) - lse_a; + kl += @exp(la) * (la - (score(rb) - lse_b)); + } + t.rows += 1; + t.kl += kl; + t.kls.append(std.heap.page_allocator, kl) catch {}; + t.kl_max = @max(t.kl_max, kl); + if (best_a == best_b) { + t.agree += 1; + } else { + var second = -std.math.inf(f64); + for (a, 0..) |ra, i| if (i != best_a) { + second = @max(second, score(ra)); + }; + t.tie_gap = @max(t.tie_gap, max_a - second); + } + if (next) |n| { + t.nll_a += lse_a - score(a[n]); + t.nll_b += lse_b - score(b[n]); + t.dnll.append(std.heap.page_allocator, (lse_b - score(b[n])) - (lse_a - score(a[n]))) catch {}; + t.scored += 1; + } + } +}; + +/// One line of a tally: agreement, score gaps, KL (mean, median, max, rows past 0.1) and each path's perplexity. +fn report(what: []const u8, t: *Tally) void { + if (t.rows == 0) return; + const n: f64 = @floatFromInt(t.rows); + const s: f64 = @floatFromInt(@max(t.scored, 1)); + std.mem.sort(f64, t.kls.items, {}, std.sort.asc(f64)); + std.mem.sort(f64, t.dnll.items, {}, std.sort.asc(f64)); + var past: usize = 0; + for (t.kls.items) |v| past += @intFromBool(v > 0.1); + const mid = if (t.dnll.items.len > 0) t.dnll.items[t.dnll.items.len / 2] else 0; + std.debug.print("{s}: {d} rows, top-1 agree {d}/{d}, max |score diff| {d:.4}, KL(captured||own) mean {e:.3} median {e:.3} max {e:.3}, {d} rows past 0.1, perplexity captured {d:.4} own {d:.4} over {d} next tokens (median NLL own - captured {e:.3}), largest top-2 gap where top-1 differs {d:.4}\n", .{ what, t.rows, t.agree, t.rows, t.max_abs, t.kl / n, t.kls.items[t.kls.items.len / 2], t.kl_max, past, @exp(t.nll_a / s), @exp(t.nll_b / s), t.scored, mid, t.tie_gap }); +} + +/// Rows from this position on attend through the indexer's picks (512 blocks of 4 keys). +const sparse_from = 2048; + +fn scores(raw: []const u8) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, raw)); +} + +fn readIds(gpa: std.mem.Allocator, io: std.Io, path: []const u8) ![]u32 { + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + errdefer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + return ids.toOwnedSlice(gpa); +} + +fn envCount(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4) { + std.debug.print("usage: flashnext-compare ... (COUNT greedy tokens, default 32)\n", .{}); + return 2; + } + const count = envCount("COUNT", 32); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const captured: flash.Tri = .{ .set = &set, .s = stream }; + // TF_FLASHNEXT_OWNED: only the entries it names run ours on this side, the rest the captured set's + const own: flash.Tri = if (glue.mixed != 0) .{ .set = &set, .own = &glue, .s = stream } else .{ .own = &glue, .s = stream }; + var e = try prompt.Engine.init(gpa, io, &driver, &stream, args[1]); + defer e.deinit(); + var out_a = try cuda.DeviceBuffer.alloc(&driver, rows_at_once * vocab * 2); + defer out_a.free(); + var out_b = try cuda.DeviceBuffer.alloc(&driver, rows_at_once * vocab * 2); + defer out_b.free(); + const host_a = try gpa.alloc(u16, rows_at_once * vocab); + defer gpa.free(host_a); + const host_b = try gpa.alloc(u16, rows_at_once * vocab); + defer gpa.free(host_b); + for (args[3..]) |path| { + const ids = try readIds(gpa, io, path); + defer gpa.free(ids); + var seq_a = try e.newSeq(ids.len + count); + defer seq_a.deinit(gpa); + var seq_b = try e.newSeq(ids.len + count); + defer seq_b.deinit(gpa); + // every prompt row's scores on both paths, the next prompt token's NLL under each; dense rows and sparse ones + var t: Tally = .{}; + var dense: Tally = .{}; + var sparse: Tally = .{}; + var at: usize = 0; + while (at < ids.len) { + const m = @min(rows_at_once, ids.len - at); + try e.forward(flash.Tri, captured, &seq_a, ids[at..][0..m], .{ .logits = .all, .logits_out = out_a.ptr }); + try e.forward(flash.Tri, own, &seq_b, ids[at..][0..m], .{ .logits = .all, .logits_out = out_b.ptr }); + try stream.synchronize(); + try out_a.download(0, std.mem.sliceAsBytes(host_a[0 .. m * vocab])); + try out_b.download(0, std.mem.sliceAsBytes(host_b[0 .. m * vocab])); + for (0..m) |r| { + const next: ?u32 = if (at + r + 1 < ids.len) ids[at + r + 1] else null; + t.row(host_a[r * vocab ..][0..vocab], host_b[r * vocab ..][0..vocab], next); + if (ids.len > sparse_from) (if (at + r < sparse_from) &dense else &sparse).row(host_a[r * vocab ..][0..vocab], host_b[r * vocab ..][0..vocab], next); + } + at += m; + } + var label: [96]u8 = undefined; + report(try std.fmt.bufPrint(&label, "rows {s}", .{std.fs.path.basename(path)}), &t); + report(try std.fmt.bufPrint(&label, "rows {s} before {d}", .{ std.fs.path.basename(path), sparse_from }), &dense); + report(try std.fmt.bufPrint(&label, "rows {s} from {d}", .{ std.fs.path.basename(path), sparse_from }), &sparse); + // each path's own greedy reply from the prompt, and where they part + const reply_a = try gpa.alloc(u32, count); + defer gpa.free(reply_a); + const reply_b = try gpa.alloc(u32, count); + defer gpa.free(reply_b); + var g: Tally = .{}; + var parted: ?usize = null; + for (0..count) |i| { + reply_a[i] = prompt.argmax(seq_a.logits); + reply_b[i] = prompt.argmax(seq_b.logits); + // the same tokens so far: the same context on both paths, until (and at) the token where they part + if (parted == null) g.row(scores(seq_a.logits), scores(seq_b.logits), null); + if (parted == null and reply_a[i] != reply_b[i]) parted = i; + if (i + 1 == count) break; + try e.predict(flash.Tri, captured, &seq_a, reply_a[i..][0..1]); + try e.predict(flash.Tri, own, &seq_b, reply_b[i..][0..1]); + } + std.debug.print("reply {s}: {d} greedy tokens, ", .{ std.fs.path.basename(path), count }); + if (parted) |p| std.debug.print("part at token {d} (top-2 gap on the captured path there {d:.4})", .{ p, g.tie_gap }) else std.debug.print("all equal", .{}); + std.debug.print(", KL up to there mean {e:.3}\n", .{g.kl / @as(f64, @floatFromInt(@max(g.rows, 1)))}); + std.debug.print("tokens captured", .{}); + for (reply_a) |v| std.debug.print(" {d}", .{v}); + std.debug.print("\ntokens own", .{}); + for (reply_b) |v| std.debug.print(" {d}", .{v}); + std.debug.print("\n", .{}); + } + return 0; +} diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig new file mode 100644 index 0000000000..e14cfab98f --- /dev/null +++ b/zig/src/families/flashnext/cuda_config.zig @@ -0,0 +1,397 @@ +//! The checkpoint's config.json against what the CUDA path is built for: another shape or setting is refused at load. + +const std = @import("std"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); + +/// One config field and the value the kernels and buffers assume. +const Want = struct { name: []const u8, value: i64 }; + +/// The text config's shapes the CUDA path's constants hold. +const shapes = [_]Want{ + .{ .name = "hidden_size", .value = wt.dims }, + .{ .name = "num_hidden_layers", .value = wt.layers_n }, + .{ .name = "full_attention_interval", .value = 4 }, + .{ .name = "vocab_size", .value = wt.vocab }, + .{ .name = "num_attention_heads", .value = wt.q_heads }, + .{ .name = "num_key_value_heads", .value = wt.kv_heads }, + .{ .name = "head_dim", .value = wt.head_dim }, + .{ .name = "num_experts", .value = wt.experts_n }, + .{ .name = "num_experts_per_tok", .value = wt.slots - 1 }, + .{ .name = "moe_intermediate_size", .value = wt.moe_w }, + .{ .name = "shared_expert_intermediate_size", .value = wt.moe_w }, + .{ .name = "hc_count", .value = wt.streams }, + .{ .name = "hc_lowrank", .value = wt.low }, + .{ .name = "linear_num_key_heads", .value = gdn.nk }, + .{ .name = "linear_num_value_heads", .value = gdn.nv }, + .{ .name = "linear_key_head_dim", .value = gdn.dk }, + .{ .name = "linear_value_head_dim", .value = gdn.dv }, + .{ .name = "linear_conv_kernel_dim", .value = gdn.taps }, + .{ .name = "indexer_n_heads", .value = wt.index_heads }, + .{ .name = "indexer_head_dim", .value = wt.index_dim }, + .{ .name = "indexer_compress_ratio", .value = wt.index_ratio }, + .{ .name = "mtp_num_hidden_layers", .value = 1 }, +}; + +/// RoPE's base: Scratch.inv's 10,000,000 (the Python engine's default when the config names none). +pub const rope_theta: f64 = 10_000_000; +/// The rotated share of a head: 64 of 256 features, `wt.half` pairs. +pub const partial_rotary: f64 = @as(f64, @floatFromInt(2 * wt.half)) / @as(f64, @floatFromInt(wt.head_dim)); + +/// Settings the kernels fix, absent meaning the value they assume: a whole number, a float, a word or a flag. +const Fixed = struct { name: []const u8, value: union(enum) { int: i64, float: f64, word: []const u8, flag: bool } }; + +/// The text config's semantic settings the evaluation hard-codes (norms, gates, activations, layer roles, n-grams). +const fixed = [_]Fixed{ + .{ .name = "output_gate_type", .value = .{ .word = "sigmoid" } }, + .{ .name = "hidden_act", .value = .{ .word = "silu" } }, + .{ .name = "attention_bias", .value = .{ .flag = false } }, + .{ .name = "attn_output_gate", .value = .{ .flag = true } }, + .{ .name = "norm_topk_prob", .value = .{ .flag = true } }, + .{ .name = "mamba_ssm_dtype", .value = .{ .word = "float32" } }, + .{ .name = "tie_word_embeddings", .value = .{ .flag = false } }, + .{ .name = "mtp_use_dedicated_embeddings", .value = .{ .flag = false } }, + .{ .name = "indexer_budget", .value = .{ .int = wt.index_top * wt.index_ratio } }, + .{ .name = "indexer_kv_heads", .value = .{ .int = 1 } }, + .{ .name = "ple_embed_dim", .value = .{ .int = wt.dims } }, + .{ .name = "ple_conv_kernel_size", .value = .{ .int = 4 } }, + .{ .name = "ngram_size", .value = .{ .int = 3 } }, +}; + +fn int(o: std.json.ObjectMap, name: []const u8) ?i64 { + const v = o.get(name) orelse return null; + return if (v == .integer) v.integer else null; +} + +fn float(v: std.json.Value) ?f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => null, + }; +} + +fn object(o: std.json.ObjectMap, name: []const u8) ?std.json.ObjectMap { + const v = o.get(name) orelse return null; + return if (v == .object) v.object else null; +} + +/// A value present and not null. +fn given(o: std.json.ObjectMap, name: []const u8) ?std.json.Value { + const v = o.get(name) orelse return null; + return if (v == .null) null else v; +} + +fn holds(v: std.json.Value, want: @FieldType(Fixed, "value")) bool { + return switch (want) { + .int => |i| v == .integer and v.integer == i, + .float => |f| float(v) == f, + .word => |w| v == .string and std.mem.eql(u8, v.string, w), + .flag => |b| v == .bool and v.bool == b, + }; +} + +/// The first config value that differs from what the CUDA path serves, or null when every one matches. +pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u8 { + const q = quant orelse return "quantization"; + if (int(q, "bits") != 4) return "quantization.bits"; + if (int(q, "group_size") != 32) return "quantization.group_size"; + if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return "quantization.mode"; + return geometry(text); +} + +fn geometry(text: std.json.ObjectMap) ?[]const u8 { + for (shapes) |w| if (int(text, w.name) != w.value) return w.name; + return semantics(text); +} + +pub const Mode = enum { affine, modelopt }; + +fn projection(q: std.json.ObjectMap, name: []const u8, algo: []const u8, group: i64) bool { + const p = object(q, name) orelse return false; + return holds(p.get("quant_algo") orelse return false, .{ .word = algo }) and int(p, "group_size") == group; +} + +fn modeloptMismatch(text: std.json.ObjectMap, quant: std.json.ObjectMap) ?[]const u8 { + if (!holds(quant.get("quant_method") orelse return "quantization_config.quant_method", .{ .word = "modelopt" })) return "quantization_config.quant_method"; + if (!holds(quant.get("quant_algo") orelse return "quantization_config.quant_algo", .{ .word = "MIXED_PRECISION" })) return "quantization_config.quant_algo"; + if (given(quant, "kv_cache_quant_algo") != null) return "quantization_config.kv_cache_quant_algo"; + if (geometry(text)) |field| return field; + const entries = object(quant, "quantized_layers") orelse return "quantization_config.quantized_layers"; + var buf: [256]u8 = undefined; + for (0..wt.layers_n) |layer| { + const routed = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts", .{layer}) catch unreachable; + if (!projection(entries, routed, "NVFP4", 16)) return "quantization_config.routed_experts"; + const p = object(entries, routed).?; + if (given(p, "weight_scale_encoding")) |encoding| if (!holds(encoding, .{ .word = "csf" })) return "quantization_config.weight_scale_encoding"; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |part| { + const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.shared_expert.{s}", .{ layer, part }) catch unreachable; + if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.shared_expert"; + } + const parts: []const []const u8 = if (wt.linearLayer(layer)) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }; + for (parts) |part| { + const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.{s}.{s}", .{ layer, if (wt.linearLayer(layer)) "linear_attn" else "self_attn", part }) catch unreachable; + if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.attention_projection"; + } + } + return null; +} + +/// The first semantic setting the kernels would ignore: eps, RoPE, gates, activations, layer roles, the head's layer. +pub fn semantics(text: std.json.ObjectMap) ?[]const u8 { + const eps = float(given(text, "rms_norm_eps") orelse return "rms_norm_eps") orelse return "rms_norm_eps"; + if (@as(f32, @floatCast(eps)) != wt.eps) return "rms_norm_eps"; // the kernels take eps as fp32 + for (fixed) |f| if (given(text, f.name)) |v| if (!holds(v, f.value)) return f.name; + if (rope(text)) |field| return field; + if (given(text, "layer_types")) |v| { + if (v != .array or v.array.items.len != wt.layers_n) return "layer_types"; + for (v.array.items, 0..) |kind, i| { + const want = if (wt.linearLayer(i)) "linear_attention" else "full_attention"; + if (!holds(kind, .{ .word = want })) return "layer_types"; + } + } + const ple = given(text, "ple_layer_ids") orelse return "ple_layer_ids"; + if (ple != .array or ple.array.items.len != 1 or !holds(ple.array.items[0], .{ .int = wt.ple_layer + 1 })) return "ple_layer_ids"; + if (object(text, "mtp")) |m| { + if (given(m, "rope_theta")) |v| if (float(v) != rope_theta) return "mtp.rope_theta"; + if (given(m, "num_hidden_layers")) |v| if (!holds(v, .{ .int = 1 })) return "mtp.num_hidden_layers"; + if (given(m, "layer_types")) |v| { + if (v != .array or v.array.items.len != 1 or !holds(v.array.items[0], .{ .word = "full_attention" })) return "mtp.layer_types"; + } + } + return null; +} + +/// RoPE as the kernels apply it: base 10,000,000 on 64 of 256 features, no scaling (text positions only). +fn rope(text: std.json.ObjectMap) ?[]const u8 { + if (given(text, "rope_scaling") != null) return "rope_scaling"; + const params = object(text, "rope_parameters"); + const theta = (if (params) |p| given(p, "rope_theta") else null) orelse given(text, "rope_theta"); + if (theta) |v| if (float(v) != rope_theta) return "rope_theta"; + const partial = (if (params) |p| given(p, "partial_rotary_factor") else null) orelse given(text, "partial_rotary_factor"); + if (partial) |v| if (float(v) != partial_rotary) return "partial_rotary_factor"; + const p = params orelse return null; + for ([_][]const u8{ "rope_type", "type" }) |name| if (given(p, name)) |v| if (!holds(v, .{ .word = "default" })) return "rope_parameters.rope_type"; + if (given(p, "factor")) |v| if (float(v) != 1) return "rope_parameters.factor"; + return null; +} + +/// Reads `dir`/config.json; error.UnsupportedCheckpoint names the field in a log line when one differs. +pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + _ = try mode(gpa, io, dir); +} + +pub fn mode(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Mode { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.UnsupportedCheckpoint; + const root = parsed.value.object; + const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; + const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; + const selected: Mode = if (quant != null) .affine else .modelopt; + var field = if (selected == .affine) mismatch(text, quant) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; + if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { + field = "tie_word_embeddings"; + }; + std.log.err("flashnext: config.json's {s} differs from the supported Flash Next CUDA layout", .{field orelse return selected}); + return error.UnsupportedCheckpoint; +} + +/// The served config's text fields, as the checked-in checkpoint has them (a test's base to change one at a time). +fn served(gpa: std.mem.Allocator, text: *std.json.ObjectMap, rope_params: *std.json.ObjectMap, mtp: *std.json.ObjectMap, kinds: *std.json.Array, ple: *std.json.Array) !void { + for (shapes) |w| try text.put(gpa, w.name, .{ .integer = w.value }); + try text.put(gpa, "rms_norm_eps", .{ .float = 1e-06 }); + try text.put(gpa, "output_gate_type", .{ .string = "sigmoid" }); + try text.put(gpa, "hidden_act", .{ .string = "silu" }); + try text.put(gpa, "attention_bias", .{ .bool = false }); + try text.put(gpa, "partial_rotary_factor", .{ .float = 0.25 }); + try text.put(gpa, "indexer_budget", .{ .integer = 2048 }); + try text.put(gpa, "ple_conv_kernel_size", .{ .integer = 4 }); + try text.put(gpa, "ngram_size", .{ .integer = 3 }); + for (0..wt.layers_n) |i| try kinds.append(.{ .string = if (wt.linearLayer(i)) "linear_attention" else "full_attention" }); + try text.put(gpa, "layer_types", .{ .array = kinds.* }); + try ple.append(.{ .integer = 2 }); + try text.put(gpa, "ple_layer_ids", .{ .array = ple.* }); + try rope_params.put(gpa, "rope_theta", .{ .integer = 10_000_000 }); + try rope_params.put(gpa, "partial_rotary_factor", .{ .float = 0.25 }); + try rope_params.put(gpa, "type", .{ .string = "default" }); + try text.put(gpa, "rope_parameters", .{ .object = rope_params.* }); + try mtp.put(gpa, "rope_theta", .{ .integer = 10_000_000 }); + try text.put(gpa, "mtp", .{ .object = mtp.* }); +} + +test "the served shapes pass and a 6-bit checkpoint or another width is named" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + var quant: std.json.ObjectMap = .empty; + try quant.put(gpa, "bits", .{ .integer = 4 }); + try quant.put(gpa, "group_size", .{ .integer = 32 }); + try std.testing.expectEqual(@as(?[]const u8, null), mismatch(text, quant)); + try quant.put(gpa, "bits", .{ .integer = 6 }); + try std.testing.expectEqualStrings("quantization.bits", mismatch(text, quant).?); + try quant.put(gpa, "bits", .{ .integer = 4 }); + try text.put(gpa, "hidden_size", .{ .integer = 4096 }); + try std.testing.expectEqualStrings("hidden_size", mismatch(text, quant).?); + try std.testing.expectEqualStrings("quantization", mismatch(text, null).?); +} + +test "a changed epsilon, RoPE, gate, activation or layer role is refused by name" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const Change = struct { field: []const u8, value: std.json.Value, refused: []const u8 }; + const changes = [_]Change{ + .{ .field = "rms_norm_eps", .value = .{ .float = 1e-5 }, .refused = "rms_norm_eps" }, + .{ .field = "output_gate_type", .value = .{ .string = "silu" }, .refused = "output_gate_type" }, + .{ .field = "hidden_act", .value = .{ .string = "gelu" }, .refused = "hidden_act" }, + .{ .field = "attention_bias", .value = .{ .bool = true }, .refused = "attention_bias" }, + .{ .field = "attn_output_gate", .value = .{ .bool = false }, .refused = "attn_output_gate" }, + .{ .field = "norm_topk_prob", .value = .{ .bool = false }, .refused = "norm_topk_prob" }, + .{ .field = "rope_scaling", .value = .{ .string = "yarn" }, .refused = "rope_scaling" }, + .{ .field = "indexer_budget", .value = .{ .integer = 4096 }, .refused = "indexer_budget" }, + .{ .field = "ngram_size", .value = .{ .integer = 4 }, .refused = "ngram_size" }, + .{ .field = "ple_layer_ids", .value = .{ .array = .init(gpa) }, .refused = "ple_layer_ids" }, + .{ .field = "layer_types", .value = .{ .array = .init(gpa) }, .refused = "layer_types" }, + }; + for (changes) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + try std.testing.expectEqual(@as(?[]const u8, null), semantics(text)); + try text.put(gpa, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, semantics(text) orelse "accepted"); + } +} + +test "each layer's label is exactly the role the kernels run: a sliding or misspelt one is refused" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const full = 3; // the first full-attention layer + try std.testing.expect(!wt.linearLayer(full) and wt.linearLayer(0)); + const Case = struct { layer: usize, label: []const u8 }; + const cases = [_]Case{ + .{ .layer = full, .label = "sliding_attention" }, + .{ .layer = full, .label = "attention" }, + .{ .layer = full, .label = "Full_Attention" }, + .{ .layer = wt.layers_n - 1, .label = "sliding_attention" }, + .{ .layer = 0, .label = "full_attention" }, + .{ .layer = 0, .label = "sliding_attention" }, + }; + for (cases) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + try std.testing.expectEqual(@as(?[]const u8, null), semantics(text)); + kinds.items[c.layer] = .{ .string = c.label }; + try text.put(gpa, "layer_types", .{ .array = kinds }); + try std.testing.expectEqualStrings("layer_types", semantics(text) orelse "accepted"); + } +} + +test "RoPE's base and type inside rope_parameters, and the MTP layer's base, are checked too" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const Case = struct { in_mtp: bool, field: []const u8, value: std.json.Value, refused: []const u8 }; + const cases = [_]Case{ + .{ .in_mtp = false, .field = "rope_theta", .value = .{ .integer = 1_000_000 }, .refused = "rope_theta" }, + .{ .in_mtp = false, .field = "partial_rotary_factor", .value = .{ .float = 0.5 }, .refused = "partial_rotary_factor" }, + .{ .in_mtp = false, .field = "type", .value = .{ .string = "yarn" }, .refused = "rope_parameters.rope_type" }, + .{ .in_mtp = false, .field = "factor", .value = .{ .float = 4 }, .refused = "rope_parameters.factor" }, + .{ .in_mtp = true, .field = "rope_theta", .value = .{ .integer = 5_000_000 }, .refused = "mtp.rope_theta" }, + .{ .in_mtp = true, .field = "layer_types", .value = .{ .array = .init(gpa) }, .refused = "mtp.layer_types" }, + }; + for (cases) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + if (c.in_mtp) { + try mtp.put(gpa, c.field, c.value); + try text.put(gpa, "mtp", .{ .object = mtp }); + } else { + try rope_params.put(gpa, c.field, c.value); + try text.put(gpa, "rope_parameters", .{ .object = rope_params }); + } + try std.testing.expectEqualStrings(c.refused, semantics(text) orelse "accepted"); + } +} + +test "mixed precision requires every routed, shared and attention declaration and refuses altered formats" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(a); + var ple: std.json.Array = .init(a); + try served(a, &text, &rope_params, &mtp, &kinds, &ple); + var quant: std.json.ObjectMap = .empty; + try quant.put(a, "quant_method", .{ .string = "modelopt" }); + try quant.put(a, "quant_algo", .{ .string = "MIXED_PRECISION" }); + var entries: std.json.ObjectMap = .empty; + for (0..wt.layers_n) |layer| { + const routed = try std.fmt.allocPrint(a, "model.language_model.layers.{d}.mlp.experts", .{layer}); + var nv: std.json.ObjectMap = .empty; + try nv.put(a, "quant_algo", .{ .string = "NVFP4" }); + try nv.put(a, "group_size", .{ .integer = 16 }); + try nv.put(a, "weight_scale_encoding", .{ .string = "csf" }); + try entries.put(a, routed, .{ .object = nv }); + const prefix = if (wt.linearLayer(layer)) "linear_attn" else "self_attn"; + const parts: []const []const u8 = if (wt.linearLayer(layer)) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }; + for (parts) |part| try addMxfp8(a, &entries, layer, prefix, part); + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |part| try addMxfp8(a, &entries, layer, "mlp.shared_expert", part); + } + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expect(modeloptMismatch(text, quant) == null); + const cases = [_]struct { name: []const u8, field: []const u8, value: std.json.Value, refused: []const u8 }{ + .{ .name = "model.language_model.layers.47.mlp.experts", .field = "group_size", .value = .{ .integer = 32 }, .refused = "quantization_config.routed_experts" }, + .{ .name = "model.language_model.layers.0.mlp.experts", .field = "weight_scale_encoding", .value = .{ .string = "other" }, .refused = "quantization_config.weight_scale_encoding" }, + .{ .name = "model.language_model.layers.17.mlp.shared_expert.down_proj", .field = "quant_algo", .value = .{ .string = "NVFP4" }, .refused = "quantization_config.shared_expert" }, + .{ .name = "model.language_model.layers.0.linear_attn.in_proj_b", .field = "group_size", .value = .{ .integer = 16 }, .refused = "quantization_config.attention_projection" }, + .{ .name = "model.language_model.layers.47.self_attn.indexer.index_qk_proj", .field = "quant_algo", .value = .{ .string = "FP8" }, .refused = "quantization_config.attention_projection" }, + }; + for (cases) |c| { + const entry = entries.getPtr(c.name).?; + const original = entry.object.get(c.field).?; + try entry.object.put(a, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, modeloptMismatch(text, quant) orelse "accepted"); + try entry.object.put(a, c.field, original); + } + try quant.put(a, "kv_cache_quant_algo", .{ .string = "FP8" }); + try std.testing.expectEqualStrings("quantization_config.kv_cache_quant_algo", modeloptMismatch(text, quant).?); + try quant.put(a, "kv_cache_quant_algo", .null); + try std.testing.expect(modeloptMismatch(text, quant) == null); + _ = entries.swapRemove("model.language_model.layers.47.self_attn.o_proj"); + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expectEqualStrings("quantization_config.attention_projection", modeloptMismatch(text, quant).?); +} + +fn addMxfp8(a: std.mem.Allocator, entries: *std.json.ObjectMap, layer: usize, prefix: []const u8, part: []const u8) !void { + var p: std.json.ObjectMap = .empty; + try p.put(a, "quant_algo", .{ .string = "MXFP8" }); + try p.put(a, "group_size", .{ .integer = 32 }); + const name = try std.fmt.allocPrint(a, "model.language_model.layers.{d}.{s}.{s}", .{ layer, prefix, part }); + try entries.put(a, name, .{ .object = p }); +} diff --git a/zig/src/families/flashnext/cuda_draft_affine.zig b/zig/src/families/flashnext/cuda_draft_affine.zig new file mode 100644 index 0000000000..c6efcf823d --- /dev/null +++ b/zig/src/families/flashnext/cuda_draft_affine.zig @@ -0,0 +1,91 @@ +//! An optional group-32 affine proposal head; the main checkpoint and its verifier remain unchanged. +const std = @import("std"); + +pub fn bf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + return @truncate((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} + +fn value(bytes: []const u8, i: usize) f32 { + return @bitCast(@as(u32, std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)) << 16); +} + +fn code(x: f32) u8 { + const bounded = std.math.clamp(x, 0, 15); + const lo = @floor(bounded); + const integer: u8 = @intFromFloat(lo); + const fraction = bounded - lo; + return integer + @intFromBool(fraction > 0.5 or (fraction == 0.5 and integer & 1 == 1)); +} + +/// Codes use the FP32 range before its BF16 table rounding, matching the Python proposal-head recipe. +pub fn row(input: []const u8, words: []u8, scales: []u8, biases: []u8) !void { + const n = input.len / 2; + if (n == 0 or n % 32 != 0 or input.len != n * 2 or words.len != n / 2 or scales.len != n / 16 or biases.len != scales.len) + return error.UnexpectedTensor; + for (0..n / 32) |group| { + const at = group * 32; + var lo = std.math.inf(f32); + var hi = -std.math.inf(f32); + for (0..32) |i| { + const v = value(input, at + i); + if (!std.math.isFinite(v)) return error.UnexpectedTensor; + lo = @min(lo, v); + hi = @max(hi, v); + } + const scale = @max((hi - lo) / 15.0, 1e-8); + if (!std.math.isFinite(scale)) return error.UnexpectedTensor; + std.mem.writeInt(u16, scales[group * 2 ..][0..2], bf16(scale), .little); + std.mem.writeInt(u16, biases[group * 2 ..][0..2], bf16(lo), .little); + for (0..16) |j| { + const low = code((value(input, at + j * 2) - lo) / scale); + const high = code((value(input, at + j * 2 + 1) - lo) / scale); + words[group * 16 + j] = low | (high << 4); + } + } +} + +test "proposal codes have known nearest-even ties and a flat row stays on its exact bias" { + var input: [128]u8 = undefined; + for (0..64) |i| std.mem.writeInt(u16, input[i * 2 ..][0..2], bf16(if (i < 32) @as(f32, @floatFromInt(@min(i, 30))) / 2 else 7), .little); + var words: [32]u8 = undefined; + var scales: [4]u8 = undefined; + var biases: [4]u8 = undefined; + try row(&input, &words, &scales, &biases); + try std.testing.expectEqual(@as(u8, 0x00), words[0]); + try std.testing.expectEqual(@as(u8, 0x21), words[1]); + try std.testing.expectEqual(@as(u8, 0x22), words[2]); + try std.testing.expectEqual(@as(u16, 0x3f80), std.mem.readInt(u16, scales[0..2], .little)); + try std.testing.expectEqual(@as(u16, 0x40e0), std.mem.readInt(u16, biases[2..4], .little)); + for (words[16..]) |b| try std.testing.expectEqual(@as(u8, 0), b); + std.mem.writeInt(u16, input[0..2], 0x7f80, .little); + try std.testing.expectError(error.UnexpectedTensor, row(&input, &words, &scales, &biases)); +} + +test "proposal approximation is bounded per value against an independent FP64 dequantizer" { + var input: [128]u8 = undefined; + for (0..64) |i| std.mem.writeInt(u16, input[i * 2 ..][0..2], bf16(@sin(@as(f32, @floatFromInt(i)) * 0.37) * 0.05), .little); + var words: [32]u8 = undefined; + var scales: [4]u8 = undefined; + var biases: [4]u8 = undefined; + try row(&input, &words, &scales, &biases); + for (0..2) |g| { + var lo = std.math.inf(f64); + var hi = -std.math.inf(f64); + for (0..32) |i| { + const v: f64 = value(&input, g * 32 + i); + lo = @min(lo, v); + hi = @max(hi, v); + } + const scale = (hi - lo) / 15; + const stored_s: f64 = value(&scales, g); + const stored_b: f64 = value(&biases, g); + const bound = scale * 0.50001 + 15 * @abs(stored_s - scale) + @abs(stored_b - lo); + for (0..32) |i| { + const at = g * 32 + i; + const q = (words[at / 2] >> @as(u3, @intCast((at % 2) * 4))) & 15; + const back = @as(f64, @floatFromInt(q)) * stored_s + stored_b; + try std.testing.expect(@abs(back - value(&input, at)) <= bound); + } + } +} diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig new file mode 100644 index 0000000000..453ba86271 --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -0,0 +1,182 @@ +//! The MLX 4-bit group-32 token table: host dequant matches glue._embed, and the same bytes go to the GPU. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const Tensor = core.safetensors.Tensor; + +pub const weight_name = "language_model.model.embed_tokens.weight"; +pub const scales_name = "language_model.model.embed_tokens.scales"; +pub const biases_name = "language_model.model.embed_tokens.biases"; + +const Triple = struct { weight: Tensor, scales: Tensor, biases: Tensor }; + +fn tensorAt(obj: std.json.ObjectMap, data: []const u8, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const o = (obj.get(name) orelse return error.MissingTensor).object; + const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + if (got != dtype) return error.UnexpectedTensor; + const shape = o.get("shape").?.array.items; + if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; + const offs = o.get("data_offsets").?.array.items; + const begin: usize = @intCast(offs[0].integer); + const end: usize = @intCast(offs[1].integer); + if (end < begin or end > data.len) return error.BadSafetensors; + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; + for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + return t; +} + +/// The three embedding tensors inside one safetensors header. Other tensors, including rank above 4, are ignored. +pub fn readTriple(gpa: std.mem.Allocator, header: []const u8, data: []const u8) !Triple { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, header, .{}); + defer parsed.deinit(); + const obj = parsed.value.object; + return .{ + .weight = try tensorAt(obj, data, weight_name, .u32), + .scales = try tensorAt(obj, data, scales_name, .bf16), + .biases = try tensorAt(obj, data, biases_name, .bf16), + }; +} + +/// The shard that holds the embedding table, mapped until `close`. +pub const Mapped = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + weight: Tensor, + scales: Tensor, + biases: Tensor, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Mapped { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + const triple = try readTriple(gpa, map.memory[8..][0..header_len], map.memory[8 + header_len ..]); + return .{ .io = io, .file = file, .map = map, .weight = triple.weight, .scales = triple.scales, .biases = triple.biases }; + } + + /// One named tensor in this shard. Rank 1 to 4. + pub fn lookup(self: *const Mapped, gpa: std.mem.Allocator, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const header_len: usize = @intCast(std.mem.readInt(u64, self.map.memory[0..8], .little)); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, self.map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + return tensorAt(parsed.value.object, self.map.memory[8 + header_len ..], name, dtype); + } + + pub fn close(self: *Mapped) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.* = undefined; + } +}; + +pub const Table = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + rows: usize, + dims: usize, + dense: bool = false, + + /// Copies the three tensors. `dims` is hidden size (eight values per stored word). + pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { + if (weight.dtype != .u32 or scales.dtype != .bf16 or biases.dtype != .bf16) return error.UnsupportedQuantization; + if (weight.rank != 2 or scales.rank != 2 or biases.rank != 2) return error.UnexpectedTensor; + const rows = weight.dim(0); + const dims = weight.dim(1) * 8; + if (dims == 0 or dims % 32 != 0 or scales.dim(0) != rows or biases.dim(0) != rows) return error.UnexpectedTensor; + if (scales.dim(1) != dims / 32 or biases.dim(1) != dims / 32) return error.UnexpectedTensor; + var w = try cuda.DeviceBuffer.fromHost(d, weight.bytes); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(d, scales.bytes); + errdefer s.free(); + var b = try cuda.DeviceBuffer.fromHost(d, biases.bytes); + errdefer b.free(); + return .{ .w = w, .s = s, .b = b, .rows = rows, .dims = dims }; + } + + /// Own a dense table without converting any BF16 values. + pub fn uploadDense(d: *const cuda.Driver, weight: Tensor) !Table { + if (weight.dtype != .bf16 or weight.rank != 2 or weight.dim(0) == 0 or weight.dim(1) == 0 or weight.dim(1) % 8 != 0) + return error.UnexpectedTensor; + const count = try std.math.mul(usize, weight.dim(0), weight.dim(1)); + if (weight.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + const zero: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = try cuda.DeviceBuffer.fromHost(d, weight.bytes), .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true }; + } + + pub fn deinit(self: *Table) void { + self.w.free(); + self.s.free(); + self.b.free(); + self.* = undefined; + } +}; + +/// fp32 to bf16, round to nearest even, the rounding Triton uses. +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn f32FromBf16(b: []const u8, i: usize) f32 { + return @bitCast(@as(u32, std.mem.readInt(u16, b[2 * i ..][0..2], .little)) << 16); +} + +/// One token row: each group of 32 is `nibble * scale + bias`, stored bf16. Same order as glue._embed. +pub fn dequant(weight: []const u8, scales: []const u8, biases: []const u8, dims: usize, token: usize, out: []u16) !void { + if (dims == 0 or dims % 32 != 0 or out.len != dims) return error.UnexpectedTensor; + const groups = dims / 32; + const words = dims / 8; + const woff = token * words * 4; + const soff = token * groups * 2; + if (woff + words * 4 > weight.len or soff + groups * 2 > scales.len or soff + groups * 2 > biases.len) return error.UnexpectedTensor; + for (0..groups) |g| { + const scale = f32FromBf16(scales[soff..], g); + const bias = f32FromBf16(biases[soff..], g); + for (0..4) |i| { + const word = std.mem.readInt(u32, weight[woff + (g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + out[g * 32 + i * 8 + j] = toBf16(q * scale + bias); + } + } + } +} + +test "embedding bytes are read when another tensor has rank above 4" { + const header = + \\{"high":{"dtype":"U8","shape":[1,1,1,1,1],"data_offsets":[0,1]}, + \\ "language_model.model.embed_tokens.weight":{"dtype":"U32","shape":[1,4],"data_offsets":[1,17]}, + \\ "language_model.model.embed_tokens.scales":{"dtype":"BF16","shape":[1,1],"data_offsets":[17,19]}, + \\ "language_model.model.embed_tokens.biases":{"dtype":"BF16","shape":[1,1],"data_offsets":[19,21]}} + ; + var data: [21]u8 = @splat(0); + data[1] = 0x21; + std.mem.writeInt(u16, data[17..19], 0x3f80, .little); + const triple = try readTriple(std.testing.allocator, header, &data); + try std.testing.expectEqual(@as(usize, 4), triple.weight.dim(1)); + var out: [32]u16 = undefined; + try dequant(triple.weight.bytes, triple.scales.bytes, triple.biases.bytes, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "one embedding group dequants nibbles with scale and bias" { + var weight: [16]u8 = @splat(0); + std.mem.writeInt(u32, weight[0..4], 0x21, .little); + var scales: [2]u8 = undefined; + std.mem.writeInt(u16, &scales, 0x3f80, .little); + const biases: [2]u8 = @splat(0); + var out: [32]u16 = undefined; + try dequant(&weight, &scales, &biases, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); + try std.testing.expectEqual(@as(u16, 0x4000), out[1]); + try std.testing.expectEqual(@as(u16, 0), out[2]); +} diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig new file mode 100644 index 0000000000..d91d104d9a --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -0,0 +1,393 @@ +//! One token through the captured _embed cubin, compared with the host dequant of the same row. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); + +const embed = flash.embed; + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(embed.weight_name).?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4) { + std.debug.print("usage: flashnext-embed \n", .{}); + return 2; + } + const token = try std.fmt.parseInt(u32, args[3], 10); + const path = try shardPath(gpa, io, args[1]); + defer gpa.free(path); + var mapped = try embed.Mapped.open(gpa, io, path); + defer mapped.close(); + const weight = mapped.weight; + const scales = mapped.scales; + const biases = mapped.biases; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var table = try embed.Table.upload(&driver, weight, scales, biases); + defer table.deinit(); + if (token >= table.rows) return error.UnexpectedTensor; + const host = try gpa.alloc(u16, table.dims); + defer gpa.free(host); + try embed.dequant(weight.bytes, scales.bytes, biases.bytes, table.dims, token, host); + + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + var ids = try cuda.DeviceBuffer.fromHost(&driver, std.mem.asBytes(&token)); + defer ids.free(); + var out = try cuda.DeviceBuffer.alloc(&driver, table.dims * 2); + defer out.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, out.ptr, table.dims, 1, 1); + + const gpu = try gpa.alloc(u8, table.dims * 2); + defer gpa.free(gpu); + try stream.synchronize(); + try out.download(0, gpu); + + var mismatch: usize = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, gpu[2 * i ..][0..2], .little); + if (got != want) mismatch += 1; + } + std.debug.print("token {d} rows {d} dims {d} mismatch {d} head", .{ token, table.rows, table.dims, mismatch }); + for (0..@min(8, table.dims)) |i| std.debug.print(" {x:0>4}", .{std.mem.readInt(u16, gpu[2 * i ..][0..2], .little)}); + std.debug.print("\n", .{}); + if (mismatch != 0) return 1; + + const streams: usize = 4; + var wide = try cuda.DeviceBuffer.alloc(&driver, streams * table.dims * 2); + defer wide.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, wide.ptr, table.dims, streams, 1); + const wide_bytes = try gpa.alloc(u8, streams * table.dims * 2); + defer gpa.free(wide_bytes); + try stream.synchronize(); + try wide.download(0, wide_bytes); + const wide_u16 = try gpa.alloc(u16, streams * table.dims); + defer gpa.free(wide_u16); + for (wide_u16, 0..) |*v, i| v.* = std.mem.readInt(u16, wide_bytes[2 * i ..][0..2], .little); + + var copies_differ = false; + for (1..streams) |s| { + if (!std.mem.eql(u16, wide_u16[0..table.dims], wide_u16[s * table.dims ..][0..table.dims])) copies_differ = true; + } + std.debug.print("embed copies {d} identical {}\n", .{ streams, !copies_differ }); + + const nc = table.dims / 256; + var pss = try cuda.DeviceBuffer.alloc(&driver, nc * streams * 4); + defer pss.free(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + try tri.hcWriteback(wide.ptr, wide.ptr, pss.ptr, wide.ptr, "*bf16", wide.ptr, wide.ptr, "*bf16", wide.ptr, "*bf16", table.dims, 1, table.dims, streams, 0, 1, 1, 1); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + try pss.download(0, pss_bytes); + const want = try gpa.alloc(f32, nc * streams); + defer gpa.free(want); + try flash.hc.partialSums(wide_u16, table.dims, streams, want); + var max_ulp: u32 = 0; + var streams_differ = false; + for (0..nc) |c| { + const first = std.mem.readInt(u32, pss_bytes[4 * c * streams ..][0..4], .little); + for (0..streams) |s| { + const bits = std.mem.readInt(u32, pss_bytes[4 * (c * streams + s) ..][0..4], .little); + if (bits != first) streams_differ = true; + const got: f32 = @bitCast(bits); + const dist = flash.hc.ulps(got, want[c * streams + s]); + if (dist > max_ulp) max_ulp = dist; + } + } + std.debug.print("writeback streams identical {} max_ulp {d}\n", .{ !streams_differ, max_ulp }); + if (copies_differ or streams_differ or max_ulp > 4) return 1; + + const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); + if (norm_mismatch != 0) return 1; + const ngram = try flash.ple_check.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); + const scored = try flash.mix.scores(flash.Tri, gpa, io, args[1], &driver, &stream, tri, wide_u16); + return if (ngram == 0 and scored == 0) 0 else 1; +} + +fn runDown(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { + const prefix = "language_model.model.layers.0.attn_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try flash.hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try flash.hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try flash.hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const sk: usize = 32; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + const tri: flash.Tri = .{ .set = set, .s = stream.* }; + try tri.hcDown(h, pss, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, part_b.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, sk); + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, sk, 1, streams, low, n, 1); + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try flash.hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try flash.hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try flash.hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + + const got = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got); + try stream.synchronize(); + try norm_b.download(0, got); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const want = try gpa.alloc(u16, streams * dims); + defer gpa.free(want); + try flash.hc.normed(wide_u16, pss_f, scale, dims, streams, 1e-6, want); + var mismatch: usize = 0; + for (want, 0..) |w, i| { + if (std.mem.readInt(u16, got[2 * i ..][0..2], .little) != w) mismatch += 1; + } + var part4: [4]u8 = undefined; + try part_b.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_part = try flash.hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = flash.hc.ulps(gpu_part, host_part); + std.debug.print("hcdown normed mismatch {d} col0 gpu {d} host {d} ulp {d}\n", .{ mismatch, gpu_part, host_part, col_ulp }); + if (col_ulp > 8) return 1; + + const part_bytes = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(part_bytes); + try part_b.download(0, part_bytes); + const part_f = try gpa.alloc(f32, sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try flash.hc.reduceAct(part_f, sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_bytes); + try xs_b.download(0, xs_bytes); + var act_mismatch: usize = 0; + for (act_want, 0..) |w, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != w) act_mismatch += 1; + } + var inj_mismatch: usize = 0; + for (inj_want, 0..) |w, i| { + if (std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little) != w) inj_mismatch += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |w, i| { + const got_xs: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = flash.hc.ulps(got_xs, w); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_mismatch, inj_mismatch, xs_ulp }); + if (act_mismatch != 0 or inj_mismatch != 0 or xs_ulp > 4) return 1; + + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, streams * dims); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try flash.hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + for (host_mixed, 0..) |host_m, i| { + const got_m = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got_m, host_m); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + const rest = try runProj(gpa, io, model_dir, driver, stream, mapped, mixed_b, xsm_b, tri, h, inj_out.ptr, wide_u16); + if (rest != 0) return 1; + return mismatch; +} + +const proj_names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" }; +const proj_rows = [_]usize{ 10240, 6144, 48, 48 }; + +fn runProj(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer, tri: flash.Tri, h: u64, inj: u64, hidden: []const u16) !u32 { + const prefix = "language_model.model.layers.0.linear_attn."; + const k8: usize = 320; + const kg: usize = 80; + var name: [160]u8 = undefined; + var n: usize = 0; + var words_len: usize = 0; + var scale_len: usize = 0; + var ws: [4]core.safetensors.Tensor = undefined; + var ss: [4]core.safetensors.Tensor = undefined; + var bs: [4]core.safetensors.Tensor = undefined; + for (proj_names, proj_rows, 0..) |part, rows, i| { + ws[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, part }), .u32); + ss[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, part }), .bf16); + bs[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, part }), .bf16); + if (ws[i].dim(0) != rows or ws[i].dim(1) != k8 or ss[i].dim(0) != rows or ss[i].dim(1) != kg) return error.UnexpectedTensor; + n += rows; + words_len += ws[i].bytes.len; + scale_len += ss[i].bytes.len; + } + if (n != 16480) return error.UnexpectedTensor; + const k = k8 * 8; + const row = try gpa.alloc(u16, k); + defer gpa.free(row); + const words = try gpa.alloc(u8, words_len); + defer gpa.free(words); + const scales = try gpa.alloc(u8, scale_len); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scale_len); + defer gpa.free(biases); + var wo: usize = 0; + var so: usize = 0; + for (ws, ss, bs) |w, s, b| { + @memcpy(words[wo..][0..w.bytes.len], w.bytes); + wo += w.bytes.len; + @memcpy(scales[so..][0..s.bytes.len], s.bytes); + @memcpy(biases[so..][0..b.bytes.len], b.bytes); + so += s.bytes.len; + } + var lane = try flash.qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + try flash.qmm.matmul(driver, stream.*, mixed.ptr, xs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, 1, n, k); + const row_bytes = try gpa.alloc(u8, k * 2); + defer gpa.free(row_bytes); + const xs_bytes = try gpa.alloc(u8, kg * 4); + defer gpa.free(xs_bytes); + const got_bytes = try gpa.alloc(u8, n * 2); + defer gpa.free(got_bytes); + try stream.synchronize(); + try mixed.download(0, row_bytes); + try xs.download(0, xs_bytes); + try out_b.download(0, got_bytes); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, kg); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host = try gpa.alloc(u16, n); + defer gpa.free(host); + try flash.qmm.dotRow(row, xs_f, words, scales, biases, n, k, host); + var off: usize = 0; + var max_steps: u32 = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got, want); + if (dist != 0) off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ n, off, max_steps, std.mem.readInt(u16, got_bytes[0..2], .little), host[0] }); + if (max_steps != 0) return max_steps; + const slices = try flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes) orelse return 1; + defer gpa.free(slices); + const wrote = try flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); + if (wrote != 0) return wrote; + const routed = try flash.moe.experts(flash.Tri, gpa, driver, stream, mapped, tri, mixed.ptr); + if (routed != 0) return routed; + return flash.attn.decode(flash.Tri, gpa, driver, stream, io, model_dir, tri, mixed.ptr, xs.ptr); +} diff --git a/zig/src/families/flashnext/cuda_gdn.zig b/zig/src/families/flashnext/cuda_gdn.zig new file mode 100644 index 0000000000..b53fb5d50d --- /dev/null +++ b/zig/src/families/flashnext/cuda_gdn.zig @@ -0,0 +1,409 @@ +//! One-row Gated DeltaNet: depthwise conv, the delta-rule state, the gated norm, then the output projection. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +pub const nk: usize = 16; +pub const nv: usize = 48; +pub const dk: usize = 128; +pub const dv: usize = 128; +pub const taps: usize = 4; +pub const conv_dim: usize = 2 * nk * dk + nv * dv; +pub const proj_width: usize = conv_dim + nv * dv + 2 * nv; +pub const value_dim: usize = nv * dv; +pub const groups: usize = value_dim / qmm.group; + +const symbol: [:0]const u8 = "flashnext_gdn_chain"; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn bf(v: f32) f32 { + return promote(toBf16(v)); +} + +fn sigmoid(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +fn softplus(x: f32) f32 { + return if (x > 20) x else @log(1.0 + @exp(x)); +} + +fn warpSum(v: *[32]f32) void { + var span: usize = 16; + while (span != 0) : (span >>= 1) { + var next: [32]f32 = undefined; + for (0..32) |lane| next[lane] = v[lane] + v[lane ^ span]; + v.* = next; + } +} + +fn convAct(proj: []const u16, cs: []const u16, cw: []const u16, row: usize, channel: usize) f32 { + var acc: f32 = 0; + for (0..taps) |tap| { + const at = row + tap; + const x: f32 = if (at < taps - 1) promote(cs[at * conv_dim + channel]) else promote(proj[(at - (taps - 1)) * proj_width + channel]); + acc = acc + promote(cw[channel * taps + tap]) * x; + } + return bf(acc / (1.0 + @exp(-acc))); +} + +fn normPair(x: []f32, scale_dk: bool) void { + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const v = x[lane * 4 + i]; + ss = ss + v * v; + } + part[lane] = ss; + } + warpSum(&part); + var inv = 1.0 / @sqrt(part[0] + 1e-6); + if (scale_dk) inv = inv * (1.0 / @sqrt(@as(f32, @floatFromInt(dk)))); + for (0..dk) |i| x[i] = x[i] * inv; +} + +fn update(state: []f32, base: usize, ks: []const f32, vs: []const f32, g: f32, beta: f32) void { + for (0..32) |warp| { + var kk: [32][4]f32 = undefined; + for (0..32) |lane| { + for (0..4) |i| kk[lane][i] = ks[lane * 4 + i]; + } + for (0..4) |j| { + var kv: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] * g; + acc = acc + state[at] * kk[lane][i]; + } + kv[lane] = acc; + } + warpSum(&kv); + const delta = (vs[row] - kv[0]) * beta; + for (0..32) |lane| { + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] + kk[lane][i] * delta; + } + } + } + } +} + +fn readout(state: []const f32, base: usize, qs: []const f32, ys: []f32) void { + for (0..32) |warp| { + for (0..4) |j| { + var part: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| acc = acc + state[base + row * dk + lane * 4 + i] * qs[lane * 4 + i]; + part[lane] = acc; + } + warpSum(&part); + ys[row] = bf(part[0]); + } + } +} + +/// One fresh or continued row. `state` is read and written in place, layout [value head, row, column]. +pub fn chain(proj: []const u16, cs: []const u16, cw: []const u16, state: []f32, a_log: []const f32, dt_bias: []const f32, norm_w: []const u16, eps: f32, rows: usize, out: []u16, xs: []f32) !void { + if (rows == 0 or proj.len < rows * proj_width or cs.len < (taps - 1) * conv_dim or cw.len < conv_dim * taps) return error.UnexpectedTensor; + if (state.len < nv * dv * dk or a_log.len < nv or dt_bias.len < nv or norm_w.len < dv) return error.UnexpectedTensor; + if (out.len < rows * value_dim or xs.len < rows * groups) return error.UnexpectedTensor; + for (0..rows) |r| { + for (0..nv) |hv| { + const hk = hv / (nv / nk); + var qs: [dk]f32 = undefined; + var ks: [dk]f32 = undefined; + var vs: [dv]f32 = undefined; + for (0..dk) |t| qs[t] = convAct(proj, cs, cw, r, hk * dk + t); + for (0..dk) |t| ks[t] = convAct(proj, cs, cw, r, nk * dk + hk * dk + t); + for (0..dv) |t| vs[t] = convAct(proj, cs, cw, r, 2 * nk * dk + hv * dv + t); + normPair(&qs, true); + normPair(&ks, false); + const b = promote(proj[r * proj_width + conv_dim + nv * dv + hv]); + const a = promote(proj[r * proj_width + conv_dim + nv * dv + nv + hv]); + const g = @exp(-@exp(a_log[hv]) * softplus(a + dt_bias[hv])); + const beta = bf(sigmoid(b)); + const base = hv * dv * dk; + update(state, base, &ks, &vs, g, beta); + var ys: [dv]f32 = undefined; + readout(state, base, &qs, &ys); + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const y = ys[lane * 4 + i]; + ss = ss + y * y; + } + part[lane] = ss; + } + warpSum(&part); + const rinv = 1.0 / @sqrt(part[0] / @as(f32, @floatFromInt(dv)) + eps); + var gated: [dv]f32 = undefined; + for (0..dv) |t| { + const yn = bf(bf(ys[t] * rinv) * promote(norm_w[t])); + const z = promote(proj[r * proj_width + conv_dim + hv * dv + t]); + gated[t] = bf(yn * sigmoid(z)); + out[r * value_dim + hv * dv + t] = toBf16(gated[t]); + } + for (0..dv / 32) |warp| { + for (0..32) |lane| part[lane] = gated[warp * 32 + lane]; + warpSum(&part); + xs[r * groups + hv * (dv / 32) + warp] = part[0]; + } + } + } +} + +fn orderedUlps(a: f32, b: f32) u32 { + const ai: i32 = @bitCast(a); + const bi: i32 = @bitCast(b); + const ao: i32 = if (ai < 0) std.math.minInt(i32) - ai else ai; + const bo: i32 = if (bi < 0) std.math.minInt(i32) - bi else bi; + const d = ao - bo; + return @intCast(if (d < 0) -d else d); +} + +/// The one-row kernel. Conv state is [3, conv channels], recurrence state is [heads, 128, 128]. +pub fn launch(d: *const cuda.Driver, stream: cuda.Stream, proj: u64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, eps: f32, out: u64, xs: u64, state_out: u64) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var module = try cuda.Module.load(d, cuda.kernels.flashnext_gdn); + defer module.unload(); + const f = try module.function(symbol); + var args: cuda.Args = .{}; + args.add(proj); + args.add(cs); + args.add(cw); + args.add(state_in); + args.add(a_log); + args.add(dt_bias); + args.add(norm_w); + args.add(eps); + args.add(@as(c_int, 1)); + args.add(out); + args.add(xs); + args.add(state_out); + try cuda.launch.launch(f, .{ .grid = .{ .x = nv }, .block = .{ .x = 1024 } }, stream, &args); +} + +fn f32FromBf16(bytes: []const u8) f32 { + return promote(std.mem.readInt(u16, bytes[0..2], .little)); +} + +/// Layer 0's conv, recurrence and output projection. A match also returns the output face's eight unreduced K slices. +pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, proj_dev: cuda.DeviceBuffer, proj_bytes: []const u8) !?[]f32 { + const prefix = "language_model.model.layers.0.linear_attn."; + var name: [160]u8 = undefined; + const conv = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}conv1d.weight", .{prefix}), .bf16); + const a_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}A_log", .{prefix}), .bf16); + const dt_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}dt_bias", .{prefix}), .bf16); + const norm = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}norm.weight", .{prefix}), .bf16); + const out_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.weight", .{prefix}), .u32); + const out_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.scales", .{prefix}), .bf16); + const out_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.biases", .{prefix}), .bf16); + if (conv.dim(0) != conv_dim or conv.dim(1) != taps or norm.dim(0) != dv) return error.UnexpectedTensor; + if (a_bits.dim(0) != nv or dt_bits.dim(0) != nv) return error.UnexpectedTensor; + if (out_w.dim(0) != 2560 or out_w.dim(1) != value_dim / 8) return error.UnexpectedTensor; + if (proj_bytes.len != proj_width * 2) return error.UnexpectedTensor; + + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + for (proj, 0..) |*o, i| o.* = std.mem.readInt(u16, proj_bytes[2 * i ..][0..2], .little); + const cs_bytes = try gpa.alloc(u8, (taps - 1) * conv_dim * 2); + defer gpa.free(cs_bytes); + @memset(cs_bytes, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + const dt_bias = try gpa.alloc(f32, nv); + defer gpa.free(dt_bias); + for (0..nv) |i| { + a_log[i] = f32FromBf16(a_bits.bytes[2 * i ..]); + dt_bias[i] = f32FromBf16(dt_bits.bytes[2 * i ..]); + } + const state_n = nv * dv * dk; + const state = try gpa.alloc(f32, state_n); + defer gpa.free(state); + @memset(state, 0); + var st_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(state)); + defer st_b.free(); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + for (cw, 0..) |*o, i| o.* = std.mem.readInt(u16, conv.bytes[2 * i ..][0..2], .little); + const norm_w = try gpa.alloc(u16, dv); + defer gpa.free(norm_w); + for (norm_w, 0..) |*o, i| o.* = std.mem.readInt(u16, norm.bytes[2 * i ..][0..2], .little); + const host_out = try gpa.alloc(u16, value_dim); + defer gpa.free(host_out); + const host_xs = try gpa.alloc(f32, groups); + defer gpa.free(host_xs); + const eps: f32 = 1e-6; + try chain(proj, cs, cw, state, a_log, dt_bias, norm_w, eps, 1, host_out, host_xs); + + var cs_b = try cuda.DeviceBuffer.fromHost(driver, cs_bytes); + defer cs_b.free(); + var cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + defer cw_b.free(); + var a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + defer a_b.free(); + var dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt_bias)); + defer dt_b.free(); + var norm_b = try cuda.DeviceBuffer.fromHost(driver, norm.bytes); + defer norm_b.free(); + var st_out = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + defer st_out.free(); + var gout = try cuda.DeviceBuffer.alloc(driver, value_dim * 2); + defer gout.free(); + var gxs = try cuda.DeviceBuffer.alloc(driver, groups * 4); + defer gxs.free(); + try launch(driver, stream.*, proj_dev.ptr, cs_b.ptr, cw_b.ptr, st_b.ptr, a_b.ptr, dt_b.ptr, norm_b.ptr, eps, gout.ptr, gxs.ptr, st_out.ptr); + + const got_bytes = try gpa.alloc(u8, value_dim * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, groups * 4); + defer gpa.free(xs_bytes); + const st_bytes = try gpa.alloc(u8, state_n * 4); + defer gpa.free(st_bytes); + try stream.synchronize(); + try gout.download(0, got_bytes); + try gxs.download(0, xs_bytes); + try st_out.download(0, st_bytes); + + var out_off: usize = 0; + var out_steps: u32 = 0; + for (host_out, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) out_off += 1; + if (dist > out_steps) out_steps = dist; + } + var xs_off: usize = 0; + var xs_ulp: u32 = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) xs_off += 1; + if (dist > xs_ulp) xs_ulp = dist; + } + var st_off: usize = 0; + var st_ulp: u32 = 0; + for (state, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, st_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) st_off += 1; + if (dist > st_ulp) st_ulp = dist; + } + const got_head = std.mem.readInt(u16, got_bytes[0..2], .little); + std.debug.print("gdn out {d} off {d} max_steps {d} xs_off {d} xs_ulp {d} state_off {d} state_ulp {d} head {x:0>4} host {x:0>4}\n", .{ value_dim, out_off, out_steps, xs_off, xs_ulp, st_off, st_ulp, got_head, host_out[0] }); + + const got_u16 = try gpa.alloc(u16, value_dim); + defer gpa.free(got_u16); + for (got_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const got_xs = try gpa.alloc(f32, groups); + defer gpa.free(got_xs); + for (got_xs, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host_proj = try gpa.alloc(u16, 2560); + defer gpa.free(host_proj); + try qmm.dotRow(got_u16, got_xs, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim, host_proj); + var lane = try qmm.pack(gpa, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim / 8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var branch = try cuda.DeviceBuffer.alloc(driver, 2560 * 2); + defer branch.free(); + try qmm.matmul(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, 1, 2560, value_dim); + const branch_bytes = try gpa.alloc(u8, 2560 * 2); + defer gpa.free(branch_bytes); + try stream.synchronize(); + try branch.download(0, branch_bytes); + var proj_off: usize = 0; + var proj_steps: u32 = 0; + for (host_proj, 0..) |want, i| { + const got = std.mem.readInt(u16, branch_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) proj_off += 1; + if (dist > proj_steps) proj_steps = dist; + } + std.debug.print("gdn out_proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ 2560, proj_off, proj_steps, std.mem.readInt(u16, branch_bytes[0..2], .little), host_proj[0] }); + if (out_off != 0 or xs_off != 0 or st_off != 0 or proj_off != 0) return null; + // Eight K slices, left unreduced for the write-back. + const sk: usize = 8; + const n: usize = 2560; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + try qmm.partials(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, part_b.ptr, 1, n, value_dim, sk); + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part_b.download(0, raw); + const slices = try gpa.alloc(f32, sk * n); + errdefer gpa.free(slices); + for (slices, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + return slices; +} + +test "warp sum is the sum of the lanes" { + var v: [32]f32 = @splat(0); + v[0] = 1; + v[31] = 3; + warpSum(&v); + try std.testing.expectEqual(@as(f32, 4), v[0]); + try std.testing.expectEqual(@as(f32, 4), v[17]); +} + +test "a zero projection and zero weights stay a zero row" { + const gpa = std.testing.allocator; + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + @memset(proj, 0); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + @memset(cw, 0); + const state = try gpa.alloc(f32, nv * dv * dk); + defer gpa.free(state); + @memset(state, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + @memset(a_log, 0); + const dt = try gpa.alloc(f32, nv); + defer gpa.free(dt); + @memset(dt, 0); + const norm = try gpa.alloc(u16, dv); + defer gpa.free(norm); + @memset(norm, 0); + const out = try gpa.alloc(u16, value_dim); + defer gpa.free(out); + const xs = try gpa.alloc(f32, groups); + defer gpa.free(xs); + try chain(proj, cs, cw, state, a_log, dt, norm, 1e-6, 1, out, xs); + try std.testing.expectEqual(@as(u16, 0), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[value_dim - 1]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); + try std.testing.expectEqual(@as(f32, 0), state[0]); +} diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig new file mode 100644 index 0000000000..30debfbb8a --- /dev/null +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -0,0 +1,343 @@ +//! Flash Next's own glue kernels (flashnext_hc.cu, flashnext_attn.cu) with the core's shared rows: the default path. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); + +/// The attention tiles' head width and query heads a KV head (flashnext_attn.cu's HD and QR). +const tile_head_dim = 256; +const tile_heads = 16; +/// fn_attn_chunks' staged rows: the query heads, `buffers` 32-key tiles of keys and values (rows padded by 8), bf16 p. +fn chunkSmem(buffers: u32) u32 { + return (tile_heads + buffers * 2 * 32) * (tile_head_dim + 8) * 2 + tile_heads * (32 + 8) * 2; +} +/// Rows up to which attention reads the next tile ahead (fewer blocks, so latency, not occupancy, sets the time). +const ahead_rows = 16; + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// The native kernels, by their extern "C" names. +const Fns = struct { + wb0: cuda.Function, + wb_moe: cuda.Function, + wb_slices: cuda.Function, + normalize: cuda.Function, + dense_mix: cuda.Function, + shared_swiglu: cuda.Function, + shared_scatter: cuda.Function, + down: cuda.Function, + reduce_act: cuda.Function, + upmix: cuda.Function, + ple_gate: cuda.Function, + ple_conv: cuda.Function, + add_streams: cuda.Function, + prep: cuda.Function, + gate: cuda.Function, + pool: cuda.Function, + scores: cuda.Function, + chunks: cuda.Function, + chunks_ahead: cuda.Function, + merge: cuda.Function, +}; + +/// The captured entry points our glue replaces, by Triton name (TF_FLASHNEXT_OWNED lists them). +pub const entries = [_][]const u8{ "_embed", "_rmsnorm", "_add_streams", "_attn_prep", "_attn_gate", "_chunks", "_merge", "_pool", "_scores", "_hc_writeback", "_qmm_hcdown", "_hc_reduce_act", "_qmm_upmix", "_ple_gate", "_ple_conv", "_router", "_topk_rows" }; + +/// The bit of entry `name` in Glue.mixed. +pub fn entryBit(comptime name: []const u8) u32 { + inline for (entries, 0..) |e, i| if (comptime std.mem.eql(u8, e, name)) return 1 << i; + @compileError("no entry point " ++ name); +} + +/// With a captured set loaded, the entries TF_FLASHNEXT_OWNED names ("all", or Triton names) run ours: an A/B switch. +fn mixedFromEnv() u32 { + const v = std.mem.span(std.c.getenv("TF_FLASHNEXT_OWNED") orelse return 0); + if (std.mem.eql(u8, v, "all")) return (1 << entries.len) - 1; + var bits: u32 = 0; + var it = std.mem.tokenizeAny(u8, v, ", "); + while (it.next()) |name| for (entries, 0..) |e, i| if (std.mem.eql(u8, e, name)) { + bits |= @as(u32, 1) << @intCast(i); + }; + return bits; +} + +pub const Glue = struct { + /// Entries ours run even with a captured set loaded (TF_FLASHNEXT_OWNED, read once at load). + mixed: u32 = 0, + hc: cuda.Module, + attn: cuda.Module, + emb: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + route: cuda.route_rows.RouteRows, + f: Fns, + + pub fn load(d: *const cuda.Driver) !Glue { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var g: Glue = undefined; + g.mixed = mixedFromEnv(); + g.hc = try cuda.Module.load(d, cuda.kernels.flashnext_hc); + errdefer g.hc.unload(); + g.attn = try cuda.Module.load(d, cuda.kernels.flashnext_attn); + errdefer g.attn.unload(); + g.emb = try cuda.embed_rows.EmbedRows.load(d); + errdefer g.emb.unload(); + g.norms = try cuda.norms.Norms.load(d); + errdefer g.norms.unload(); + g.route = try cuda.route_rows.RouteRows.load(d); + errdefer g.route.unload(); + g.f = .{ + .wb0 = try g.hc.function("fn_hc_writeback0"), + .wb_moe = try g.hc.function("fn_hc_writeback_moe"), + .wb_slices = try g.hc.function("fn_hc_writeback_slices"), + .normalize = try g.hc.function("fn_hc_normalize"), + .dense_mix = try g.hc.function("fn_hc_dense_mix"), + .shared_swiglu = try g.hc.function("fn_shared_swiglu"), + .shared_scatter = try g.hc.function("fn_shared_scatter"), + .down = try g.hc.function("fn_hc_down"), + .reduce_act = try g.hc.function("fn_hc_reduce_act"), + .upmix = try g.hc.function("fn_hc_upmix"), + .ple_gate = try g.hc.function("fn_ple_gate"), + .ple_conv = try g.hc.function("fn_ple_conv"), + .add_streams = try g.hc.function("fn_add_streams"), + .prep = try g.attn.function("fn_attn_prep"), + .gate = try g.attn.function("fn_attn_gate"), + .pool = try g.attn.function("fn_qsa_pool"), + .scores = try g.attn.function("fn_qsa_scores"), + .chunks = try g.attn.function("fn_attn_chunks"), + .chunks_ahead = try g.attn.function("fn_attn_chunks_ahead"), + .merge = try g.attn.function("fn_attn_merge"), + }; + try g.f.chunks.allowDynamicShared(chunkSmem(1)); + try g.f.chunks_ahead.allowDynamicShared(chunkSmem(2)); + return g; + } + + pub fn deinit(g: *Glue) void { + g.route.unload(); + g.norms.unload(); + g.emb.unload(); + g.attn.unload(); + g.hc.unload(); + } + + fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: u32, shared: u32, a: *cuda.Args) !void { + if (grid[0] == 0 or grid[1] == 0 or grid[2] == 0) return; + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = block }, .shared = shared }; + try cuda.launch.launch(f, cfg, s, a); + } + + fn ptrs(a: *cuda.Args, xs: []const u64) void { + for (xs) |x| a.add(x); + } + + fn ints(a: *cuda.Args, xs: []const usize) void { + for (xs) |x| a.add(int(x)); + } + + /// Token rows of the MLX 4-bit group-32 table into `copies` streams each (the core's embed rows). + pub fn embed(g: *const Glue, s: cuda.Stream, ids: u64, w: u64, sc: u64, b: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try g.emb.affine4Groups(s, ids, w, sc, b, out, rows, dims, 32, copies); + } + + pub fn embedDense(g: *const Glue, s: cuda.Stream, ids: u64, w: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try g.emb.bf16Copies(s, ids, w, out, rows, dims, copies); + } + + /// One RMS norm over each row's `dims` features (rows `x_stride` apart), bf16 out and 32-group sums (core norms). + pub fn rmsnorm(g: *const Glue, s: cuda.Stream, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { + if (group != dims) return error.RowShape; // the head norms each row whole + try g.norms.rmsSums(s, x, x_stride, w, out, xs, rows, dims, 32, eps); + } + + pub fn addStreams(g: *const Glue, s: cuda.Stream, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ e, hs, out }); + ints(&a, &.{ dims, streams }); + try go(g.f.add_streams, s, .{ rows, cdiv(dims, 256), 1 }, 256, 0, &a); + } + + pub fn attnGate(g: *const Glue, s: cuda.Stream, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + if (head_dim != tile_head_dim) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ o, proj, out, xs }); + ints(&a, &.{ pair_width, q_heads }); + try go(g.f.gate, s, .{ rows, q_heads, 1 }, tile_head_dim, 0, &a); + } + + /// Text positions, a bf16 cache: the model's heads (cuda_prompt_wt.zig) normed, roped and stored. + pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + if (wt.head_dim != tile_head_dim) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); + a.add(eps); + ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); + try go(g.f.prep, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); + } + + /// Attention over 512-key chunks; `qsa` the sparse form (NCH 5, rows' key lists), else dense (NCH 3). + pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + const group = wt.q_heads / kv; + if (group > tile_heads) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ q, kc, vc, pos, po, pm, pl, ids, nk, sparse }); + ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); + a.add(@as(f32, 0.0625)); + ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); + const ahead = rows <= ahead_rows; + try go(if (ahead) g.f.chunks_ahead else g.f.chunks, s, .{ rows, kv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &a); + } + + pub fn attnMerge(g: *const Glue, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ po, pm, pl, pos, out, nk, sparse }); + ints(&a, &.{ wt.q_heads, kv, wt.q_heads / kv, 512, if (qsa) wt.nch else wt.dense_nch, @intFromBool(qsa) }); + try go(g.f.merge, s, .{ rows, wt.q_heads, 1 }, tile_head_dim, 0, &a); + } + + pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ ikc, pooled, pos, k_scale, inv }); + a.add(eps); + ints(&a, &.{ rows, wt.index_dim, wt.half, wt.index_ratio }); + try go(g.f.pool, s, .{ rows / wt.index_ratio + 2, 1, 1 }, wt.index_dim, 0, &a); + } + + pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + if (wt.index_dim != 128 or wt.index_heads > 8) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ iq, pooled, pos, sc }); + ints(&a, &.{ nb, wt.index_heads, wt.index_ratio, wt.index_top }); + try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); + } + + pub fn hcNormalize(g: *const Glue, s: cuda.Stream, h: u64, pss: u64, scale: u64, out: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + var args: cuda.Args = .{}; + for ([_]u64{ h, pss, scale, out }) |v| args.add(v); + args.add(eps); + args.add(int(dims)); + args.add(int(streams)); + try cuda.launch.launch(g.f.normalize, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(streams) }, .block = .{ .x = 256 } }, s, &args); + } + + pub fn hcDenseMix(g: *const Glue, s: cuda.Stream, up: u64, norm: u64, mixed: u64, sums: u64, rows: usize, dims: usize, streams: usize) !void { + var args: cuda.Args = .{}; + for ([_]u64{ up, norm, mixed, sums }) |v| args.add(v); + args.add(int(dims)); + args.add(int(streams)); + try cuda.launch.launch(g.f.dense_mix, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(dims / 32) }, .block = .{ .x = 32 } }, s, &args); + } + + pub fn sharedSwiGLU(g: *const Glue, s: cuda.Stream, gate: u64, up: u64, act: u64, rows: usize, width: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ gate, up, act }); + ints(&a, &.{width}); + try go(g.f.shared_swiglu, s, .{ rows, cdiv(width, 256), 1 }, 256, 0, &a); + } + + pub fn sharedScatter(g: *const Glue, s: cuda.Stream, shared: u64, out: u64, rows: usize, dims: usize, slots: usize) !void { + if (slots == 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ shared, out }); + ints(&a, &.{ dims, slots }); + try go(g.f.shared_scatter, s, .{ rows, cdiv(dims, 256), 1 }, 256, 0, &a); + } + + pub fn hcReduceAct(g: *const Glue, s: cuda.Stream, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + const cols = low + has_inj * streams; + if (cols > 1024 or sk > 32) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ part, act, xs, inj }); + ints(&a, &.{ sk, rows, streams, low, has_inj, ndn }); + try go(g.f.reduce_act, s, .{ rows, 1, 1 }, @intCast(cdiv(@max(cols, low / 32), 32) * 32), 0, &a); + } + + pub fn hcUpmix(g: *const Glue, s: cuda.Stream, act: u64, xs: u64, w: u64, sc: u64, b: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + if (streams != 4 or dims % 32 != 0 or k % 32 != 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ act, xs, w, sc, b, normed, mixed, xs_mixed }); + ints(&a, &.{ rows, n, k, dims, streams }); + const shared: u32 = @intCast(16 * (k + 8) * 2 + streams * 16 * 32 * 4); + try go(g.f.upmix, s, .{ cdiv(rows, 16), dims / 32, 1 }, 256, shared, &a); + } + + /// The normed streams times the down rows, each K slice's fp32 sums to `part` [sk, rows, n] (sk > 1). + pub fn hcDown(g: *const Glue, s: cuda.Stream, h: u64, pss: u64, scale: u64, normed: u64, w: u64, sc: u64, b: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + _ = out; + const per = k / 32 / sk; + if (sk < 2 or k % (32 * sk) != 0 or dims % (per * 32) != 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ h, pss, scale, normed, w, sc, b, part }); + a.add(int(rows)); + a.add(eps); + ints(&a, &.{ n, k, dims, nc, streams, sk }); + const shared: u32 = @intCast(16 * (per * 32 + 8) * 2 + 16 * per * 4); + try go(g.f.down, s, .{ cdiv(rows, 16), cdiv(n, 64), sk }, 128, shared, &a); + } + + /// Mode 0: the streams' sums; 2: the MoE slots (fp32 y and weights) written back; 4: `world` fp32 K slices. + pub fn hcWriteback(g: *const Glue, s: cuda.Stream, h: u64, hout: u64, pss: u64, branch: u64, inj: u64, y: u64, wts: u64, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + if (dims % 256 != 0 or streams > 8) return error.RowShape; + var a: cuda.Args = .{}; + const grid: [3]usize = .{ rows, dims / 256, 1 }; + switch (mode) { + 0 => { + ptrs(&a, &.{ h, pss }); + ints(&a, &.{ dims, streams }); + try go(g.f.wb0, s, grid, 256, 0, &a); + }, + 2 => { + ptrs(&a, &.{ h, hout, pss, inj, y, wts }); + ints(&a, &.{ dims, streams, topk, slots }); + try go(g.f.wb_moe, s, grid, 256, 0, &a); + }, + 4 => { + ptrs(&a, &.{ h, hout, pss, branch, inj }); + a.add(@as(i64, @intCast(rs))); + ints(&a, &.{ dims, streams, world }); + try go(g.f.wb_slices, s, grid, 256, 0, &a); + }, + else => return error.RowShape, + } + } + + pub fn pleGate(g: *const Glue, s: cuda.Stream, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ keys, vals, h, norm_key, norm_query, gated, pss }); + a.add(eps); + ints(&a, &.{ dims, streams }); + try go(g.f.ple_gate, s, .{ rows, 1, 1 }, 512, 0, &a); + } + + pub fn pleConv(g: *const Glue, s: cuda.Stream, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ gated, pss, norm_conv, tail, conv_w, h, hout, nrow }); + a.add(eps); + ints(&a, &.{ dims, streams, taps, dilation }); + try go(g.f.ple_conv, s, .{ rows, cdiv(streams * dims, 512), 1 }, 512, 0, &a); + } + + /// Router logits [rows, experts + 1] fp32: the experts' rows, then the shared expert's gate row (core route rows). + pub fn routerMany(g: *const Glue, s: cuda.Stream, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { + try g.route.logits(s, x, x_stride, w, out, rows, wt.dims, wt.experts_n + 1); + } + + /// Each row's ten experts and the shared one: picks and weights [rows, 11] (core route rows). + pub fn topkMany(g: *const Glue, s: cuda.Stream, logits: u64, pick: u64, wts: u64, rows: usize) !void { + try g.route.topk(s, logits, wt.experts_n + 1, wt.experts_n, wt.slots - 1, pick, wts, wt.slots, rows, true); + } +}; + +test "every entry point has its own bit" { + try std.testing.expectEqual(@as(u32, 1), entryBit("_embed")); + try std.testing.expectEqual(@as(u32, 1) << 16, entryBit("_topk_rows")); +} + +test "the attention tiles hold the model's heads" { + try std.testing.expect(wt.head_dim == tile_head_dim and wt.q_heads / wt.kv_heads <= tile_heads); +} diff --git a/zig/src/families/flashnext/cuda_hc.zig b/zig/src/families/flashnext/cuda_hc.zig new file mode 100644 index 0000000000..65b661904b --- /dev/null +++ b/zig/src/families/flashnext/cuda_hc.zig @@ -0,0 +1,307 @@ +//! Layer-0 hyper-connection host checks: mode 0 writes each 256-wide chunk's sum of squares. + +const std = @import("std"); + +/// glue._hc_writeback mode 0: pss[c, s] is a pairwise sum of squares of stream s, chunk c. +pub fn partialSums(h: []const u16, dims: usize, streams: usize, out: []f32) !void { + const block: usize = 256; + if (dims == 0 or dims % block != 0 or h.len != streams * dims) return error.UnexpectedTensor; + const nc = dims / block; + if (out.len != nc * streams) return error.UnexpectedTensor; + for (0..nc) |c| { + for (0..streams) |s| { + var buf: [256]f32 = undefined; + for (0..block) |i| { + const bits: u32 = h[s * dims + c * block + i]; + const v: f32 = @bitCast(bits << 16); + buf[i] = v * v; + } + var n: usize = block; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[c * streams + s] = buf[0]; + } + } +} + +test "mode 0 partial sums are the squares of one 256-wide chunk" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + var out: [1]f32 = undefined; + try partialSums(&h, 256, 1, &out); + try std.testing.expectEqual(@as(f32, 1), out[0]); +} + +/// Steps between two positive finite floats. Triton's block sum lands a couple of steps from a pairwise sum. +pub fn ulps(a: f32, b: f32) u32 { + const xa: i32 = @bitCast(a); + const xb: i32 = @bitCast(b); + const d = xa - xb; + return @intCast(if (d < 0) -d else d); +} + +test "ulps counts steps between positive floats" { + const a: f32 = 1; + const b: f32 = @bitCast(@as(u32, @bitCast(a)) + 2); + try std.testing.expectEqual(@as(u32, 2), ulps(a, b)); +} + +pub const tile_bn: usize = 64; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// MLX (n, k8) words to [n/64][k/32][64][4], n padded up to 64. qmm.tile_words. +pub fn tileWords(gpa: std.mem.Allocator, words: []const u8, n: usize, k8: usize) ![]u8 { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const npad = (n + tile_bn - 1) / tile_bn * tile_bn; + const kg = k8 / 4; + const out = try gpa.alloc(u8, (npad / tile_bn) * kg * tile_bn * 4 * 4); + @memset(out, 0); + for (0..n) |r| { + const tile = r / tile_bn; + const local = r % tile_bn; + for (0..kg) |g| { + for (0..4) |w| { + const src = ((r * k8) + g * 4 + w) * 4; + const dst = (((tile * kg + g) * tile_bn + local) * 4 + w) * 4; + @memcpy(out[dst..][0..4], words[src..][0..4]); + } + } + } + return out; +} + +/// (rows, cols) bf16 to (cols, rows), the tiled Q4 scale layout. +pub fn transposeBf16(gpa: std.mem.Allocator, src: []const u8, rows: usize, cols: usize) ![]u8 { + if (src.len != rows * cols * 2) return error.UnexpectedTensor; + const out = try gpa.alloc(u8, src.len); + for (0..rows) |r| { + for (0..cols) |c| { + const from = (r * cols + c) * 2; + const to = (c * rows + r) * 2; + @memcpy(out[to..][0..2], src[from..][0..2]); + } + } + return out; +} + +/// The normed row _qmm_hcdown writes: bf16(h * rinv * scale), one rinv per stream from that stream's partial sums. +pub fn normed(h: []const u16, pss: []const f32, scale: []const f32, dims: usize, streams: usize, eps: f32, out: []u16) !void { + const nc = dims / 256; + if (dims == 0 or dims % 256 != 0 or h.len != streams * dims or pss.len != nc * streams or scale.len != h.len or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + var total: f32 = 0; + for (0..nc) |c| total += pss[c * streams + s]; + const rinv = 1.0 / @sqrt(total / @as(f32, @floatFromInt(dims)) + eps); + for (0..dims) |i| { + const hv: f32 = @bitCast(@as(u32, h[s * dims + i]) << 16); + out[s * dims + i] = toBf16(hv * rinv * scale[s * dims + i]); + } + } +} + +/// fp32 dot of the first K-slice against output column 0, groups of 32, from the MLX row (not the tiled copy). +pub fn firstPartial(row_normed: []const u16, row_words: []const u8, scale_row: []const u8, bias_row: []const u8, groups: usize) !f32 { + if (row_normed.len < groups * 32 or row_words.len < groups * 16 or scale_row.len < groups * 2 or bias_row.len < groups * 2) return error.UnexpectedTensor; + var acc: f32 = 0; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, row_words[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + const xv: f32 = @bitCast(@as(u32, row_normed[g * 32 + i * 8 + j]) << 16); + buf[i * 8 + j] = xv; + p += xv * q; + } + } + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + const s: f32 = @bitCast(@as(u32, std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)) << 16); + const b: f32 = @bitCast(@as(u32, std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)) << 16); + acc += p * s + buf[0] * b; + } + return acc; +} + +test "tiled words keep row 0 in the first tile's local row 0" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + const out = try tileWords(std.testing.allocator, &words, 1, 4); + defer std.testing.allocator.free(out); + try std.testing.expectEqual(@as(u32, 0x21), std.mem.readInt(u32, out[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, out[16..20], .little)); +} + +test "the down norm writes bf16 of h times rinv times scale" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + const pss = [_]f32{1}; + var scale: [256]f32 = @splat(1); + var out: [256]u16 = undefined; + try normed(&h, &pss, &scale, 256, 1, 1e-6, &out); + const rinv = 1.0 / @sqrt(1.0 / 256.0 + 1e-6); + try std.testing.expectEqual(toBf16(rinv), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +fn roundBf16(v: f32) f32 { + return @bitCast(@as(u32, toBf16(v)) << 16); +} + +/// glue._hc_reduce_act: sum the K slices in order, round, divide by the stream count, SiLU, and the inject gates. +pub fn reduceAct(part: []const f32, sk: usize, ndn: usize, low: usize, streams: usize, act: []u16, inj: []u16, xs: []f32) !void { + if (low == 0 or low % 32 != 0 or ndn < low + streams or part.len != sk * ndn or act.len != low or inj.len != streams or xs.len != low / 32) return error.UnexpectedTensor; + for (0..low) |i| { + var sum: f32 = 0; + for (0..sk) |s| sum += part[s * ndn + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + var buf: [32]f32 = undefined; + for (0..low / 32) |g| { + for (0..32) |j| buf[j] = @bitCast(@as(u32, act[g * 32 + j]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + xs[g] = buf[0]; + } + for (0..streams) |s| { + var sum: f32 = 0; + for (0..sk) |k| sum += part[k * ndn + low + s]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + const sig = roundBf16(1.0 / (1.0 + @exp(-y))); + inj[s] = toBf16(2.0 * sig); + } +} + +test "a zero down-projection sum is SiLU zero and an inject gate of one" { + var part: [33]f32 = @splat(0); + var act: [32]u16 = undefined; + var inj: [1]u16 = undefined; + var xs: [1]f32 = undefined; + try reduceAct(&part, 1, 33, 32, 1, &act, &inj, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), inj[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One mixed hidden row: up-dot each stream, sigmoid-gate it by that stream's normed value, then average. +pub fn mix(act: []const u16, xs: []const f32, normed_row: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize, dims: usize, streams: usize, out: []u16) !void { + const kg = k8 / 4; + if (dims == 0 or k8 == 0 or k8 % 4 != 0 or act.len < kg * 32 or xs.len < kg or normed_row.len != streams * dims or out.len != dims) return error.UnexpectedTensor; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..dims) |col| { + var total: f32 = 0; + for (0..streams) |s| { + const row = s * dims + col; + var acc: f32 = 0; + const word_row = words[row * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[row * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[row * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(act[g * 32 + i * 8 + j]) * q; + } + } + const sc = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const bias = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc += p * sc + xs[g] * bias; + } + const up = roundBf16(acc); + const sig = roundBf16(1.0 / (1.0 + @exp(-up))); + total += roundBf16(sig * promote(normed_row[row])); + } + out[col] = toBf16(total / @as(f32, @floatFromInt(streams))); + } +} + +fn bf16Steps(a: u16, b: u16) u32 { + const oa: i32 = if (a >= 0x8000) ~@as(i32, a) else a; + const ob: i32 = if (b >= 0x8000) ~@as(i32, b) else b; + const d = oa - ob; + return @intCast(if (d < 0) -d else d); +} + +pub fn mixedSteps(gpu: u16, host: u16) u32 { + return bf16Steps(gpu, host); +} + +/// Mode 4's branch: the K slices added in slice order, then one bf16 round. +pub fn sumSlices(part: []const f32, sk: usize, n: usize, out: []u16) !void { + if (sk == 0 or n == 0 or part.len != sk * n or out.len != n) return error.UnexpectedTensor; + for (0..n) |i| { + var acc = part[i]; + for (1..sk) |k| acc += part[k * n + i]; + out[i] = toBf16(acc); + } +} + +/// Mode 4's hidden update: one bf16 rounding of hidden + branch * inject, the write-back's fused multiply-add. +pub fn applyBranch(h: []const u16, branch: []const u16, inj: []const u16, dims: usize, streams: usize, out: []u16) !void { + if (dims == 0 or h.len != streams * dims or branch.len != dims or inj.len != streams or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + const gate = promote(inj[s]); + for (0..dims) |i| { + out[s * dims + i] = toBf16(promote(h[s * dims + i]) + promote(branch[i]) * gate); + } + } +} + +test "slices add in order and round once" { + const part = [_]f32{ 1, 0, 0.5, 0 }; + var out: [2]u16 = undefined; + try sumSlices(&part, 2, 2, &out); + try std.testing.expectEqual(toBf16(1.5), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +test "a branch of one times an inject gate of one adds one" { + const h = [_]u16{0}; + const branch = [_]u16{0x3f80}; + const inj = [_]u16{0x3f80}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "the fused product is rounded once with the hidden value" { + const h = [_]u16{0x3b34}; + const branch = [_]u16{0xbe16}; + const inj = [_]u16{0x3c9a}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0xb6f0), out[0]); +} + +test "a zero activation mixes to a zero hidden value" { + var act: [32]u16 = @splat(0); + const xs = [_]f32{0}; + var row: [1]u16 = @splat(0); + var words: [16]u8 = @splat(0); + var scale: [2]u8 = @splat(0); + var bias: [2]u8 = @splat(0); + var mixed: [1]u16 = undefined; + try mix(&act, &xs, &row, &words, &scale, &bias, 1, 4, 1, 1, &mixed); + try std.testing.expectEqual(@as(u16, 0), mixed[0]); +} diff --git a/zig/src/families/flashnext/cuda_head.zig b/zig/src/families/flashnext/cuda_head.zig new file mode 100644 index 0000000000..4793a99b60 --- /dev/null +++ b/zig/src/families/flashnext/cuda_head.zig @@ -0,0 +1,48 @@ +//! Optional vocabulary-column partitioning: each rank retains the checkpoint bits for its own output rows. +const std = @import("std"); +const Tensor = @import("core").safetensors.Tensor; + +pub fn setting(value: ?[]const u8) !bool { + const v = value orelse return false; + if (std.mem.eql(u8, v, "0")) return false; + if (std.mem.eql(u8, v, "1")) return true; + return error.BadHeadSharding; +} + +pub fn enabled() !bool { + return setting(if (std.c.getenv("TF_FLASHNEXT_SHARD_HEAD")) |v| std.mem.span(v) else null); +} + +pub fn shard(t: Tensor, world: usize, rank: usize) !Tensor { + if (world != 2 or rank >= world or t.rank != 2 or t.dtype != .bf16 or t.dim(0) == 0 or t.dim(0) % 256 != 0) + return error.BadHeadSharding; + const bytes = try std.math.mul(usize, try std.math.mul(usize, t.dim(0), t.dim(1)), 2); + if (t.bytes.len != bytes) return error.UnexpectedTensor; + var out = t; + out.shape[0] /= world; + out.bytes = t.bytes[rank * (bytes / world) ..][0 .. bytes / world]; + return out; +} + +test "head partitions preserve every BF16 byte in global vocabulary order" { + var bytes: [256 * 64 * 2]u8 = undefined; + for (&bytes, 0..) |*v, i| v.* = @truncate(i * 17 + i / 251); + const t: Tensor = .{ .dtype = .bf16, .rank = 2, .shape = .{ 256, 64, 0, 0, 0 }, .bytes = &bytes }; + const a = try shard(t, 2, 0); + const b = try shard(t, 2, 1); + try std.testing.expectEqual(@as(usize, 128), a.dim(0)); + try std.testing.expectEqualSlices(u8, bytes[0 .. bytes.len / 2], a.bytes); + try std.testing.expectEqualSlices(u8, bytes[bytes.len / 2 ..], b.bytes); + try std.testing.expectError(error.BadHeadSharding, shard(t, 1, 0)); + try std.testing.expectError(error.BadHeadSharding, shard(t, 2, 2)); + var bad = t; + bad.shape[0] = 255; + try std.testing.expectError(error.BadHeadSharding, shard(bad, 2, 0)); + bad = t; + bad.bytes = bytes[1..]; + try std.testing.expectError(error.UnexpectedTensor, shard(bad, 2, 0)); + try std.testing.expect(!try setting(null)); + try std.testing.expect(!try setting("0")); + try std.testing.expect(try setting("1")); + try std.testing.expectError(error.BadHeadSharding, setting("2")); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig new file mode 100644 index 0000000000..6fd33fdafc --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -0,0 +1,593 @@ +//! Flash Next's lane backend on the native CUDA server: chunked prompts, MTP draft windows, and shared rounds. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); +const weights = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); +const ranks = @import("cuda_rank_engine.zig"); +const config = @import("cuda_config.zig"); +const Glue = @import("cuda_glue.zig").Glue; + +const be = lanes.backend; +const ring = 1024; +/// Streams a round takes at most (a round with more live streams takes them in turns). +const round_streams = 32; + +/// The Triton entry points the pass launches and the head's two more: a captured set replaces our glue with all. +pub const captured_entries = [_][]const u8{ "_embed", "_hc_writeback", "_qmm_hcdown", "_hc_reduce_act", "_qmm_upmix", "_ple_gate", "_ple_conv", "_attn_prep", "_attn_gate", "_pool", "_scores", "_chunks", "_merge", "_router", "_topk_rows" }; +pub const captured_head_entries = [_][]const u8{ "_rmsnorm", "_add_streams" }; + +/// `Tri` launches our glue, or a captured set's. `open` loads the checkpoint and serves greedy or keyed draws. +pub fn Namespace(comptime Tri: type) type { + return struct { + pub const model_type = "qwen4_exp"; + pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed" }; + // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) + pub const default_context: i64 = @intCast(prompt.max_tokens); + // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments + pub const max_segments: u32 = 1; + pub const prompt_rows: u32 = @intCast(prompt.batch.default_chunk); + + /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. + pub const snap_bytes = snapBytes; + pub const snap_save = snapSave; + pub const snap_restore = snapRestore; + pub const snap_drop = snapDrop; + /// The pass ignores its cuts, so states keep at any token; the MTP head reads one token on (lookahead 1). + pub const cache_rules = struct { + pub const lookahead: u32 = 1; + pub const planned = false; + pub const min_gap: u32 = 256; + pub const min_prompt: u32 = 512; + }; + + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + + /// One stream's prompt-and-reply driver, as native/cuda.zig types it; this family has none. + 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: be.Backend, + facts: lanes.Model, + rows: u32, + prompt_rows: u32, + /// Device bytes each admitted stream allocates: its caches (and with the head, its cache and window). + stream_bytes: usize, + /// Every stream allocates its own sequence; none rides on the engine's. + free_streams: u32 = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + learns: ?*anyopaque = null, + learn_bytes: usize = 0, + /// What the prompt cache's copies are taken from: the backend itself (snap_* look up the owner stream). + target: *anyopaque, + /// Whether our own glue runs (no captured set given, or one without this family's entry points). + own_glue: bool = true, + follow: ?*const fn (*anyopaque) anyerror!void = null, + }; + + /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. + pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + const main = weights.deviceBytes(gpa, io, dir) catch return 0; + return main + mtp.headBytes(); + } + + pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const mode = try config.mode(gpa, io, dir); + if (mode == .modelopt) { + const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); + // Dense draft projections and vocabulary rows exceed the affine head's estimate. + return main + if (drafts and rank == 0) mtp.headBytes() + (1 << 30) else @as(u64, 0); + } + if (world != 1) return error.UnsupportedQuantization; + const main = try weights.deviceBytes(gpa, io, dir); + return main + if (drafts) mtp.headBytes() else @as(u64, 0); + } + + pub fn healthy(ptr: *anyopaque) bool { + if (of(ptr).ranks) |r| return r.healthy(); + return true; + } + + pub fn rankRoom(ptr: *anyopaque, local: u64) !u64 { + if (of(ptr).ranks) |r| return r.room(local); + return local; + } + + /// A request or load this engine refuses, in words; null: none of its own. + pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", + error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", + error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", + error.LearnedPocket => "this folder holds Living Weights pockets, which Flash Next on CUDA does not serve yet", + else => null, + }; + } + + /// A stream's caches; with the head, its cache, a window's buffers and the drafts it holds for the next round. + const Lane = struct { + seq: prompt.Seq, + draft: ?mtp.Draft = null, + chain: mtp.Chain = .{}, + /// The prompt's next row while its pass runs a chunk a call; null once the prompt is in. + filled: ?usize = 0, + }; + + const Ops = mtp.Ops(Tri); + + const Cuda = struct { + gpa: std.mem.Allocator, + context: usize, + engine: *prompt.Engine, + ranks: ?*ranks.Session = null, + head: ?*mtp.Head = null, + depth: usize = mtp.default_depth, + confidence: f64 = mtp.default_confidence, + /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. + set: ?cuda.aot.Set = null, + glue: Glue, + stream: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, + /// A round's scores, every row [round_streams, vocab] bf16. + logits: cuda.DeviceBuffer, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + fn tri(self: *Cuda) Tri { + return .{ .set = if (self.set) |*x| x else null, .own = &self.glue, .s = self.stream }; + } + + fn ops(self: *Cuda) Ops { + return .{ .e = self.engine, .tri = self.tri(), .h = self.head.?, .main = if (self.ranks != null) .{ + .ptr = self, + .forward = mainForward, + .verify = mainVerify, + .rollback = mainRollback, + } else null }; + } + + fn mainForward(ptr: *anyopaque, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) anyerror!void { + const self = of(ptr); + return self.ranks.?.forward(Tri, self.tri(), q, ids, hooks); + } + + fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) anyerror!void { + const self = of(ptr); + return self.ranks.?.verify(Tri, self.tri(), q, d, ids); + } + + fn mainRollback(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, keep: usize) anyerror!void { + const self = of(ptr); + return self.ranks.?.keep(Tri, self.tri(), q, d, keep); + } + + fn backend(self: *Cuda) be.Backend { + return .{ .ptr = self, .vtable = &.{ + .prefill = prefillFn, + .prefill_step = prefillStepFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + .tree = treeFn, + .probabilities = probabilitiesFn, + } }; + } + + /// With the head: windows of up to 16 rows cut by its confidence; else shared one-row rounds. + fn facts(self: *const Cuda) lanes.Model { + const drafting = self.head != null; + return .{ + .exact_width = if (drafting) mtp.max_rows else 2, + .gpu_tokens = false, + .mtp = drafting, + .speculate = drafting, + .speculate_early = false, + .drafts = if (drafting) @intCast(self.depth) else 0, + .draft_probabilities = drafting, + .streams_exact = true, + .hidden_rows = true, + .batch_rows = if (drafting) 4 * mtp.max_rows else round_streams, + .max_streams = round_streams, + }; + } + + fn take(self: *Cuda, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + fn value(self: *const Cuda, feed: be.Feed) u32 { + return switch (feed) { + .handle => |h| self.drawn[h % ring], + .value => |v| v, + }; + } + + /// The token after the sequence's last row (seq.logits) at `position`: the draws serial decoding makes. + fn tokenAt(self: *Cuda, seq: *prompt.Seq, s: *lanes.Stream, position: u64) !u32 { + if (s.sampling == null) return prompt.argmax(seq.logits); + const row = try self.gpa.alloc(u16, prompt.vocabulary); + defer self.gpa.free(row); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + return mtp.drawRow(self.gpa, row, null, position, s.sampling); + } + + fn lane(self: *Cuda, s: *lanes.Stream) !*Lane { + return self.lanes.getPtr(s) orelse return error.UnknownStream; + } + + /// The stream's caches, from the kept state the cache names on a restore; its prompt pass starts there. + fn open(self: *Cuda, s: *lanes.Stream) !*Lane { + const ids = s.prompt(); + if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.gpa, s); + if (gop.found_existing) dropLane(self, gop.value_ptr); + // Admit shared plain rounds when four or more streams are active. + const drafting = self.head != null and s.drafts and self.lanes.count() < 4; + // without the head's drafts, one row a stream a round: only a drafted window can rewind its states + if (!drafting) { + s.drafts = false; + s.proposer = null; + } + // a window's rows past the reply's last token (its drafts) still need cache slots + const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else self.engine.newSeq(cap)) catch |err| { + self.lanes.removeByPtr(gop.key_ptr); + return err; + } }; + const l = gop.value_ptr; + errdefer { + dropLane(self, l); + self.lanes.removeByPtr(gop.key_ptr); + } + if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + s.cached = 0; + const saved = s.reuse.saved orelse return l; + const kept: *snap.Snap = @ptrCast(@alignCast(saved)); + if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { + s.reuse_failed = true; + } else if (self.ranks != null and l.draft != null) { + s.reuse_failed = true; + } else if (if (self.ranks) |r| r.restore(&l.seq, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { + l.filled = kept.at; + s.cached = @intCast(kept.at); + } else |err| switch (err) { + error.SequenceTooShort, error.NoHeadState => s.reuse_failed = true, + else => return err, + } + return l; + } + + /// One prompt chunk, cut at the next cache mark; true once the last row's scores are in seq.logits. + fn step(self: *Cuda, s: *lanes.Stream, l: *Lane) !bool { + const ids = s.prompt(); + const at = l.filled orelse return true; + if (s.isCancelled()) return error.Cancelled; + var stop = @min(ids.len, at + self.engine.chunk); + for (s.reuse.marks) |m| if (m > at and m < stop) { + stop = m; + break; + }; + const last = stop == ids.len; + if (l.draft) |*d| { + try self.ops().prefillChunk(&l.seq, d, ids, at, stop); + } else if (self.ranks) |r| { + try r.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); + } else try self.engine.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); + l.filled = if (last) null else stop; + if (!last and std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(stop)) != null) { + if (s.reuse.hook) |k| k.at(k.ptr, s, @intCast(stop)); + } + return last; + } + }; + + fn envUsize(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; + } + + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: Options) !Loaded { + if (o.context > prompt.max_tokens) return error.ContextOverWindow; + if (o.segments != 1 or o.slide) return error.UnsupportedOption; + if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; + const mode = try config.mode(gpa, io, dir); + if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; + const driver: *cuda.Driver = @constCast(ctx.d); + const own = try gpa.create(Cuda); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .context = o.context, .engine = undefined, .glue = undefined, .stream = undefined, .logits = undefined }; + own.stream = try cuda.Stream.init(driver, true); + errdefer own.stream.deinit(); + own.glue = try Glue.load(driver); + errdefer own.glue.deinit(); + // a captured set replaces our glue, for comparison, when it holds every entry point this engine launches + if (captured) |kernels| own.set = try comparison(gpa, io, driver, ctx, kernels, o.drafts); + errdefer if (own.set) |*x| x.deinit(); + if (o.tp == 2) { + const configured_chunk = envUsize("TF_FLASHNEXT_CHUNK", prompt.batch.default_chunk); + var id = try ranks.identity(gpa, io, dir, o.context, if (configured_chunk == 0) prompt.batch.default_chunk else configured_chunk); + id.mtp = @intFromBool(o.drafts); + own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id); + } + errdefer if (own.ranks) |r| { + r.poison(); + r.deinit(); + }; + own.engine = try prompt.Engine.initRank(gpa, io, driver, &own.stream, dir, o.tp, o.rank); + if (own.ranks) |r| { + r.engine = own.engine; + own.engine.collective = &r.collective; + } + errdefer own.engine.deinit(); + // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) + _ = try own.engine.reserve(@min(o.context, own.engine.chunk)); + // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) + if (o.drafts and (o.tp == 1 or o.rank == 0)) { + own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); + own.depth = @max(1, @min(mtp.max_chain, envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); + if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; + } + errdefer if (own.head) |h| h.deinit(); + const window = @min(o.context + mtp.max_rows, prompt.max_tokens); + // and an attention block's block scores at the window (64 MiB at 262,144 tokens) + _ = try own.engine.reserveScores(window, @min(o.context, own.engine.chunk)); + own.logits = try cuda.DeviceBuffer.alloc(driver, round_streams * prompt.vocabulary * 2); + errdefer own.logits.free(); + own.lanes = .empty; + own.next = 0; + var stream_bytes = prompt.seqBytes(window); + if (own.head) |h| { + stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); + } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); + return .{ + .backend = own.backend(), + .facts = own.facts(), + .rows = if (own.head != null) mtp.max_rows else 2, + .prompt_rows = @intCast(own.engine.chunk), + .stream_bytes = stream_bytes, + .ctx = own, + .deinit = release, + .target = own, + .own_glue = own.set == null, + .follow = if (o.rank != 0) follow else null, + }; + } + + /// The captured set at `dir` when it holds every entry point (the head's too when drafting), else null. + fn comparison(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, ctx: *const cuda.Context, dir: []const u8, drafts: bool) !?cuda.aot.Set { + var set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, dir); + const holds = for (captured_entries) |name| { + if (!set.has(name)) break false; + } else for (captured_head_entries) |name| { + if (drafts and !set.has(name)) break false; + } else true; + if (holds) return set; + set.deinit(); + std.log.info("flashnext: the captured set at {s} lacks Flash Next's kernels; its own glue runs", .{dir}); + return null; + } + + fn follow(ptr: *anyopaque) anyerror!void { + const self = of(ptr); + return self.ranks.?.follow(Tri, self.tri(), self.logits.ptr); + } + + fn release(p: *anyopaque) void { + const own: *Cuda = @ptrCast(@alignCast(p)); + var it = own.lanes.valueIterator(); + while (it.next()) |l| dropLane(own, l); + own.lanes.deinit(own.gpa); + own.engine.round_graph.clear(); + own.logits.free(); + if (own.head) |h| h.deinit(); + if (own.ranks) |r| r.deinit(); + own.engine.deinit(); + if (own.set) |*x| x.deinit(); + own.glue.deinit(); + own.stream.deinit(); + own.gpa.destroy(own); + } + + fn dropLane(self: *Cuda, l: *Lane) void { + if (l.draft) |*d| d.deinit(); + if (self.ranks) |r| r.drop(&l.seq) else l.seq.deinit(self.gpa); + } + + fn of(ptr: *anyopaque) *Cuda { + return @ptrCast(@alignCast(ptr)); + } + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const l = try self.open(s); + while (!try self.step(s, l)) {} + } + + /// One chunk of the stream's prompt between the other streams' rounds; the first call opens its caches. + fn prefillStepFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!bool { + const self = of(ptr); + const l = if (self.lanes.getPtr(s)) |x| (if (x.filled != null) x else try self.open(s)) else try self.open(s); + return self.step(s, l); + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + const l = try self.lane(s); + if (l.filled != null or position != l.seq.fed) return error.PositionMismatch; + return self.take(try self.tokenAt(&l.seq, s, position)); + } + + 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 drafting stream's window in one batched pass; streams without drafts share rounds, each row its own bits. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + const seqs = try self.gpa.alloc(*prompt.Seq, windows.len); + defer self.gpa.free(seqs); + const ids = try self.gpa.alloc(u32, windows.len); + defer self.gpa.free(ids); + const at_window = try self.gpa.alloc(usize, windows.len); + defer self.gpa.free(at_window); + var plain: usize = 0; + for (windows, out, 0..) |w, *o, wi| { + if (w.held != 0 or w.parents != null) return error.HeldDraftsNotBuilt; + const l = try self.lane(w.stream); + const rows = w.rows(); + if (l.draft) |*d| { + if (rows > mtp.max_rows) return error.WindowTooWide; + for (w.positions, 0..) |p, r| if (p != l.seq.fed + 1 + r) return error.PositionMismatch; + var row_ids: [mtp.max_rows]u32 = undefined; + row_ids[0] = w.pending; + @memcpy(row_ids[1..rows], w.tokens); + try self.ops().verify(&l.seq, d, row_ids[0..rows]); + try self.ops().draws(d, w.stream.sampling, o.sampled[0..rows]); + } else { + if (rows != 1) return error.CannotRewind; + if (w.positions[0] != l.seq.fed + 1) return error.PositionMismatch; + seqs[plain] = &l.seq; + ids[plain] = w.pending; + at_window[plain] = wi; + plain += 1; + } + if (w.tokens.len > 0) @memcpy(o.drafts[0..w.tokens.len], w.tokens); + } + var at: usize = 0; + while (at < plain) : (at += round_streams) { + const n = @min(round_streams, plain - at); + if (self.ranks) |r| try r.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr) else try self.engine.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr); + } + for (seqs[0..plain], at_window[0..plain]) |q, wi| out[wi].sampled[0] = try self.tokenAt(q, windows[wi].stream, windows[wi].positions[0]); + } + + /// Each window keeps its first rows (a chain's prefix): the DeltaNet states replayed over them alone. + 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 or path.len > w.rows()) return error.CannotRewind; + for (path, 0..) |r, i| if (r != i) return error.CannotRewind; + const l = try self.lane(w.stream); + if (l.draft) |*d| try self.ops().rollback(&l.seq, d, path.len) else if (path.len != w.rows()) return error.CannotRewind; + } + } + + /// The head absorbs the kept rows and chains the next window's drafts (Python decode.draft). + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + const self = of(ptr); + if (self.head == null) return error.NoDraftHead; + for (requests) |r| { + if (r.lanes != null) return error.TreesNotBuilt; + const l = try self.lane(r.stream); + const d = if (l.draft) |*x| x else return error.NoDraftHead; + // the window's kept rows: rollback is idempotent, so a later keep finds it done + if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else { + if (d.open) try self.ops().rollback(&l.seq, d, d.rows); + } + if (r.position != l.seq.fed + 1) return error.PositionMismatch; + // the window's rows past the cache's room cannot run: fewer drafts there + const room = l.seq.cap - l.seq.fed - 1; + const count = @min(@min(r.depth, self.depth), room); + if (r.rows) |path| { + for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + if (r.follow.len < path.len) return error.PositionMismatch; + try self.ops().draft(d, r.follow[0..path.len], d.streams.ptr, r.position, count, r.stream.sampling, self.confidence, &l.chain); + } else { + const first = self.value(r.first orelse return error.NoFirstToken); + try self.ops().draft(d, &.{first}, d.prompt_last, r.position, count, r.stream.sampling, self.confidence, &l.chain); + } + } + } + + /// The held chain as host tokens: the round loop verifies them as it would copied ones. + fn treeFn(ptr: *anyopaque, s: *lanes.Stream, gpa: std.mem.Allocator) anyerror!?lanes.stream.Held { + const self = of(ptr); + const l = try self.lane(s); + if (l.draft == null) return null; + return .{ .count = @intCast(l.chain.n), .tokens = try gpa.dupe(u32, l.chain.slice()) }; + } + + /// Each held draft's temperature-1 probability under the head. + fn probabilitiesFn(ptr: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { + const self = of(ptr); + const l = try self.lane(s); + if (l.draft == null or out.len > l.chain.n) return false; + @memcpy(out, l.chain.probs[0..out.len]); + return true; + } + + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + if (self.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + self.stream.synchronize() catch {}; + dropLane(self, &l); + } + } + + test { + _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane); + } + + // -- the prompt cache's copies ------------------------------------------------------------------------ + + fn snapBytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(at, of(ptr).head != null); + } + + /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). + fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (l.seq.fed != at) return error.NotAtMark; + if (self.ranks) |r| { + if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq); + } + return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); + } + + fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (self.ranks) |r| { + if (l.draft != null) return error.NoHeadState; + return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + } + try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); + } + + fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { + const self = of(ptr); + if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); + snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig new file mode 100644 index 0000000000..4d2734942b --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -0,0 +1,144 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); +const ranks = @import("cuda_rank_engine.zig"); +const wire = @import("cuda_rank_wire.zig"); + +// Virtual device addresses; these callbacks never open CUDA or dereference device memory. +const Device = struct { + var next: u64 = 0x100000; + var copies: usize = 0; + + fn alloc(out: *u64, n: usize) callconv(.c) c_int { + out.* = next; + next += n + 256; + return 0; + } + fn free(_: u64) callconv(.c) c_int { + return 0; + } + fn fill(_: u64, _: u8, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn copy(_: u64, _: u64, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + copies += 1; + return 0; + } + fn sync(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn driver() cuda.Driver { + var d: cuda.Driver = undefined; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = free; + d.api.cuMemsetD8Async = fill; + d.api.cuMemcpyDtoDAsync_v2 = copy; + d.api.cuStreamSynchronize = sync; + return d; + } +}; + +/// Exercise the actual lane open and explicit snapshot entry points with a CPU control peer. +pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytype, comptime drop: anytype, explicit: bool, drafting: bool) !void { + const a = std.testing.allocator; + Device.next = 0x100000; + Device.copies = 0; + var driver = Device.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.gpa = a; + engine.driver = &driver; + engine.stream = &stream; + engine.eos = 7; + engine.chunk = 16; + engine.round_graph = .{}; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = std.c.close(pair[0]); + defer _ = std.c.close(pair[1]); + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const peer: wire.Link = .{ .fd = pair[1] }; + // A successful create, followed by a successful restore if the lane attempts it. + try peer.sendHeader(.{ .kind = .ack, .sequence = 1, .object = 1 }, wire.deadline(1000)); + try peer.sendHeader(.{ .kind = .ack, .sequence = 2, .object = 1 }, wire.deadline(1000)); + var rank: ranks.Session = .{ .gpa = a, .rank = 0, .link = .{ .fd = pair[0] }, .collective = undefined, .engine = &engine }; + rank.collective.failed = false; + var ids = [_]u32{1}; + var head: mtp.Head = undefined; + head.ids = &ids; + var self: Cuda = undefined; + self.gpa = a; + self.context = 1024; + self.engine = &engine; + self.stream = stream; + self.ranks = &rank; + self.head = if (drafting) &head else null; + self.lanes = .empty; + defer self.lanes.deinit(a); + var kept: snap.Snap = .{ .rank_id = 99, .buf = undefined, .at = 2, .history = .{ 8, 9 }, .head = false }; + kept.buf.ptr = 0x100; + var s = try lanes.Stream.init(a, .{ .id = "restore", .prompt = &.{ 1, 2, 3, 4 }, .max_new = 1, .drafts = drafting, .reuse = if (explicit) .{} else .{ .saved = &kept, .at = 2 } }); + defer s.deinit(a); + const lane = try open(&self, &s); + defer { + rank.stopped = true; + drop(&self, lane); + } + if (explicit) { + if (drafting) { + try std.testing.expectError(error.NoHeadState, restore(&self, &s, &kept)); + } else try restore(&self, &s, &kept); + } + const created = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectEqual(wire.Kind.create, created.kind); + try std.testing.expectEqual(@as(u64, 1), created.object); + if (drafting) { + try std.testing.expectEqual(@as(usize, 0), lane.seq.fed); + try std.testing.expectEqual(@as(usize, 0), lane.draft.?.len); + try std.testing.expectEqual(@as(usize, 0), lane.filled.?); + try std.testing.expectEqual(@as(u32, 0), s.cached); + try std.testing.expectEqual(!explicit, s.reuse_failed); + try std.testing.expectEqual(@as(u64, 2), rank.next); + try std.testing.expectEqual(@as(usize, 0), Device.copies); + try std.testing.expectEqual([2]i64{ 7, 7 }, lane.seq.history); + } else { + const resumed = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectEqual(wire.Kind.restore, resumed.kind); + try std.testing.expectEqual(@as(u64, 1), resumed.object); + try std.testing.expectEqual(kept.rank_id, resumed.arg); + try std.testing.expectEqual(@as(usize, 2), lane.seq.fed); + try std.testing.expect(lane.draft == null); + try std.testing.expect(!s.reuse_failed); + try std.testing.expectEqual(@as(u64, 3), rank.next); + try std.testing.expect(Device.copies > 0); + try std.testing.expectEqual([2]i64{ 8, 9 }, lane.seq.history); + } +} + +/// Each case runs independently in the CUDA family's standard host test root. +pub fn Cases(comptime Cuda: type, comptime open: anytype, comptime restore: anytype, comptime drop: anytype) type { + return struct { + test "rank cache restore: drafted open refuses a plain snapshot before rank restore" { + try run(Cuda, open, restore, drop, false, true); + } + + test "rank cache restore: explicit drafted restore refuses a plain snapshot" { + try run(Cuda, open, restore, drop, true, true); + } + + test "rank cache restore: plain open retains the rank restore command" { + try run(Cuda, open, restore, drop, false, false); + } + + test "rank cache restore: explicit plain restore retains the rank command" { + try run(Cuda, open, restore, drop, true, false); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig new file mode 100644 index 0000000000..05143c866f --- /dev/null +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -0,0 +1,397 @@ +//! Flash Next's glue launches: our own kernels (cuda_glue.zig), or the captured Triton cubins where a set is loaded. + +const std = @import("std"); +const cuda = @import("cuda"); +const aot = cuda.aot; + +pub const embed = @import("cuda_embed.zig"); +pub const projection = @import("cuda_projection.zig"); +pub const rank_engine = @import("cuda_rank_engine.zig"); +pub const collective = @import("cuda_collective.zig"); +pub const rank_wire = @import("cuda_rank_wire.zig"); +pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); +pub const hc = @import("cuda_hc.zig"); +pub const draft_affine = @import("cuda_draft_affine.zig"); +pub const qmm = @import("cuda_qmm.zig"); +pub const gdn = @import("cuda_gdn.zig"); +pub const mlp = @import("cuda_mlp.zig"); +pub const moe = @import("cuda_moe.zig"); +pub const attn = @import("cuda_attn.zig"); +pub const ple = @import("cuda_ple.zig"); +pub const ple_check = @import("cuda_ple_check.zig"); +pub const mix = @import("cuda_mix.zig"); +pub const prompt = @import("cuda_prompt.zig"); +/// The row-by-row reference path, for flashnext-prompt +pub const row_path = @import("cuda_rows.zig"); +/// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback +pub const mtp = @import("cuda_mtp.zig"); +pub const native = @import("cuda_lanes.zig").Namespace(Tri); +/// Our own glue kernels, which every launch takes when no captured set is loaded +pub const Glue = glue_.Glue; +const glue_ = @import("cuda_glue.zig"); + +const p = aot.ptr; + +fn int(name: []const u8, v: usize) aot.Arg { + return aot.int(name, @intCast(v)); +} + +fn ci(name: []const u8, v: usize) aot.Const { + return aot.ci(name, @intCast(v)); +} + +fn u(x: usize) u32 { + return @intCast(x); +} + +fn cdiv(a: usize, b: usize) u32 { + return u((a + b - 1) / b); +} + +/// One program per token row and per group of 32 hidden features. +pub fn embedGrid(rows: usize, dims: usize) [3]u32 { + return .{ u(rows), u(dims / 32), 1 }; +} + +/// RMSNorm's BLOCK is the next power of two at or above the group width. +pub fn rmsBlock(group: usize) usize { + return std.math.ceilPowerOfTwo(usize, group) catch unreachable; +} + +/// Each launch's grid and constexprs match the Python wrapper's when a captured set is loaded; without one, ours run. +pub const Tri = struct { + set: ?*const aot.Set = null, + own: ?*const Glue = null, + s: cuda.Stream, + + fn run(t: Tri, name: []const u8, grid: [3]u32, args: []const aot.Arg, consts: []const aot.Const) !void { + try t.set.?.run(t.s, name, grid, args, consts); + } + + /// Ours for `name` without a captured set, or with one when TF_FLASHNEXT_OWNED names it; else null: captured. + fn glue(t: Tri, comptime name: []const u8) ?*const Glue { + if (t.set == null) return t.own.?; + const g = t.own orelse return null; + return if (g.mixed & glue_.entryBit(name) != 0) g else null; + } + + /// glue.embed: ids (R,) int32 into (R, copies * dims) bf16, MLX 4-bit group 32. + pub fn embed(t: Tri, ids: u64, weight: u64, scales: u64, biases: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + if (t.glue("_embed")) |g| return g.embed(t.s, ids, weight, scales, biases, out, dims, copies, rows); + try t.run("_embed", embedGrid(rows, dims), &.{ + p("IDS", "*i32", ids), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + }, &.{ ci("D", dims), ci("S_COPIES", copies) }); + } + + pub fn embedDense(t: Tri, ids: u64, w: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + try g.embedDense(t.s, ids, w, out, dims, copies, rows); + } + + /// glue.rmsnorm per row and `group` features: bf16 out [rows, dims] and 32-group sums; x rows `x_stride` apart. + pub fn rmsnorm(t: Tri, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { + if (t.glue("_rmsnorm")) |g| return g.rmsnorm(t.s, x, w, out, xs, eps, x_stride, rows, dims, group); + try t.run("_rmsnorm", .{ u(rows), u(dims / group), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*fp32", w), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + aot.float("eps", eps), + int("x_stride", x_stride), + }, &.{ ci("D", dims), ci("G", group), ci("BLOCK", rmsBlock(group)) }); + } + + /// glue.add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. Program (row, 256 columns). + pub fn addStreams(t: Tri, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { + if (t.glue("_add_streams")) |g| return g.addStreams(t.s, e, hs, out, rows, dims, streams); + try t.run("_add_streams", .{ u(rows), u(dims / 256), 1 }, &.{ + p("E", "*bf16", e), + p("HS", "*bf16", hs), + p("OUT", "*bf16", out), + }, &.{ ci("D", dims), ci("S", streams), ci("BLOCK", 256) }); + } + + /// glue.attn_gate: one program per row and per query head. + pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + if (t.glue("_attn_gate")) |g| return g.attnGate(t.s, o, proj, out, xs, rows, pair_width, q_heads, head_dim); + try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ + p("O", "*bf16", o), + p("P", "*bf16", proj), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("PW", pair_width), ci("NQ", q_heads), ci("HD", head_dim) }); + } + + /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { + if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); + try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ + p("P", "*bf16", proj), + p("POS0", "*i32", pos), + p("QW", "*fp32", q_scale), + p("KW", "*fp32", k_scale), + p("IW", "*fp32", i_scale), + p("INV", "*fp32", inv), + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KS", "*fp16", rope), + p("VS", "*fp16", rope), + p("IQ", "*bf16", iq), + p("IKC", "*bf16", ikc), + p("ROPE", "*i32", rope), + p("DELTA", "*i32", delta), + int("length", length), + aot.float("eps", eps), + }, &.{ ci("PW", 13952), ci("NQ", 24), ci("NKV", 2), ci("HD", 256), ci("NI", 4), ci("IHD", 128), ci("HALF", 32), ci("BITS", 0), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); + } + + /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); + try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KSC", "*fp16", sparse), + p("VSC", "*fp16", sparse), + p("POS0", "*i32", pos), + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("IDS", "*i32", ids), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", if (qsa) 5 else 3), aot.cf("SCALE", 0.0625), ci("IDW", 2052), ci("QSA", @intFromBool(qsa)), ci("BITS", 0) }); + } + + /// attention._merge: one program per row and per KV head, in attnChunks' form. + pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { + if (t.glue("_merge")) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, kv, qsa); + try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("POS0", "*i32", pos), + p("OUT", "*bf16", out), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", if (qsa) 5 else 3), ci("QSA", @intFromBool(qsa)), ci("BITS", 0) }); + } + + /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. + pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); + try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ + p("IKC", "*bf16", ikc), + p("POOLED", "*bf16", pooled), + p("POS0", "*i32", pos), + p("W", "*fp32", k_scale), + p("INV", "*fp32", inv), + aot.float("eps", eps), + int("R", rows), + int("length", 0), + }, &.{ ci("DI", 128), ci("HALF", 32), ci("RATIO", 4), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); + } + + /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. + pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); + try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ + p("IQ", "*bf16", iq), + p("POOLED", "*bf16", pooled), + p("POS0", "*i32", pos), + p("SC", "*fp32", sc), + int("NB", nb), + }, &.{ ci("HI", 4), ci("DI", 128), ci("RATIO", 4), ci("TOP", 512), ci("BB", 64) }); + } + + /// Dense connections always use native glue; a captured affine projection cannot replace these operators. + pub fn hcNormalize(t: Tri, h: u64, pss: u64, scale: u64, out: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + return g.hcNormalize(t.s, h, pss, scale, out, eps, rows, dims, streams); + } + + pub fn hcDenseMix(t: Tri, up: u64, norm: u64, mixed: u64, sums: u64, rows: usize, dims: usize, streams: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + return g.hcDenseMix(t.s, up, norm, mixed, sums, rows, dims, streams); + } + + /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. + pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + if (t.glue("_hc_reduce_act")) |g| return g.hcReduceAct(t.s, part, act, xs, inj, sk, rows, streams, low, ndn, has_inj); + try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ + p("PART", "*fp32", part), + p("ACT", "*bf16", act), + p("XS", "*fp32", xs), + p("INJ", "*bf16", inj), + int("M", rows), + }, &.{ ci("SK", sk), ci("S", streams), ci("LOW", low), ci("LOWP", rmsBlock(low)), ci("HAS_INJ", has_inj), ci("NDN", ndn) }); + } + + /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. + pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + if (t.glue("_qmm_upmix")) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); + try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ + p("X", "*bf16", act), + p("XS", "*fp32", xs), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("NORMED", "*bf16", normed), + p("MIXED", "*bf16", mixed), + p("XSM", "*fp32", xs_mixed), + int("M", rows), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("SS", streams), ci("BM", 16), ci("DB", 32), ci("GPI", 2), ci("SBN", 64) }); + } + + /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. + pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + if (t.glue("_qmm_hcdown")) |g| return g.hcDown(t.s, h, pss, scale, normed, weight, scales, biases, out, part, eps, rows, n, k, dims, nc, streams, sk); + try t.run("_qmm_hcdown", .{ cdiv(rows, 16), cdiv(n, 64), u(sk) }, &.{ + p("H", "*bf16", h), + p("PSS", "*fp32", pss), + p("SCALE", "*fp32", scale), + p("NORMED", "*bf16", normed), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("PART", "*fp32", part), + int("M", rows), + aot.float("eps", eps), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("NC", nc), ci("SS", streams), ci("SK", sk), ci("BM", 16), ci("BLOCK_N", 64), ci("GPI", 2), ci("SBN", 64) }); + } + + /// glue.hc_writeback. Pointer types follow the tensor dtypes Triton specialized (mode 0..4). + pub fn hcWriteback(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, branch_ty: []const u8, inj: u64, y: u64, y_ty: []const u8, wts: u64, wts_ty: []const u8, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + if (t.glue("_hc_writeback")) |g| return g.hcWriteback(t.s, h, hout, pss, branch, inj, y, wts, rs, rows, dims, streams, mode, topk, slots, world); + try t.run("_hc_writeback", .{ u(rows), u(dims / 256), 1 }, &.{ + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("PSS", "*fp32", pss), + p("BR", branch_ty, branch), + p("INJ", "*bf16", inj), + p("Y", y_ty, y), + p("WTS", wts_ty, wts), + int("RS", rs), + }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); + } + + /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. + pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { + if (t.set == null) return error.CapturedKernelsRequired; // the row path's n-gram check only + try t.run("_ple_embed", .{ u(rows), u(heads), 1 }, &.{ + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("HEADS", heads), ci("DH", dh), aot.cf("SCALE", scale) }); + } + + /// glue.ple_gate: one program per row. + pub fn pleGate(t: Tri, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + if (t.glue("_ple_gate")) |g| return g.pleGate(t.s, keys, vals, h, norm_key, norm_query, gated, pss, eps, rows, dims, streams); + try t.run("_ple_gate", .{ u(rows), 1, 1 }, &.{ + p("KEYS", "*bf16", keys), + p("VALS", "*bf16", vals), + p("H", "*bf16", h), + p("NK", "*fp32", norm_key), + p("NQ", "*fp32", norm_query), + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + aot.float("eps", eps), + }, &.{ ci("D", dims), ci("S", streams), ci("BLOCK", 512) }); + } + + /// glue.ple_conv. R folds to a constexpr when it is 1. + pub fn pleConv(t: Tri, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + if (t.glue("_ple_conv")) |g| return g.pleConv(t.s, gated, pss, norm_conv, tail, conv_w, h, hout, nrow, eps, rows, dims, streams, taps, dilation); + try t.run("_ple_conv", .{ u(rows), u(streams * dims / 512), 1 }, &.{ + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + p("NC", "*fp32", norm_conv), + p("TAIL", "*bf16", tail), + p("CW", "*bf16", conv_w), + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("NROW", "*bf16", nrow), + aot.float("eps", eps), + int("R", rows), + }, &.{ ci("D", dims), ci("S", streams), ci("TAPS", taps), ci("DIL", dilation), ci("BLOCK", 512) }); + } + + /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. + pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { + try t.routerMany(x, w, out, x_stride, 1); + } + + /// moe._topk_rows: one program, the row's 10 experts and the shared expert. + pub fn topkRows(t: Tri, logits: u64, pick: u64, wts: u64) !void { + try t.topkMany(logits, pick, wts, 1); + } + + /// moe._topk_rows over `rows` rows of [rows, 513] logits: picks and weights [rows, 11]. + pub fn topkMany(t: Tri, logits: u64, pick: u64, wts: u64, rows: usize) !void { + if (t.glue("_topk_rows")) |g| return g.topkMany(t.s, logits, pick, wts, rows); + try t.run("_topk_rows", .{ u(rows), 1, 1 }, &.{ + p("L", "*fp32", logits), + p("PICK", "*i32", pick), + p("WTS", "*fp32", wts), + }, &.{ ci("NE", 512), ci("NL", 513), ci("TOPK", 10), ci("SLOTS", 11), ci("BLOCK", 1024), ci("SLOTP", 16) }); + } + + /// moe.router over `rows` rows at the one-row tiles (BM 16, BLOCK_E 32, BK 256), each row keeping its own bits. + pub fn routerMany(t: Tri, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { + if (t.glue("_router")) |g| return g.routerMany(t.s, x, w, out, x_stride, rows); + try t.run("_router", .{ cdiv(rows, 16), cdiv(513, 32), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*bf16", w), + p("OUT", "*fp32", out), + int("M", rows), + int("x_stride", x_stride), + }, &.{ ci("D", 2560), ci("NE", 513), ci("BM", 16), ci("BLOCK_E", 32), ci("BK", 256) }); + } +}; + +test { + _ = embed; + _ = hc; + _ = qmm; + _ = gdn; + _ = mlp; + _ = moe; + _ = attn; + _ = ple; + _ = mix; + _ = prompt; + _ = prompt.batch; + _ = mtp; + _ = native; + _ = @import("cuda_snap.zig"); + _ = @import("cuda_config.zig"); + _ = @import("cuda_projection_load.zig"); + _ = @import("cuda_rank_engine.zig"); + _ = @import("cuda_round_graph.zig"); + _ = @import("cuda_rank_memory.zig"); + _ = @import("cuda_rank_wire.zig"); + _ = @import("cuda_modelopt_experts.zig"); + _ = row_path; + _ = @import("cuda_mtp_check.zig"); + _ = @import("cuda_draft_affine.zig"); + _ = @import("cuda_head.zig"); +} + +test "embed grid is one program per row and per group of 32" { + try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); +} + +test "rmsnorm block is the next power of two at or above the group" { + try std.testing.expectEqual(@as(usize, 4096), rmsBlock(2560)); + try std.testing.expectEqual(@as(usize, 16384), rmsBlock(10240)); +} diff --git a/zig/src/families/flashnext/cuda_mix.zig b/zig/src/families/flashnext/cuda_mix.zig new file mode 100644 index 0000000000..ce050ef0af --- /dev/null +++ b/zig/src/families/flashnext/cuda_mix.zig @@ -0,0 +1,289 @@ +//! The final mixer and the word scores for one decode token: four embedding copies, no inject gate. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const streams: usize = 4; +const low: usize = 320; +const wide: usize = streams * dims; +const sk: usize = 32; +const per: usize = 10; +const eps: f32 = 1e-6; +const vocab: usize = 248320; +const prefix = "language_model.model.hyper_connection_mixer."; + +const Tensor = core.safetensors.Tensor; +const Triple = struct { w: Tensor, s: Tensor, b: Tensor }; + +const Gap = struct { off: usize = 0, steps: u32 = 0, at: usize = 0 }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundBf16(v: f32) f32 { + return promote(toBf16(v)); +} + +fn tree32(buf: *[32]f32) f32 { + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn readU16(raw: []const u8, i: usize) u16 { + return std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "hc_norm.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +/// Split-K sum, one bf16 round, divide by the stream count, SiLU. No inject columns. +fn siluRow(part: []const f32, slices: usize, n: usize, nstreams: usize, act: []u16, xs: []f32) !void { + if (n == 0 or n % 32 != 0 or slices == 0 or part.len != slices * n or act.len != n or xs.len != n / 32) return error.UnexpectedTensor; + for (0..n) |i| { + var sum = part[i]; + for (1..slices) |s| sum += part[s * n + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(nstreams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + for (xs, 0..) |*o, g| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(act[g * 32 + j]); + o.* = tree32(&buf); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var g = Gap{}; + for (got, want, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) g.off += 1; + if (d > g.steps) { + g.steps = d; + g.at = i; + } + } + return g; +} + +fn u16s(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = readU16(raw, i); +} + +test "a zero down sum without an inject gate is SiLU zero" { + var part: [32]f32 = @splat(0); + var act: [32]u16 = undefined; + var xs: [1]f32 = undefined; + try siluRow(&part, 1, 32, 1, &act, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +/// Mixer read-out of `hidden` (four stream copies) and the head's scores of that mixed row. +pub fn scores(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, hidden: []const u16) !u32 { + if (hidden.len != wide) return error.UnexpectedTensor; + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const gamma = file.get(prefix ++ "hc_norm.weight") orelse return error.MissingTensor; + if (!gamma.is(.bf16, &.{wide})) return error.UnexpectedTensor; + const down = try triple(&file, prefix ++ "input_mix_weight_down", low, wide / 8); + const up = try triple(&file, prefix ++ "input_mix_weight_up", wide, low / 8); + const head = try triple(&file, "language_model.lm_head", vocab, dims / 8); + + const scale = try gpa.alloc(f32, wide); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = promote(readU16(gamma.bytes, i)); + const pss = try gpa.alloc(f32, (dims / 256) * streams); + defer gpa.free(pss); + try hc.partialSums(hidden, dims, streams, pss); + const host_norm = try gpa.alloc(u16, wide); + defer gpa.free(host_norm); + try hc.normed(hidden, pss, scale, dims, streams, eps, host_norm); + + var h_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(hidden)); + defer h_b.free(); + var pss_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(pss)); + defer pss_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + const down_w = try hc.tileWords(gpa, down.w.bytes, low, wide / 8); + defer gpa.free(down_w); + const down_s = try hc.transposeBf16(gpa, down.s.bytes, low, wide / 32); + defer gpa.free(down_s); + const down_b = try hc.transposeBf16(gpa, down.b.bytes, low, wide / 32); + defer gpa.free(down_b); + var dw = try cuda.DeviceBuffer.fromHost(driver, down_w); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(driver, down_s); + defer ds.free(); + var db = try cuda.DeviceBuffer.fromHost(driver, down_b); + defer db.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, wide * 2); + defer norm_b.free(); + var dn_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer dn_b.free(); + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * low * 4); + defer part_b.free(); + try tri.hcDown(h_b.ptr, pss_b.ptr, scale_b.ptr, norm_b.ptr, dw.ptr, ds.ptr, db.ptr, dn_b.ptr, part_b.ptr, eps, 1, low, wide, dims, dims / 256, streams, sk); + + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, act_b.ptr, sk, 1, streams, low, low, 0); + + const up_w = try hc.tileWords(gpa, up.w.bytes, wide, low / 8); + defer gpa.free(up_w); + const up_s = try hc.transposeBf16(gpa, up.s.bytes, wide, low / 32); + defer gpa.free(up_s); + const up_b = try hc.transposeBf16(gpa, up.b.bytes, wide, low / 32); + defer gpa.free(up_b); + var uw = try cuda.DeviceBuffer.fromHost(driver, up_w); + defer uw.free(); + var us = try cuda.DeviceBuffer.fromHost(driver, up_s); + defer us.free(); + var ub = try cuda.DeviceBuffer.fromHost(driver, up_b); + defer ub.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, uw.ptr, us.ptr, ub.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, wide, low, dims, streams); + + const norm_raw = try gpa.alloc(u8, wide * 2); + defer gpa.free(norm_raw); + const part_raw = try gpa.alloc(u8, sk * low * 4); + defer gpa.free(part_raw); + const act_raw = try gpa.alloc(u8, low * 2); + defer gpa.free(act_raw); + const xs_raw = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_raw); + const mixed_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_raw); + const xsm_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xsm_raw); + try stream.synchronize(); + try norm_b.download(0, norm_raw); + try part_b.download(0, part_raw); + try act_b.download(0, act_raw); + try xs_b.download(0, xs_raw); + try mixed_b.download(0, mixed_raw); + try xsm_b.download(0, xsm_raw); + + var norm_off: usize = 0; + for (host_norm, 0..) |want, i| { + if (readU16(norm_raw, i) != want) norm_off += 1; + } + const gpu_norm = try gpa.alloc(u16, per * 32); + defer gpa.free(gpu_norm); + u16s(gpu_norm, norm_raw); + const host_part = try hc.firstPartial(gpu_norm, down.w.bytes[0 .. per * 16], down.s.bytes[0 .. per * 2], down.b.bytes[0 .. per * 2], per); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, part_raw[0..4], .little)); + const col_ulp = hc.ulps(gpu_part, host_part); + + const part_f = try gpa.alloc(f32, sk * low); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_raw[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try siluRow(part_f, sk, low, streams, act_want, xs_want); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (readU16(act_raw, i) != want) act_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const act_u = try gpa.alloc(u16, low); + defer gpa.free(act_u); + u16s(act_u, act_raw); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const norm_u = try gpa.alloc(u16, wide); + defer gpa.free(norm_u); + u16s(norm_u, norm_raw); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u, xs_f, norm_u, up.w.bytes, up.s.bytes, up.b.bytes, wide, low / 8, dims, streams, host_mixed); + const got_mixed = try gpa.alloc(u16, dims); + defer gpa.free(got_mixed); + u16s(got_mixed, mixed_raw); + const mixed_gap = gap(got_mixed, host_mixed); + std.debug.print("mix norm {d} col_ulp {d} act {d} xs_ulp {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ norm_off, col_ulp, act_off, xs_ulp, mixed_gap.off, mixed_gap.steps, mixed_gap.at, got_mixed[mixed_gap.at], host_mixed[mixed_gap.at], got_mixed[0] }); + + var lane = try qmm.pack(gpa, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims / 8); + defer lane.deinit(gpa); + var hw = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer hw.free(); + var hs = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer hs.free(); + var hb = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer hb.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + defer logits_b.free(); + try qmm.matmul(driver, stream.*, mixed_b.ptr, xsm_b.ptr, hw.ptr, hs.ptr, hb.ptr, logits_b.ptr, 1, vocab, dims); + const xsm_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xsm_f); + for (xsm_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xsm_raw[4 * i ..][0..4], .little)); + const host_logits = try gpa.alloc(u16, vocab); + defer gpa.free(host_logits); + try qmm.dotRow(got_mixed, xsm_f, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims, host_logits); + const logits_raw = try gpa.alloc(u8, vocab * 2); + defer gpa.free(logits_raw); + try stream.synchronize(); + try logits_b.download(0, logits_raw); + const got_logits = try gpa.alloc(u16, vocab); + defer gpa.free(got_logits); + u16s(got_logits, logits_raw); + const score_gap = gap(got_logits, host_logits); + std.debug.print("scores n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ vocab, score_gap.off, score_gap.steps, score_gap.at, got_logits[score_gap.at], host_logits[score_gap.at], got_logits[0] }); + if (score_gap.off > 0 and score_gap.off <= 8) { + for (got_logits, host_logits, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) std.debug.print("score {d} gpu {x:0>4} host {x:0>4} steps {d}\n", .{ i, a, b, d }); + } + } + + if (norm_off != 0 or col_ulp > 8 or act_off != 0 or xs_ulp > 4 or mixed_gap.steps > 2 or score_gap.steps > 1 or score_gap.off > 3) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_mixed_experts.zig b/zig/src/families/flashnext/cuda_mixed_experts.zig new file mode 100644 index 0000000000..b26225731d --- /dev/null +++ b/zig/src/families/flashnext/cuda_mixed_experts.zig @@ -0,0 +1,36 @@ +//! Routed NVFP4 intermediate partitions and a separately owned MXFP8 shared expert. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const routed = @import("cuda_modelopt_experts.zig"); +const projections = @import("cuda_projection_load.zig"); + +pub const Owned = struct { + experts: routed.Owned, + shared: [3]wt.Q4, + rank: usize, + world: usize, + + pub fn free(self: *Owned) void { + self.experts.free(); + for (&self.shared) |*p| p.free(); + } +}; + +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { + var e = try routed.load(a, d, store, prefix, world, rank); + errdefer e.free(); + var shared: [3]wt.Q4 = @splat(wt.Q4.empty(d)); + errdefer for (&shared) |*p| p.free(); + // The shared branch participates exactly once in the sum across ranks. + if (rank == 0) { + var buf: [256]u8 = undefined; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |name, i| { + const n = if (i == 2) wt.dims else wt.moe_w; + const k = if (i == 2) wt.moe_w else wt.dims; + shared[i] = try projections.load(a, d, store, try std.fmt.bufPrint(&buf, "{s}.shared_expert.{s}", .{ prefix, name }), n, k); + if (shared[i].native == null) return error.UnsupportedQuantization; + } + } + return .{ .experts = e, .shared = shared, .rank = rank, .world = world }; +} diff --git a/zig/src/families/flashnext/cuda_mlp.zig b/zig/src/families/flashnext/cuda_mlp.zig new file mode 100644 index 0000000000..bd7e875d72 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mlp.zig @@ -0,0 +1,255 @@ +//! Layer 0's second hyper-connection: the mixer branch written back, then the same read-out on the mlp weights. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const branch_sk: usize = 8; + +/// `slices` are the output face's unreduced K slices. `h` is the four streams, `inj` the first connection's gates. +pub fn connect(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, h: u64, inj: u64, hidden: []const u16, slices: []const f32) !u32 { + if (hidden.len == 0 or hidden.len % 4 != 0) return error.UnexpectedTensor; + const streams: usize = 4; + const dims = hidden.len / streams; + if (slices.len != branch_sk * dims) return error.UnexpectedTensor; + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + try stream.synchronize(); + var inj_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = inj, .len = inj_bytes.len }; + try inj_buf.download(0, inj_bytes); + const inj_u16 = try gpa.alloc(u16, streams); + defer gpa.free(inj_u16); + for (inj_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little); + const branch = try gpa.alloc(u16, dims); + defer gpa.free(branch); + try hc.sumSlices(slices, branch_sk, dims, branch); + const host_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(host_h); + try hc.applyBranch(hidden, branch, inj_u16, dims, streams, host_h); + + var part = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(slices)); + defer part.free(); + const nc = dims / 256; + var pss_b = try cuda.DeviceBuffer.alloc(driver, nc * streams * 4); + defer pss_b.free(); + try tri.hcWriteback(h, h, pss_b.ptr, part.ptr, "*fp32", inj, h, "*bf16", h, "*bf16", dims, 1, dims, streams, 4, 1, 1, branch_sk); + const got_h = try gpa.alloc(u8, hidden.len * 2); + defer gpa.free(got_h); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + var h_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = h, .len = got_h.len }; + try h_buf.download(0, got_h); + try pss_b.download(0, pss_bytes); + var h_off: usize = 0; + var h_steps: u32 = 0; + for (host_h, 0..) |want, i| { + const got = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) h_off += 1; + if (dist > h_steps) h_steps = dist; + } + const want_pss = try gpa.alloc(f32, nc * streams); + defer gpa.free(want_pss); + try hc.partialSums(host_h, dims, streams, want_pss); + var pss_ulp: u32 = 0; + for (want_pss, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > pss_ulp) pss_ulp = dist; + } + std.debug.print("mlp writeback off {d} max_steps {d} pss_ulp {d} head {x:0>4} host {x:0>4} branch {x:0>4}\n", .{ h_off, h_steps, pss_ulp, std.mem.readInt(u16, got_h[0..2], .little), host_h[0], branch[0] }); + if (h_steps != 0 or pss_ulp > 4) return 1; + + const prefix = "language_model.model.layers.0.mlp_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const down_sk: usize = 32; + var down_part = try cuda.DeviceBuffer.alloc(driver, down_sk * n * 4); + defer down_part.free(); + try tri.hcDown(h, pss_b.ptr, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, down_part.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, down_sk); + const got_norm = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got_norm); + try stream.synchronize(); + try norm_b.download(0, got_norm); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const gpu_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(gpu_h); + for (gpu_h, 0..) |*o, i| o.* = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const want_norm = try gpa.alloc(u16, hidden.len); + defer gpa.free(want_norm); + try hc.normed(gpu_h, pss_f, scale, dims, streams, 1e-6, want_norm); + var norm_off: usize = 0; + var norm_steps: u32 = 0; + var norm_at: usize = 0; + for (want_norm, 0..) |want, i| { + const got = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + if (got != want) { + const dist = hc.mixedSteps(got, want); + if (norm_off == 0 or dist > norm_steps) { + norm_steps = dist; + norm_at = i; + } + norm_off += 1; + } + } + var part4: [4]u8 = undefined; + try down_part.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_part = try hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = hc.ulps(gpu_part, host_part); + std.debug.print("mlp hcdown normed off {d} max_steps {d} at {d} col0 gpu {d} host {d} ulp {d}\n", .{ norm_off, norm_steps, norm_at, gpu_part, host_part, col_ulp }); + if (norm_steps > 1 or col_ulp > 8) return 1; + + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(down_part.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, down_sk, 1, streams, low, n, 1); + const part_bytes = try gpa.alloc(u8, down_sk * n * 4); + defer gpa.free(part_bytes); + try stream.synchronize(); + try down_part.download(0, part_bytes); + const part_f = try gpa.alloc(f32, down_sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try hc.reduceAct(part_f, down_sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_out_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_out_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_out_bytes); + try xs_b.download(0, xs_bytes); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != want) act_off += 1; + } + var inj_off: usize = 0; + for (inj_want, 0..) |want, i| { + if (std.mem.readInt(u16, inj_out_bytes[2 * i ..][0..2], .little) != want) inj_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("mlp reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_off, inj_off, xs_ulp }); + if (act_off != 0 or inj_off != 0 or xs_ulp > 4) return 1; + + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try stream.synchronize(); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, hidden.len); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + var mixed_at: usize = 0; + for (host_mixed, 0..) |want, i| { + const got = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) { + max_steps = dist; + mixed_at = i; + } + } + std.debug.print("mlp upmix dims {d} off {d} max_steps {d} at {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, mixed_at, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_modelopt_experts.zig b/zig/src/families/flashnext/cuda_modelopt_experts.zig new file mode 100644 index 0000000000..40a3b6a35e --- /dev/null +++ b/zig/src/families/flashnext/cuda_modelopt_experts.zig @@ -0,0 +1,132 @@ +//! Rank-local intermediate-channel TP for NVFP4 routed experts; the shared expert retains its own projection format. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const ex = cuda.experts.Nvfp4Experts; +const slice = core.nvfp4_slice; + +pub const Cut = struct { + from: usize, + to: usize, + + pub fn init(total_width: usize, world: usize, rank: usize) !Cut { + if (world == 0 or rank >= world or total_width == 0 or total_width % world != 0) return error.UnevenExpertSplit; + const part = total_width / world; + if (part % 32 != 0) return error.UnevenExpertSplit; + return .{ .from = rank * part, .to = (rank + 1) * part }; + } + + pub fn width(self: Cut) usize { + return self.to - self.from; + } +}; + +pub const Owned = struct { + up: cuda.DeviceBuffer, + down: cuda.DeviceBuffer, + up_scale: cuda.DeviceBuffer, + down_scale: cuda.DeviceBuffer, + width: usize, + dims: usize, + count: usize, + + pub fn free(self: *Owned) void { + self.up.free(); + self.down.free(); + self.up_scale.free(); + self.down_scale.free(); + self.* = undefined; + } + + pub fn view(self: *const Owned) cuda.experts.Layer { + return .{ .format = .nvfp4, .up = self.up.ptr, .down = self.down.ptr, .up_scale = self.up_scale.ptr, .down_scale = self.down_scale.ptr, .width = self.width, .dims = self.dims, .experts = self.count }; + } +}; + +/// Own only the selected rectangle; the exclusive loader store releases every borrowed plane before return. +pub fn readPart(a: std.mem.Allocator, store: anytype, name: []const u8, cut: Cut, down: bool, dims: usize, full_width: usize) !slice.Linear { + defer store.release(); + var buf: [512]u8 = undefined; + const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})); + const n = if (down) dims else full_width; + const k = if (down) full_width else dims; + if (!w.is(.u8, &.{ n, k / 2 })) return error.UnexpectedTensor; + var linear = try slice.read(store, a, name, .{ .from = if (down) 0 else cut.from, .to = if (down) dims else cut.to }, .{ .from = if (down) cut.from else 0, .to = if (down) cut.to else dims }); + errdefer linear.deinit(a); + if (!std.math.isFinite(linear.global) or linear.global <= 0) return error.UnsupportedQuantization; + return linear; +} + +/// Pack one expert at a time with BF16 activations (W4A16); input_scale is validated as metadata. +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { + return loadShape(a, d, store, prefix, world, rank, wt.dims, wt.moe_w, wt.experts_n); +} + +pub fn loadShape(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, world: usize, rank: usize, dims: usize, full_width: usize, count: usize) !Owned { + const cut = try Cut.init(full_width, world, rank); + if (dims == 0 or dims % 32 != 0 or count == 0) return error.UnexpectedTensor; + const width = cut.width(); + const up_bytes = try packedBytes(width, dims, 2); + const down_bytes = try packedBytes(dims, width, 1); + const up_total = std.math.mul(usize, count, up_bytes) catch return error.UnexpectedTensor; + const down_total = std.math.mul(usize, count, down_bytes) catch return error.UnexpectedTensor; + const globals_count = std.math.mul(usize, count, 2) catch return error.UnexpectedTensor; + var up = try cuda.DeviceBuffer.alloc(d, up_total); + errdefer up.free(); + var down = try cuda.DeviceBuffer.alloc(d, down_total); + errdefer down.free(); + const up_globals = try a.alloc(f32, globals_count); + defer a.free(up_globals); + const down_globals = try a.alloc(f32, count); + defer a.free(down_globals); + const packed_up = try a.alloc(u32, up_bytes / 4); + defer a.free(packed_up); + const packed_down = try a.alloc(u32, down_bytes / 4); + defer a.free(packed_down); + var buf: [512]u8 = undefined; + for (0..count) |i| { + var gate = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.gate_proj", .{ prefix, i }), cut, false, dims, full_width); + defer gate.deinit(a); + var value = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.up_proj", .{ prefix, i }), cut, false, dims, full_width); + defer value.deinit(a); + ex.packGateUp(packed_up, .{ gate.codes, gate.scales }, .{ value.codes, value.scales }, width, dims); + try up.upload(i * up_bytes, std.mem.sliceAsBytes(packed_up)); + up_globals[i * 2] = gate.global; + up_globals[i * 2 + 1] = value.global; + var output = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.down_proj", .{ prefix, i }), cut, true, dims, full_width); + defer output.deinit(a); + ex.packOne(packed_down, output.codes, output.scales, dims, width); + try down.upload(i * down_bytes, std.mem.sliceAsBytes(packed_down)); + down_globals[i] = output.global; + } + var us = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(up_globals)); + errdefer us.free(); + const ds = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(down_globals)); + return .{ .up = up, .down = down, .up_scale = us, .down_scale = ds, .width = width, .dims = dims, .count = count }; +} + +test "intermediate-channel cuts cover the source once and preserve 32-input kernel groups" { + const a = try Cut.init(640, 2, 0); + const b = try Cut.init(640, 2, 1); + try std.testing.expectEqual(@as(usize, 0), a.from); + try std.testing.expectEqual(@as(usize, 320), a.to); + try std.testing.expectEqual(a.to, b.from); + try std.testing.expectEqual(@as(usize, 640), b.to); + try std.testing.expectEqual(a.width(), b.width()); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 0, 0)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 2, 2)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 3, 0)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 32, 0)); +} + +fn packedBytes(n: usize, k: usize, parts: usize) !usize { + var size = std.math.mul(usize, n / ex.cols, k / 32) catch return error.UnexpectedTensor; + for ([_]usize{ parts, ex.words, 4 }) |factor| size = std.math.mul(usize, size, factor) catch return error.UnexpectedTensor; + return size; +} + +test "expert packed extents reject overflow before device allocation" { + try std.testing.expectEqual(ex.bytes(1, 320, 2560, 2), try packedBytes(320, 2560, 2)); + try std.testing.expectError(error.UnexpectedTensor, packedBytes(std.math.maxInt(usize), std.math.maxInt(usize), 2)); +} diff --git a/zig/src/families/flashnext/cuda_moe.zig b/zig/src/families/flashnext/cuda_moe.zig new file mode 100644 index 0000000000..a59407876f --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe.zig @@ -0,0 +1,512 @@ +//! Layer 0's experts on one mixed row: router, top 10 plus the shared expert, then each slot's gate, up and down. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const count: usize = 512; +const top_k: usize = 10; +const slots: usize = top_k + 1; +const dims: usize = 2560; +const width: usize = 640; +const routed_rows: usize = count + 1; +const tile: usize = 16; +const block_bytes: usize = 160 * 4; + +const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Proj = struct { + routed_w: []const u8, + routed_s: []const u8, + routed_b: []const u8, + shared_w: []const u8, + shared_s: []const u8, + shared_b: []const u8, + n: usize, + k: usize, +}; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundTrip(v: f32) f32 { + return promote(toBf16(v)); +} + +/// The expert kernel's group sum: each lane adds its eight inputs, then the quad adds as (0+1)+(2+3). +fn groupSum(x: []const u16, g: usize) f32 { + var pair: [2]f32 = .{ 0, 0 }; + for (0..4) |t| { + var s: f32 = 0; + const base = g * 32 + t * 8; + for (0..8) |j| s += promote(x[base + j]); + if (t % 2 == 0) pair[t / 2] = s else pair[t / 2] += s; + } + return pair[0] + pair[1]; +} + +/// One expert matrix: groups in order, `p * scale + xs * bias`, fp32. `xs` is `groupSum` of `x`. +fn project(x: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []f32) void { + const kg = k / 32; + const k8 = k / 8; + var xs: [80]f32 = undefined; + for (0..kg) |g| xs[g] = groupSum(x, g); + for (0..n) |col| { + var acc: f32 = 0; + const wb = col * k8 * 4; + const sb = col * kg * 2; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, words[wb + (g * 4 + i) * 4 ..][0..4], .little); + const base = g * 32 + i * 8; + inline for (0..8) |j| { + const q: f32 = @floatFromInt((word >> (j * 4)) & 0xF); + p += promote(x[base + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scales[sb + g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, biases[sb + g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = acc; + } +} + +/// bf16(SiLU(bf16(gate)) * bf16(up)), the SwiGLU epilogue with no clip. +fn swiglu(gate: f32, up: f32) u16 { + const gv = roundTrip(gate); + const uv = roundTrip(up); + const silu = roundTrip(gv / (1 + @exp(-gv))); + return toBf16(silu * uv); +} + +/// Router rows in K tiles of 256, the same tile the captured program uses, added in tile order. +fn routerDots(x: []const u16, w: []const u16, out: []f32) void { + const k = x.len; + for (out, 0..) |*o, e| { + var acc: f32 = 0; + const row = w[e * k ..][0..k]; + var t: usize = 0; + while (t < k) : (t += 256) { + var part: f32 = 0; + for (0..256) |i| part += promote(x[t + i]) * promote(row[t + i]); + acc += part; + } + o.* = acc; + } +} + +/// Largest `k` of the first `n` logits (lower id on a tie), softmax weights rounded to bf16, then the shared gate. +fn choose(logits: []const f32, n: usize, k: usize, picks: []i32, wts: []f32) void { + var scratch: [count]f32 = undefined; + for (0..n) |i| scratch[i] = logits[i]; + var top = scratch[0]; + for (scratch[1..n]) |v| top = @max(top, v); + var total: f32 = 0; + var exs: [16]f32 = undefined; + for (0..k) |slot| { + var m = scratch[0]; + var idx: usize = 0; + for (scratch[1..n], 1..) |v, i| if (v > m) { + m = v; + idx = i; + }; + const ex = @exp(m - top); + picks[slot] = @intCast(idx); + exs[slot] = ex; + total += ex; + scratch[idx] = -std.math.inf(f32); + } + for (0..k) |slot| wts[slot] = roundTrip(exs[slot] / total); + const shared = roundTrip(logits[n]); + picks[k] = @intCast(n); + wts[k] = roundTrip(1 / (1 + @exp(-shared))); +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn loadProj(mapped: *const embed.Mapped, gpa: std.mem.Allocator, name: []const u8, n: usize, k: usize) !Proj { + const prefix = "language_model.model.layers.0.mlp."; + var buf: [180]u8 = undefined; + const k8 = k / 8; + const kg = k / 32; + const rw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.weight", .{ prefix, name }), .u32); + const rs = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.scales", .{ prefix, name }), .bf16); + const rb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.biases", .{ prefix, name }), .bf16); + const sw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.weight", .{ prefix, name }), .u32); + const ss = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.scales", .{ prefix, name }), .bf16); + const sb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.biases", .{ prefix, name }), .bf16); + if (rw.dim(0) != count or rw.dim(1) != n or rw.dim(2) != k8) return error.UnexpectedTensor; + if (rs.dim(0) != count or rs.dim(1) != n or rs.dim(2) != kg) return error.UnexpectedTensor; + if (rb.dim(0) != count or rb.dim(1) != n or rb.dim(2) != kg) return error.UnexpectedTensor; + if (sw.dim(0) != n or sw.dim(1) != k8 or ss.dim(0) != n or ss.dim(1) != kg) return error.UnexpectedTensor; + if (sb.dim(0) != n or sb.dim(1) != kg) return error.UnexpectedTensor; + return .{ .routed_w = rw.bytes, .routed_s = rs.bytes, .routed_b = rb.bytes, .shared_w = sw.bytes, .shared_s = ss.bytes, .shared_b = sb.bytes, .n = n, .k = k }; +} + +fn copyExpert(dst_w: []u8, dst_s: []u8, dst_b: []u8, p: Proj, id: i32, slot: usize) void { + const k8 = p.k / 8; + const kg = p.k / 32; + const wb = p.n * k8 * 4; + const sb = p.n * kg * 2; + const routed = id != count; + const wsrc = if (routed) p.routed_w[@as(usize, @intCast(id)) * wb ..][0..wb] else p.shared_w; + const ssrc = if (routed) p.routed_s[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_s; + const bsrc = if (routed) p.routed_b[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_b; + @memcpy(dst_w[slot * wb ..][0..wb], wsrc); + @memcpy(dst_s[slot * sb ..][0..sb], ssrc); + @memcpy(dst_b[slot * sb ..][0..sb], bsrc); +} + +fn packDownload(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) ![]u8 { + const kg = k / 32; + const nb = n / 32; + const out_len = e * nb * kg * block_bytes; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, out_len); + defer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream, &a); + const out = try gpa.alloc(u8, out_len); + errdefer gpa.free(out); + try stream.synchronize(); + try out_b.download(0, out); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +fn maxUlp(got: []const f32, want: []const f32) struct { off: usize, ulp: u32 } { + var off: usize = 0; + var ulp: u32 = 0; + for (got, want) |g, w| { + const dist = hc.ulps(g, w); + if (dist != 0) off += 1; + if (dist > ulp) ulp = dist; + } + return .{ .off = off, .ulp = ulp }; +} + +/// The 512 routed experts and the shared expert, on the mixed row already at `mixed`. +pub fn experts(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, mixed: u64) !u32 { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const prefix = "language_model.model.layers.0.mlp."; + var name: [180]u8 = undefined; + const gate = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}gate.weight", .{prefix}), .bf16); + const sg_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.weight", .{prefix}), .u32); + const sg_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.scales", .{prefix}), .bf16); + const sg_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.biases", .{prefix}), .bf16); + if (gate.dim(0) != count or gate.dim(1) != dims) return error.UnexpectedTensor; + if (sg_w.dim(0) != 1 or sg_w.dim(1) != dims / 8) return error.UnexpectedTensor; + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + try stream.synchronize(); + var mixed_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + try mixed_buf.download(0, row_raw); + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_raw[2 * i ..][0..2], .little); + + const router = try gpa.alloc(u16, routed_rows * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. count * dims]), gate.bytes); + try embed.dequant(sg_w.bytes, sg_s.bytes, sg_b.bytes, dims, 0, router[count * dims ..][0..dims]); + const host_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(host_l); + routerDots(row, router, host_l); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + defer w_b.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, routed_rows * 4); + defer logits_b.free(); + var pick_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer pick_b.free(); + var wts_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer wts_b.free(); + try tri.router(mixed, w_b.ptr, logits_b.ptr, dims); + try tri.topkRows(logits_b.ptr, pick_b.ptr, wts_b.ptr); + + const logit_raw = try gpa.alloc(u8, routed_rows * 4); + defer gpa.free(logit_raw); + const pick_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(pick_raw); + const wts_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(wts_raw); + try stream.synchronize(); + try logits_b.download(0, logit_raw); + try pick_b.download(0, pick_raw); + try wts_b.download(0, wts_raw); + const gpu_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(gpu_l); + for (gpu_l, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, logit_raw[4 * i ..][0..4], .little)); + var gpu_pick: [slots]i32 = undefined; + var gpu_w: [slots]f32 = undefined; + for (0..slots) |i| { + gpu_pick[i] = std.mem.readInt(i32, pick_raw[4 * i ..][0..4], .little); + gpu_w[i] = @bitCast(std.mem.readInt(u32, wts_raw[4 * i ..][0..4], .little)); + } + const route = maxUlp(gpu_l, host_l); + var from_gpu: [slots]i32 = undefined; + var w_gpu: [slots]f32 = undefined; + choose(gpu_l, count, top_k, &from_gpu, &w_gpu); + var from_host: [slots]i32 = undefined; + var w_host: [slots]f32 = undefined; + choose(host_l, count, top_k, &from_host, &w_host); + var pick_off: usize = 0; + var host_pick_off: usize = 0; + for (gpu_pick, from_gpu, from_host) |g, a, b| { + if (g != a) pick_off += 1; + if (g != b) host_pick_off += 1; + } + const wts_same = maxUlp(&gpu_w, &w_gpu); + var host_wts_off: usize = 0; + var host_wts_steps: u32 = 0; + for (gpu_w, w_host) |g, h| { + const steps = hc.mixedSteps(toBf16(g), toBf16(h)); + if (steps != 0) host_wts_off += 1; + if (steps > host_wts_steps) host_wts_steps = steps; + } + std.debug.print("router off {d} max_ulp {d} head {x:0>8} host {x:0>8}\n", .{ route.off, route.ulp, @as(u32, @bitCast(gpu_l[0])), @as(u32, @bitCast(host_l[0])) }); + std.debug.print("topk pick_off {d} host_pick_off {d} wts_ulp {d} host_wts_off {d} host_wts_steps {d} picks", .{ pick_off, host_pick_off, wts_same.ulp, host_wts_off, host_wts_steps }); + for (gpu_pick) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + + var packed_ids: [slots]i32 = undefined; + var npack: usize = 0; + var local: [slots]i32 = undefined; + for (gpu_pick, 0..) |id, s| { + if (id < 0 or id > count) return error.UnexpectedTensor; + var found: ?usize = null; + for (packed_ids[0..npack], 0..) |prev, i| if (prev == id) { + found = i; + }; + if (found) |i| { + local[s] = @intCast(i); + } else { + packed_ids[npack] = id; + local[s] = @intCast(npack); + npack += 1; + } + } + + const gate_p = try loadProj(mapped, gpa, "gate_proj", width, dims); + const up_p = try loadProj(mapped, gpa, "up_proj", width, dims); + const down_p = try loadProj(mapped, gpa, "down_proj", dims, width); + const gathered_w = try gpa.alloc(u8, npack * width * (dims / 8) * 4); + defer gpa.free(gathered_w); + const gathered_s = try gpa.alloc(u8, npack * width * (dims / 32) * 2); + defer gpa.free(gathered_s); + const gathered_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(gathered_b); + const up_w = try gpa.alloc(u8, gathered_w.len); + defer gpa.free(up_w); + const up_s = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_s); + const up_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_b); + const down_w = try gpa.alloc(u8, npack * dims * (width / 8) * 4); + defer gpa.free(down_w); + const down_s = try gpa.alloc(u8, npack * dims * (width / 32) * 2); + defer gpa.free(down_s); + const down_b = try gpa.alloc(u8, down_s.len); + defer gpa.free(down_b); + for (packed_ids[0..npack], 0..) |id, slot| { + copyExpert(gathered_w, gathered_s, gathered_b, gate_p, id, slot); + copyExpert(up_w, up_s, up_b, up_p, id, slot); + copyExpert(down_w, down_s, down_b, down_p, id, slot); + } + + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + var exp_mod = try cuda.Module.load(driver, cuda.kernels.experts); + defer exp_mod.unload(); + const pack_fn = try pack_mod.function(pack_symbol); + const packed_g = try packDownload(gpa, driver, stream.*, pack_fn, gathered_w, gathered_s, gathered_b, npack, width, dims); + defer gpa.free(packed_g); + const packed_u = try packDownload(gpa, driver, stream.*, pack_fn, up_w, up_s, up_b, npack, width, dims); + defer gpa.free(packed_u); + const packed_d = try packDownload(gpa, driver, stream.*, pack_fn, down_w, down_s, down_b, npack, dims, width); + defer gpa.free(packed_d); + const stacked = try stackHalves(gpa, packed_g, packed_u); + defer gpa.free(stacked); + + var up_wb = try cuda.DeviceBuffer.fromHost(driver, stacked); + defer up_wb.free(); + var down_wb = try cuda.DeviceBuffer.fromHost(driver, packed_d); + defer down_wb.free(); + var local_pick: [slots]i32 = undefined; + for (0..slots) |s| local_pick[s] = local[s]; + var picks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&local_pick)); + defer picks_b.free(); + var members_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer members_b.free(); + var items_b = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + defer items_b.free(); + var counts_b = try cuda.DeviceBuffer.alloc(driver, 8); + defer counts_b.free(); + var act_b = try cuda.DeviceBuffer.alloc(driver, slots * width * 2); + defer act_b.free(); + var y_b = try cuda.DeviceBuffer.alloc(driver, slots * dims * 4); + defer y_b.free(); + const plan_fn = try exp_mod.function(plan_symbol); + var plan_a: cuda.Args = .{}; + plan_a.add(picks_b.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(npack)); + plan_a.add(cint(tile)); + plan_a.add(members_b.ptr); + plan_a.add(items_b.ptr); + plan_a.add(counts_b.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + + const up_fn = try exp_mod.function(up_symbol); + const down_fn = try exp_mod.function(down_symbol); + try launchExpert(up_fn, stream.*, mixed, dims, slots, up_wb.ptr, dims / 32, width / 32, items_b.ptr, counts_b.ptr, members_b.ptr, act_b.ptr, width, slots * (width / 32)); + try launchExpert(down_fn, stream.*, act_b.ptr, width, 0, down_wb.ptr, width / 32, dims / 32, items_b.ptr, counts_b.ptr, members_b.ptr, y_b.ptr, dims, slots * (dims / 32)); + + const act_raw = try gpa.alloc(u8, slots * width * 2); + defer gpa.free(act_raw); + const y_raw = try gpa.alloc(u8, slots * dims * 4); + defer gpa.free(y_raw); + try stream.synchronize(); + try act_b.download(0, act_raw); + try y_b.download(0, y_raw); + + var act_off: usize = 0; + var act_steps: u32 = 0; + var y_off: usize = 0; + var y_ulp: u32 = 0; + var y_steps: u32 = 0; + var head_g: u32 = 0; + var head_h: u32 = 0; + const gate_col = try gpa.alloc(f32, width); + defer gpa.free(gate_col); + const up_col = try gpa.alloc(f32, width); + defer gpa.free(up_col); + const act_row = try gpa.alloc(u16, width); + defer gpa.free(act_row); + const down_col = try gpa.alloc(f32, dims); + defer gpa.free(down_col); + const gw = width * (dims / 8) * 4; + const gs = width * (dims / 32) * 2; + const dw = dims * (width / 8) * 4; + const ds = dims * (width / 32) * 2; + for (0..slots) |s| { + const src: usize = @intCast(local[s]); + project(row, gathered_w[src * gw ..][0..gw], gathered_s[src * gs ..][0..gs], gathered_b[src * gs ..][0..gs], width, dims, gate_col); + project(row, up_w[src * gw ..][0..gw], up_s[src * gs ..][0..gs], up_b[src * gs ..][0..gs], width, dims, up_col); + for (act_row, gate_col, up_col) |*o, gv, uv| o.* = swiglu(gv, uv); + for (0..width) |c| { + const got = std.mem.readInt(u16, act_raw[(s * width + c) * 2 ..][0..2], .little); + const dist = hc.mixedSteps(got, act_row[c]); + if (dist != 0) act_off += 1; + if (dist > act_steps) act_steps = dist; + } + project(act_row, down_w[src * dw ..][0..dw], down_s[src * ds ..][0..ds], down_b[src * ds ..][0..ds], dims, width, down_col); + for (0..dims) |c| { + const got: f32 = @bitCast(std.mem.readInt(u32, y_raw[(s * dims + c) * 4 ..][0..4], .little)); + const dist = hc.ulps(got, down_col[c]); + const steps = hc.mixedSteps(toBf16(got), toBf16(down_col[c])); + if (dist != 0) y_off += 1; + if (dist > y_ulp) y_ulp = dist; + if (steps > y_steps) y_steps = steps; + if (s == 0 and c == 0) { + head_g = @bitCast(got); + head_h = @bitCast(down_col[c]); + } + } + } + std.debug.print("experts slots {d} act_off {d} act_steps {d} y_off {d} y_ulp {d} y_steps {d} head {x:0>8} host {x:0>8}\n", .{ slots, act_off, act_steps, y_off, y_ulp, y_steps, head_g, head_h }); + if (pick_off != 0 or host_pick_off != 0 or wts_same.ulp != 0 or host_wts_steps > 1 or act_steps != 0 or y_steps != 0) return 1; + return 0; +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + const grid = (units + 3) / 4; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "a tied logit keeps the lower expert id and the shared gate rounds to one half" { + const logits = [_]f32{ 1, 3, 3, 0 }; + var picks: [3]i32 = undefined; + var wts: [3]f32 = undefined; + choose(&logits, 3, 2, &picks, &wts); + try std.testing.expectEqual([3]i32{ 1, 2, 3 }, picks); + try std.testing.expectEqual(@as(f32, 0.5), wts[0]); + try std.testing.expectEqual(@as(f32, 0.5), wts[1]); + try std.testing.expectEqual(@as(f32, 0.5), wts[2]); +} + +test "each lane of the quad contributes to the group sum" { + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[8] = 0x3f80; + x[16] = 0x3f80; + x[24] = 0x3f80; + try std.testing.expectEqual(@as(f32, 4), groupSum(&x, 0)); +} + +test "scale times the nibble dot plus the group bias is five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + var out: [1]f32 = undefined; + project(&x, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(f32, 5), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig new file mode 100644 index 0000000000..8c3d03e9ac --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -0,0 +1,277 @@ +//! Flash Next's MTP drafts on CUDA (Python qwen4_exp/cuda/mtp.py): the head, its draft chains and verify windows. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; + +pub const Sampling = lanes.sampling.Sampling; + +/// Drafts a verify window can hold beside its pending token. +pub const max_chain: usize = max_rows - 1; +/// Preserve the existing recipe unless an explicit depth is selected. +pub const default_depth: usize = 6; +/// The live recipe's --mtp-confidence (Python's CONFIDENCE default is 0.7). +pub const default_confidence: f64 = 0.6; +/// A verify window's rows at most: the pending token and a chain, or a copied continuation. +pub const max_rows: usize = 16; +/// qmm.split_for(2560, 2560): fc_embedding and fc_hidden sum K in four slices, in order. +pub const fc_sk: usize = 4; +/// Head rows one launch takes at most (a prompt's rows go through in pieces). +pub const piece_rows: usize = 256; +/// Normed-stream partial sums a row (cuda_batch.zig's pss_n). +pub const pss_n = (dims / 256) * streams; +pub const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; +pub const proj_bytes = gdn.proj_width * 2; + +const vocab_text = @embedFile("draft_vocab_default.txt"); + +fn cint(v: usize) c_int { + return @intCast(v); +} + +/// The draft vocabulary's ids, sorted (Python draft_token_ids("default")): column j of the draft head is ids[j]. +pub fn draftIds(gpa: std.mem.Allocator) ![]u32 { + var out: std.ArrayList(u32) = .empty; + errdefer out.deinit(gpa); + var it = std.mem.tokenizeAny(u8, vocab_text, " \t\r\n,"); + while (it.next()) |part| { + const id = try std.fmt.parseInt(u32, part, 10); + if (id >= vocab) continue; + try out.append(gpa, id); + } + std.mem.sort(u32, out.items, {}, std.sort.asc(u32)); + var n: usize = 0; + for (out.items, 0..) |id, i| { + if (i > 0 and id == out.items[n - 1]) continue; + out.items[n] = id; + n += 1; + } + out.shrinkRetainingCapacity(n); + return out.toOwnedSlice(gpa); +} + +// ---- sampling ------------------------------------------------------------------------------------------------------ + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// A row's first maximum (argmax's choice): its column. +pub fn argmaxRow(row: []const u16) usize { + var best = promote(row[0]); + var at: usize = 0; + for (row[1..], 1..) |b, i| { + const v = promote(b); + if (v > best) { + best = v; + at = i; + } + } + return at; +} + +/// The keyed draw over one bf16 score row, column j being token ids[j] (or j), as lanes.sampling.choose draws. +pub fn drawRow(gpa: std.mem.Allocator, row: []const u16, ids: ?[]const u32, position: u64, sm: ?Sampling) !u32 { + const id = struct { + fn of(map: ?[]const u32, col: usize) u64 { + return if (map) |m| m[col] else col; + } + }.of; + const s = sm orelse return @intCast(id(ids, argmaxRow(row))); + if (s.temperature <= 0) return @intCast(id(ids, argmaxRow(row))); + const n = row.len; + const k = if (s.top_k == 0 or s.top_k >= n) n else @as(usize, s.top_k); + const vals = try gpa.alloc(f64, k); + defer gpa.free(vals); + const keys = try gpa.alloc(u64, k); + defer gpa.free(keys); + if (k == n) { + for (vals, keys, row, 0..) |*v, *key, b, col| { + v.* = promote(b); + key.* = id(ids, col); + } + } else { + // the k best columns, best first; a later column with an equal value ranks below (the id order) + const top_v = try gpa.alloc(f32, k); + defer gpa.free(top_v); + const top_c = try gpa.alloc(usize, k); + defer gpa.free(top_c); + var held: usize = 0; + for (row, 0..) |b, col| { + const v = promote(b); + if (held == k and !(v > top_v[k - 1])) continue; + var at: usize = if (held < k) held else k - 1; + if (held < k) held += 1; + while (at > 0 and v > top_v[at - 1]) : (at -= 1) { + top_v[at] = top_v[at - 1]; + top_c[at] = top_c[at - 1]; + } + top_v[at] = v; + top_c[at] = col; + } + for (vals, keys, top_v, top_c) |*v, *key, tv, c| { + v.* = tv; + key.* = id(ids, c); + } + } + return @intCast(try lanes.sampling.choose(gpa, vals, keys, position, s)); +} + +// ---- the head's weights -------------------------------------------------------------------------------------------- + +pub const Head = struct { + gpa: std.mem.Allocator, + norm_e: cuda.DeviceBuffer, // [D] fp32 + norm_h: cuda.DeviceBuffer, // [S D] fp32 + fc_e: Q4, + fc_h: Q4, + layer: load.Layer, + mixer: HcW, + /// lm_head's rows at `ids`, in their order. + draft: Q4, + ids: []u32, + module: cuda.Module, + best_fn: cuda.Function, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, dir: []const u8) !*Head { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const h = try gpa.create(Head); + errdefer gpa.destroy(h); + h.gpa = gpa; + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + try store.detectNorms(); + var buf: [120]u8 = undefined; + const ne = try store.normScale(gpa, try store.tensor("language_model.mtp.pre_fc_norm_embedding.weight")); + defer gpa.free(ne); + const nh = try store.normScale(gpa, try store.tensor("language_model.mtp.pre_fc_norm_hidden.weight")); + defer gpa.free(nh); + if (ne.len != dims or nh.len != wide) return error.UnexpectedTensor; + h.norm_e = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(ne)); + errdefer h.norm_e.free(); + h.norm_h = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nh)); + errdefer h.norm_h.free(); + h.fc_e = try fc(gpa, driver, &store, &buf, "language_model.mtp.fc_embedding"); + errdefer h.fc_e.free(); + h.fc_h = try fc(gpa, driver, &store, &buf, "language_model.mtp.fc_hidden"); + errdefer h.fc_h.free(); + h.mixer = try wt.loadHc(gpa, driver, &store, "language_model.mtp.hyper_connection_mixer", false); + errdefer h.mixer.free(); + { + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const pack_fn = try pack_mod.function(wt.pack_symbol); + h.layer = try load.loadLayerAt(gpa, driver, stream, &store, pack_fn, "language_model.mtp.layers.0", false, false); + } + errdefer h.layer.free(gpa); + h.ids = try draftIds(gpa); + errdefer gpa.free(h.ids); + h.draft = try @import("cuda_mtp_load.zig").draft(gpa, driver, &store, h.ids); + errdefer h.draft.free(); + h.module = try cuda.Module.load(driver, cuda.kernels.flashnext_mtp); + errdefer h.module.unload(); + h.best_fn = try h.module.function("fn_row_best"); + return h; + } + + fn fc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *wt.Store, buf: []u8, name: []const u8) !Q4 { + _ = buf; + return @import("cuda_projection_load.zig").load(gpa, driver, store, name, dims, dims); + } + + pub fn deinit(h: *Head) void { + h.module.unload(); + h.draft.free(); + h.gpa.free(h.ids); + h.layer.free(h.gpa); + h.mixer.free(); + h.fc_h.free(); + h.fc_e.free(); + h.norm_h.free(); + h.norm_e.free(); + h.gpa.destroy(h); + } + + /// fn_row_best over `rows` rows of `v` scores `ld` apart: [rows, 4] fp32 at `out` (best, log-sum-exp, id bits). + pub fn best(h: *const Head, stream: cuda.Stream, scores: u64, v: usize, ld: usize, rows: usize, out: u64) !void { + var a: cuda.Args = .{}; + a.add(scores); + a.add(cint(v)); + a.add(@as(c_longlong, @intCast(ld))); + a.add(out); + try cuda.launch.launch(h.best_fn, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, stream, &a); + } +}; + +/// Device bytes the head's weights take beyond the main model's (for the server's admission estimate). +pub fn headBytes() usize { + const experts = (wt.experts_n + 1) * 3 * wt.moe_w * dims / 2 * 5 / 4; + const attn = (wt.proj_n + dims) * dims / 2 * 5 / 4; + const draft = 79_591 * dims / 2 * 5 / 4; + return experts + attn + draft + 2 * dims * dims / 2 * 5 / 4; +} + +pub const Pick = @import("cuda_mtp_draft.zig").Pick; +pub const Chain = @import("cuda_mtp_draft.zig").Chain; +pub const Draft = @import("cuda_mtp_draft.zig").Draft; +pub const Ops = @import("cuda_mtp_ops.zig").Ops; +pub const check = @import("cuda_mtp_check.zig").check; + +test "the draft vocabulary is 79,591 sorted ids under the vocabulary" { + const ids = try draftIds(std.testing.allocator); + defer std.testing.allocator.free(ids); + try std.testing.expectEqual(@as(usize, 79_591), ids.len); + for (ids[1..], ids[0 .. ids.len - 1]) |b, a| try std.testing.expect(b > a); +} + +test "a keyed draw over the top-k columns is the draw over every column" { + const gpa = std.testing.allocator; + var row: [300]u16 = undefined; + var prng = std.Random.DefaultPrng.init(3); + for (&row) |*v| v.* = @as(u16, 0x3f00) + prng.random().int(u8); // values near 0.5..1, many ties + const all = struct { + fn draw(a: std.mem.Allocator, r: []const u16, position: u64, s: Sampling) !u64 { + const vals = try a.alloc(f64, r.len); + defer a.free(vals); + const ids = try a.alloc(u64, r.len); + defer a.free(ids); + for (vals, ids, r, 0..) |*v, *id, b, i| { + v.* = promote(b); + id.* = i; + } + return lanes.sampling.choose(a, vals, ids, position, s); + } + }.draw; + for (0..64) |position| { + const s: Sampling = .{ .seed = 11, .temperature = 1.0, .top_k = 20, .top_p = 0.95 }; + try std.testing.expectEqual(try all(gpa, &row, position, s), try drawRow(gpa, &row, null, position, s)); + } + try std.testing.expectEqual(@as(u32, @intCast(argmaxRow(&row))), try drawRow(gpa, &row, null, 0, null)); +} + +test "draft chains fit one verify window while the default stays six" { + try std.testing.expectEqual(@as(usize, 16), max_chain + 1); + try std.testing.expectEqual(@as(usize, 6), default_depth); + try std.testing.expect(default_depth <= max_chain); +} diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig new file mode 100644 index 0000000000..435cf6c9da --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -0,0 +1,171 @@ +//! The drafted check (flashnext-prompt with TF_FLASHNEXT_DRAFTS): serial and drafted decoding, token for token. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); + +const vocab = wt.vocab; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const default_confidence = mtp.default_confidence; +const Head = mtp.Head; +const Draft = mtp.Draft; +const Chain = mtp.Chain; +const Ops = mtp.Ops; +const Sampling = mtp.Sampling; +const drawRow = mtp.drawRow; + +/// What a drafted decode did: rounds, drafts offered and kept, and its wall time. +pub const Stats = struct { + rounds: usize = 0, + drafted: usize = 0, + accepted: usize = 0, + ms: f64 = 0, + /// Wall time in the windows (forward and draws) and in the head (rollback, absorb, chain). + verify_ms: f64 = 0, + draft_ms: f64 = 0, +}; + +/// Python mtp_decode over a prefilled sequence: `out` holds the first token; decode until `count` tokens. +pub fn decode(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, d: *Draft, out: *std.ArrayList(u32), count: usize, depth: usize, confidence: f64, sm: ?Sampling, stop: []const u32) !Stats { + const gpa = o.e.gpa; + var st: Stats = .{}; + var ch: Chain = .{}; + var sampled: [max_rows]u32 = undefined; + var ids: [max_rows]u32 = undefined; + const t0 = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + const isStop = struct { + fn f(list: []const u32, t: u32) bool { + return std.mem.indexOfScalar(u32, list, t) != null; + } + }.f; + try o.draft(d, out.items[out.items.len - 1 ..], d.prompt_last, seq.fed + 1, @min(depth, count - out.items.len), sm, confidence, &ch); + while (out.items.len < count and !isStop(stop, out.items[out.items.len - 1])) { + ids[0] = out.items[out.items.len - 1]; + @memcpy(ids[1..][0..ch.n], ch.slice()); + const rows = 1 + ch.n; + const tv = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + try o.verify(seq, d, ids[0..rows]); + try o.draws(d, sm, sampled[0..rows]); + const td = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + st.verify_ms += @as(f64, @floatFromInt(td - tv)) / 1e6; + defer st.draft_ms += @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - td)) / 1e6; + var keep: usize = 1; + for (ch.slice(), 0..) |t, i| { + if (sampled[i] != t or isStop(stop, sampled[i])) break; + keep += 1; + } + try o.rollback(seq, d, keep); + st.rounds += 1; + st.drafted += ch.n; + st.accepted += keep - 1; + const take = @min(keep, count - out.items.len); + try out.appendSlice(gpa, sampled[0..take]); + if (out.items.len >= count or isStop(stop, out.items[out.items.len - 1])) break; + try o.draft(d, sampled[0..keep], d.streams.ptr, seq.fed + 1, @min(depth, count - out.items.len), sm, confidence, &ch); + } + st.ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - t0)) / 1e6; + return st; +} + +/// Serial decoding through the same pass and draws: one row a forward (the reference drafted decoding must equal). +pub fn serial(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, out: *std.ArrayList(u32), count: usize, sm: ?Sampling, stop: []const u32) !f64 { + const gpa = o.e.gpa; + const t0 = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + const row = try gpa.alloc(u16, vocab); + defer gpa.free(row); + while (out.items.len < count and std.mem.indexOfScalar(u32, stop, out.items[out.items.len - 1]) == null) { + var one = [_]u32{out.items[out.items.len - 1]}; + try o.e.predict(Tri, o.tri, seq, &one); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try out.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + } + return @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - t0)) / 1e6; +} + +fn envUsize(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} + +fn envF64(name: [:0]const u8, default: f64) f64 { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch default; +} + +/// The drafted check (TF_FLASHNEXT_DRAFTS): serial and drafted decoding of `count` tokens a prompt, token for token. +pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, prompts: []const []const u32, count: usize) !bool { + var e = try prompt.Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + const h = try Head.init(gpa, io, driver, stream, model_dir); + defer h.deinit(); + const o: Ops(Tri) = .{ .e = e, .tri = tri, .h = h }; + const depth = envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth); + const confidence = envF64("TF_FLASHNEXT_CONFIDENCE", default_confidence); + const sm: ?Sampling = if (std.c.getenv("TF_FLASHNEXT_SEED")) |seed| .{ + .seed = std.fmt.parseInt(u64, std.mem.span(seed), 10) catch 0, + .temperature = envF64("TF_FLASHNEXT_TEMPERATURE", 1.0), + .top_p = envF64("TF_FLASHNEXT_TOP_P", 0.95), + .top_k = @intCast(envUsize("TF_FLASHNEXT_TOP_K", 20)), + } else null; + const repeat = envUsize("TF_FLASHNEXT_REPEAT", 1); + const stop: []const u32 = &.{}; + var same = true; + for (prompts, 0..) |ids, pi| { + const cap = ids.len + count + max_rows; + var seq = try e.newSeq(cap); + defer seq.deinit(gpa); + var d = try Draft.init(gpa, driver, stream.*, cap, h.ids.len, e.chunk); + defer d.deinit(); + for (0..repeat) |rep| { + // serial: the plain prompt pass, then one row a forward + seq.fed = 0; + seq.history = .{ e.eos, e.eos }; + try resetSeq(&seq, stream.*); + var ref: std.ArrayList(u32) = .empty; + defer ref.deinit(gpa); + try e.predict(Tri, tri, &seq, ids); + const row = try gpa.alloc(u16, vocab); + defer gpa.free(row); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try ref.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + const serial_ms = try serial(Tri, o, &seq, &ref, count, sm, stop); + // drafted: the prompt pass with its streams and the head's cache, then windows + seq.fed = 0; + seq.history = .{ e.eos, e.eos }; + try resetSeq(&seq, stream.*); + var got: std.ArrayList(u32) = .empty; + defer got.deinit(gpa); + try o.prefill(&seq, &d, ids); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try got.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + const st = try decode(Tri, o, &seq, &d, &got, count, depth, confidence, sm, stop); + const equal = std.mem.eql(u32, ref.items, got.items); + same = same and equal; + const n_ref: f64 = @floatFromInt(ref.items.len - 1); + const n_got: f64 = @floatFromInt(got.items.len - 1); + std.debug.print("prompt {d} ({d} ids) run {d}: serial {d:.2} tok/s, drafted {d:.2} tok/s ({d:.2}x), {d} rounds, {d} drafts, {d} kept ({d:.1}%), {d:.2} tokens a round: {s}\n", .{ + pi, ids.len, rep, + n_ref / (serial_ms / 1e3), n_got / (st.ms / 1e3), (n_got / (st.ms / 1e3)) / (n_ref / (serial_ms / 1e3)), + st.rounds, st.drafted, st.accepted, + if (st.drafted > 0) 100.0 * @as(f64, @floatFromInt(st.accepted)) / @as(f64, @floatFromInt(st.drafted)) else 0.0, if (st.rounds > 0) n_got / @as(f64, @floatFromInt(st.rounds)) else 0.0, if (equal) "same tokens" else "DIFFER", + }); + std.debug.print(" drafted rounds: {d:.2} ms a window ({d:.2} rows), {d:.2} ms in the head a round\n", .{ st.verify_ms / @as(f64, @floatFromInt(@max(st.rounds, 1))), @as(f64, @floatFromInt(st.rounds + st.drafted)) / @as(f64, @floatFromInt(@max(st.rounds, 1))), st.draft_ms / @as(f64, @floatFromInt(@max(st.rounds, 1))) }); + std.debug.print("serial ", .{}); + for (ref.items) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\ndrafted", .{}); + for (got.items) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + } + } + return same; +} + +/// A sequence's caches back to empty (the check reruns a prompt on one allocation). +fn resetSeq(seq: *prompt.Seq, stream: cuda.Stream) !void { + for (seq.conv) |*b| try b.fill8(0, stream.handle); + for (seq.state) |*b| try b.fill8(0, stream.handle); + try seq.tail.fill8(0, stream.handle); +} diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig new file mode 100644 index 0000000000..7f354f48de --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -0,0 +1,176 @@ +//! A drafting stream's state: the head's cache, a window's buffers and what a partly kept window rolls back to. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); +const mtp = @import("cuda_mtp.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const conv_bytes = mtp.conv_bytes; +const proj_bytes = mtp.proj_bytes; + +/// One draft and its temperature-1 probability under the head. +pub const Pick = struct { token: u32, prob: f64 }; + +/// A chain of drafts the head holds for the next window. +pub const Chain = struct { + tokens: [max_chain]u32 = undefined, + probs: [max_chain]f64 = undefined, + n: usize = 0, + + pub fn slice(c: *const Chain) []const u32 { + return c.tokens[0..c.n]; + } +}; + +/// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. +pub const Draft = struct { + gpa: std.mem.Allocator, + kc: cuda.DeviceBuffer, + vc: cuda.DeviceBuffer, + ikc: cuda.DeviceBuffer, + /// Its pooled indexer keys [cap / 4, 128] bf16, for sparse attention past 2,048 keys (Python State.mtp_pooled). + pooled: cuda.DeviceBuffer, + /// Entries the head's cache holds (Python mtp_len), the last `drafted` of them a chain's. + len: usize = 0, + drafted: usize = 0, + cap: usize, + /// The last pass's rows of main-model streams [rows_cap, wide]: a prompt chunk's, then each window's. + streams: cuda.DeviceBuffer, + rows_cap: usize, + /// The prompt's last row of streams in `streams`, which the first draft reads with the first token. + prompt_last: u64 = 0, + /// The head's last output row (its streams after the MoE write-back): the next chained draft's input. + out: cuda.DeviceBuffer, + /// A window's scores [max_rows, vocab], the draft head's [ids], best values [max_rows, 4]. + logits: cuda.DeviceBuffer, + dlogits: cuda.DeviceBuffer, + best: cuda.DeviceBuffer, + /// Before a window: each DeltaNet layer's conv window and (from the pass) projection rows; the n-gram tail. + saved: cuda.DeviceBuffer, + /// A window's DeltaNet end states: swapped in when every row is kept, else the kept rows replay on the sequence's. + spare: [linear_n]cuda.DeviceBuffer = undefined, + /// A window ran with the states written to `spare` and is not yet kept or rolled back. + open: bool = false, + keep: batch.Keep = .{}, + history: [2]i64 = .{ 0, 0 }, + unused_history: [2]i64 = .{ 0, 0 }, + fed0: usize = 0, + rows: usize = 0, + ids: [max_rows]u32 = undefined, + /// The scores a sampled draw reads on the host (a window's rows, or the draft head's row). + host: []u16, + best_host: [max_rows * 4]f32 = undefined, + + const hist_at: usize = 0; + const proj_at: usize = hist_at + linear_n * conv_bytes; + pub const tail_at: usize = proj_at + linear_n * max_rows * proj_bytes; + /// The window's n-gram rows (the pass's ple_nrow), saved for a rollback to shift the kept ones into the tail. + pub const nrow_at: usize = tail_at + ple_tail_n * wide * 2; + const saved_bytes: usize = nrow_at + max_rows * wide * 2; + + /// Device bytes `init(cap, draft_n, rows)` allocates. + pub fn bytes(cap: usize, draft_n: usize, rows: usize) usize { + return cap * (2 * prompt.kv_row_bytes + index_dim * 2) + prompt.pooledBytes(cap) + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + max_rows * 16 + saved_bytes + linear_n * state_n * 4; + } + + /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { + var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.kc.free(); + d.vc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.vc.free(); + d.ikc = try cuda.DeviceBuffer.alloc(driver, cap * index_dim * 2); + errdefer d.ikc.free(); + d.pooled = try cuda.DeviceBuffer.alloc(driver, prompt.pooledBytes(cap)); + errdefer d.pooled.free(); + d.streams = try cuda.DeviceBuffer.alloc(driver, d.rows_cap * wide * 2); + errdefer d.streams.free(); + d.out = try cuda.DeviceBuffer.alloc(driver, wide * 2); + errdefer d.out.free(); + d.logits = try cuda.DeviceBuffer.alloc(driver, max_rows * vocab * 2); + errdefer d.logits.free(); + d.dlogits = try cuda.DeviceBuffer.alloc(driver, draft_n * 2); + errdefer d.dlogits.free(); + d.best = try cuda.DeviceBuffer.alloc(driver, max_rows * 16); + errdefer d.best.free(); + d.saved = try cuda.DeviceBuffer.alloc(driver, saved_bytes); + errdefer d.saved.free(); + var made: usize = 0; + errdefer for (d.spare[0..made]) |*b| b.free(); + while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); + errdefer gpa.free(d.host); + // zeroed on the compute stream, ahead of the passes that read them + for ([_]cuda.DeviceBuffer{ d.kc, d.vc, d.ikc, d.pooled }) |b| try b.fill8(0, stream.handle); + var li: usize = 0; + for (&d.keep.layers, 0..) |*lk, layer| { + if (!wt.linearLayer(layer)) continue; + lk.hist = d.saved.ptr + hist_at + li * conv_bytes; + lk.proj = d.saved.ptr + proj_at + li * max_rows * proj_bytes; + li += 1; + } + d.aim(); + return d; + } + + /// Each DeltaNet layer's window end state into its spare buffer. + fn aim(d: *Draft) void { + var li: usize = 0; + for (&d.keep.layers, 0..) |*lk, layer| { + if (!wt.linearLayer(layer)) continue; + lk.state_out = d.spare[li].ptr; + li += 1; + } + } + + /// An open window kept whole: its end states become the sequence's (buffers swapped, nothing copied). + pub fn settle(d: *Draft, seq: *prompt.Seq) void { + if (!d.open) return; + for (seq.state, &d.spare) |*a, *b| std.mem.swap(cuda.DeviceBuffer, a, b); + d.aim(); + d.open = false; + } + + pub fn deinit(d: *Draft) void { + d.gpa.free(d.host); + for (&d.spare) |*b| b.free(); + d.saved.free(); + d.best.free(); + d.dlogits.free(); + d.logits.free(); + d.out.free(); + d.streams.free(); + d.pooled.free(); + d.ikc.free(); + d.vc.free(); + d.kc.free(); + } + + /// Drop the head's cache (a new prompt). + pub fn reset(d: *Draft) void { + d.len = 0; + d.drafted = 0; + d.rows = 0; + } +}; diff --git a/zig/src/families/flashnext/cuda_mtp_load.zig b/zig/src/families/flashnext/cuda_mtp_load.zig new file mode 100644 index 0000000000..dc1a226d31 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_load.zig @@ -0,0 +1,55 @@ +//! The draft vocabulary borrows exact rows of the checkpoint's dense or affine language head. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); + +pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []const u32) !wt.Q4 { + const mode = if (std.c.getenv("TF_FLASHNEXT_DRAFT_HEAD")) |v| std.mem.span(v) else "bf16"; + const affine = std.mem.eql(u8, mode, "affine4"); + if (!affine and !std.mem.eql(u8, mode, "bf16")) return error.BadDraftHead; + const tensor = try store.tensor("language_model.lm_head.weight"); + if (tensor.dtype == .bf16) { + if (!tensor.is(.bf16, &.{ wt.vocab, wt.dims })) return error.UnexpectedTensor; + const rows = try a.alloc(u8, ids.len * wt.dims * 2); + defer a.free(rows); + for (ids, 0..) |id, i| { + if (id >= wt.vocab) return error.InvalidToken; + @memcpy(rows[i * wt.dims * 2 ..][0 .. wt.dims * 2], tensor.bytes[id * wt.dims * 2 ..][0 .. wt.dims * 2]); + } + if (affine) return quantized(a, d, rows, ids.len); + return projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ ids.len, wt.dims, 1, 1, 1 }, .bytes = rows }); + } + var buf: [120]u8 = undefined; + const head = try wt.join3(a, store, &buf, "language_model.lm_head", ""); + defer a.free(head.w); + defer a.free(head.s); + defer a.free(head.b); + const wrow = wt.dims / 2; + const grow = wt.dims / 32 * 2; + if (head.w.len != wt.vocab * wrow or head.s.len != wt.vocab * grow) return error.UnexpectedTensor; + const w = try a.alloc(u8, ids.len * wrow); + defer a.free(w); + const s = try a.alloc(u8, ids.len * grow); + defer a.free(s); + const b = try a.alloc(u8, ids.len * grow); + defer a.free(b); + for (ids, 0..) |id, i| { + if (id >= wt.vocab) return error.InvalidToken; + @memcpy(w[i * wrow ..][0..wrow], head.w[id * wrow ..][0..wrow]); + @memcpy(s[i * grow ..][0..grow], head.s[id * grow ..][0..grow]); + @memcpy(b[i * grow ..][0..grow], head.b[id * grow ..][0..grow]); + } + return wt.uploadQ4(a, d, w, s, b, ids.len, wt.dims / 8); +} + +fn quantized(a: std.mem.Allocator, d: *cuda.Driver, rows: []const u8, n: usize) !wt.Q4 { + const words = try a.alloc(u8, n * wt.dims / 2); + defer a.free(words); + const scales = try a.alloc(u8, n * wt.dims / 16); + defer a.free(scales); + const biases = try a.alloc(u8, scales.len); + defer a.free(biases); + for (0..n) |i| try @import("cuda_draft_affine.zig").row(rows[i * wt.dims * 2 ..][0 .. wt.dims * 2], words[i * wt.dims / 2 ..][0 .. wt.dims / 2], scales[i * wt.dims / 16 ..][0 .. wt.dims / 16], biases[i * wt.dims / 16 ..][0 .. wt.dims / 16]); + return wt.uploadQ4(a, d, words, scales, biases, n, wt.dims / 8); +} diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig new file mode 100644 index 0000000000..0e36c57869 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -0,0 +1,322 @@ +//! The head's work over a sequence: chunked prompt passes, its forward and chains, verify windows and rollbacks. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); +const mtp = @import("cuda_mtp.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const conv_bytes = mtp.conv_bytes; +const proj_bytes = mtp.proj_bytes; +const fc_sk = mtp.fc_sk; +const piece_rows = mtp.piece_rows; +const pss_n = mtp.pss_n; +const Head = mtp.Head; +const Draft = mtp.Draft; +const Chain = mtp.Chain; +const Pick = mtp.Pick; +const Sampling = mtp.Sampling; +const drawRow = mtp.drawRow; +const promote = mtp.promote; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn copyAsync(stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { + if (n == 0) return; + const d = stream.d; + try d.check(d.api.cuMemcpyDtoDAsync_v2(dst, src, n, stream.handle), "cuMemcpyDtoDAsync"); +} + +fn download(stream: cuda.Stream, dst: []u8, src: u64) !void { + const d = stream.d; + try stream.synchronize(); + try d.check(d.api.cuMemcpyDtoH_v2(dst.ptr, src, dst.len), "cuMemcpyDtoH"); +} + +/// An owner can mirror main-model mutations to a peer without exposing device pointers on the control wire. +pub const Main = struct { + ptr: *anyopaque, + forward: *const fn (*anyopaque, *prompt.Seq, []const u32, batch.Hooks) anyerror!void, + verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32) anyerror!void, + rollback: *const fn (*anyopaque, *prompt.Seq, *Draft, usize) anyerror!void, +}; + +/// The head's work over a sequence: prompt pass, verify windows, rollback, forward and chains; `Tri` launches kernels. +pub fn Ops(comptime Tri: type) type { + return struct { + e: *prompt.Engine, + tri: Tri, + h: *Head, + main: ?Main = null, + + const Self = @This(); + + fn forwardMain(o: Self, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (o.main) |m| return m.forward(m.ptr, seq, ids, hooks); + return o.e.forward(Tri, o.tri, seq, ids, hooks); + } + + fn stream(o: Self) cuda.Stream { + return o.e.stream.*; + } + + /// Prompt rows [at, stop): the pass with its streams kept, then the head over each row with a known next token. + pub fn prefillChunk(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, at: usize, stop: usize) !void { + if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; + if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; + const last = stop == ids.len; + try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); + const head_to = @min(stop, ids.len - 1); + if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); + if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; + } + + /// The whole prompt a chunk of the engine's rows at a time, from an empty head cache. + pub fn prefill(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + d.reset(); + var at: usize = 0; + while (at < ids.len) { + const stop = @min(ids.len, at + @min(o.e.chunk, d.rows_cap)); + try o.prefillChunk(seq, d, ids, at, stop); + at = stop; + } + } + + /// The head over `next.len` rows of streams `rows`: cache at d.len, last row in d.out, scores in d.dlogits. + pub fn forwardHead(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + const e = o.e; + const tri = o.tri; + const h = o.h; + if (next.len == 0) return; + if (d.len + next.len > d.cap) return error.PromptTooLong; + const want = @max(@as(usize, 32), 2 * @min(next.len, piece_rows)); + const s = try e.reserve(want); + const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); + // the engine's block scores (the passes share one stream), for rows past 2,048 keys of the head's cache + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, + .seq = .{ .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .hooks = .{ .logits = .none }, + }; + var at: usize = 0; + while (at < next.len) { + const m = @min(piece, next.len - at); + p.rows = m; + p.pos0 = d.len; + p.pending = null; + try p.stage(next[at..][0..m]); + // the embedding branch: the token's row (stream 0 of the four copies), normed, fc_embedding + if (e.table.dense) { + try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, m); + } else { + try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, m); + } + try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, m, dims, dims); + try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); + try p.sumSlices(s.part, s.branch, true, m * dims, fc_sk); + // the hidden branch: the four streams normed as one row, fc_hidden on each stream + try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, s.norm, s.gq, eps, wide, m, wide, wide); + try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * m, fc_sk); + try p.sumSlices(s.part, s.ple_gated, true, streams * m * dims, fc_sk); + try tri.addStreams(s.branch, s.ple_gated, s.h, m, dims, streams); + // the layer: attention over the head's own cache, its experts + try p.writeback(0, m); + try p.readout(h.layer.attn_hc, 0, m); + try p.attnBlock(&h.layer, 0); + try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(h.layer.mlp_hc, 0, m); + try p.moeBlock(&h.layer); + p.pending = 0; + try p.writeback(0, m); + try p.copy(d.out.ptr, s.h + (m - 1) * wide * 2, wide * 2); + d.len += m; + at += m; + if (at == next.len and scores) { + // the last row's streams and sums at row 0, where the captured kernels' pointers stay aligned + if (m > 1) { + try p.copy(s.h, s.h + (m - 1) * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + (m - 1) * pss_n * 4, pss_n * 4); + } + try p.readout(h.mixer, 0, 1); + try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); + } + } + } + + /// Python draft.absorb: the chain's entries trimmed, then the kept rows (main streams `rows`, next tokens). + pub fn absorb(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + if (d.drafted > 0) { + d.len -= d.drafted; + d.drafted = 0; + } + try o.forwardHead(d, next, rows, scores); + } + + /// The draft at `position` and its temperature-1 probability (Python Engine.sample_draft). + pub fn pick(o: Self, d: *Draft, position: u64, sm: ?Sampling) !Pick { + const h = o.h; + const n = h.ids.len; + try h.best(o.stream(), d.dlogits.ptr, n, n, 1, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0..4]), d.best.ptr); + const top = d.best_host[0]; + const lse = d.best_host[1]; + const greedy = if (sm) |s| s.temperature <= 0 else true; + if (greedy) { + const col: usize = @intCast(@as(i32, @bitCast(d.best_host[2]))); + return .{ .token = h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) }; + } + const row = d.host[0..n]; + try download(o.stream(), std.mem.sliceAsBytes(row), d.dlogits.ptr); + const tok = try drawRow(o.e.gpa, row, h.ids, position, sm); + const col = std.sort.binarySearch(u32, h.ids, tok, struct { + fn order(t: u32, x: u32) std.math.Order { + return std.math.order(t, x); + } + }.order) orelse return .{ .token = tok, .prob = 0 }; + return .{ .token = tok, .prob = @exp(@as(f64, promote(row[col])) - @as(f64, lse)) }; + } + + /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts, the first always kept. + pub fn draft(o: Self, d: *Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?Sampling, confidence: f64, out: *Chain) !void { + out.n = 0; + try o.absorb(d, next, rows, count > 0); + var j: usize = 0; + while (j < @min(count, max_chain)) : (j += 1) { + const got = try o.pick(d, position + j, sm); + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + if (j + 1 < count) { + try o.forwardHead(d, &.{got.token}, d.out.ptr, true); + d.drafted += 1; + } + } + } + + /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams. + pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + if (o.main) |m| return m.verify(m.ptr, seq, d, ids); + o.e.round_graph.clear(); + const n = ids.len; + if (n == 0 or n > max_rows) return error.WindowTooWide; + if (seq.fed + n > seq.cap) return error.PromptTooLong; + d.settle(seq); + d.fed0 = seq.fed; + d.rows = n; + d.history = seq.history; + @memcpy(d.ids[0..n], ids); + const keeps = n > 1; + if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); + try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); + d.open = keeps; + } + + /// The window's rows drawn at their keyed positions (seq position fed0 + 1 + r), as serial decoding draws. + pub fn draws(o: Self, d: *Draft, sm: ?Sampling, out: []u32) !void { + const n = d.rows; + const greedy = if (sm) |s| s.temperature <= 0 else true; + if (greedy) { + try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); + return; + } + const rows = d.host[0 .. n * vocab]; + try download(o.stream(), std.mem.sliceAsBytes(rows), d.logits.ptr); + for (out[0..n], 0..) |*t, r| t.* = try drawRow(o.e.gpa, rows[r * vocab ..][0..vocab], null, d.fed0 + 1 + r, sm); + } + + /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. + pub fn rollback(o: Self, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + if (o.main) |m| return m.rollback(m.ptr, seq, d, keep); + o.e.round_graph.clear(); + if (keep == 0 or keep > d.rows) return error.CannotRewind; + if (keep == d.rows) return d.settle(seq); + if (!d.open) return error.CannotRewind; + const e = o.e; + const st = o.stream(); + const s = if (e.batch_s) |*bs| bs else return error.CannotRewind; + const k = &e.batch_k; + seq.fed = d.fed0 + keep; + var hist = d.history; + for (d.ids[0..keep]) |t| hist = .{ hist[1], t }; + seq.history = hist; + try copyAsync(st, seq.tail.ptr, d.saved.ptr + Draft.tail_at, ple_tail_n * wide * 2); + try shift(k, st, seq.tail.ptr, d.saved.ptr + Draft.nrow_at, wide, keep, ple_tail_n, wide); + var li: usize = 0; + for (e.layers[0..e.n_layers], 0..) |*lw, layer| { + if (!lw.linear) continue; + defer li += 1; + // the sequence's state is still the window's start (the pass wrote its end state to the spare) + const lk = d.keep.layers[layer]; + try copyAsync(st, seq.conv[li].ptr, lk.hist, conv_bytes); + var f: cuda.Args = .{}; + f.add(lk.proj); + f.add(seq.conv[li].ptr); + f.add(lw.cw.ptr); + f.add(lw.a_log.ptr); + f.add(lw.dt.ptr); + f.add(s.gq); + f.add(s.gk); + f.add(s.gv); + f.add(s.ggb); + try cuda.launch.launch(k.front, .{ .grid = .{ .x = @intCast(keep), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, st, &f); + var a: cuda.Args = .{}; + a.add(s.gq); + a.add(s.gk); + a.add(s.gv); + a.add(s.ggb); + a.add(seq.state[li].ptr); + a.add(cint(keep)); + a.add(s.ys); + try cuda.launch.launch(k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); + try shift(k, st, seq.conv[li].ptr, lk.proj, gdn.proj_width, keep, gdn.taps - 1, gdn.conv_dim); + } + d.rows = keep; + d.open = false; + } + + fn shift(k: *const batch.Kernels, st: cuda.Stream, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { + var a: cuda.Args = .{}; + a.add(win); + a.add(rows_ptr); + a.add(@as(c_longlong, @intCast(stride))); + a.add(cint(n)); + a.add(cint(taps)); + a.add(cint(c)); + try cuda.launch.launch(k.shift, .{ .grid = .{ .x = @intCast((c + 255) / 256) }, .block = .{ .x = 256 } }, st, &a); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_norms.zig b/zig/src/families/flashnext/cuda_norms.zig new file mode 100644 index 0000000000..f257ac56ca --- /dev/null +++ b/zig/src/families/flashnext/cuda_norms.zig @@ -0,0 +1,27 @@ +//! Classify the checkpoint-wide centered or direct RMSNorm convention, rejecting ambiguous mixtures. +const std = @import("std"); + +pub fn offset(comptime layers_n: usize, input: [layers_n]f64) !f32 { + comptime std.debug.assert(layers_n >= 2 and layers_n % 2 == 0); + var means = input; + var above: usize = 0; + for (means) |mean| { + if (!std.math.isFinite(mean)) return error.InvalidNormConvention; + if (mean > 0.5) above += 1; + } + std.mem.sort(f64, &means, {}, std.sort.asc(f64)); + const median = (means[layers_n / 2 - 1] + means[layers_n / 2]) / 2; + if (above * 10 >= layers_n * 9 and median >= 0.75 and median <= 1.5) return 0; + if (above * 10 <= layers_n and median >= -0.5 and median <= 0.25) return 1; + return error.InvalidNormConvention; +} + +test "centered and direct norms have distinct offsets; mixed and nonfinite conventions fail" { + try std.testing.expectEqual(@as(f32, 1), try offset(48, @splat(-0.1))); + try std.testing.expectEqual(@as(f32, 0), try offset(48, @splat(1.05))); + var mixed: [48]f64 = @splat(1.0); + @memset(mixed[0 .. 48 / 2], 0); + try std.testing.expectError(error.InvalidNormConvention, offset(48, mixed)); + mixed[0] = std.math.nan(f64); + try std.testing.expectError(error.InvalidNormConvention, offset(48, mixed)); +} diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig new file mode 100644 index 0000000000..459f56905d --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -0,0 +1,520 @@ +//! Layer 2's n-gram table: the hash of each token's row ids, and the mapped table a pass gathers its rows from. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const golden: u64 = 0x9E3779B97F4A7C15; +const mix1: u64 = 0xBF58476D1CE4E5B9; +const mix2: u64 = 0x94D049BB133111EB; +const prime_step: u64 = 10007; +pub const max_n: usize = 8; +pub const max_heads: usize = 64; +const max_shards: usize = 128; +pub const layer: usize = 1; + +const Tensor = core.safetensors.Tensor; + +pub const Hash = struct { + n: usize, + per: usize, + heads: usize, + eos: i64, + dh: usize, + shards: usize, + sizes: [max_heads]i64 = undefined, + offsets: [max_heads]i64 = undefined, + mult: [max_n]i64 = undefined, + starts: [max_shards + 1]i64 = undefined, + + pub fn init(args: struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, + }) !Hash { + if (args.ngram < 2 or args.ngram > max_n or args.per == 0 or args.shards == 0 or args.shards > max_shards) return error.UnexpectedTensor; + if (args.base < 2 or args.divisor == 0) return error.UnexpectedTensor; + const heads = (args.ngram - 1) * args.per; + if (heads > max_heads or args.ple_dim == 0 or args.ple_dim % heads != 0) return error.UnexpectedTensor; + var h = Hash{ .n = args.ngram, .per = args.per, .heads = heads, .eos = args.eos, .dh = args.ple_dim / heads, .shards = args.shards }; + var total: i64 = 0; + for (0..heads) |head| { + const size: i64 = @intCast(nthPrimeAfter(args.base - 1, args.ple_index * heads + head + 1)); + h.sizes[head] = size; + h.offsets[head] = total; + total += size; + } + const divisor: i64 = @intCast(args.divisor); + const shard_n: i64 = @intCast(args.shards); + const padded = divCeil(total, divisor) * divisor; + const base: i64 = @divFloor(padded, shard_n); + const extra: usize = @intCast(@mod(padded, shard_n)); + var acc: i64 = 0; + for (0..args.shards) |i| { + h.starts[i] = acc; + acc += base + @as(i64, @intFromBool(i < extra)); + } + h.starts[args.shards] = acc; + const half: u64 = @max(1, (((1 << 63) - 1) / @max(args.vocab, 1)) / 2); + const seed = args.seed +% prime_step *% @as(u64, @intCast(args.ple_index)); + for (0..args.ngram) |i| { + const mixed = splitmix64(seed +% golden *% @as(u64, @intCast(i + 1))); + h.mult[i] = @intCast(2 * (mixed % half) + 1); + } + return h; + } + + pub fn locate(h: Hash, row: i64) !struct { shard: usize, local: usize } { + if (row < 0 or row >= h.starts[h.shards]) return error.UnexpectedTensor; + var lo: usize = 0; + var hi = h.shards; + while (lo + 1 < hi) { + const mid = (lo + hi) / 2; + if (h.starts[mid] <= row) lo = mid else hi = mid; + } + return .{ .shard = lo, .local = @intCast(row - h.starts[lo]) }; + } +}; + +fn splitmix64(value: u64) u64 { + var v = value +% golden; + v = (v ^ (v >> 30)) *% mix1; + v = (v ^ (v >> 27)) *% mix2; + return v ^ (v >> 31); +} + +fn isPrime(value: u64) bool { + if (value < 2) return false; + if (value % 2 == 0) return value == 2; + var d: u64 = 3; + while (d <= value / d) : (d += 2) if (value % d == 0) return false; + return true; +} + +fn nthPrimeAfter(start: u64, count: usize) u64 { + var prime = start; + for (0..count) |_| { + prime += 1; + while (!isPrime(prime)) prime += 1; + } + return prime; +} + +fn divCeil(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +/// Row ids of one new token after `history` (n - 1 ids). A fresh decode window is n - 1 copies of eos. +pub fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { + if (history.len != h.n - 1 or out.len != h.heads) return error.UnexpectedTensor; + var seq: [max_n]i64 = undefined; + for (history, 0..) |v, i| seq[i] = v; + seq[history.len] = token; + const width = h.n; + var eos_at: [max_n]i64 = undefined; + for (0..width) |i| eos_at[i] = if (seq[i] == h.eos) @intCast(i) else -1; + var acc: [max_n]i64 = undefined; + acc[0] = eos_at[0]; + for (1..width) |i| acc[i] = @max(acc[i - 1], eos_at[i]); + var before: [max_n]i64 = undefined; + before[0] = -1; + for (1..width) |i| before[i] = acc[i - 1]; + var shifted: [max_n]i64 = undefined; + const pos: i64 = @intCast(width - 1); + const in_segment = pos - (before[width - 1] + 1); + for (0..h.n) |shift| { + const source = pos - @as(i64, @intCast(shift)); + const taken = seq[@intCast(@max(source, 0))]; + const keep = in_segment >= @as(i64, @intCast(shift)) and source >= 0; + shifted[shift] = if (keep) taken else h.eos; + } + var at: usize = 0; + var ngram: usize = 2; + while (ngram <= h.n) : (ngram += 1) { + var mixed = shifted[0] *% h.mult[0]; + for (1..ngram) |p| mixed = mixed ^ (shifted[p] *% h.mult[p]); + for (0..h.per) |_| { + out[at] = @mod(mixed, h.sizes[at]) + h.offsets[at]; + at += 1; + } + } +} + +pub fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +pub fn sumGroups(row: []const u16, out: []f32) void { + const groups = row.len / 32; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @bitCast(@as(u32, row[g * 32 + i]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +/// One mapped shard (core.safetensors.File) with a typed lookup. +pub const Open = struct { + file: core.safetensors.File, + io: std.Io, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Open { + return .{ .file = try core.safetensors.File.open(gpa, io, path), .io = io }; + } + + pub fn close(self: *Open, gpa: std.mem.Allocator) void { + _ = gpa; + self.file.close(self.io); + } + + pub fn tensor(self: *const Open, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const t = self.file.get(name) orelse return error.MissingTensor; + if (t.dtype != dtype) return error.UnexpectedTensor; + return t; + } +}; + +pub fn readI64(t: Tensor, out: []i64) !void { + if (t.dtype != .i64 or t.rank != 1 or t.dim(0) != out.len or t.bytes.len != out.len * 8) return error.UnexpectedTensor; + for (out, 0..) |*o, i| o.* = std.mem.readInt(i64, t.bytes[8 * i ..][0..8], .little); +} + +pub fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { + return if (nested) + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.weight", .{ layer, shard }) + else + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.weight", .{ layer, shard }); +} + +pub fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { + return partNameFor(buf, shard, nested, part, false); +} + +fn partNameFor(buf: []u8, shard: usize, nested: bool, part: []const u8, modelopt: bool) ![]const u8 { + const prefix = if (modelopt) "model.language_model" else "language_model.model"; + return if (nested) + std.fmt.bufPrint(buf, "{s}.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ prefix, layer, shard, part }) + else + std.fmt.bufPrint(buf, "{s}.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ prefix, layer, shard, part }); +} + +pub fn findOpen(opened: []Open, path: []const u8) ?*Open { + for (opened) |*o| if (std.mem.eql(u8, o.file.path, path)) return o; + return null; +} + +pub fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usize, words: []u8, scales: []u8, biases: []u8) !void { + var name: [180]u8 = undefined; + const w = try file.tensor(try partName(&name, shard, nested, "weight"), .u32); + const s = try file.tensor(try partName(&name, shard, nested, "scales"), .bf16); + const b = try file.tensor(try partName(&name, shard, nested, "biases"), .bf16); + const words_w = dh / 8; + const groups = dh / 32; + if (w.dim(1) != words_w or s.dim(1) != groups or b.dim(1) != groups or local >= w.dim(0)) return error.UnexpectedTensor; + const wb = words_w * 4; + const gb = groups * 2; + @memcpy(words, w.bytes[local * wb ..][0..wb]); + @memcpy(scales, s.bytes[local * gb ..][0..gb]); + @memcpy(biases, b.bytes[local * gb ..][0..gb]); +} + +pub fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { + return readScale(gpa, io, dir, wm, "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale", false); +} + +fn readScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap, name: []const u8, required: bool) !f32 { + const file = wm.get(name) orelse return if (required) error.MissingTensor else 1; + const path = try std.fs.path.join(gpa, &.{ dir, file.string }); + defer gpa.free(path); + var opened = try Open.load(gpa, io, path); + defer opened.close(gpa); + const t = opened.file.get(name) orelse return error.MissingTensor; + if (t.numel() != 1) return error.UnexpectedTensor; + return switch (t.dtype) { + .f32 => @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)), + .f64 => @floatCast(@as(f64, @bitCast(std.mem.readInt(u64, t.bytes[0..8], .little)))), + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[0..2], .little)) << 16), + else => error.UnexpectedTensor, + }; +} + +fn jsonU64(v: std.json.Value) !u64 { + if (v != .integer or v.integer < 0) return error.UnexpectedTensor; + return @intCast(v.integer); +} + +fn optU64(o: std.json.ObjectMap, name: []const u8, default: u64) !u64 { + const v = o.get(name) orelse return default; + if (v == .null) return default; + return jsonU64(v); +} + +const NgramSpec = struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, +}; + +/// `ple_layer_ids` value 2 is decoder layer 1. The eos id is the text config's, which the hash mixes in. +pub fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const text = if (root.get("text_config")) |v| if (v == .object and v.object.count() > 0) v.object else root else root; + const layers = try jsonU64(text.get("num_hidden_layers") orelse return error.MissingTensor); + if (layers <= layer) return error.UnexpectedTensor; + const hidden = try jsonU64(text.get("hidden_size") orelse return error.MissingTensor); + const ids = text.get("ple_layer_ids") orelse return error.MissingTensor; + if (ids != .array) return error.UnexpectedTensor; + var seen: [512]bool = @splat(false); + var ple_index: usize = 0; + var found = false; + for (ids.array.items) |item| { + const id = try jsonU64(item); + if (id == 0 or id > layers or id > seen.len) return error.UnexpectedTensor; + if (id == layer + 1) found = true; + if (!seen[id - 1] and id - 1 < layer) ple_index += 1; + seen[id - 1] = true; + } + if (!found) return error.UnexpectedTensor; + const eos: i64 = if (text.get("eos_token_id")) |v| blk: { + if (v == .null) break :blk 0; + if (v == .array) { + if (v.array.items.len == 0) return error.UnexpectedTensor; + break :blk @intCast(try jsonU64(v.array.items[0])); + } + break :blk @intCast(try jsonU64(v)); + } else 0; + return .{ + .vocab = try jsonU64(text.get("vocab_size") orelse return error.MissingTensor), + .ngram = @intCast(try optU64(text, "ngram_size", 3)), + .per = @intCast(try optU64(text, "heads_per_ngram", 8)), + .base = try optU64(text, "ngram_vocab_size_base", 20_000_000), + .divisor = try optU64(text, "make_ngram_vocab_size_divisible_by", 128), + .shards = @intCast(try optU64(text, "split_ngram_parts", 128)), + .seed = try optU64(text, "seed", 1234), + .eos = eos, + .ple_dim = @intCast(try optU64(text, "ple_embed_dim", hidden)), + .ple_index = ple_index, + }; +} + +/// The text config's eos id, which a fresh n-gram window repeats. +pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { + const cfg = try readSpec(gpa, io, model_dir); + return cfg.eos; +} + +/// The n-gram table held open for an engine's life: the hash, the scale, and each shard's slices of its mapped file. +pub const Ngram = struct { + gpa: std.mem.Allocator, + hash: Hash, + scale: f32, + files: std.ArrayList(Open) = .empty, + w: [max_shards][]const u8 = undefined, + s: [max_shards][]const u8 = undefined, + b: [max_shards][]const u8 = undefined, + format: enum { mlx, nvfp4 } = .mlx, + tables: [max_shards]core.nvfp4_rows.Table = undefined, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { + const cfg = try readSpec(gpa, io, model_dir); + var g: Ngram = .{ .gpa = gpa, .hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }), .scale = 1 }; + errdefer g.deinit(); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + const first = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight"; + const modelopt = wm.contains(first); + g.format = if (modelopt) .nvfp4 else .mlx; + g.scale = if (modelopt) + try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale_2", true) + else + try tableScale(gpa, io, model_dir, wm); + const dh = g.hash.dh; + var key: [180]u8 = undefined; + for (0..g.hash.shards) |shard| { + var nested = false; + const mapped = if (wm.get(try partNameFor(&key, shard, false, "weight", modelopt))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try partNameFor(&key, shard, true, "weight", modelopt)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(g.files.items, path) orelse file: { + var o = try Open.load(gpa, io, path); + errdefer o.close(gpa); + std.posix.madvise(o.file.map.memory.ptr, o.file.map.memory.len, std.posix.MADV.RANDOM) catch {}; + try g.files.append(gpa, o); + break :file &g.files.items[g.files.items.len - 1]; + }; + const w = try file.tensor(try partNameFor(&key, shard, nested, "weight", modelopt), if (modelopt) .u8 else .u32); + const expected_rows: usize = @intCast(g.hash.starts[shard + 1] - g.hash.starts[shard]); + if (w.rank != 2 or w.dim(0) != expected_rows) return error.UnexpectedTensor; + if (modelopt) { + const sc = try file.tensor(try partNameFor(&key, shard, nested, "weight_scale", true), .f8_e4m3); + const table = try core.nvfp4_rows.Table.init(w, sc, g.scale); + if (table.width != dh) return error.UnexpectedTensor; + g.tables[shard] = table; + } else { + const sc = try file.tensor(try partNameFor(&key, shard, nested, "scales", false), .bf16); + const bias = try file.tensor(try partNameFor(&key, shard, nested, "biases", false), .bf16); + if (!w.is(.u32, &.{ expected_rows, dh / 8 }) or !sc.is(.bf16, &.{ expected_rows, dh / 32 }) or + !bias.is(.bf16, &.{ expected_rows, dh / 32 })) return error.UnexpectedTensor; + g.w[shard] = w.bytes; + g.s[shard] = sc.bytes; + g.b[shard] = bias.bytes; + } + } + return g; + } + + pub fn deinit(g: *Ngram) void { + for (g.files.items) |*o| o.close(g.gpa); + g.files.deinit(g.gpa); + } + + /// Values a row's embedding holds (heads x head width). + pub fn width(g: *const Ngram) usize { + return g.hash.heads * g.hash.dh; + } + + /// Rows [lo, hi) of a chunk's token-major row ids `ids`, each dequantized into its head's slice of `out`. + fn fill(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize) !void { + const dh = g.hash.dh; + for (lo..hi) |i| { + const loc = try g.hash.locate(ids[i]); + const row = out[i * dh ..][0..dh]; + switch (g.format) { + .mlx => try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, row), + .nvfp4 => try g.tables[loc.shard].gather(&.{loc.local}, row), + } + } + } + + fn worker(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize, failed: *?anyerror) void { + g.fill(ids, out, lo, hi) catch |err| { + failed.* = err; + }; + } + + /// `embedding` for each token after `history` (moved past the chunk) into out and xs, read by up to 128 threads. + pub fn rows(g: *const Ngram, history: []i64, tokens: []const u32, out: []u16, xs: []f32) !void { + const heads = g.hash.heads; + const w = g.width(); + if (history.len != g.hash.n - 1 or out.len != tokens.len * w or xs.len != out.len / 32) return error.UnexpectedTensor; + const items = tokens.len * heads; + const ids = try g.gpa.alloc(i64, items); + defer g.gpa.free(ids); + var hist: [max_n]i64 = undefined; + @memcpy(hist[0..history.len], history); + for (tokens, 0..) |token, t| { + try lastIds(g.hash, hist[0..history.len], @intCast(token), ids[t * heads ..][0..heads]); + for (0..history.len - 1) |i| hist[i] = hist[i + 1]; + hist[history.len - 1] = token; + } + const max_threads = 128; + const n_threads = @min(max_threads, items / 4); + if (n_threads <= 1) { + try g.fill(ids, out, 0, items); + } else { + var threads: [max_threads]std.Thread = undefined; + var failed: [max_threads]?anyerror = @splat(null); + var spawned: usize = 0; + defer for (threads[0..spawned]) |th| th.join(); + const per = (items + n_threads - 1) / n_threads; + for (0..n_threads) |i| { + const lo = @min(items, i * per); + const hi = @min(items, lo + per); + threads[i] = try std.Thread.spawn(.{}, worker, .{ g, ids, out, lo, hi, &failed[i] }); + spawned += 1; + } + for (threads[0..spawned]) |th| th.join(); + spawned = 0; + for (failed[0..n_threads]) |f| if (f) |err| return err; + } + for (0..tokens.len) |t| { + const row = out[t * w ..][0..w]; + if (g.format == .mlx and g.scale != 1) { + for (row) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * g.scale); + } + } + sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); + } + @memcpy(history, hist[0..history.len]); + } +}; + +test "a fresh window hashes token 7 into sixteen rows" { + const h = try Hash.init(.{ + .vocab = 248320, + .ngram = 3, + .per = 8, + .base = 20_000_000, + .divisor = 128, + .shards = 128, + .seed = 1234, + .eos = 248044, + .ple_dim = 2560, + .ple_index = 0, + }); + try std.testing.expectEqual(@as(i64, 23703573157769), h.mult[0]); + const history = [_]i64{ 248044, 248044 }; + var ids: [16]i64 = undefined; + try lastIds(h, &history, 7, &ids); + const want = [_]i64{ 2927653, 34980843, 54748278, 66612378, 97814964, 109013870, 126560393, 151352333, 167888935, 182235580, 215170017, 237467519, 247510681, 278779700, 296141806, 304994522 }; + try std.testing.expectEqualSlices(i64, &want, &ids); + const loc = try h.locate(ids[0]); + try std.testing.expectEqual(@as(usize, 1), loc.shard); + try std.testing.expectEqual(@as(usize, 427641), loc.local); +} + +test { + _ = @import("cuda_ple_nvfp4_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_ple_check.zig b/zig/src/families/flashnext/cuda_ple_check.zig new file mode 100644 index 0000000000..b31417259a --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_check.zig @@ -0,0 +1,235 @@ +//! The n-gram table's host checks: one token's rows dequantized on the host (the row path) and against the GPU gather. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const ple = @import("cuda_ple.zig"); + +const max_n = ple.max_n; +const max_heads = ple.max_heads; +const layer = ple.layer; +const Hash = ple.Hash; +const Open = ple.Open; +const Tensor = core.safetensors.Tensor; +const lastIds = ple.lastIds; +const toBf16 = ple.toBf16; +const sumGroups = ple.sumGroups; +const readI64 = ple.readI64; +const weightName = ple.weightName; +const partName = ple.partName; +const findOpen = ple.findOpen; +const copyRow = ple.copyRow; +const tableScale = ple.tableScale; +const readSpec = ple.readSpec; + +/// Host dequant of one token's sixteen n-gram rows. `history` is the previous n - 1 ids. +pub fn embedding(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, history: []const i64, token: u32, out: []u16, xs: []f32) !void { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + if (history.len != hash.n - 1 or out.len != hash.heads * hash.dh or xs.len != out.len / 32) return error.UnexpectedTensor; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history, @intCast(token), ids[0..hash.heads]); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + const scale = try tableScale(gpa, io, model_dir, wm); + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, out[h * dh ..][0..dh]); + if (scale != 1) { + for (out) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + sumGroups(out, xs); +} + +/// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. +pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + var history: [max_n]i64 = undefined; + for (0..hash.n - 1) |i| history[i] = hash.eos; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history[0 .. hash.n - 1], token, ids[0..hash.heads]); + + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + + var cname: [160]u8 = undefined; + const mult_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.layer_multipliers", .{layer}); + const const_file = wm.get(mult_name) orelse return error.MissingTensor; + const const_path = try std.fs.path.join(gpa, &.{ model_dir, const_file.string }); + defer gpa.free(const_path); + var constants = try Open.load(gpa, io, const_path); + defer constants.close(gpa); + var got_m: [max_n]i64 = undefined; + var got_o: [max_heads]i64 = undefined; + var got_s: [max_heads]i64 = undefined; + try readI64(try constants.tensor(mult_name, .i64), got_m[0..hash.n]); + const off_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_offsets", .{layer}); + try readI64(try constants.tensor(off_name, .i64), got_o[0..hash.heads]); + const size_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_vocab_sizes", .{layer}); + try readI64(try constants.tensor(size_name, .i64), got_s[0..hash.heads]); + if (!std.mem.eql(i64, got_m[0..hash.n], hash.mult[0..hash.n]) or !std.mem.eql(i64, got_o[0..hash.heads], hash.offsets[0..hash.heads]) or !std.mem.eql(i64, got_s[0..hash.heads], hash.sizes[0..hash.heads])) { + std.debug.print("ngram constants differ\n", .{}); + return 1; + } + const scale = try tableScale(gpa, io, model_dir, wm); + + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + + const host = try gpa.alloc(u16, heads * dh); + defer gpa.free(host); + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, host[h * dh ..][0..dh]); + if (scale != 1) { + for (host) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + const host_xs = try gpa.alloc(f32, heads * dh / 32); + defer gpa.free(host_xs); + sumGroups(host, host_xs); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, host.len * 2); + defer out_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, host_xs.len * 4); + defer xs_b.free(); + try tri.pleEmbed(w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, xs_b.ptr, 1, heads, dh, scale); + const got_bytes = try gpa.alloc(u8, host.len * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, host_xs.len * 4); + defer gpa.free(xs_bytes); + try stream.synchronize(); + try out_b.download(0, got_bytes); + try xs_b.download(0, xs_bytes); + const gpu = try gpa.alloc(u16, host.len); + defer gpa.free(gpu); + for (gpu, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (gpu, host, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + var xs_ulp: u32 = 0; + var xs_at: usize = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) { + xs_ulp = dist; + xs_at = i; + } + } + std.debug.print("ngram token {d} heads {d} values {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} xs_ulp {d} xs_at {d} head {x:0>4}\n", .{ token, heads, host.len, off, steps, at, gpu[at], host[at], xs_ulp, xs_at, gpu[0] }); + if (off != 0 or steps != 0 or xs_ulp != 0) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig b/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig new file mode 100644 index 0000000000..ee81d7a7b4 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig @@ -0,0 +1,111 @@ +//! Synthetic mapped shards exercise the production gather, threading and chunk-history boundaries. +const std = @import("std"); +const core = @import("core"); +const ple = @import("cuda_ple.zig"); +const a = std.testing.allocator; +const io = std.testing.io; +const base = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"; + +fn fixture(dir: std.Io.Dir, include_scale: bool) !void { + const config = + \\{"text_config":{"num_hidden_layers":2,"hidden_size":32,"vocab_size":100,"ple_layer_ids":[2], + \\ "ngram_size":2,"heads_per_ngram":2,"ngram_vocab_size_base":13,"make_ngram_vocab_size_divisible_by":2, + \\ "split_ngram_parts":2,"ple_embed_dim":32,"eos_token_id":99}} + ; + try dir.writeFile(io, .{ .sub_path = "config.json", .data = config }); + const h = try ple.Hash.init(.{ .vocab = 100, .ngram = 2, .per = 2, .base = 13, .divisor = 2, .shards = 2, .seed = 1234, .eos = 99, .ple_dim = 32, .ple_index = 0 }); + var header: std.Io.Writer.Allocating = .init(a); + defer header.deinit(); + var index: std.Io.Writer.Allocating = .init(a); + defer index.deinit(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(a); + try header.writer.writeByte('{'); + try index.writer.writeAll("{\"weight_map\":{"); + for (0..2) |shard| { + const rows: usize = @intCast(h.starts[shard + 1] - h.starts[shard]); + for ([_][]const u8{ "weight", "weight_scale" }, 0..) |part, component| { + const cols: usize = if (component == 0) 8 else 1; + const dtype = if (component == 0) "U8" else "F8_E4M3"; + const begin = data.items.len; + for (0..rows) |row| { + const id: usize = @intCast(h.starts[shard] + @as(i64, @intCast(row))); + const code: u8 = if (component == 0) @intCast(id % 16 + ((id + 3) % 16) * 16) else 0x38; + try data.appendNTimes(a, code, cols); + } + if (begin > 0) { + try header.writer.writeByte(','); + try index.writer.writeByte(','); + } + try header.writer.print("\"{s}.shard_{d}.{s}\":{{\"dtype\":\"{s}\",\"shape\":[{d},{d}],\"data_offsets\":[{d},{d}]}}", .{ base, shard, part, dtype, rows, cols, begin, data.items.len }); + try index.writer.print("\"{s}.shard_{d}.{s}\":\"table.safetensors\"", .{ base, shard, part }); + } + } + if (include_scale) { + const begin = data.items.len; + const bytes = std.mem.toBytes(@as(u32, @bitCast(@as(f32, 0.375)))); + try data.appendSlice(a, &bytes); + try header.writer.print(",\"{s}.weight_scale_2\":{{\"dtype\":\"F32\",\"shape\":[],\"data_offsets\":[{d},{d}]}}", .{ base, begin, data.items.len }); + try index.writer.print(",\"{s}.weight_scale_2\":\"table.safetensors\"", .{base}); + } + try header.writer.writeByte('}'); + try index.writer.writeAll("}}"); + var image: std.Io.Writer.Allocating = .init(a); + defer image.deinit(); + const size = std.mem.toBytes(std.mem.nativeToLittle(u64, header.written().len)); + try image.writer.writeAll(&size); + try image.writer.writeAll(header.written()); + try image.writer.writeAll(data.items); + try dir.writeFile(io, .{ .sub_path = "table.safetensors", .data = image.written() }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.written() }); +} + +test "NVFP4 mapped PLE rows equal token decode across prompt chunks and threaded gathers" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, true); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + var g = try ple.Ngram.init(a, io, path); + defer g.deinit(); + try std.testing.expectEqual(@as(usize, 1), g.files.items.len); + var tokens: [64]u32 = undefined; + for (&tokens, 0..) |*v, i| v.* = @intCast((i * 7) % 100); + var full: [64 * 32]u16 = undefined; + var xs_full: [64]f32 = undefined; + var history = [_]i64{99}; + try g.rows(&history, &tokens, &full, &xs_full); + var chunked: [64 * 32]u16 = undefined; + var xs_chunked: [64]f32 = undefined; + var chunk_hist = [_]i64{99}; + try g.rows(&chunk_hist, tokens[0..17], chunked[0 .. 17 * 32], xs_chunked[0..17]); + try g.rows(&chunk_hist, tokens[17..], chunked[17 * 32 ..], xs_chunked[17..]); + try std.testing.expectEqualSlices(u16, &full, &chunked); + try std.testing.expectEqualSlices(f32, &xs_full, &xs_chunked); + var single_hist = [_]i64{99}; + const values = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (tokens, 0..) |token, row| { + var ids: [2]i64 = undefined; + try ple.lastIds(g.hash, &single_hist, token, &ids); + var one: [32]u16 = undefined; + var xs: [1]f32 = undefined; + try g.rows(&single_hist, &.{token}, &one, &xs); + try std.testing.expectEqualSlices(u16, full[row * 32 ..][0..32], &one); + for (ids, 0..) |id, head| for (0..16) |col| { + const code: usize = @intCast(@mod(id + @as(i64, if (col % 2 == 0) 0 else 3), 16)); + const v = values[code % 8] * @as(f32, if (code < 8) 1 else -1) * 0.375; + try std.testing.expectEqual(core.quant.convert.bf16(v), one[head * 16 + col]); + }; + } + try std.testing.expectEqualSlices(i64, &history, &single_hist); + try std.testing.expectEqualSlices(i64, &history, &chunk_hist); +} + +test "NVFP4 PLE requires its global table scale" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, false); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + try std.testing.expectError(error.MissingTensor, ple.Ngram.init(a, io, path)); +} diff --git a/zig/src/families/flashnext/cuda_projection.zig b/zig/src/families/flashnext/cuda_projection.zig new file mode 100644 index 0000000000..c0b96adb3b --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection.zig @@ -0,0 +1,71 @@ +//! Owned affine or native lane projections; packed device views borrow these exact buffers. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +pub const Projection = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + n: usize, + k: usize, + native: ?cuda.qmmf.Weight = null, + + pub fn free(self: *Projection) void { + self.w.free(); + self.s.free(); + self.b.free(); + } + + pub fn empty(d: *const cuda.Driver) Projection { + const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; + } +}; + +/// Repack the checkpoint's BF16 bits without rounding or widening the resident weight. +pub fn dense(gpa: std.mem.Allocator, driver: *const cuda.Driver, t: core.safetensors.Tensor) !Projection { + if (t.rank != 2 or t.dtype != .bf16 or t.dim(0) == 0 or t.dim(1) == 0 or t.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = t.dim(0); + const k = t.dim(1); + const count = try std.math.mul(usize, n, k); + if (t.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + const input = try gpa.alloc(u16, count); + defer gpa.free(input); + for (input, 0..) |*v, i| v.* = std.mem.readInt(u16, t.bytes[i * 2 ..][0..2], .little); + const tiled = try gpa.alloc(u16, cuda.nvfp4.padded(n) * k); + defer gpa.free(tiled); + cuda.qmmf.packBf16(tiled, input, n, k); + var p = Projection.empty(driver); + p.w = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(tiled)); + p.n = n; + p.k = k; + p.native = cuda.modelopt.bf16Weight(p.w.ptr, n, k); + return p; +} + +/// E4M3 codes and UE8M0 group-32 scales are permuted only; calibration metadata is not applied as another scale. +pub fn mxfp8(gpa: std.mem.Allocator, driver: *const cuda.Driver, w: core.safetensors.Tensor, s: core.safetensors.Tensor) !Projection { + if (w.rank != 2 or w.dtype != .f8_e4m3 or w.dim(0) == 0 or w.dim(1) == 0 or w.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = w.dim(0); + const k = w.dim(1); + const count = try std.math.mul(usize, n, k); + if (w.bytes.len != count or !s.is(.u8, &.{ n, k / 32 }) or s.bytes.len != count / 32) + return error.UnexpectedTensor; + const codes = try gpa.alloc(u8, cuda.mxfp8.padded(n) * k); + defer gpa.free(codes); + const scales = try gpa.alloc(u8, codes.len / 32); + defer gpa.free(scales); + cuda.mxfp8.packCodes(codes, w.bytes, n, k); + cuda.mxfp8.packScales(scales, s.bytes, n, k); + var p = Projection.empty(driver); + p.w = try cuda.DeviceBuffer.fromHost(driver, codes); + errdefer p.w.free(); + p.s = try cuda.DeviceBuffer.fromHost(driver, scales); + p.n = n; + p.k = k; + p.native = cuda.mxfp8.weight(p.w.ptr, p.s.ptr, n, k); + return p; +} diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig new file mode 100644 index 0000000000..c0c37f1b79 --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -0,0 +1,132 @@ +//! Projection loading preserves dense BF16 bits and MXFP8 codes/scales; affine layouts keep their existing packer. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); + +pub const Micro = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + + pub fn deinit(self: *Micro, a: std.mem.Allocator) void { + a.free(self.codes); + a.free(self.scales); + self.* = undefined; + } +}; + +/// Own each plane before the store advances to a possibly different mapped shard. +pub fn readMicro(a: std.mem.Allocator, store: anytype, prefix: []const u8) !Micro { + var buf: [512]u8 = undefined; + const w = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); + if (w.dtype != .f8_e4m3 or w.rank != 2 or w.dim(0) == 0 or w.dim(1) == 0 or w.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = w.dim(0); + const k = w.dim(1); + const count = try std.math.mul(usize, n, k); + if (w.bytes.len != count) return error.UnexpectedTensor; + const codes = try a.dupe(u8, w.bytes); + errdefer a.free(codes); + const s = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{prefix})); + if (!s.is(.u8, &.{ n, k / 32 }) or s.bytes.len != count / 32) return error.UnexpectedTensor; + return .{ .n = n, .k = k, .codes = codes, .scales = try a.dupe(u8, s.bytes) }; +} + +/// Concatenate output rows of compatible projections without converting or recalibrating either plane. +pub fn readStitched(a: std.mem.Allocator, store: anytype, prefix: []const u8, names: []const []const u8) !Micro { + if (names.len == 0 or names.len > 6) return error.UnexpectedTensor; + var parts: [6]Micro = undefined; + var held: usize = 0; + defer for (parts[0..held]) |*part| part.deinit(a); + var n: usize = 0; + var k: usize = 0; + var buf: [512]u8 = undefined; + for (names) |name| { + parts[held] = try readMicro(a, store, try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name })); + held += 1; + const part = parts[held - 1]; + if (k != 0 and k != part.k) return error.UnexpectedTensor; + k = part.k; + n = try std.math.add(usize, n, part.n); + } + const count = try std.math.mul(usize, n, k); + const codes = try a.alloc(u8, count); + errdefer a.free(codes); + const scales = try a.alloc(u8, count / 32); + var row: usize = 0; + for (parts[0..held]) |part| { + @memcpy(codes[row * k ..][0..part.codes.len], part.codes); + @memcpy(scales[row * (k / 32) ..][0..part.scales.len], part.scales); + row += part.n; + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales }; +} + +fn uploadMicro(a: std.mem.Allocator, d: *const cuda.Driver, value: Micro) !wt.Q4 { + return projection.mxfp8(a, d, .{ .dtype = .f8_e4m3, .rank = 2, .shape = .{ value.n, value.k, 1, 1, 1 }, .bytes = value.codes }, .{ .dtype = .u8, .rank = 2, .shape = .{ value.n, value.k / 32, 1, 1, 1 }, .bytes = value.scales }); +} + +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, n: usize, k: usize) !wt.Q4 { + var buf: [512]u8 = undefined; + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); + switch (t.dtype) { + .bf16 => { + if (!t.is(.bf16, &.{ n, k })) return error.UnexpectedTensor; + return projection.dense(a, d, t); + }, + .f8_e4m3 => { + var value = try readMicro(a, store, prefix); + defer value.deinit(a); + if (value.n != n or value.k != k) return error.UnexpectedTensor; + return uploadMicro(a, d, value); + }, + .u32 => { + if (!t.is(.u32, &.{ n, k / 8 })) return error.UnexpectedTensor; + const planes = try wt.join3(a, store, &buf, prefix, ""); + defer a.free(planes.w); + defer a.free(planes.s); + defer a.free(planes.b); + return wt.uploadQ4(a, d, planes.w, planes.s, planes.b, n, k / 8); + }, + else => return error.UnsupportedQuantization, + } +} + +pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !wt.Q4 { + if (names.len == 0) return error.UnexpectedTensor; + var buf: [512]u8 = undefined; + const first = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, names[0] })); + if (first.dtype == .f8_e4m3) { + var value = try readStitched(a, store, prefix, names); + defer value.deinit(a); + if (value.n != n or value.k != k) return error.UnexpectedTensor; + return uploadMicro(a, d, value); + } + if (first.dtype == .bf16) { + const bytes = try a.alloc(u8, try std.math.mul(usize, try std.math.mul(usize, n, k), 2)); + defer a.free(bytes); + var at: usize = 0; + for (names) |name| { + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, name })); + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k or t.bytes.len != t.dim(0) * k * 2 or t.bytes.len > bytes.len - at) + return error.UnexpectedTensor; + @memcpy(bytes[at..][0..t.bytes.len], t.bytes); + at += t.bytes.len; + } + if (at != bytes.len) return error.UnexpectedTensor; + return projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = bytes }); + } + if (first.dtype != .u32) return error.UnsupportedQuantization; + const planes = try wt.stitch(a, store, prefix, names); + defer a.free(planes.w); + defer a.free(planes.s); + defer a.free(planes.b); + if (planes.w.len != n * (k / 8) * 4) return error.UnexpectedTensor; + return wt.uploadQ4(a, d, planes.w, planes.s, planes.b, n, k / 8); +} + +test { + _ = @import("cuda_projection_load_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_projection_load_test.zig b/zig/src/families/flashnext/cuda_projection_load_test.zig new file mode 100644 index 0000000000..a76e08c304 --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection_load_test.zig @@ -0,0 +1,63 @@ +//! A store that invalidates its previous view on every read, like advancing mapped shards. +const std = @import("std"); +const core = @import("core"); +const source = @import("cuda_projection_load.zig"); +const Tensor = core.safetensors.Tensor; + +const Fake = struct { + scratch: [1024]u8 = undefined, + bad: enum { none, weight, scale, extent, width, missing } = .none, + + pub fn tensor(self: *Fake, name: []const u8) !Tensor { + const part: usize = if (std.mem.startsWith(u8, name, "root.a.")) 0 else if (std.mem.startsWith(u8, name, "root.b.")) 1 else return error.MissingTensor; + const n: usize = if (part == 0) 3 else 5; + const k: usize = if (self.bad == .width and part == 1) 128 else 64; + const scale = std.mem.endsWith(u8, name, ".weight_scale"); + if (scale and self.bad == .missing) return error.MissingTensor; + if (!scale and !std.mem.endsWith(u8, name, ".weight")) return error.MissingTensor; + const cols = if (scale) k / 32 else k; + const len = n * cols; + @memset(&self.scratch, 0xee); + for (self.scratch[0..len], 0..) |*v, i| v.* = @truncate(i * 17 + part * 53 + @as(usize, if (scale) 129 else 7)); + return .{ + .dtype = if (scale) (if (self.bad == .scale) .f8_e4m3 else .u8) else (if (self.bad == .weight) .bf16 else .f8_e4m3), + .rank = 2, + .shape = .{ n, cols, 1, 1, 1 }, + .bytes = self.scratch[0 .. len - @as(usize, if (self.bad == .extent) 1 else 0)], + }; + } +}; + +fn trial(a: std.mem.Allocator) !void { + var fake: Fake = .{}; + var value = try source.readStitched(a, &fake, "root", &.{ ".a", ".b" }); + defer value.deinit(a); + try std.testing.expectEqual(@as(usize, 8), value.n); + try std.testing.expectEqual(@as(usize, 64), value.k); + for (0..2) |part| { + const start: usize = if (part == 0) 0 else 3; + const n: usize = if (part == 0) 3 else 5; + for (0..n * 64) |i| try std.testing.expectEqual(@as(u8, @truncate(i * 17 + part * 53 + 7)), value.codes[start * 64 + i]); + for (0..n * 2) |i| try std.testing.expectEqual(@as(u8, @truncate(i * 17 + part * 53 + 129)), value.scales[start * 2 + i]); + } +} + +test "MXFP8 fusion owns each shard plane and preserves exact output row order" { + try trial(std.testing.allocator); +} + +test "MXFP8 fusion cleans every partial allocation on failure" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, trial, .{}); +} + +test "MXFP8 source refuses incompatible shapes, dtypes, extents and missing scales" { + const a = std.testing.allocator; + for ([_]@FieldType(Fake, "bad"){ .weight, .scale, .extent, .width }) |bad| { + var fake: Fake = .{ .bad = bad }; + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{ ".a", ".b" })); + } + var fake: Fake = .{ .bad = .missing }; + try std.testing.expectError(error.MissingTensor, source.readMicro(a, &fake, "root.a")); + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{})); + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{ ".a", ".a", ".a", ".a", ".a", ".a", ".a" })); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig new file mode 100644 index 0000000000..71a6ae74cf --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -0,0 +1,464 @@ +//! The loaded engine and one stream's caches: prompts and windows on the batched pass, decode rounds over streams. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const weights = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); +pub const batch = @import("cuda_batch.zig"); + +const layers_n = weights.layers_n; +const linear_n = weights.linear_n; +const attn_n = weights.attn_n; +const dims = weights.dims; +const streams = weights.streams; +const wide = weights.wide; +const vocab = weights.vocab; +const experts_n = weights.experts_n; +const slots = weights.slots; +const moe_w = weights.moe_w; +const head_dim = weights.head_dim; +const index_dim = weights.index_dim; +const proj_n = weights.proj_n; +const out_k = weights.out_k; +const ple_layer = weights.ple_layer; +const inj_stride = weights.inj_stride; +const wts_stride = weights.wts_stride; +const ple_tail_n = weights.ple_tail_n; +const state_n = weights.state_n; +const kv_heads = weights.kv_heads; +const Store = weights.Store; +const Face = weights.Face; +const Q4 = weights.Q4; +const HcW = weights.HcW; +const linearLayer = weights.linearLayer; +const linearIndex = weights.linearIndex; +const attnIndex = weights.attnIndex; +const promote = weights.promote; +const asF32 = weights.asF32; +const uploadQ4 = weights.uploadQ4; +const join3 = weights.join3; +const loadHc = weights.loadHc; +const stitch = weights.stitch; +const Layer = load.Layer; +const loadLayer = load.loadLayer; + +pub fn argmax(raw: []const u8) u32 { + var best = promote(std.mem.readInt(u16, raw[0..2], .little)); + var at: usize = 0; + var i: usize = 1; + while (i < vocab) : (i += 1) { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v > best) { + best = v; + at = i; + } + } + return @intCast(at); +} + +/// `n` zeroed buffers, the fill on the compute stream: a legacy-stream memset returns before it lands. +fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, n: usize, len: usize) ![]cuda.DeviceBuffer { + const Zeroed = struct { + driver: *cuda.Driver, + stream: cuda.Stream, + len: usize, + fn make(z: @This()) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.alloc(z.driver, z.len); + } + fn fill(z: @This(), b: *cuda.DeviceBuffer) !void { + try b.fill8(0, z.stream.handle); + } + }; + return made(cuda.DeviceBuffer, gpa, n, Zeroed{ .driver = driver, .stream = stream, .len = len }); +} + +/// `n` buffers from `src.make`, each through `src.fill`; on a failure every buffer made so far is freed, the last too. +fn made(comptime B: type, gpa: std.mem.Allocator, n: usize, src: anytype) ![]B { + const out = try gpa.alloc(B, n); + var held: usize = 0; + errdefer { + for (out[0..held]) |*b| b.free(); + gpa.free(out); + } + while (held < n) { + out[held] = try src.make(); + held += 1; // counted before its fill, so a failed fill frees it with the rest + try src.fill(&out[held - 1]); + } + return out; +} + +pub const vocabulary: usize = vocab; +/// Prompt plus reply tokens a sequence can hold: the model's window. +pub const max_tokens: usize = weights.window; + +/// One position's keys and values in one attention layer, bf16. +pub const kv_row_bytes: usize = kv_heads * head_dim * 2; + +/// Device bytes `Engine.newSeq(cap)` allocates: conv tails, recurrent states, KV with indexer keys, the n-gram tail. +pub fn seqBytes(cap: usize) usize { + const linear = linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + state_n * 4); + const attn = attn_n * (cap * (2 * kv_row_bytes + index_dim * 2) + pooledBytes(cap)); + return linear + attn + ple_tail_n * wide * 2; +} + +/// One attention layer's pooled indexer keys for a cache of `cap` tokens: a bf16 row of 128 a block of 4. +pub fn pooledBytes(cap: usize) usize { + return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; +} + +const round_graph = @import("cuda_round_graph.zig"); + +/// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + store: Store, + model_dir: []u8, + table: embed.Table, + mixer: HcW, + head_q: Q4, + eos: i64, + layers: [layers_n]Layer, + n_layers: usize, + batch_k: batch.Kernels, + /// The n-gram table, open for the engine's life (the batched pass gathers its rows from it). + ngram: ple.Ngram, + /// The pass's scratch, grown to the widest chunk run. + batch_s: ?batch.Scratch = null, + /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. + scores: ?cuda.DeviceBuffer = null, + /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). + chunk: usize, + round_graph: round_graph.Cache = .{}, + collective: ?*@import("cuda_collective.zig").Collective = null, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { + return initRank(gpa, io, driver, stream, model_dir, 1, 0); + } + + pub fn initRank(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8, world: usize, rank: usize) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.io = io; + e.driver = driver; + e.stream = stream; + e.n_layers = 0; + e.collective = null; + e.round_graph = .{}; + errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); + e.model_dir = try gpa.dupe(u8, model_dir); + errdefer gpa.free(e.model_dir); + e.store = try Store.open(gpa, io, e.model_dir); + errdefer e.store.deinit(); + try e.store.detectNorms(); + const weight = try e.store.tensor(embed.weight_name); + if (weight.dtype == .bf16) { + e.table = try embed.Table.uploadDense(driver, weight); + } else { + const scales = try e.store.tensor(embed.scales_name); + const biases = try e.store.tensor(embed.biases_name); + e.table = try embed.Table.upload(driver, weight, scales, biases); + } + errdefer e.table.deinit(); + e.eos = try ple.eosOf(gpa, io, e.model_dir); + e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); + errdefer e.mixer.free(); + const head_tensor = try e.store.tensor("language_model.lm_head.weight"); + if (head_tensor.dtype == .bf16) { + if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; + e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, if (try @import("cuda_head.zig").enabled()) try @import("cuda_head.zig").shard(head_tensor, world, rank) else head_tensor); + } else { + if (try @import("cuda_head.zig").enabled()) return error.BadHeadSharding; + var buf: [80]u8 = undefined; + const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; + e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + } + errdefer e.head_q.free(); + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const pack_fn = try pack_mod.function(weights.pack_symbol); + for (0..layers_n) |i| { + e.layers[i] = try load.loadLayerRank(gpa, driver, stream, &e.store, pack_fn, i, world, rank); + e.n_layers += 1; + std.log.info("flashnext: rank {d}/{d} loaded layer {d}/{d}", .{ rank + 1, world, i + 1, layers_n }); + } + e.store.release(); + e.store.closeFile(); + e.ngram = try ple.Ngram.init(gpa, io, e.model_dir); + errdefer e.ngram.deinit(); + e.batch_k = try batch.Kernels.load(driver); + e.batch_s = null; + e.scores = null; + e.chunk = batch.default_chunk; + if (std.c.getenv("TF_FLASHNEXT_CHUNK")) |v| e.chunk = std.fmt.parseInt(usize, std.mem.span(v), 10) catch batch.default_chunk; + if (e.chunk == 0) e.chunk = batch.default_chunk; + return e; + } + + pub fn deinit(e: *Engine) void { + e.round_graph.clear(); + if (e.batch_s) |*bs| bs.free(e.gpa); + if (e.scores) |*b| b.free(); + e.batch_k.deinit(); + e.ngram.deinit(); + for (e.layers[0..e.n_layers]) |*l| l.free(e.gpa); + e.head_q.free(); + e.mixer.free(); + e.table.deinit(); + e.store.deinit(); + e.gpa.free(e.model_dir); + e.gpa.destroy(e); + } + + /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. + pub fn newSeq(e: *Engine, cap: usize) !Seq { + e.round_graph.clear(); + if (cap == 0) return error.UnexpectedTensor; + const logits = try e.gpa.alloc(u8, vocab * 2); + errdefer e.gpa.free(logits); + const conv = try zeros(e.gpa, e.driver, e.stream.*, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + errdefer freeOwned(e.gpa, conv); + const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); + errdefer freeOwned(e.gpa, state); + const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, kc); + const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, vc); + const ikc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); + errdefer freeOwned(e.gpa, ikc); + const pooled = try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); + errdefer freeOwned(e.gpa, pooled); + var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); + errdefer tail.free(); + try tail.fill8(0, e.stream.handle); + return .{ + .conv = conv, + .state = state, + .kc = kc, + .vc = vc, + .ikc = ikc, + .pooled = pooled, + .tail = tail, + .history = .{ e.eos, e.eos }, + .fed = 0, + .logits = logits, + .cap = cap, + }; + } + + /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. + pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { + return e.forward(Tri, tri, seq, ids, .{}); + } + + fn complete(e: *Engine) !void { + if (e.collective) |c| return c.finish(e.stream.*); + try e.stream.synchronize(); + } + + /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. + pub fn reserveScores(e: *Engine, cap: usize, rows: usize) !u64 { + const need = batch.scoreBytes(cap, rows); + if (e.scores) |*b| { + if (b.len >= need) return b.ptr; + try e.complete(); + e.round_graph.clear(); + b.free(); + e.scores = null; + } + e.scores = try cuda.DeviceBuffer.alloc(e.driver, need); + return e.scores.?.ptr; + } + + /// Scratch for chunks of `rows` rows, reallocated only to grow. + pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { + if (e.batch_s) |*bs| { + if (bs.cap >= rows) return bs; + try e.complete(); + e.round_graph.clear(); + bs.free(e.gpa); + e.batch_s = null; + } + e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); + return &e.batch_s.?; + } + + /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. + pub fn forward(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; + const s = try e.reserve(@min(ids.len, e.chunk)); + const scores = try e.reserveScores(seq.cap, s.cap); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, + .seq = seq.view(scores), + .hooks = hooks, + }; + var at: usize = 0; + while (at < ids.len) { + const n = @min(s.cap, ids.len - at); + const last = at + n == ids.len; + p.hooks.logits = if (last) hooks.logits else .none; + p.hooks.streams = if (hooks.streams != 0) hooks.streams + at * wide * 2 else 0; + try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none) seq.logits else null); + at += n; + } + if (hooks.logits == .none) try e.complete(); + seq.fed += ids.len; + } + + /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone. + pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { + if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor; + var cap: usize = 0; + for (seqs) |q| { + if (q.fed + 1 > q.cap) return error.UnexpectedTensor; + cap = @max(cap, q.cap); + } + const rows = try e.gpa.alloc(batch.Lane, seqs.len); + defer e.gpa.free(rows); + const s = try e.reserve(seqs.len); + // each row's attention runs alone on the block scores + const scores = try e.reserveScores(cap, 1); + for (rows, seqs) |*l, q| l.* = .{ .seq = q.view(scores), .pos = q.fed }; + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, + .seq = rows[0].seq, + .lanes = rows, + .hooks = .{ .logits = .all, .logits_out = out }, + }; + if (e.collective != null and round_graph.enabled()) { + const key = try round_graph.Key.init(seqs, s.ids, scores, out); + try e.round_graph.run(key, &p, ids); + } else try p.chunk(ids, 0, null); + try e.complete(); + const d = e.stream.d; + for (seqs, 0..) |q, r| { + try d.check(d.api.cuMemcpyDtoH_v2(q.logits.ptr, out + r * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + q.fed += 1; + } + } +}; + +/// One stream's caches. `fed` tokens have already gone through the model. +pub const Seq = struct { + rank_id: u64 = 0, + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + pooled: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + history: [2]i64, + fed: usize, + logits: []u8, + cap: usize, + + /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). + pub fn view(s: *Seq, scores: u64) batch.SeqState { + return .{ .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores }; + } + + pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { + freeBufs(s.conv); + freeBufs(s.state); + freeBufs(s.kc); + freeBufs(s.vc); + freeBufs(s.ikc); + freeBufs(s.pooled); + gpa.free(s.conv); + gpa.free(s.state); + gpa.free(s.kc); + gpa.free(s.vc); + gpa.free(s.ikc); + gpa.free(s.pooled); + if (s.tail.ptr != 0) s.tail.free(); + gpa.free(s.logits); + } +}; + +fn freeBufs(bs: []cuda.DeviceBuffer) void { + for (bs) |*b| b.free(); +} + +fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { + freeBufs(bs); + gpa.free(bs); +} + +test "linear layers skip every fourth index and attention starts at 3" { + try std.testing.expect(linearLayer(0)); + try std.testing.expect(!linearLayer(3)); + try std.testing.expect(!linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); +} + +test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { + const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); + try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); + try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); + try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); +} + +test "a failed fill frees the buffer it was filling and every one before it" { + const Fake = struct { + live: *usize, + pub fn free(b: *@This()) void { + b.live.* -= 1; + } + }; + const Src = struct { + live: *usize, + fills: *usize, + fail_at: usize, + fn make(s: @This()) !Fake { + s.live.* += 1; + return .{ .live = s.live }; + } + fn fill(s: @This(), _: *Fake) !void { + s.fills.* += 1; + if (s.fills.* == s.fail_at) return error.FillFailed; + } + }; + for ([_]usize{ 1, 2, 4 }) |fail_at| { + var live: usize = 0; + var fills: usize = 0; + try std.testing.expectError(error.FillFailed, made(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); + try std.testing.expectEqual(@as(usize, 0), live); + } + var live: usize = 0; + var fills: usize = 0; + const all = try made(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); + defer std.testing.allocator.free(all); + try std.testing.expectEqual(@as(usize, 3), live); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig new file mode 100644 index 0000000000..94df19ae5c --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -0,0 +1,413 @@ +//! One resident copy of every decoder layer. The prompt forward reads these, it does not reload them. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const mixed = @import("cuda_mixed_experts.zig"); +const projections = @import("cuda_projection_load.zig"); + +const layers_n = wt.layers_n; +const dims = wt.dims; +const wide = wt.wide; +const experts_n = wt.experts_n; +const moe_w = wt.moe_w; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const ple_layer = wt.ple_layer; +const block_bytes = wt.block_bytes; +const cint = wt.cint; +const Store = wt.Store; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const linearLayer = wt.linearLayer; +const asF32 = wt.asF32; +const join3 = wt.join3; +const loadHc = wt.loadHc; + +const OwnedFace = struct { + face: Face, + w: []u8, + s: []u8, + b: []u8, + + fn free(self: *OwnedFace, gpa: std.mem.Allocator) void { + gpa.free(self.w); + gpa.free(self.s); + gpa.free(self.b); + } +}; + +pub fn faceOf(store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !Face { + const w = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + const sc = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + const b = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + const k8 = k / 8; + const kg = k / 32; + if (routed) { + if (!w.is(.u32, &.{ experts_n, n, k8 }) or !sc.is(.bf16, &.{ experts_n, n, kg }) or !b.is(.bf16, &.{ experts_n, n, kg })) return error.UnexpectedTensor; + } else if (!w.is(.u32, &.{ n, k8 }) or !sc.is(.bf16, &.{ n, kg }) or !b.is(.bf16, &.{ n, kg })) return error.UnexpectedTensor; + return .{ .w = w.bytes, .s = sc.bytes, .b = b.bytes, .n = n, .k = k, .routed = routed }; +} + +pub const Layer = struct { + attn_hc: HcW, + mlp_hc: HcW, + proj: Q4, + out: Q4, + router: cuda.DeviceBuffer, + linear: bool, + cw: cuda.DeviceBuffer, + a_log: cuda.DeviceBuffer, + dt: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + q_scale: cuda.DeviceBuffer, + k_scale: cuda.DeviceBuffer, + i_scale: cuda.DeviceBuffer, + /// The indexer's k_layernorm (fp32), the pooled keys' RMSNorm scale. + ik_scale: cuda.DeviceBuffer, + key: Q4, + val: Q4, + nk: cuda.DeviceBuffer, + nq: cuda.DeviceBuffer, + nc: cuda.DeviceBuffer, + ple_cw: cuda.DeviceBuffer, + /// Packed gate and up for every routed expert plus the shared expert at index 512. + experts_up: cuda.DeviceBuffer, + experts_down: cuda.DeviceBuffer, + native_experts: ?mixed.Owned = null, + + pub fn free(self: *Layer, gpa: std.mem.Allocator) void { + _ = gpa; + self.attn_hc.free(); + self.mlp_hc.free(); + self.proj.free(); + self.out.free(); + self.router.free(); + self.cw.free(); + self.a_log.free(); + self.dt.free(); + self.norm.free(); + self.q_scale.free(); + self.k_scale.free(); + self.i_scale.free(); + self.ik_scale.free(); + self.key.free(); + self.val.free(); + self.nk.free(); + self.nq.free(); + self.nc.free(); + self.ple_cw.free(); + if (self.native_experts) |*e| e.free(); + self.experts_up.free(); + self.experts_down.free(); + } +}; + +fn deadBuf(d: *cuda.Driver) cuda.DeviceBuffer { + return .{ .d = d, .ptr = 0, .len = 0 }; +} + +fn blankFace() OwnedFace { + return .{ .face = .{ .w = &.{}, .s = &.{}, .b = &.{}, .n = 0, .k = 0, .routed = false }, .w = &.{}, .s = &.{}, .b = &.{} }; +} + +/// Copy one expert tensor out of the checkpoint so the shard can be closed. +fn keepFace(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !OwnedFace { + const view = try faceOf(store, buf, prefix, name, routed, n, k); + const w = try gpa.dupe(u8, view.w); + errdefer gpa.free(w); + const s = try gpa.dupe(u8, view.s); + errdefer gpa.free(s); + const b = try gpa.dupe(u8, view.b); + return .{ .face = .{ .w = w, .s = s, .b = b, .n = n, .k = k, .routed = routed }, .w = w, .s = s, .b = b }; +} + +fn packDevice(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) !cuda.DeviceBuffer { + const kg = k / 32; + const nb = n / 32; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, e * nb * kg * block_bytes); + errdefer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream.*, &a); + try stream.synchronize(); + return out_b; +} + +/// Routed experts, then the shared expert as index `experts_n`, in the kernel's packed layout. +fn packAll(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, routed: OwnedFace, shared: OwnedFace) !cuda.DeviceBuffer { + const n = routed.face.n; + const k = routed.face.k; + var bulk = try packDevice(driver, stream, pack_fn, routed.w, routed.s, routed.b, experts_n, n, k); + defer bulk.free(); + var one = try packDevice(driver, stream, pack_fn, shared.w, shared.s, shared.b, 1, n, k); + defer one.free(); + if (bulk.len != experts_n * one.len) return error.UnexpectedTensor; + var out = try cuda.DeviceBuffer.alloc(driver, bulk.len + one.len); + errdefer out.free(); + try out.copyFrom(0, bulk.ptr, bulk.len, stream.handle); + try out.copyFrom(bulk.len, one.ptr, one.len, stream.handle); + try stream.synchronize(); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +/// Gate and up blocks side by side, which is the SwiGLU weight the expert kernel reads. +fn stackUp(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, gate: cuda.DeviceBuffer, up: cuda.DeviceBuffer) !cuda.DeviceBuffer { + const graw = try gpa.alloc(u8, gate.len); + defer gpa.free(graw); + const uraw = try gpa.alloc(u8, up.len); + defer gpa.free(uraw); + try stream.synchronize(); + try gate.download(0, graw); + try up.download(0, uraw); + const stacked = try stackHalves(gpa, graw, uraw); + defer gpa.free(stacked); + return cuda.DeviceBuffer.fromHost(driver, stacked); +} + +pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize) !Layer { + return loadLayerRank(gpa, driver, stream, store, pack_fn, layer, 1, 0); +} + +pub fn loadLayerRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize, world: usize, rank: usize) !Layer { + var buf: [80]u8 = undefined; + const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); + return loadLayerAtRank(gpa, driver, stream, store, pack_fn, base, linearLayer(layer), layer == ple_layer, world, rank); +} + +/// The decoder layer under `base` (the MTP head's "language_model.mtp.layers.0": full attention, no n-gram branch). +pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool) !Layer { + return loadLayerAtRank(gpa, driver, stream, store, pack_fn, base, linear, with_ple, 1, 0); +} + +fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool, world: usize, rank: usize) !Layer { + var buf: [200]u8 = undefined; + var base_buf: [80]u8 = undefined; + @memcpy(base_buf[0..base.len], base); + const root = base_buf[0..base.len]; + var attn_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.attn_hyper_connection", .{root}), true); + errdefer attn_hc.free(); + var mlp_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.mlp_hyper_connection", .{root}), true); + errdefer mlp_hc.free(); + const mlp = try std.fmt.bufPrint(&buf, "{s}.mlp", .{root}); + var mlp_buf: [96]u8 = undefined; + @memcpy(mlp_buf[0..mlp.len], mlp); + const mlp_root = mlp_buf[0..mlp.len]; + const gate = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.gate.weight", .{mlp_root})); + if (!gate.is(.bf16, &.{ experts_n, dims })) return error.UnexpectedTensor; + const gate_copy = try gpa.alloc(u8, gate.bytes.len); + defer gpa.free(gate_copy); + @memcpy(gate_copy, gate.bytes); + const router = try gpa.alloc(u16, (experts_n + 1) * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); + const sg_tensor = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.shared_expert_gate.weight", .{mlp_root})); + if (sg_tensor.dtype == .bf16) { + if (!sg_tensor.is(.bf16, &.{ 1, dims }) or sg_tensor.bytes.len != dims * 2) return error.UnexpectedTensor; + @memcpy(std.mem.sliceAsBytes(router[experts_n * dims ..]), sg_tensor.bytes); + } else { + const sg = try join3(gpa, store, &buf, mlp_root, ".shared_expert_gate"); + defer gpa.free(sg.w); + defer gpa.free(sg.s); + defer gpa.free(sg.b); + try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + } + var router_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + errdefer router_b.free(); + + var proj_q = Q4.empty(driver); + var out_q = Q4.empty(driver); + errdefer proj_q.free(); + errdefer out_q.free(); + var cw_b = deadBuf(driver); + var a_b = deadBuf(driver); + var dt_b = deadBuf(driver); + var nw_b = deadBuf(driver); + var qs_b = deadBuf(driver); + var ks_b = deadBuf(driver); + var is_b = deadBuf(driver); + var iks_b = deadBuf(driver); + var key_q = Q4.empty(driver); + var val_q = Q4.empty(driver); + errdefer key_q.free(); + errdefer val_q.free(); + var nk_b = deadBuf(driver); + var nq_b = deadBuf(driver); + var nc_b = deadBuf(driver); + var pcw_b = deadBuf(driver); + errdefer cw_b.free(); + errdefer a_b.free(); + errdefer dt_b.free(); + errdefer nw_b.free(); + errdefer qs_b.free(); + errdefer ks_b.free(); + errdefer is_b.free(); + errdefer iks_b.free(); + errdefer nk_b.free(); + errdefer nq_b.free(); + errdefer nc_b.free(); + errdefer pcw_b.free(); + if (linear) { + const pre = try std.fmt.bufPrint(&buf, "{s}.linear_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const lp = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; + proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); + out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); + if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); + defer gpa.free(a_log); + const dt = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.dt_bias", .{lp}))); + defer gpa.free(dt); + if (a_log.len != gdn.nv or dt.len != gdn.nv) return error.UnexpectedTensor; + a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt)); + const norm = try store.copyOf(try std.fmt.bufPrint(&buf, "{s}.norm.weight", .{lp})); + defer gpa.free(norm); + nw_b = try cuda.DeviceBuffer.fromHost(driver, norm); + } else { + const pre = try std.fmt.bufPrint(&buf, "{s}.self_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const ap = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; + proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); + out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); + const qn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); + defer gpa.free(qn); + const kn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); + defer gpa.free(kn); + const iqn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); + defer gpa.free(iqn); + const ikn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.k_layernorm.weight", .{ap}))); + defer gpa.free(ikn); + if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim or ikn.len != index_dim) return error.UnexpectedTensor; + qs_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(qn)); + ks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(kn)); + is_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(iqn)); + iks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(ikn)); + } + if (with_ple) { + const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); + var pre_buf: [96]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const pp = pre_buf[0..pre.len]; + key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); + val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); + const nk = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); + defer gpa.free(nk); + const nq = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); + defer gpa.free(nq); + const nc = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); + defer gpa.free(nc); + if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; + nk_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nk)); + nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); + nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); + if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + } + + var native_experts: ?mixed.Owned = null; + errdefer if (native_experts) |*e| e.free(); + var experts_up = deadBuf(driver); + errdefer experts_up.free(); + var experts_down = deadBuf(driver); + errdefer experts_down.free(); + if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { + native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); + } else { + if (world != 1 or rank != 0) return error.UnsupportedQuantization; + const legacy = legacy_block: { + var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&shared) |*f| f.free(gpa); + // a folder Living Weights taught holds a wider shared expert (its 640 neurons, then the pocket's) + const width = (try store.hold(try std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.weight", .{mlp_root}))).dim(0); + if (width != moe_w) return error.LearnedPocket; + shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); + shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); + shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); + var routed: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&routed) |*f| f.free(gpa); + routed[0] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); + routed[1] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); + routed[2] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); + var gate_p = try packAll(driver, stream, pack_fn, routed[0], shared[0]); + errdefer gate_p.free(); + var up_p = try packAll(driver, stream, pack_fn, routed[1], shared[1]); + errdefer up_p.free(); + var down_p = try packAll(driver, stream, pack_fn, routed[2], shared[2]); + errdefer down_p.free(); + for (&shared) |*f| f.free(gpa); + for (&routed) |*f| f.free(gpa); + shared = .{ blankFace(), blankFace(), blankFace() }; + routed = .{ blankFace(), blankFace(), blankFace() }; + var combined_up = try stackUp(gpa, driver, stream, gate_p, up_p); + errdefer combined_up.free(); + gate_p.free(); + up_p.free(); + break :legacy_block .{ combined_up, down_p }; + }; + experts_up = legacy[0]; + experts_down = legacy[1]; + } + store.release(); + return .{ + .attn_hc = attn_hc, + .mlp_hc = mlp_hc, + .proj = proj_q, + .out = out_q, + .router = router_b, + .linear = linear, + .cw = cw_b, + .a_log = a_b, + .dt = dt_b, + .norm = nw_b, + .q_scale = qs_b, + .k_scale = ks_b, + .i_scale = is_b, + .ik_scale = iks_b, + .key = key_q, + .val = val_q, + .nk = nk_b, + .nq = nq_b, + .nc = nc_b, + .ple_cw = pcw_b, + .experts_up = experts_up, + .experts_down = experts_down, + .native_experts = native_experts, + }; +} diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig new file mode 100644 index 0000000000..9ba52f395f --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -0,0 +1,87 @@ +//! Prompt tokens in, greedy token ids out. The ids file is whitespace or comma separated. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const want = [_]u32{ 1596, 1144, 310, 5707, 310, 1156, 25, 328 }; + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 5) { + std.debug.print("usage: flashnext-prompt \n", .{}); + return 2; + } + const count = try std.fmt.parseInt(usize, args[4], 10); + // TF_FLASHNEXT_DRAFTS: the MTP check (cuda_mtp.zig) over every comma-separated ids file, one load + if (std.c.getenv("TF_FLASHNEXT_DRAFTS") != null) return drafted(gpa, io, args[1], args[2], args[3], count); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 24)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + if (ids.items.len == 0 or count == 0) return 2; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set: ?cuda.aot.Set = if (std.mem.eql(u8, args[2], "native")) null else try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer if (set) |*x| x.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const tri: flash.Tri = .{ .set = if (set) |*x| x else null, .own = &glue, .s = stream }; + const got = try gpa.alloc(u32, count); + defer gpa.free(got); + try flash.row_path.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); + std.debug.print("tokens", .{}); + for (got) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + if (count == want.len and ids.items.len == 58) { + var same = true; + for (got, &want) |g, w| if (g != w) { + same = false; + }; + std.debug.print("{s}\n", .{if (same) "match" else "differ"}); + if (!same) return 1; + } + return 0; +} + +/// Serial and drafted decoding of each prompt compared token for token (flash.mtp.check), with their speeds. +fn drafted(gpa: std.mem.Allocator, io: std.Io, model: []const u8, kernels: []const u8, files: []const u8, count: usize) !u8 { + var prompts: std.ArrayList([]u32) = .empty; + defer { + for (prompts.items) |p| gpa.free(p); + prompts.deinit(gpa); + } + var names = std.mem.tokenizeScalar(u8, files, ','); + while (names.next()) |name| { + const text = try std.Io.Dir.cwd().readFileAlloc(io, name, gpa, .limited(1 << 20)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + errdefer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + try prompts.append(gpa, try ids.toOwnedSlice(gpa)); + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set: ?cuda.aot.Set = if (std.mem.eql(u8, kernels, "native")) null else try cuda.aot.Set.load(gpa, io, &driver, ctx.device, kernels); + defer if (set) |*x| x.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const tri: flash.Tri = .{ .set = if (set) |*x| x else null, .own = &glue, .s = stream }; + const same = try flash.mtp.check(flash.Tri, gpa, io, &driver, &stream, tri, model, prompts.items, count); + std.debug.print("{s}\n", .{if (same) "drafted == serial on every prompt" else "drafted DIFFERS from serial"}); + return if (same) 0 else 1; +} diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig new file mode 100644 index 0000000000..36d364152f --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -0,0 +1,321 @@ +//! One fused decode row: hyper-connection, linear or full attention, experts, and the n-gram step. + +const std = @import("std"); +const cuda = @import("cuda"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const block_bytes = wt.block_bytes; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const inj_stride = wt.inj_stride; +const wts_stride = wt.wts_stride; +const vocab = wt.vocab; +const cint = wt.cint; +const toBf16 = wt.toBf16; +const pack_symbol = wt.pack_symbol; +const plan_symbol = wt.plan_symbol; +const up_symbol = wt.up_symbol; +const down_symbol = wt.down_symbol; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; + +pub const Scratch = struct { + h: cuda.DeviceBuffer, + y: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + inj: cuda.DeviceBuffer, + pss: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + dn: cuda.DeviceBuffer, + hc_part: cuda.DeviceBuffer, + act: cuda.DeviceBuffer, + xs_act: cuda.DeviceBuffer, + mixed: cuda.DeviceBuffer, + xsm: cuda.DeviceBuffer, + proj: cuda.DeviceBuffer, + part: cuda.DeviceBuffer, + gout: cuda.DeviceBuffer, + gxs: cuda.DeviceBuffer, + st: cuda.DeviceBuffer, + ple_emb: cuda.DeviceBuffer, + ple_xs: cuda.DeviceBuffer, + ple_keys: cuda.DeviceBuffer, + ple_vals: cuda.DeviceBuffer, + ple_gated: cuda.DeviceBuffer, + ple_pss: cuda.DeviceBuffer, + ple_nrow: cuda.DeviceBuffer, + ple_part: cuda.DeviceBuffer, + pa: cuda.DeviceBuffer, + q: cuda.DeviceBuffer, + iq: cuda.DeviceBuffer, + po: cuda.DeviceBuffer, + pm: cuda.DeviceBuffer, + pl: cuda.DeviceBuffer, + ao: cuda.DeviceBuffer, + gated: cuda.DeviceBuffer, + gxs_attn: cuda.DeviceBuffer, + dummy: cuda.DeviceBuffer, + inv: cuda.DeviceBuffer, + pos: cuda.DeviceBuffer, + ids: cuda.DeviceBuffer, + logits: cuda.DeviceBuffer, + router_l: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + moe_members: cuda.DeviceBuffer, + moe_items: cuda.DeviceBuffer, + moe_counts: cuda.DeviceBuffer, + moe_act: cuda.DeviceBuffer, + + pub fn alloc(driver: *cuda.Driver, rows: usize) !Scratch { + var s: Scratch = undefined; + s.h = try cuda.DeviceBuffer.alloc(driver, rows * wide * 2); + s.y = try cuda.DeviceBuffer.alloc(driver, rows * slots * dims * 4); + s.wts = try cuda.DeviceBuffer.alloc(driver, rows * wts_stride); + s.inj = try cuda.DeviceBuffer.alloc(driver, rows * inj_stride); + s.pss = try cuda.DeviceBuffer.alloc(driver, (dims / 256) * streams * 4); + s.norm = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.dn = try cuda.DeviceBuffer.alloc(driver, (low + 4) * 2); + s.hc_part = try cuda.DeviceBuffer.alloc(driver, hc_sk * (low + 4) * 4); + s.act = try cuda.DeviceBuffer.alloc(driver, low * 2); + s.xs_act = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + s.mixed = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.xsm = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.proj = try cuda.DeviceBuffer.alloc(driver, gdn.proj_width * 2); + s.part = try cuda.DeviceBuffer.alloc(driver, out_sk * dims * 4); + s.gout = try cuda.DeviceBuffer.alloc(driver, gdn.value_dim * 2); + s.gxs = try cuda.DeviceBuffer.alloc(driver, (gdn.value_dim / 32) * 4); + s.st = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + s.ple_emb = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_xs = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.ple_keys = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_vals = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_gated = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_pss = try cuda.DeviceBuffer.alloc(driver, streams * 4); + s.ple_nrow = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_part = try cuda.DeviceBuffer.alloc(driver, 4 * dims * 4); + s.pa = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + s.q = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + s.iq = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + s.po = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + s.pm = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.pl = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.ao = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gated = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gxs_attn = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + s.dummy = try cuda.DeviceBuffer.alloc(driver, 4096); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + s.pos = try cuda.DeviceBuffer.alloc(driver, 4); + s.ids = try cuda.DeviceBuffer.alloc(driver, 4); + s.logits = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + s.router_l = try cuda.DeviceBuffer.alloc(driver, (experts_n + 1) * 4); + s.picks = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_members = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_items = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + s.moe_counts = try cuda.DeviceBuffer.alloc(driver, 8); + s.moe_act = try cuda.DeviceBuffer.alloc(driver, slots * moe_w * 2); + return s; + } + + pub fn free(self: *Scratch) void { + self.h.free(); + self.y.free(); + self.wts.free(); + self.inj.free(); + self.pss.free(); + self.norm.free(); + self.dn.free(); + self.hc_part.free(); + self.act.free(); + self.xs_act.free(); + self.mixed.free(); + self.xsm.free(); + self.proj.free(); + self.part.free(); + self.gout.free(); + self.gxs.free(); + self.st.free(); + self.ple_emb.free(); + self.ple_xs.free(); + self.ple_keys.free(); + self.ple_vals.free(); + self.ple_gated.free(); + self.ple_pss.free(); + self.ple_nrow.free(); + self.ple_part.free(); + self.pa.free(); + self.q.free(); + self.iq.free(); + self.po.free(); + self.pm.free(); + self.pl.free(); + self.ao.free(); + self.gated.free(); + self.gxs_attn.free(); + self.dummy.free(); + self.inv.free(); + self.pos.free(); + self.ids.free(); + self.logits.free(); + self.router_l.free(); + self.picks.free(); + self.moe_members.free(); + self.moe_items.free(); + self.moe_counts.free(); + self.moe_act.free(); + } +}; + +pub fn writeback(tri: anytype, s: *Scratch, h: u64, mode: usize, inj: u64, y: u64, wts: u64) !void { + const dummy = s.dummy.ptr; + if (mode == 0) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 0, 1, 1, 1); + } else if (mode == 2) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", dims, 1, dims, streams, 2, 10, 11, 1); + } else if (mode == 4) { + try tri.hcWriteback(h, h, s.pss.ptr, s.part.ptr, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + } else return error.UnexpectedTensor; +} + +pub fn readout(tri: anytype, s: *Scratch, hc_w: HcW, h: u64, inj: u64) !void { + try tri.hcDown(h, s.pss.ptr, hc_w.scale.ptr, s.norm.ptr, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn.ptr, s.hc_part.ptr, eps, 1, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else s.act.ptr; + try tri.hcReduceAct(s.hc_part.ptr, s.act.ptr, s.xs_act.ptr, inj_ptr, hc_sk, 1, streams, low, hc_w.n_down, hc_w.has_inj); + try tri.hcUpmix(s.act.ptr, s.xs_act.ptr, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.norm.ptr, s.mixed.ptr, s.xsm.ptr, 1, wide, low, dims, streams); +} + +fn shiftTail(gpa: std.mem.Allocator, stream: *cuda.Stream, window: cuda.DeviceBuffer, taps_n: usize, row_bytes: usize, fresh: []const u8) !void { + const bytes = taps_n * row_bytes; + const raw = try gpa.alloc(u8, bytes); + defer gpa.free(raw); + try stream.synchronize(); + try window.download(0, raw); + const next = try gpa.alloc(u8, bytes); + defer gpa.free(next); + const keep = taps_n - 1; + @memcpy(next[0 .. keep * row_bytes], raw[row_bytes..]); + @memcpy(next[keep * row_bytes ..], fresh); + try window.upload(0, next); +} + +pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s: *Scratch, proj_q: Q4, out_q: Q4, conv: cuda.DeviceBuffer, cw: u64, a_log: u64, dt: u64, norm: u64, state: *cuda.DeviceBuffer) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.proj.ptr, 1, proj_q.n, proj_q.k); + try gdn.launch(driver, stream.*, s.proj.ptr, conv.ptr, cw, state.ptr, a_log, dt, norm, eps, s.gout.ptr, s.gxs.ptr, s.st.ptr); + try state.copyFrom(0, s.st.ptr, state_n * 4, stream.handle); + const row = try gpa.alloc(u8, gdn.conv_dim * 2); + defer gpa.free(row); + try stream.synchronize(); + try s.proj.download(0, row); + try shiftTail(gpa, stream, conv, gdn.taps - 1, gdn.conv_dim * 2, row); + try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); +} + +pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); + // stream-ordered (queued kernels still read s.pos), then waited for: `pos` lives on this frame + try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); + try stream.synchronize(); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) + const keys: usize = @as(usize, @intCast(pos)) + 1; + if (keys > wt.dense_nch * 512) return error.UnexpectedTensor; + try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false); + try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, false); + try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); + try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, stream, &a); +} + +pub fn moeStep(comptime Tri: type, tri: Tri, s: *Scratch, stream: *cuda.Stream, plan_fn: cuda.Function, up_fn: cuda.Function, down_fn: cuda.Function, router_w: u64, up: u64, down: u64, y: u64, wts: u64) !void { + try tri.router(s.mixed.ptr, router_w, s.router_l.ptr, dims); + try tri.topkRows(s.router_l.ptr, s.picks.ptr, wts); + var plan_a: cuda.Args = .{}; + plan_a.add(s.picks.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(experts_n + 1)); + plan_a.add(cint(tile)); + plan_a.add(s.moe_members.ptr); + plan_a.add(s.moe_items.ptr); + plan_a.add(s.moe_counts.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + try launchExpert(up_fn, stream.*, s.mixed.ptr, dims, slots, up, dims / 32, moe_w / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, s.moe_act.ptr, moe_w, slots * (moe_w / 32)); + try launchExpert(down_fn, stream.*, s.moe_act.ptr, moe_w, 0, down, moe_w / 32, dims / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, y, dims, slots * (dims / 32)); +} + +fn reduceBf16(gpa: std.mem.Allocator, stream: *cuda.Stream, part: cuda.DeviceBuffer, sk: usize, n: usize, dst: cuda.DeviceBuffer) !void { + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part.download(0, raw); + const row = try gpa.alloc(u16, n); + defer gpa.free(row); + for (0..n) |i| { + var acc: f32 = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + for (1..sk) |s| acc += @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * (s * n + i) ..][0..4], .little))); + row[i] = toBf16(acc); + } + try dst.upload(0, std.mem.sliceAsBytes(row)); +} + +pub fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, model_dir: []const u8, history: *[2]i64, token: u32, h: u64, key_q: Q4, val_q: Q4, nk: u64, nq: u64, nc: u64, cw: u64, tail: *cuda.DeviceBuffer) !void { + const emb = try gpa.alloc(u16, dims); + defer gpa.free(emb); + const xs = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs); + try @import("cuda_ple_check.zig").embedding(gpa, io, model_dir, history, token, emb, xs); + try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); + try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); + try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); + try qmm.partials(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, val_q.w.ptr, val_q.s.ptr, val_q.b.ptr, s.ple_vals.ptr, s.ple_part.ptr, 1, dims, dims, 4); + try reduceBf16(gpa, stream, s.ple_part, 4, dims, s.ple_vals); + try tri.pleGate(s.ple_keys.ptr, s.ple_vals.ptr, h, nk, nq, s.ple_gated.ptr, s.ple_pss.ptr, eps, 1, dims, streams); + try tri.pleConv(s.ple_gated.ptr, s.ple_pss.ptr, nc, tail.ptr, cw, h, h, s.ple_nrow.ptr, eps, 1, dims, streams, 4, 3); + const nrow = try gpa.alloc(u8, wide * 2); + defer gpa.free(nrow); + try stream.synchronize(); + try s.ple_nrow.download(0, nrow); + try shiftTail(gpa, stream, tail.*, ple_tail_n, wide * 2, nrow); + history[0] = history[1]; + history[1] = token; +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig new file mode 100644 index 0000000000..6e54d694ed --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -0,0 +1,599 @@ +//! Checkpoint bytes for one drafts-off prompt: the safetensors reader and the quantized linears. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); + +pub const layers_n: usize = 48; +pub const linear_n: usize = 36; +pub const attn_n: usize = 12; +pub const dims: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * dims; +pub const low: usize = 320; +pub const eps: f32 = 1e-6; +pub const vocab: usize = 248320; +pub const experts_n: usize = 512; +pub const slots: usize = 11; +pub const moe_w: usize = 640; +pub const block_bytes: usize = 160 * 4; +pub const tile: usize = 16; +pub const q_heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const half: usize = 32; +pub const proj_n: usize = 13952; +pub const out_k: usize = 6144; +/// Chunks of 512 keys a row's attention reads at most: 2,048 selected keys plus a tail, or 2,051 dense keys. +pub const nch: usize = 5; +pub const dense_nch: usize = 3; +/// The block indexer (QSA): keys pooled in blocks of 4, a row past 512 complete blocks reads its 512 best. +pub const index_ratio: usize = 4; +pub const index_top: usize = 512; +/// A row's key list: the budget's 2,048 keys plus a tail of up to ratio - 1 (attention.AttnScratch.idw). +pub const index_idw: usize = index_top * index_ratio + index_ratio; +/// Keys a dense row reads at most before selection starts: (top + 1) x ratio - 1. +pub const dense_keys: usize = (index_top + 1) * index_ratio - 1; +/// The model's window (config max_position_embeddings). +pub const window: usize = 262144; +pub const ple_layer: usize = 1; +pub const inj_stride: usize = 16; +pub const wts_stride: usize = 48; +pub const ple_tail_n: usize = 9; +pub const state_n: usize = gdn.nv * gdn.dv * gdn.dk; +pub const hc_sk: usize = 32; +pub const out_sk: usize = 8; + +pub const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +pub const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +pub const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Tensor = core.safetensors.Tensor; + +const Held = struct { rel: []u8, file: Shard }; + +/// One mapped shard: core.safetensors.File. +const Shard = struct { + f: core.safetensors.File, + io: std.Io, + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { + return .{ .f = try core.safetensors.File.open(gpa, io, path), .io = io }; + } + + fn close(self: *Shard) void { + self.f.close(self.io); + } + + fn get(self: *const Shard, name: []const u8) ?Tensor { + return self.f.get(name); + } +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: std.Io, + dir: []const u8, + parsed: std.json.Parsed(std.json.Value), + path: []u8 = &.{}, + rel: []u8 = &.{}, + file: ?Shard = null, + held: [12]Held = undefined, + nheld: usize = 0, + released: [2]Held = undefined, + nreleased: usize = 0, + norm_offset: f32 = 0, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; + } + + /// Infer one checkpoint-wide convention from every decoder HC norm, never from the final mixer's outlier. + pub fn detectNorms(self: *Store) !void { + var means: [layers_n]f64 = undefined; + var buf: [160]u8 = undefined; + for (0..layers_n) |i| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.attn_hyper_connection.hc_norm.weight", .{i}); + const values = try asF32(self.gpa, try self.tensor(name)); + defer self.gpa.free(values); + if (values.len != wide) return error.UnexpectedTensor; + var sum: f64 = 0; + for (values) |v| { + if (!std.math.isFinite(v)) return error.InvalidNormConvention; + sum += v; + } + means[i] = sum / @as(f64, @floatFromInt(values.len)); + } + self.norm_offset = try @import("cuda_norms.zig").offset(layers_n, means); + } + + pub fn normScale(self: *const Store, gpa: std.mem.Allocator, t: Tensor) ![]f32 { + const values = try asF32(gpa, t); + errdefer gpa.free(values); + for (values) |*v| { + if (!std.math.isFinite(v.*)) return error.InvalidNormConvention; + v.* += self.norm_offset; + } + return values; + } + + pub fn deinit(self: *Store) void { + self.release(); + self.closeReleased(); + self.closeFile(); + self.parsed.deinit(); + } + + /// Borrowers expire here; two indexed shards cover separate weight and calibration planes. + pub fn release(self: *Store) void { + for (self.held[0..self.nheld]) |h| { + if (self.nreleased == self.released.len) { + self.released[0].file.close(); + self.gpa.free(self.released[0].rel); + self.released[0] = self.released[1]; + self.nreleased -= 1; + } + self.released[self.nreleased] = h; + self.nreleased += 1; + } + self.nheld = 0; + } + + fn closeReleased(self: *Store) void { + for (self.released[0..self.nreleased]) |*h| { + h.file.close(); + self.gpa.free(h.rel); + } + self.nreleased = 0; + } + + pub fn hold(self: *Store, requested: []const u8) !Tensor { + var mapped: [512]u8 = undefined; + const key = try self.resolveKey(requested, &mapped); + const file = try self.shard(key); + for (self.held[0..self.nheld]) |*h| { + if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + } + if (self.nheld == self.held.len) return error.UnexpectedTensor; + for (self.released[0..self.nreleased], 0..) |cached, i| { + if (std.mem.eql(u8, cached.rel, file)) { + self.held[self.nheld] = cached; + self.nreleased -= 1; + if (i < self.nreleased) self.released[i] = self.released[self.nreleased]; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse return error.MissingTensor; + } + } + const rel = try self.gpa.dupe(u8, file); + errdefer self.gpa.free(rel); + const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + defer self.gpa.free(path); + self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + pub fn closeFile(self: *Store) void { + if (self.file) |*f| f.close(); + self.file = null; + if (self.path.len != 0) self.gpa.free(self.path); + self.path = &.{}; + if (self.rel.len != 0) self.gpa.free(self.rel); + self.rel = &.{}; + } + + pub fn resolveKey(self: *const Store, requested: []const u8, buf: []u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + if (wm.object.contains(requested)) return requested; + const aliases = .{ + .{ "language_model.model.", "model.language_model." }, + .{ "language_model.lm_head.", "lm_head." }, + .{ "language_model.mtp.", "mtp." }, + }; + inline for (aliases) |pair| { + if (std.mem.startsWith(u8, requested, pair[0])) { + const name = try std.fmt.bufPrint(buf, "{s}{s}", .{ pair[1], requested[pair[0].len..] }); + if (wm.object.contains(name)) return name; + } + } + return error.MissingTensor; + } + + fn shard(self: *Store, key: []const u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + const v = wm.object.get(key) orelse { + std.debug.print("missing {s}\n", .{key}); + return error.MissingTensor; + }; + return v.string; + } + + fn use(self: *Store, key: []const u8) !void { + const file = try self.shard(key); + if (self.file != null and std.mem.eql(u8, file, self.rel)) return; + self.closeFile(); + self.rel = try self.gpa.dupe(u8, file); + self.path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + self.file = try Shard.load(self.gpa, self.io, self.path); + } + + pub fn tensor(self: *Store, requested: []const u8) !Tensor { + var mapped: [512]u8 = undefined; + const key = try self.resolveKey(requested, &mapped); + try self.use(key); + return self.file.?.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + /// Check the exact or compatibility name without mapping its shard. + pub fn has(self: *const Store, requested: []const u8) bool { + var mapped: [512]u8 = undefined; + _ = self.resolveKey(requested, &mapped) catch return false; + return true; + } + + /// Checkpoint-reader adapter: the borrowed tensor stays mapped until release(). + pub fn get(self: *Store, key: []const u8) !Tensor { + return self.hold(key); + } + + pub fn copyOf(self: *Store, key: []const u8) ![]u8 { + const t = try self.tensor(key); + const out = try self.gpa.alloc(u8, t.bytes.len); + @memcpy(out, t.bytes); + return out; + } +}; + +pub const Face = struct { + w: []const u8, + s: []const u8, + b: []const u8, + n: usize, + k: usize, + routed: bool, +}; + +pub const Q4 = @import("cuda_projection.zig").Projection; + +pub const HcW = struct { + scale: cuda.DeviceBuffer, + down: Q4, + up: Q4, + n_down: usize, + has_inj: usize, + + pub fn free(self: *HcW) void { + self.scale.free(); + self.down.free(); + self.up.free(); + } +}; + +pub fn linearLayer(layer: usize) bool { + return layer % 4 != 3; +} + +pub fn linearIndex(layer: usize) usize { + return layer - layer / 4; +} + +pub fn attnIndex(layer: usize) usize { + return layer / 4; +} + +pub fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +pub fn cint(v: usize) c_int { + return @intCast(v); +} + +pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { + var n: usize = 0; + for (parts) |p| n += p.len; + const out = try gpa.alloc(u8, n); + var at: usize = 0; + for (parts) |p| { + @memcpy(out[at..][0..p.len], p); + at += p.len; + } + return out; +} + +/// A tensor the engine leaves off the device: the n-gram tables, the draft head and the vision tower. +pub fn hostOnly(name: []const u8) bool { + return std.mem.indexOf(u8, name, ".ple_embedding.ngram_embedding.") != null or + std.mem.startsWith(u8, name, "language_model.mtp.") or + std.mem.startsWith(u8, name, "mtp.") or + std.mem.startsWith(u8, name, "model.visual.") or + std.mem.startsWith(u8, name, "vision_tower."); +} + +/// The checkpoint's tensor bytes the engine uploads, every shard the index names read for its header alone. +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { + var store = try Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = 0; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var sh = try Shard.load(gpa, io, path); + defer sh.close(); + var ni = sh.f.names.iterator(); + while (ni.next()) |e| { + if (!hostOnly(e.key_ptr.*)) total += e.value_ptr.end - e.value_ptr.begin; + } + } + return total; +} + +pub fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { + if (t.rank != 1) return error.UnexpectedTensor; + const out = try gpa.alloc(f32, t.dim(0)); + errdefer gpa.free(out); + if (t.dtype == .f32 and t.bytes.len == out.len * 4) { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); + } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); + } else return error.UnexpectedTensor; + return out; +} + +pub fn uploadQ4(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + var lane = try qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn uploadTiled(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + const kg = k8 / 4; + const tiled = try hc.tileWords(gpa, words, n, k8); + defer gpa.free(tiled); + const ts = try hc.transposeBf16(gpa, scales, n, kg); + defer gpa.free(ts); + const tb = try hc.transposeBf16(gpa, biases, n, kg); + defer gpa.free(tb); + var w = try cuda.DeviceBuffer.fromHost(driver, tiled); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, ts); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, tb); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8) !struct { w: []u8, s: []u8, b: []u8 } { + const w = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + errdefer gpa.free(w); + const s = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + errdefer gpa.free(s); + const b = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + return .{ .w = w, .s = s, .b = b }; +} + +pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + var buf: [180]u8 = undefined; + const down_name = try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_down.weight", .{prefix}); + const down_tensor = try store.tensor(down_name); + if (down_tensor.dtype == .bf16) return loadDenseHc(gpa, driver, store, prefix, inject); + const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); + defer gpa.free(down.w); + defer gpa.free(down.s); + defer gpa.free(down.b); + const up = try join3(gpa, store, &buf, prefix, ".input_mix_weight_up"); + defer gpa.free(up.w); + defer gpa.free(up.s); + defer gpa.free(up.b); + const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); + const scale = try store.normScale(gpa, gamma); + defer gpa.free(scale); + if (scale.len != wide) return error.UnexpectedTensor; + var n: usize = low; + var words: []u8 = undefined; + var scales: []u8 = undefined; + var biases: []u8 = undefined; + var inj_w: []u8 = &.{}; + var inj_s: []u8 = &.{}; + var inj_b: []u8 = &.{}; + if (inject) { + const inj = try join3(gpa, store, &buf, prefix, ".block_inject_weight"); + inj_w = inj.w; + inj_s = inj.s; + inj_b = inj.b; + n = low + 4; + words = try cat(gpa, &.{ down.w, inj.w }); + scales = try cat(gpa, &.{ down.s, inj.s }); + biases = try cat(gpa, &.{ down.b, inj.b }); + } else { + words = try gpa.dupe(u8, down.w); + scales = try gpa.dupe(u8, down.s); + biases = try gpa.dupe(u8, down.b); + } + defer gpa.free(words); + defer gpa.free(scales); + defer gpa.free(biases); + defer if (inject) gpa.free(inj_w); + defer if (inject) gpa.free(inj_s); + defer if (inject) gpa.free(inj_b); + const k8 = wide / 8; + if (words.len != n * k8 * 4) return error.UnexpectedTensor; + var down_q = try uploadTiled(gpa, driver, words, scales, biases, n, k8); + errdefer down_q.free(); + var up_q = try uploadTiled(gpa, driver, up.w, up.s, up.b, wide, low / 8); + errdefer up_q.free(); + const scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; +} + +fn loadDenseHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + const projection = @import("cuda_projection.zig"); + var buf: [200]u8 = undefined; + const n = low + if (inject) streams else @as(usize, 0); + const data = try gpa.alloc(u8, n * wide * 2); + defer gpa.free(data); + const down = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_down.weight", .{prefix})); + if (!down.is(.bf16, &.{ low, wide }) or down.bytes.len != low * wide * 2) return error.UnexpectedTensor; + @memcpy(data[0..down.bytes.len], down.bytes); + if (inject) { + const inj = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.block_inject_weight.weight", .{prefix})); + if (!inj.is(.bf16, &.{ streams, wide }) or inj.bytes.len != streams * wide * 2) return error.UnexpectedTensor; + @memcpy(data[low * wide * 2 ..], inj.bytes); + } + var dq = try projection.dense(gpa, driver, .{ .dtype = .bf16, .rank = 2, .shape = .{ n, wide, 1, 1, 1 }, .bytes = data }); + errdefer dq.free(); + const up = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_up.weight", .{prefix})); + if (!up.is(.bf16, &.{ wide, low })) return error.UnexpectedTensor; + var uq = try projection.dense(gpa, driver, up); + errdefer uq.free(); + const gamma = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix}))); + defer gpa.free(gamma); + if (gamma.len != wide) return error.UnexpectedTensor; + const scale = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(gamma)); + return .{ .scale = scale, .down = dq, .up = uq, .n_down = n, .has_inj = @intFromBool(inject) }; +} + +pub fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { + var buf: [180]u8 = undefined; + var ws: [6][]u8 = undefined; + var ss: [6][]u8 = undefined; + var bs: [6][]u8 = undefined; + var n: usize = 0; + for (names, 0..) |part, i| { + const got = try join3(gpa, store, &buf, prefix, part); + ws[i] = got.w; + ss[i] = got.s; + bs[i] = got.b; + n += 1; + } + defer for (0..n) |i| { + gpa.free(ws[i]); + gpa.free(ss[i]); + gpa.free(bs[i]); + }; + return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; +} + +test "the n-gram tables, the draft head and the vision tower stay off the device" { + try std.testing.expect(hostOnly("language_model.model.layers.2.ple.ple_embedding.ngram_embedding.shards.7.weight")); + try std.testing.expect(hostOnly("language_model.mtp.layers.0.mlp.switch_mlp.up_proj.weight")); + try std.testing.expect(hostOnly("vision_tower.blocks.0.attn.qkv.weight")); + try std.testing.expect(!hostOnly("language_model.model.layers.2.ple.ple_embedding.ngram_heads_offsets")); + try std.testing.expect(!hostOnly("language_model.model.layers.0.mlp.switch_mlp.up_proj.weight")); + try std.testing.expect(!hostOnly("language_model.lm_head.weight")); +} + +test "canonical checkpoint aliases preserve exact stored namespace and prefer a direct key" { + const a = std.testing.allocator; + const parsed = try std.json.parseFromSlice(std.json.Value, a, + \\{"weight_map":{"model.language_model.layers.0.x.weight":"a","lm_head.weight":"b","mtp.x":"c","language_model.model.direct":"d","model.language_model.direct":"e"}} + , .{}); + var store: Store = .{ .gpa = a, .io = std.testing.io, .dir = "", .parsed = parsed }; + defer store.deinit(); + var buf: [512]u8 = undefined; + try std.testing.expectEqualStrings("model.language_model.layers.0.x.weight", try store.resolveKey("language_model.model.layers.0.x.weight", &buf)); + try std.testing.expectEqualStrings("lm_head.weight", try store.resolveKey("language_model.lm_head.weight", &buf)); + try std.testing.expectEqualStrings("mtp.x", try store.resolveKey("language_model.mtp.x", &buf)); + try std.testing.expectEqualStrings("language_model.model.direct", try store.resolveKey("language_model.model.direct", &buf)); + try std.testing.expectError(error.MissingTensor, store.resolveKey("language_model.model.missing", &buf)); + try std.testing.expect(hostOnly("model.visual.weight")); + try std.testing.expect(hostOnly("mtp.layers.0.weight")); +} + +test "released expert planes reuse two indexed shards without invalidating held peers" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + const header = "{\"x\":{\"dtype\":\"U8\",\"shape\":[1],\"data_offsets\":[0,1]}}"; + var data: [8 + header.len + 1]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + data[data.len - 1] = 17; + try tmp.dir.writeFile(io, .{ .sub_path = "a", .data = &data }); + data[data.len - 1] = 29; + try tmp.dir.writeFile(io, .{ .sub_path = "b", .data = &data }); + try tmp.dir.writeFile(io, .{ .sub_path = "c", .data = &data }); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{\"x\":\"a\"}}" }); + var store = try Store.open(a, io, dir); + defer store.deinit(); + const first = try store.hold("x"); + const address = first.bytes.ptr; + store.release(); + try std.testing.expectEqual(@as(usize, 0), store.nheld); + try std.testing.expectEqual(@as(usize, 1), store.nreleased); + // Removing the path proves the next borrow uses the existing mapping and parsed header. + try tmp.dir.deleteFile(io, "a"); + const again = try store.hold("x"); + try std.testing.expectEqual(address, again.bytes.ptr); + try std.testing.expectEqual(@as(u8, 17), again.bytes[0]); + const wm = &store.parsed.value.object.getPtr("weight_map").?.object; + wm.getPtr("x").?.* = .{ .string = "b" }; + const other = try store.hold("x"); + try std.testing.expectEqual(@as(u8, 29), other.bytes[0]); + try std.testing.expectEqual(@as(u8, 17), again.bytes[0]); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + try tmp.dir.deleteFile(io, "b"); + wm.getPtr("x").?.* = .{ .string = "a" }; + try std.testing.expectEqual(@as(u8, 17), (try store.hold("x")).bytes[0]); + wm.getPtr("x").?.* = .{ .string = "b" }; + try std.testing.expectEqual(@as(u8, 29), (try store.hold("x")).bytes[0]); + store.release(); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + wm.getPtr("x").?.* = .{ .string = "c" }; + _ = try store.hold("x"); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + for (store.released[0..store.nreleased]) |h| try std.testing.expect(!std.mem.eql(u8, "a", h.rel)); +} + +test "centered scales add one in FP32 while direct scales preserve the original bits" { + _ = @import("cuda_norms.zig"); + const raw = [_]f32{ -0.25, 0.0001, 2.75 }; + const t: Tensor = .{ .dtype = .f32, .rank = 1, .shape = .{ 3, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&raw) }; + var store: Store = undefined; + store.norm_offset = 1; + const scales = try store.normScale(std.testing.allocator, t); + defer std.testing.allocator.free(scales); + for (raw, scales) |v, scale| try std.testing.expectEqual(v + @as(f32, 1), scale); + store.norm_offset = 0; + const direct = try store.normScale(std.testing.allocator, t); + defer std.testing.allocator.free(direct); + try std.testing.expectEqualSlices(f32, &raw, direct); +} diff --git a/zig/src/families/flashnext/cuda_qmm.zig b/zig/src/families/flashnext/cuda_qmm.zig new file mode 100644 index 0000000000..50df60a80f --- /dev/null +++ b/zig/src/families/flashnext/cuda_qmm.zig @@ -0,0 +1,211 @@ +//! Group-32 lane matmul: MLX words packed as qmm.pack, then the shared CUDA kernel's 16-row tile. + +const std = @import("std"); +const cuda = @import("cuda"); + +pub const group: usize = 32; + +const offsets = [8]u32{ 0, 8, 16, 24, 1, 9, 17, 25 }; + +/// LaneTile<32, 16, 64, 1, 4, 4>::SMEM. One by four warps, four stages. +pub const smem: u32 = 9472; +pub const threads: u32 = 128; + +const symbol: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; + +pub const Packed = struct { + weight: []u8, + scales: []u8, + biases: []u8, + n: usize, + k: usize, + npad: usize, + + pub fn deinit(self: *Packed, gpa: std.mem.Allocator) void { + gpa.free(self.weight); + gpa.free(self.scales); + gpa.free(self.biases); + self.* = undefined; + } +}; + +fn nibble(words: []const u8, k8: usize, col: usize, pos: usize) u32 { + const word_i = pos / 8; + const shift: u5 = @intCast((pos % 8) * 4); + const off = (col * k8 + word_i) * 4; + const word = std.mem.readInt(u32, words[off..][0..4], .little); + return (word >> shift) & 0xF; +} + +/// MLX (n, k/8) words and (n, k/32) scales and biases -> the lane layout. n is padded to 128. +pub fn pack(gpa: std.mem.Allocator, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Packed { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const k = k8 * 8; + if (k % group != 0) return error.UnexpectedTensor; + const kg = k / group; + if (scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + const npad = (n + 127) / 128 * 128; + const tiles = npad / 64; + const weight = try gpa.alloc(u8, tiles * kg * 8 * 32 * 4); + errdefer gpa.free(weight); + @memset(weight, 0); + for (0..n) |col| { + const tile = col / 64; + const inner = col % 64; + const j = inner / 8; + const r = inner % 8; + for (0..kg) |g| { + for (0..4) |c| { + var bits: u32 = 0; + for (0..8) |p| { + const pos = g * group + 2 * c + offsets[p]; + bits |= nibble(words, k8, col, pos) << @intCast(p * 4); + } + const idx = r * 4 + c; + const dst = (((tile * kg + g) * 8 + j) * 32 + idx) * 4; + std.mem.writeInt(u32, weight[dst..][0..4], bits, .little); + } + } + } + const scales_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(scales_out); + const biases_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(biases_out); + @memset(scales_out, 0); + @memset(biases_out, 0); + for (0..kg) |g| { + for (0..n) |col| { + const src = (col * kg + g) * 2; + const dst = (g * npad + col) * 2; + @memcpy(scales_out[dst..][0..2], scales[src..][0..2]); + @memcpy(biases_out[dst..][0..2], biases[src..][0..2]); + } + } + return .{ .weight = weight, .scales = scales_out, .biases = biases_out, .n = n, .k = k, .npad = npad }; +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One row: each group is `dot(x, q) * scale + xs * bias`, groups added in order, then rounded to bf16. +pub fn dotRow(x: []const u16, xs: []const f32, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []u16) !void { + if (k == 0 or k % group != 0 or x.len < k or xs.len < k / group or out.len < n) return error.UnexpectedTensor; + const kg = k / group; + const k8 = k / 8; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..n) |col| { + var acc: f32 = 0; + const word_row = words[col * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[col * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[col * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(x[g * group + i * 8 + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = toBf16(acc); + } +} + +/// qmm.cu's 16-row launch for one K slice: x (m, k) bf16 and xs (m, k/32) fp32 into out (m, n) bf16. +pub fn matmul(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, m: usize, n: usize, k: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or k == 0 or k % group != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(@as(u64, 0)); + for ([_]usize{ m, n, k, 1, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t) }, .block = .{ .x = threads }, .shared = smem }, stream, &args); +} + +/// The same kernel as `matmul` with one block per K slice. `part` receives the unreduced [sk, m, n] fp32 slices. +pub fn partials(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, m: usize, n: usize, k: usize, sk: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or sk == 0 or k == 0 or k % group != 0 or (k / group) % sk != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(part); + for ([_]usize{ m, n, k, sk, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ + .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, + .block = .{ .x = threads }, + .shared = smem, + }, stream, &args); +} + +test "one group's first nibble pair lands in the lane word" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + const bias: [2]u8 = @splat(0); + var got = try pack(std.testing.allocator, &words, &scale, &bias, 1, 4); + defer got.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 128), got.npad); + try std.testing.expectEqual(@as(u32, 0x20001), std.mem.readInt(u32, got.weight[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, got.weight[4..8], .little)); + try std.testing.expectEqual(@as(u16, 0x3f80), std.mem.readInt(u16, got.scales[0..2], .little)); +} + +test "scale times the nibble dot plus the group bias rounds to five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + const xs = [_]f32{2}; + var out: [1]u16 = undefined; + try dotRow(&x, &xs, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(u16, 0x40a0), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig new file mode 100644 index 0000000000..ff7db8e5ae --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -0,0 +1,471 @@ +//! Leader-owned operations replayed on the other CUDA rank; sequence and cache identities survive map relocation. +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const snaps = @import("cuda_snap.zig"); +const wire = @import("cuda_rank_wire.zig"); +const Collective = @import("cuda_collective.zig").Collective; +const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); + +pub const Identity = extern struct { + config: [32]u8, + index: [32]u8, + program: [32]u8, + checkpoint: [32]u8, + context: u64, + chunk: u64, + decode_graph: u64 = 0, + mtp: u64 = 0, + head_sharded: u64 = 0, +}; + +fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 28)); + defer a.free(bytes); + var out: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(bytes, &out, .{}); + return out; +} + +pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usize, chunk: usize) !Identity { + const digest = std.mem.span(std.c.getenv("TENSORFOLD_CHECKPOINT_SHA256") orelse return error.NoCheckpointIdentity); + if (digest.len != 64) return error.NoCheckpointIdentity; + var checkpoint: [32]u8 = undefined; + _ = std.fmt.hexToBytes(&checkpoint, digest) catch return error.NoCheckpointIdentity; + const config = try std.fs.path.join(a, &.{ dir, "config.json" }); + defer a.free(config); + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + defer a.free(index); + const program = try std.process.executablePathAlloc(io, a); + defer a.free(program); + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()) }; +} + +pub const Session = struct { + gpa: std.mem.Allocator, + rank: usize, + link: wire.Link, + collective: Collective, + engine: *prompt.Engine = undefined, + next: u64 = 1, + next_object: u64 = 1, + seqs: std.AutoHashMapUnmanaged(u64, *prompt.Seq) = .empty, + saved: std.AutoHashMapUnmanaged(u64, *snaps.Snap) = .empty, + windows: std.AutoHashMapUnmanaged(u64, mtp.Draft) = .empty, + drafting: bool = false, + stopped: bool = false, + available: std.atomic.Value(bool) = .init(true), + + pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity) !*Session { + const s = try a.create(Session); + errdefer a.destroy(s); + const link = try wire.Link.open(io, host, port, rank, 60_000); + errdefer link.close(); + const until = wire.deadline(60_000); + var peer: Identity = undefined; + const hello: wire.Header = .{ .kind = .hello, .arg = rank, .bytes = @sizeOf(Identity) }; + if (rank == 0) { + try link.sendHeader(hello, until); + try link.send(std.mem.asBytes(&mine), until); + const h = try link.receiveHeader(until); + if (h.kind != .hello or h.sequence != 0 or h.arg != 1 or h.bytes != @sizeOf(Identity) or h.flags != 0 or h.object != 0) return error.RankIdentity; + try link.receive(std.mem.asBytes(&peer), until); + } else { + const h = try link.receiveHeader(until); + if (h.kind != .hello or h.sequence != 0 or h.arg != 0 or h.bytes != @sizeOf(Identity) or h.flags != 0 or h.object != 0) return error.RankIdentity; + try link.receive(std.mem.asBytes(&peer), until); + try link.sendHeader(hello, until); + try link.send(std.mem.asBytes(&mine), until); + } + if (!sameIdentity(mine, peer)) return error.RankIdentity; + s.* = .{ .gpa = a, .rank = rank, .link = link, .collective = try Collective.init(link, rank, 30_000), .drafting = mine.mtp == 1 }; + return s; + } + + pub fn healthy(s: *const Session) bool { + if (!s.available.load(.acquire)) return false; + s.link.checkPeer() catch return false; + return true; + } + + pub fn poison(s: *Session) void { + s.stopped = true; + s.available.store(false, .release); + s.collective.abort(); + _ = std.c.shutdown(s.link.fd, std.posix.SHUT.RDWR); + } + + pub fn deinit(s: *Session) void { + if (s.rank == 0 and !s.stopped) { + const h = s.command(.{ .kind = .stop }, &.{}) catch null; + if (h) |sent| s.ack(sent) catch s.poison(); + } + var wi = s.windows.valueIterator(); + while (wi.next()) |d| d.deinit(); + s.windows.deinit(s.gpa); + s.collective.deinit(); + var qi = s.seqs.valueIterator(); + while (qi.next()) |q| { + q.*.deinit(s.gpa); + s.gpa.destroy(q.*); + } + s.seqs.deinit(s.gpa); + var si = s.saved.valueIterator(); + while (si.next()) |saved| snaps.drop(s.gpa, s.engine.stream.*, saved.*); + s.saved.deinit(s.gpa); + s.link.close(); + s.gpa.destroy(s); + } + + /// Every rank's host budgets the same number of streams and cache bytes, using the least local room. + pub fn room(s: *Session, local: u64) !u64 { + const until = wire.deadline(600_000); + errdefer s.poison(); + const mine: wire.Header = .{ .kind = .ready, .arg = local }; + const peer = if (s.rank == 0) peer: { + try s.link.sendHeader(mine, until); + break :peer try s.link.receiveHeader(until); + } else peer: { + const h = try s.link.receiveHeader(until); + try s.link.sendHeader(mine, until); + break :peer h; + }; + if (peer.kind != .ready or peer.sequence != 0 or peer.object != 0 or peer.bytes != 0 or peer.flags != 0) return error.BadRankFrame; + return @min(local, peer.arg); + } + + fn command(s: *Session, input: wire.Header, payload: []const u8) !wire.Header { + if (s.rank != 0 or s.stopped or s.collective.failed) return error.RankDisconnected; + if (payload.len > wire.max_payload or input.bytes != payload.len) return error.BadRankFrame; + errdefer s.poison(); + var h = input; + h.sequence = s.next; + s.next = std.math.add(u64, s.next, 1) catch return error.BadRankFrame; + const until = wire.deadline(30_000); + try s.link.sendHeader(h, until); + try s.link.send(payload, until); + return h; + } + + fn ack(s: *Session, sent: wire.Header) !void { + errdefer s.poison(); + const h = try s.link.receiveHeader(wire.deadline(30_000)); + if (h.kind != .ack or h.sequence != sent.sequence or h.object != sent.object or h.bytes != 0 or h.arg != 0) return error.BadRankFrame; + if (h.flags == 1) return error.OutOfDeviceMemory; + if (h.flags != 0) return error.RemoteRank; + } + + pub fn create(s: *Session, cap: usize) !prompt.Seq { + var q = try s.engine.newSeq(cap); + errdefer q.deinit(s.gpa); + const id = s.next_object; + s.next_object = std.math.add(u64, id, 1) catch return error.BadRankFrame; + const h = try s.command(.{ .kind = .create, .object = id, .arg = cap }, &.{}); + try s.ack(h); + q.rank_id = id; + return q; + } + + pub fn drop(s: *Session, q: *prompt.Seq) void { + s.engine.round_graph.clear(); + if (!s.stopped) { + const h = s.command(.{ .kind = .drop_seq, .object = q.rank_id }, &.{}) catch { + s.poison(); + q.deinit(s.gpa); + return; + }; + s.ack(h) catch s.poison(); + } + q.deinit(s.gpa); + } + + pub fn forward(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) !void { + if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.logits == .all) return error.UnsupportedRankOperation; + try tokenIds(ids); + const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); + errdefer s.poison(); + try s.engine.forward(Tri, tri, q, ids, hooks); + try s.ack(h); + } + + pub fn round(s: *Session, comptime Tri: type, tri: Tri, seqs: []const *prompt.Seq, ids: []const u32, out: u64) !void { + if (seqs.len == 0 or seqs.len != ids.len or seqs.len > 32) return error.UnsupportedRankOperation; + try tokenIds(ids); + var raw: [32 * 12]u8 = undefined; + for (seqs, ids, 0..) |q, token, i| { + if (q.rank_id == 0 or q.fed >= q.cap) return error.UnexpectedTensor; + for (seqs[0..i]) |prior| if (prior.rank_id == q.rank_id) return error.DuplicateRankSequence; + std.mem.writeInt(u64, raw[i * 12 ..][0..8], q.rank_id, .little); + std.mem.writeInt(u32, raw[i * 12 + 8 ..][0..4], token, .little); + } + const bytes = raw[0 .. seqs.len * 12]; + const h = try s.command(.{ .kind = .round, .bytes = @intCast(bytes.len) }, bytes); + errdefer s.poison(); + try s.engine.round(Tri, tri, seqs, ids, out); + try s.ack(h); + } + + /// Both ranks keep their own recurrent window; only the leader consumes draft scores and hidden rows. + pub fn verify(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) !void { + if (!s.drafting or ids.len == 0 or ids.len > mtp.max_rows or q.fed + ids.len > q.cap) return error.UnsupportedRankOperation; + try tokenIds(ids); + const h = try s.command(.{ .kind = .verify, .object = q.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); + errdefer s.poison(); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verify(q, d, ids); + try s.ack(h); + } + + pub fn keep(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, count: usize) !void { + if (!s.drafting or count == 0 or count > d.rows) return error.CannotRewind; + const h = try s.command(.{ .kind = .keep, .object = q.rank_id, .arg = count }, &.{}); + errdefer s.poison(); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.rollback(q, d, count); + try s.collective.finish(s.engine.stream.*); + try s.ack(h); + } + + fn window(s: *Session, q: *prompt.Seq) !*mtp.Draft { + const item = try s.windows.getOrPut(s.gpa, q.rank_id); + if (!item.found_existing) { + errdefer _ = s.windows.remove(q.rank_id); + item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk); + } + return item.value_ptr; + } + + pub fn save(s: *Session, q: *prompt.Seq) !*snaps.Snap { + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); + kept.rank_id = s.next_object; + s.next_object = std.math.add(u64, s.next_object, 1) catch return error.BadRankFrame; + const h = try s.command(.{ .kind = .save, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + try s.ack(h); + return kept; + } + + pub fn restore(s: *Session, q: *prompt.Seq, kept: *const snaps.Snap) !void { + if (kept.at > q.cap) return error.SequenceTooShort; + const h = try s.command(.{ .kind = .restore, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + errdefer s.poison(); + try snaps.restore(s.engine.stream.*, q, null, kept); + try s.ack(h); + } + + pub fn dropSaved(s: *Session, kept: *snaps.Snap) void { + if (!s.stopped) { + const h = s.command(.{ .kind = .drop_snap, .object = kept.rank_id }, &.{}) catch { + s.poison(); + snaps.drop(s.gpa, s.engine.stream.*, kept); + return; + }; + s.ack(h) catch s.poison(); + } + snaps.drop(s.gpa, s.engine.stream.*, kept); + } + + pub fn follow(s: *Session, comptime Tri: type, tri: Tri, logits: u64) !void { + if (s.rank != 1) return error.BadRank; + errdefer s.poison(); + const payload = try s.gpa.alloc(u8, wire.max_payload); + defer s.gpa.free(payload); + while (true) { + const h = try s.link.receiveHeader(null); + if (h.sequence != s.next) return error.BadRankFrame; + s.next = std.math.add(u64, s.next, 1) catch return error.BadRankFrame; + try s.link.receive(payload[0..h.bytes], wire.deadline(30_000)); + const stop = s.apply(Tri, tri, logits, h, payload[0..h.bytes]) catch |err| { + s.link.sendHeader(.{ .kind = .ack, .sequence = h.sequence, .object = h.object, .flags = 2 }, wire.deadline(30_000)) catch {}; + return err; + }; + try s.link.sendHeader(.{ .kind = .ack, .sequence = h.sequence, .object = h.object }, wire.deadline(30_000)); + if (stop) { + s.stopped = true; + s.available.store(false, .release); + return; + } + } + } + + fn apply(s: *Session, comptime Tri: type, tri: Tri, logits: u64, h: wire.Header, payload: []const u8) !bool { + try shape(h); + switch (h.kind) { + .create => { + if (s.seqs.contains(h.object)) return error.BadRankFrame; + const q = try s.gpa.create(prompt.Seq); + errdefer s.gpa.destroy(q); + q.* = try s.engine.newSeq(@intCast(h.arg)); + errdefer q.deinit(s.gpa); + q.rank_id = h.object; + try s.seqs.put(s.gpa, h.object, q); + }, + .forward => { + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const ids = try s.gpa.alloc(u32, payload.len / 4); + defer s.gpa.free(ids); + for (ids, 0..) |*v, i| v.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); + try tokenIds(ids); + try s.engine.forward(Tri, tri, q, ids, .{ .logits = @fromBackingInt(@intCast(h.flags)) }); + }, + .round => { + var seqs: [32]*prompt.Seq = undefined; + var ids: [32]u32 = undefined; + const n = payload.len / 12; + try roundPayload(payload); + for (0..n) |i| { + seqs[i] = s.seqs.get(std.mem.readInt(u64, payload[i * 12 ..][0..8], .little)) orelse return error.UnknownRankSequence; + ids[i] = std.mem.readInt(u32, payload[i * 12 + 8 ..][0..4], .little); + if (seqs[i].fed >= seqs[i].cap) return error.UnexpectedTensor; + } + try s.engine.round(Tri, tri, seqs[0..n], ids[0..n], logits); + }, + .verify => { + if (!s.drafting) return error.UnsupportedRankOperation; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + var ids: [mtp.max_rows]u32 = undefined; + const n = payload.len / 4; + for (ids[0..n], 0..) |*id, i| id.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); + try tokenIds(ids[0..n]); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verify(q, try s.window(q), ids[0..n]); + }, + .keep => { + if (!s.drafting) return error.UnsupportedRankOperation; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const d = s.windows.getPtr(h.object) orelse return error.CannotRewind; + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.rollback(q, d, @intCast(h.arg)); + try s.collective.finish(s.engine.stream.*); + }, + .drop_seq => { + s.engine.round_graph.clear(); + if (s.windows.fetchRemove(h.object)) |entry| { + var d = entry.value; + d.deinit(); + } + const entry = s.seqs.fetchRemove(h.object) orelse return error.UnknownRankSequence; + entry.value.deinit(s.gpa); + s.gpa.destroy(entry.value); + }, + .save => { + if (s.saved.contains(h.arg)) return error.BadRankFrame; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); + try s.saved.put(s.gpa, h.arg, kept); + }, + .restore => { + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; + try snaps.restore(s.engine.stream.*, q, null, kept); + }, + .drop_snap => { + const entry = s.saved.fetchRemove(h.object) orelse return error.UnknownRankSnapshot; + snaps.drop(s.gpa, s.engine.stream.*, entry.value); + }, + .stop => return true, + else => return error.BadRankFrame, + } + return false; + } +}; + +fn tokenIds(ids: []const u32) !void { + for (ids) |id| if (id >= wt.vocab) return error.InvalidToken; +} + +fn roundPayload(payload: []const u8) !void { + if (payload.len == 0 or payload.len % 12 != 0 or payload.len / 12 > 32) return error.BadRankFrame; + var seen: [32]u64 = undefined; + for (0..payload.len / 12) |i| { + const id = std.mem.readInt(u64, payload[i * 12 ..][0..8], .little); + const token = std.mem.readInt(u32, payload[i * 12 + 8 ..][0..4], .little); + if (id == 0) return error.BadRankFrame; + for (seen[0..i]) |prior| if (prior == id) return error.DuplicateRankSequence; + try tokenIds(&.{token}); + seen[i] = id; + } +} + +fn sameIdentity(a: Identity, b: Identity) bool { + return std.mem.eql(u8, std.mem.asBytes(&a), std.mem.asBytes(&b)); +} + +/// Refuse malformed operations before looking up an owner or touching the GPU. +pub fn shape(h: wire.Header) !void { + switch (h.kind) { + .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, + .verify => if (h.object == 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 4 != 0 or h.bytes / 4 > mtp.max_rows) return error.BadRankFrame, + .keep => if (h.object == 0 or h.arg == 0 or h.arg > mtp.max_rows or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, + .save, .restore => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .drop_seq, .drop_snap => if (h.object == 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .stop => if (h.object != 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + else => return error.BadRankFrame, + } +} + +test "rank identity binds checkpoint, config, index, program and execution settings" { + const base: Identity = .{ .config = @splat(1), .index = @splat(2), .program = @splat(3), .checkpoint = @splat(4), .context = 4096, .chunk = 64 }; + try std.testing.expect(sameIdentity(base, base)); + inline for (.{ "config", "index", "program", "checkpoint" }) |name| { + var other = base; + @field(other, name)[0] ^= 1; + try std.testing.expect(!sameIdentity(base, other)); + } + var other = base; + other.context += 1; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.chunk += 1; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.decode_graph = 1; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.mtp = 1; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.head_sharded = 1; + try std.testing.expect(!sameIdentity(base, other)); +} + +test "rank operations refuse invalid geometry and unsupported hooks before execution" { + try shape(.{ .kind = .create, .object = 1, .arg = 4096 }); + try shape(.{ .kind = .forward, .object = 1, .bytes = 4 }); + try shape(.{ .kind = .round, .bytes = 24 }); + try shape(.{ .kind = .verify, .object = 1, .bytes = mtp.max_rows * 4 }); + try shape(.{ .kind = .keep, .object = 1, .arg = mtp.max_rows }); + for ([_]wire.Header{ + .{ .kind = .verify, .object = 1, .bytes = (mtp.max_rows + 1) * 4 }, + .{ .kind = .verify, .object = 1, .bytes = 3 }, + .{ .kind = .verify, .bytes = 4 }, + .{ .kind = .verify, .object = 1, .bytes = 4, .flags = 1 }, + .{ .kind = .keep, .object = 1 }, + .{ .kind = .keep, .object = 1, .arg = mtp.max_rows + 1 }, + .{ .kind = .keep, .object = 1, .arg = 1, .bytes = 4 }, + }) |bad| try std.testing.expectError(error.BadRankFrame, shape(bad)); + for ([_]wire.Header{ + .{ .kind = .hello }, .{ .kind = .create, .arg = 4096 }, .{ .kind = .create, .object = 1, .arg = prompt.max_tokens + 1 }, + .{ .kind = .forward, .object = 1, .bytes = 3 }, .{ .kind = .forward, .object = 1, .bytes = 4, .flags = 1 }, .{ .kind = .round, .bytes = 33 * 12 }, + .{ .kind = .round, .bytes = 13 }, .{ .kind = .stop, .object = 1 }, .{ .kind = .save, .object = 1 }, + .{ .kind = .drop_seq, .object = 1, .bytes = 4 }, + }) |h| try std.testing.expectError(error.BadRankFrame, shape(h)); +} + +test "round payload rejects duplicate owners and out-of-vocabulary tokens before GPU execution" { + var bytes: [24]u8 = @splat(0); + std.mem.writeInt(u64, bytes[0..8], 9, .little); + std.mem.writeInt(u64, bytes[12..20], 10, .little); + try roundPayload(&bytes); + std.mem.writeInt(u64, bytes[12..20], 9, .little); + try std.testing.expectError(error.DuplicateRankSequence, roundPayload(&bytes)); + std.mem.writeInt(u64, bytes[12..20], 10, .little); + std.mem.writeInt(u32, bytes[20..24], wt.vocab, .little); + try std.testing.expectError(error.InvalidToken, roundPayload(&bytes)); + try std.testing.expectError(error.InvalidToken, tokenIds(&.{ 0, wt.vocab })); + try std.testing.expectError(error.BadRankFrame, roundPayload(bytes[0..13])); +} diff --git a/zig/src/families/flashnext/cuda_rank_memory.zig b/zig/src/families/flashnext/cuda_rank_memory.zig new file mode 100644 index 0000000000..642510a680 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_memory.zig @@ -0,0 +1,78 @@ +//! Conservative rank-local checkpoint bytes: replicated dense operators, partitioned routed weights, one shared branch. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const Entry = core.safetensors.Entry; + +pub fn estimate(a: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + var store = try wt.Store.open(a, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(a); + var wi = wm.object.iterator(); + while (wi.next()) |item| try files.put(a, item.value_ptr.string, {}); + var total: u64 = 0; + var fi = files.keyIterator(); + while (fi.next()) |name| { + const path = try std.fs.path.join(a, &.{ dir, name.* }); + defer a.free(path); + var file = try core.safetensors.File.open(a, io, path); + defer file.close(io); + var tensors = file.names.iterator(); + while (tensors.next()) |tensor| { + total = try std.math.add(u64, total, try tensorBytes(tensor.key_ptr.*, tensor.value_ptr.*, world, rank)); + } + } + return total; +} + +fn tensorBytes(name: []const u8, e: Entry, world: usize, rank: usize) !u64 { + const bytes: u64 = e.end - e.begin; + if (wt.hostOnly(name)) return 0; + if (rank != 0 and std.mem.indexOf(u8, name, ".mlp.shared_expert.") != null) return 0; + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { + if (std.mem.endsWith(u8, name, ".weight_scale_2")) return bytes; + if (!std.mem.endsWith(u8, name, ".weight")) return 0; + if (e.dtype != .u8 or e.rank != 2 or bytes % (8 * world) != 0) return error.UnexpectedTensor; + // NVFP4 device blocks hold 128 code words and 16 scale words; global scales stay replicated. + return try std.math.mul(u64, bytes / (8 * world), 9); + } + if (std.mem.indexOf(u8, name, ".weight_scale") != null or std.mem.endsWith(u8, name, ".input_scale")) return 0; + if (e.dtype == .f8_e4m3 and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + const codes = try std.math.mul(u64, cuda.mxfp8.padded(e.dim(0)), e.dim(1)); + return try std.math.add(u64, codes, codes / 32); + } + if (e.dtype == .bf16 and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + if (std.mem.indexOf(u8, name, ".mlp.gate.") != null or std.mem.indexOf(u8, name, ".shared_expert_gate.") != null) return bytes; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(e.dim(0)), e.dim(1)), 2); + } + if (e.dtype == .bf16 and centered(name)) return try std.math.mul(u64, bytes, 2); + return bytes; +} + +fn centered(name: []const u8) bool { + for ([_][]const u8{ ".hc_norm.weight", ".q_norm.weight", ".k_norm.weight", ".q_layernorm.weight", ".k_layernorm.weight", ".norm_key.weight", ".norm_query.weight", ".norm_conv.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return true; + } + return false; +} + +test "rank estimates partition routed storage while retaining replicated and rank-zero shared operators" { + const e: Entry = .{ .dtype = .u8, .rank = 2, .shape = .{ 640, 1280, 1, 1, 1 }, .begin = 0, .end = 640 * 1280 }; + const name = "model.language_model.layers.0.mlp.experts.0.gate_proj.weight"; + try std.testing.expectEqual(@as(u64, cuda.experts.Nvfp4Experts.bytes(1, 320, 2560, 1)), try tensorBytes(name, e, 2, 0)); + try std.testing.expectEqual(try tensorBytes(name, e, 2, 0), try tensorBytes(name, e, 2, 1)); + var shared = e; + shared.dtype = .f8_e4m3; + shared.shape = .{ 640, 2560, 1, 1, 1 }; + shared.end = 640 * 2560; + const shared_name = "model.language_model.layers.0.mlp.shared_expert.gate_proj.weight"; + try std.testing.expectEqual(@as(u64, 640 * 2560 * 33 / 32), try tensorBytes(shared_name, shared, 2, 0)); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared_name, shared, 2, 1)); + const norm: Entry = .{ .dtype = .bf16, .rank = 1, .shape = .{ 256, 1, 1, 1, 1 }, .begin = 0, .end = 512 }; + try std.testing.expectEqual(@as(u64, 1024), try tensorBytes("model.language_model.layers.3.self_attn.q_norm.weight", norm, 2, 1)); + try std.testing.expectEqual(@as(u64, 512), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", norm, 2, 1)); +} diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig new file mode 100644 index 0000000000..afad725e94 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -0,0 +1,218 @@ +//! Bounded Linux TCP control channel. Tensor bytes travel over NCCL; this channel owns command order and peer EOF. +const std = @import("std"); +const builtin = @import("builtin"); +const p = std.posix; +const net = std.Io.net; + +pub const max_payload = 1 << 20; +pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep }; +pub const Header = struct { + kind: Kind, + sequence: u64 = 0, + object: u64 = 0, + arg: u64 = 0, + bytes: u32 = 0, + flags: u32 = 0, + + pub fn encode(h: Header) ![40]u8 { + if (h.bytes > max_payload) return error.BadRankFrame; + var out: [40]u8 = @splat(0); + @memcpy(out[0..4], "TFR1"); + std.mem.writeInt(u32, out[4..8], @backingInt(h.kind), .little); + std.mem.writeInt(u64, out[8..16], h.sequence, .little); + std.mem.writeInt(u64, out[16..24], h.object, .little); + std.mem.writeInt(u64, out[24..32], h.arg, .little); + std.mem.writeInt(u32, out[32..36], h.bytes, .little); + std.mem.writeInt(u32, out[36..40], h.flags, .little); + return out; + } + + pub fn decode(raw: *const [40]u8) !Header { + if (!std.mem.eql(u8, raw[0..4], "TFR1")) return error.BadRankFrame; + const h: Header = .{ + .kind = std.enums.fromInt(Kind, std.mem.readInt(u32, raw[4..8], .little)) orelse return error.BadRankFrame, + .sequence = std.mem.readInt(u64, raw[8..16], .little), + .object = std.mem.readInt(u64, raw[16..24], .little), + .arg = std.mem.readInt(u64, raw[24..32], .little), + .bytes = std.mem.readInt(u32, raw[32..36], .little), + .flags = std.mem.readInt(u32, raw[36..40], .little), + }; + if (h.bytes > max_payload) return error.BadRankFrame; + return h; + } +}; + +pub fn now() u64 { + var ts: std.c.timespec = undefined; + if (std.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0) return 0; + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_ms * 1000 + @as(u64, @intCast(ts.nsec)); +} + +pub fn deadline(ms: u64) u64 { + return now() +| (ms *| std.time.ns_per_ms); +} + +pub fn pause() void { + const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = std.c.nanosleep(&ts, null); +} + +fn wait(fd: p.socket_t, events: i16, until: ?u64) !void { + while (true) { + const current = now(); + if (until) |end| if (current >= end) return error.RankTimeout; + const ms: i32 = if (until) |end| @intCast(@min(1000, (end - current + std.time.ns_per_ms - 1) / std.time.ns_per_ms)) else 1000; + var fds = [_]p.pollfd{.{ .fd = fd, .events = events, .revents = 0 }}; + const count = try p.poll(&fds, ms); + if (count == 0) continue; + if (fds[0].revents & events != 0) return; + if (fds[0].revents & (p.POLL.ERR | p.POLL.HUP | p.POLL.NVAL) != 0) return error.RankDisconnected; + } +} + +pub const Link = struct { + fd: p.socket_t, + + pub fn close(l: Link) void { + _ = std.c.shutdown(l.fd, p.SHUT.RDWR); + _ = std.c.close(l.fd); + } + + /// Exactly two rank processes. Retry only connection establishment, never an accepted command. + pub fn open(io: std.Io, host: []const u8, port: u16, rank: usize, timeout_ms: u64) !Link { + if (builtin.os.tag != .linux) return error.RankPlatform; + if (rank > 1 or port == 0) return error.BadRank; + const address = net.IpAddress.parse(host, port) catch try net.IpAddress.resolve(io, host, port); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&address, &storage); + const family: u32 = if (address == .ip4) p.AF.INET else p.AF.INET6; + const until = deadline(timeout_ms); + if (rank == 0) { + const fd = try socket(family); + defer _ = std.c.close(fd); + const one: c_int = 1; + try p.setsockopt(fd, p.SOL.SOCKET, p.SO.REUSEADDR, std.mem.asBytes(&one)); + if (std.c.bind(fd, &storage.any, len) != 0) return error.RankBind; + if (std.c.listen(fd, 1) != 0) return error.RankListen; + while (true) { + try wait(fd, p.POLL.IN, until); + const peer = std.c.accept4(fd, null, null, p.SOCK.NONBLOCK | p.SOCK.CLOEXEC); + switch (p.errno(peer)) { + .SUCCESS => { + errdefer _ = std.c.close(peer); + try noDelay(peer); + return .{ .fd = peer }; + }, + .INTR, .AGAIN => continue, + else => return error.RankAccept, + } + } + } + while (now() < until) { + const fd = try socket(family); + const rc = std.c.connect(fd, &storage.any, len); + if (rc == 0) return .{ .fd = fd }; + const e = p.errno(rc); + if (e == .INPROGRESS or e == .INTR) { + wait(fd, p.POLL.OUT, until) catch |err| { + _ = std.c.close(fd); + if (err == error.RankTimeout) return err; + pause(); + continue; + }; + var status: c_int = 0; + var size: p.socklen_t = @sizeOf(c_int); + if (std.c.getsockopt(fd, p.SOL.SOCKET, p.SO.ERROR, &status, &size) == 0 and status == 0) return .{ .fd = fd }; + } + _ = std.c.close(fd); + pause(); + } + return error.RankTimeout; + } + + pub fn send(l: Link, bytes: []const u8, until: u64) !void { + var at: usize = 0; + while (at < bytes.len) { + if (now() >= until) return error.RankTimeout; + const rc = std.c.send(l.fd, bytes.ptr + at, bytes.len - at, p.MSG.NOSIGNAL); + switch (p.errno(rc)) { + .SUCCESS => { + if (rc == 0) return error.RankDisconnected; + at += @intCast(rc); + }, + .INTR => continue, + .AGAIN => try wait(l.fd, p.POLL.OUT, until), + else => return error.RankDisconnected, + } + } + } + + pub fn receive(l: Link, bytes: []u8, until: ?u64) !void { + var at: usize = 0; + while (at < bytes.len) { + if (until) |end| if (now() >= end) return error.RankTimeout; + const rc = std.c.recv(l.fd, bytes.ptr + at, bytes.len - at, 0); + switch (p.errno(rc)) { + .SUCCESS => { + if (rc == 0) return error.RankDisconnected; + at += @intCast(rc); + }, + .INTR => continue, + .AGAIN => try wait(l.fd, p.POLL.IN, until), + else => return error.RankDisconnected, + } + } + } + + pub fn sendHeader(l: Link, h: Header, until: u64) !void { + const raw = try h.encode(); + try l.send(&raw, until); + } + + pub fn receiveHeader(l: Link, until: ?u64) !Header { + var raw: [40]u8 = undefined; + // An idle follower can wait indefinitely; a partial frame cannot. + try l.receive(raw[0..1], until); + try l.receive(raw[1..], until orelse deadline(30_000)); + return Header.decode(&raw); + } + + /// A non-consuming EOF check while a GPU collective progresses. A future control frame stays queued. + pub fn checkPeer(l: Link) !void { + var byte: u8 = 0; + const rc = std.c.recv(l.fd, &byte, 1, p.MSG.PEEK | p.MSG.DONTWAIT); + if (rc == 0) return error.RankDisconnected; + if (rc < 0) switch (p.errno(rc)) { + .AGAIN, .INTR => {}, + else => return error.RankDisconnected, + }; + } +}; + +fn socket(family: u32) !p.socket_t { + const fd = std.c.socket(family, p.SOCK.STREAM | p.SOCK.NONBLOCK | p.SOCK.CLOEXEC, 0); + if (p.errno(fd) != .SUCCESS) return error.RankSocket; + errdefer _ = std.c.close(fd); + try noDelay(fd); + return fd; +} + +fn noDelay(fd: p.socket_t) !void { + const one: c_int = 1; + try p.setsockopt(fd, p.IPPROTO.TCP, p.TCP.NODELAY, std.mem.asBytes(&one)); +} + +test "rank frames are versioned, bounded and exact" { + const h: Header = .{ .kind = .forward, .sequence = 700, .object = 19, .arg = 42, .bytes = max_payload, .flags = 3 }; + var raw = try h.encode(); + try std.testing.expectEqualDeep(h, try Header.decode(&raw)); + raw[0] = 0; + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + raw = try h.encode(); + @memset(raw[4..8], 0xff); + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + raw = try h.encode(); + @memset(raw[32..36], 0xff); + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + try std.testing.expectError(error.BadRankFrame, (Header{ .kind = .forward, .bytes = max_payload + 1 }).encode()); +} diff --git a/zig/src/families/flashnext/cuda_round_graph.zig b/zig/src/families/flashnext/cuda_round_graph.zig new file mode 100644 index 0000000000..e367e31dd9 --- /dev/null +++ b/zig/src/families/flashnext/cuda_round_graph.zig @@ -0,0 +1,98 @@ +//! A single bounded decode graph, recaptured when sequence owners, scratch or attention geometry change. +const std = @import("std"); +const cuda = @import("cuda"); +const dims = 2560; + +pub fn enabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_DECODE_GRAPH") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub const Key = struct { + count: usize, + owners: [32]u64 = @splat(0), + chunks: [32]u32 = @splat(0), + score_tiles: [32]u32 = @splat(0), + scratch: u64, + scores: u64, + output: u64, + + pub fn init(seqs: anytype, scratch: u64, scores: u64, output: u64) !Key { + if (seqs.len == 0 or seqs.len > 32) return error.UnexpectedTensor; + var key: Key = .{ .count = seqs.len, .scratch = scratch, .scores = scores, .output = output }; + for (seqs, 0..) |q, i| { + if (q.rank_id == 0) return error.UnsupportedRankOperation; + key.owners[i] = q.rank_id; + const end = q.fed + 1; + key.chunks[i] = @intCast((@min(end, 2051) + 511) / 512); + key.score_tiles[i] = if (end / 4 > 512) @intCast(((end + 3) / 4 + 63) / 64) else 0; + } + return key; + } +}; + +pub const Cache = struct { + exec: ?cuda.graph.Exec = null, + key: ?Key = null, + + pub fn clear(c: *Cache) void { + if (c.exec) |*e| e.deinit(); + c.* = .{}; + } + + /// CPU histories and pinned inputs advance once per round, outside stream capture and replay. + fn prepare(pass: anytype, ids: []const u32) !void { + const s = pass.s; + try s.ready(); + for (s.hostIds(ids.len), ids) |*out, token| out.* = @intCast(token); + const positions = s.hostPos(ids.len * 4); + @memset(positions, 0); + for (pass.lanes, 0..) |lane, i| { + positions[i * 4] = @intCast(lane.pos); + try pass.m.ngram.rows(lane.seq.history[0..], ids[i..][0..1], s.hostEmb(ids.len)[i * dims ..][0..dims], s.hostXs(ids.len)[i * (dims / 32) ..][0 .. dims / 32]); + } + } + + pub fn run(c: *Cache, key: Key, pass: anytype, ids: []const u32) !void { + try prepare(pass, ids); + if (c.key == null or !std.meta.eql(c.key.?, key)) { + c.clear(); + try cuda.graph.beginCapture(pass.stream, .thread_local); + pass.prepared = true; + pass.chunk(ids, 0, null) catch |err| { + if (cuda.graph.endCapture(pass.stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return err; + }; + var graph = try cuda.graph.endCapture(pass.stream); + defer graph.deinit(); + c.exec = try graph.instantiate(); + errdefer c.clear(); + try c.exec.?.upload(pass.stream); + c.key = key; + } + try c.exec.?.launchOn(pass.stream); + } +}; + +test "decode graph keys invalidate on owner order, attention seams and score tile seams" { + const Seq = struct { rank_id: u64, fed: usize }; + var first: Seq = .{ .rank_id = 1, .fed = 10 }; + var second: Seq = .{ .rank_id = 2, .fed = 10 }; + const base = try Key.init(@as([]const *Seq, &.{ &first, &second }), 1, 2, 3); + try std.testing.expect(!std.meta.eql(base, try Key.init(@as([]const *Seq, &.{ &second, &first }), 1, 2, 3))); + first.fed = 511; + const before = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 512; + try std.testing.expect(!std.meta.eql(before, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); + first.fed = 2048; + const dense = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 2051; + try std.testing.expect(!std.meta.eql(dense, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); + first.fed = 2303; + const score = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 2304; + try std.testing.expect(!std.meta.eql(score, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); +} diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig new file mode 100644 index 0000000000..89095b53ee --- /dev/null +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -0,0 +1,224 @@ +//! The row-by-row reference: each prompt row one fused decode step, the batched pass's check in flashnext-prompt. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const qmm = @import("cuda_qmm.zig"); +const weights = @import("cuda_prompt_wt.zig"); +const step = @import("cuda_prompt_step.zig"); +const load = @import("cuda_prompt_load.zig"); +const prompt = @import("cuda_prompt.zig"); + +const dims = weights.dims; +const streams = weights.streams; +const wide = weights.wide; +const vocab = weights.vocab; +const slots = weights.slots; +const ple_layer = weights.ple_layer; +const inj_stride = weights.inj_stride; +const wts_stride = weights.wts_stride; +const Q4 = weights.Q4; +const HcW = weights.HcW; +const Scratch = step.Scratch; +const Layer = load.Layer; +const promote = weights.promote; + +/// The row path's own scratch and expert kernels, made for a check and dropped after it. +pub const Rows = struct { + gpa: std.mem.Allocator, + driver: *cuda.Driver, + scratch: Scratch, + scratch_rows: usize, + mode: []u8, + experts: cuda.Module, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, + + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream) !Rows { + var r: Rows = .{ .gpa = gpa, .driver = driver, .scratch = try Scratch.alloc(driver, 1), .scratch_rows = 1, .mode = undefined, .experts = undefined, .plan_fn = undefined, .up_fn = undefined, .down_fn = undefined }; + errdefer r.scratch.free(); + try r.scratch.dummy.fill8(0, stream.handle); + r.mode = try gpa.alloc(u8, 1); + errdefer gpa.free(r.mode); + r.experts = try cuda.Module.load(driver, cuda.kernels.experts); + errdefer r.experts.unload(); + r.plan_fn = try r.experts.function(weights.plan_symbol); + r.up_fn = try r.experts.function(weights.up_symbol); + r.down_fn = try r.experts.function(weights.down_symbol); + return r; + } + + pub fn deinit(r: *Rows) void { + r.experts.unload(); + r.gpa.free(r.mode); + r.scratch.free(); + } + + fn grow(r: *Rows, stream: cuda.Stream, rows: usize) !void { + if (rows <= r.scratch_rows) return; + var next = try Scratch.alloc(r.driver, rows); + errdefer next.free(); + try next.dummy.fill8(0, stream.handle); + const mode = try r.gpa.alloc(u8, rows); + r.scratch.free(); + r.gpa.free(r.mode); + r.scratch = next; + r.mode = mode; + r.scratch_rows = rows; + } + + /// Every row of `ids` from seq.fed one fused decode step; the last row's scores land in seq.logits. + pub fn predict(r: *Rows, e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq, ids: []const u32) !void { + if (e.table.dense or e.head_q.native != null) return error.UnsupportedQuantization; + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + // dense attention only: the reference stops where the indexer would start choosing + if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; + try r.grow(e.stream.*, ids.len); + var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; + @memset(pass.mode, 0); + for (ids, 0..) |token, t| try embedRow(Tri, &pass, tri, token, pass.s.h.ptr + @as(u64, t) * wide * 2); + for (e.layers[0..e.n_layers], 0..) |*lw, layer| try runLayer(Tri, &pass, tri, lw, layer, ids, seq.fed); + try head(Tri, &pass, tri, ids.len - 1); + seq.fed += ids.len; + } +}; + +const Pass = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + model_dir: []const u8, + s: *Scratch, + table: embed.Table, + mode: []u8, + history: *[2]i64, + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + mixer: HcW, + head_q: Q4, + logits: []u8, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, +}; + +fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { + const id: i32 = @intCast(token); + // on the compute stream (queued _embed launches still read ids), then waited for: `id` lives on this frame + try p.s.ids.uploadAsync(0, std.mem.asBytes(&id), p.stream.handle); + try p.stream.synchronize(); + try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); +} + +fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, ids: []const u32, pos0: usize) !void { + for (ids, 0..) |token, t| { + const h = p.s.h.ptr + @as(u64, t) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; + const y = p.s.y.ptr + @as(u64, t) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, t) * wts_stride; + if (layer == ple_layer and p.mode[t] == 2) { + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } + if (layer == ple_layer) try step.pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, lw.key, lw.val, lw.nk.ptr, lw.nq.ptr, lw.nc.ptr, lw.ple_cw.ptr, &p.tail); + if (p.mode[t] == 2) { + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } else try step.writeback(tri, p.s, h, 0, inj, y, wt); + try step.readout(tri, p.s, lw.attn_hc, h, inj); + if (lw.linear) { + const li = weights.linearIndex(layer); + try step.gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); + } else { + const ai = weights.attnIndex(layer); + try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + } + try step.writeback(tri, p.s, h, 4, inj, y, wt); + try step.readout(tri, p.s, lw.mlp_hc, h, inj); + try step.moeStep(Tri, tri, p.s, p.stream, p.plan_fn, p.up_fn, p.down_fn, lw.router.ptr, lw.experts_up.ptr, lw.experts_down.ptr, y, wt); + p.mode[t] = 2; + } +} + +fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !void { + if (p.mode[row] != 2) return error.UnexpectedTensor; + const h = p.s.h.ptr + @as(u64, row) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; + const y = p.s.y.ptr + @as(u64, row) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, row) * wts_stride; + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[row] = 0; + try step.readout(tri, p.s, p.mixer, h, inj); + try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); + try p.stream.synchronize(); + try p.s.logits.download(0, p.logits); +} + +/// The five best scores of one row, best first, as "id:score" pairs on stderr (the capture's top5 field). +fn printTop(raw: []const u8) void { + var ids: [5]usize = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (0..vocab) |i| { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v <= vals[4]) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = i; + } + std.debug.print("top5", .{}); + for (ids, vals) |id, v| std.debug.print(" {d}:{d:.4}", .{ id, v }); + std.debug.print("\n", .{}); +} + +fn ms(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +/// `out` gets `new_tokens` greedy ids: the prompt's argmax, then one fed back at a time (TF_FLASHNEXT_PROMPT=rows). +pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > prompt.max_tokens) return error.UnexpectedTensor; + var e = try prompt.Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + const by_rows = if (std.c.getenv("TF_FLASHNEXT_PROMPT")) |v| std.mem.eql(u8, std.mem.span(v), "rows") else false; + var rows: ?Rows = if (by_rows) try Rows.init(gpa, driver, stream) else null; + defer if (rows) |*r| r.deinit(); + var seq = try e.newSeq(tokens.len + new_tokens); + defer seq.deinit(gpa); + const loaded = cuda.usage(true); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + if (rows) |*r| try r.predict(e, Tri, tri, &seq, tokens) else try e.predict(Tri, tri, &seq, tokens); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); + out[0] = prompt.argmax(seq.logits); + printTop(seq.logits); + // TF_FLASHNEXT_LOGITS=path: the prompt's last row of scores (bf16), to compare the two paths byte for byte + if (std.c.getenv("TF_FLASHNEXT_LOGITS")) |path| try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(path), .data = seq.logits }); + std.debug.print("token {d}\n", .{out[0]}); + var tok = out[0]; + for (1..new_tokens) |i| { + var one = [_]u32{tok}; + if (rows) |*r| try r.predict(e, Tri, tri, &seq, &one) else try e.predict(Tri, tri, &seq, &one); + tok = prompt.argmax(seq.logits); + if (i < 4) printTop(seq.logits); + out[i] = tok; + std.debug.print("token {d}\n", .{tok}); + } + const t2 = std.Io.Clock.awake.now(io).toNanoseconds(); + const per = if (new_tokens > 1) ms(t2 - t1) / @as(f64, @floatFromInt(new_tokens - 1)) else 0; + std.debug.print("timing prompt {d} rows {d:.1} ms, decode {d} tokens {d:.2} ms/token\n", .{ tokens.len, ms(t1 - t0), new_tokens - 1, per }); + const used = cuda.usage(false); + std.debug.print("memory: {d:.2} GiB on the device after the load and the sequence, peak {d:.2} GiB during the prompt and decode (+{d:.2} GiB)\n", .{ gib(loaded.device), gib(used.peak), gib(used.peak - loaded.device) }); +} diff --git a/zig/src/families/flashnext/cuda_select_main.zig b/zig/src/families/flashnext/cuda_select_main.zig new file mode 100644 index 0000000000..80b69a4d6e --- /dev/null +++ b/zig/src/families/flashnext/cuda_select_main.zig @@ -0,0 +1,113 @@ +//! fn_qsa_select (flashnext_batch.cu) against the lists Triton's _select and _select_tiles made (capture_sparse.sh). + +const std = @import("std"); +const cuda = @import("cuda"); + +const Case = struct { name: []const u8, rows: usize, pos0: i32, nb: usize, idw: usize }; + +fn readAll(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, ext: []const u8, len: usize) ![]u8 { + const path = try std.fmt.allocPrint(gpa, "{s}/{s}.{s}", .{ dir, name, ext }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 30)); + if (bytes.len != len) return error.BadFixture; + return bytes; +} + +fn int(bytes: []const u8, i: usize) i32 { + return std.mem.readInt(i32, bytes[4 * i ..][0..4], .little); +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: flashnext-select \n", .{}); + return 2; + } + const dir = args[1]; + const index = try std.fmt.allocPrint(gpa, "{s}/cases.json", .{dir}); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 20)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice([]Case, gpa, text, .{ .ignore_unknown_fields = true }); + defer parsed.deinit(); + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + var module = try cuda.Module.load(&driver, cuda.kernels.flashnext_batch); + defer module.unload(); + const select = try module.function("fn_qsa_select"); + + var failed = false; + for (parsed.value) |c| { + const scores = try readAll(gpa, io, dir, c.name, "scores", c.rows * c.nb * 4); + defer gpa.free(scores); + const want_ids = try readAll(gpa, io, dir, c.name, "ids", c.rows * c.idw * 4); + defer gpa.free(want_ids); + const want_nk = try readAll(gpa, io, dir, c.name, "nk", c.rows * 4); + defer gpa.free(want_nk); + const want_sp = try readAll(gpa, io, dir, c.name, "sparse", c.rows * 4); + defer gpa.free(want_sp); + + var sc_b = try cuda.DeviceBuffer.fromHost(&driver, scores); + defer sc_b.free(); + var pos_b = try cuda.DeviceBuffer.fromHost(&driver, std.mem.asBytes(&c.pos0)); + defer pos_b.free(); + var ids_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * c.idw * 4); + defer ids_b.free(); + try ids_b.fill8(0xFF, null); + var nk_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * 4); + defer nk_b.free(); + var sp_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * 4); + defer sp_b.free(); + var a: cuda.Args = .{}; + a.add(sc_b.ptr); + a.add(pos_b.ptr); + a.add(ids_b.ptr); + a.add(nk_b.ptr); + a.add(sp_b.ptr); + a.add(@as(c_int, @intCast(c.nb))); + a.add(@as(c_int, @intCast(c.idw))); + a.add(@as(c_int, 512)); + a.add(@as(c_int, 4)); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try cuda.launch.launch(select, .{ .grid = .{ .x = @intCast(c.rows) }, .block = .{ .x = 1024 } }, stream, &a); + try stream.synchronize(); + const us = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e3; + const got_ids = try gpa.alloc(u8, c.rows * c.idw * 4); + defer gpa.free(got_ids); + const got_nk = try gpa.alloc(u8, c.rows * 4); + defer gpa.free(got_nk); + const got_sp = try gpa.alloc(u8, c.rows * 4); + defer gpa.free(got_sp); + try ids_b.download(0, got_ids); + try nk_b.download(0, got_nk); + try sp_b.download(0, got_sp); + var bad: usize = 0; + var sparse_rows: usize = 0; + for (0..c.rows) |r| { + const nk = int(want_nk, r); + if (int(got_nk, r) != nk or int(got_sp, r) != int(want_sp, r)) { + bad += 1; + continue; + } + if (int(want_sp, r) == 0) continue; + sparse_rows += 1; + for (0..@intCast(nk)) |i| { + if (int(got_ids, r * c.idw + i) != int(want_ids, r * c.idw + i)) { + bad += 1; + break; + } + } + } + std.debug.print("{s}: {d} rows at {d}, {d} blocks a row, {d} sparse: {d} differ ({d:.0} us)\n", .{ c.name, c.rows, c.pos0, c.nb, sparse_rows, bad, us }); + if (bad != 0) failed = true; + } + std.debug.print("{s}\n", .{if (failed) "differ" else "match"}); + return if (failed) 1 else 0; +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig new file mode 100644 index 0000000000..0e551809b3 --- /dev/null +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -0,0 +1,112 @@ +//! Flash Next prompt states for the shared prompt cache: one stream's caches after `at` tokens, copied on the device. + +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const Draft = @import("cuda_mtp_draft.zig").Draft; + +const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; +const state_bytes = wt.state_n * 4; +const index_row = wt.index_dim * 2; +const tail_bytes = wt.ple_tail_n * wt.wide * 2; + +/// One attention cache's first `at` key, value and indexer-key rows, and the pooled blocks they complete. +fn attnBytes(at: usize) usize { + return at * (2 * prompt.kv_row_bytes + index_row) + at / wt.index_ratio * index_row; +} + +/// Device bytes a state after `at` tokens takes: the DeltaNet's are fixed, each attention cache grows by the token. +pub fn bytes(at: usize, head: bool) usize { + return wt.linear_n * (conv_bytes + state_bytes) + (wt.attn_n + @intFromBool(head)) * attnBytes(at) + tail_bytes; +} + +/// One kept state; `head`: it holds the MTP head's cache over the same rows, which a drafting stream needs. +pub const Snap = struct { + rank_id: u64 = 0, + buf: cuda.DeviceBuffer, + at: usize, + history: [2]i64, + head: bool, +}; + +const Way = enum { save, restore }; + +const Piece = struct { bufs: []cuda.DeviceBuffer, n: usize }; + +/// An attention cache's pieces for `at` rows: complete pooled blocks only (the rest pool again as later rows land). +fn attnPieces(kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, at: usize) [4]Piece { + return .{ + .{ .bufs = kc, .n = at * prompt.kv_row_bytes }, + .{ .bufs = vc, .n = at * prompt.kv_row_bytes }, + .{ .bufs = ikc, .n = at * index_row }, + .{ .bufs = pooled, .n = at / wt.index_ratio * index_row }, + }; +} + +/// Each piece of the sequence (and the head's cache, with `d`) to or from its place in the buffer at `base`. +fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, way: Way) !void { + const dr = stream.d; + var off: u64 = 0; + var pieces: [11]Piece = undefined; + pieces[0] = .{ .bufs = seq.conv, .n = conv_bytes }; + pieces[1] = .{ .bufs = seq.state, .n = state_bytes }; + pieces[2..6].* = attnPieces(seq.kc, seq.vc, seq.ikc, seq.pooled, at); + pieces[6] = .{ .bufs = (&seq.tail)[0..1], .n = tail_bytes }; + var count: usize = 7; + if (d) |h| { + pieces[7..11].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); + count = 11; + } + for (pieces[0..count]) |p| for (p.bufs) |b| { + if (p.n == 0) continue; + const dst, const src = if (way == .save) .{ base + off, b.ptr } else .{ b.ptr, base + off }; + try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, p.n, stream.handle), "cuMemcpyDtoDAsync"); + off += p.n; + }; + std.debug.assert(off == bytes(at, d != null)); +} + +/// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. +pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft) !*Snap { + const at = seq.fed; + if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; + const snap = try gpa.create(Snap); + errdefer gpa.destroy(snap); + snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; + errdefer snap.buf.free(); + try move(stream, seq, d, snap.buf.ptr, at, .save); + try stream.synchronize(); + return snap; +} + +/// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. +pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { + if (snap.at > seq.cap) return error.SequenceTooShort; + if (d) |h| if (!snap.head or snap.at > h.cap) return error.NoHeadState; + // a state with the head's rows serves a stream without drafts too: those rows stay in the buffer + try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); + try stream.synchronize(); + seq.history = snap.history; + seq.fed = snap.at; + if (d) |h| { + h.reset(); + h.len = snap.at; + } +} + +pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { + stream.synchronize() catch {}; + snap.buf.free(); + gpa.destroy(snap); +} + +test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { + try std.testing.expectEqual(@as(usize, 12 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, false) - bytes(0, false)); + // a complete block of 4 adds its pooled row + try std.testing.expectEqual(@as(usize, 12 * (4 * (2 * 2 * 256 * 2 + 128 * 2) + 128 * 2)), bytes(4, false) - bytes(0, false)); + // the head's cache is one more attention layer's rows + try std.testing.expectEqual(@as(usize, 13 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, true) - bytes(0, true)); + try std.testing.expect(bytes(0, false) > 100 << 20 and bytes(0, false) < 120 << 20); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..d425d0d166 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -4,6 +4,9 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const flashnext = @import("flashnext"); +const family_memory = @import("cuda_family_memory.zig"); +pub const follow = @import("cuda_follow.zig").follow; const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); const Allocator = std.mem.Allocator; @@ -15,7 +18,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, flashnext.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; @@ -37,25 +40,8 @@ pub const deviceName = cuda_device.deviceName; const getenv = cuda_device.getenv; const setting = cuda_device.setting; -/// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. -fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { - const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; - const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; - const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; - if (doc != .object) return 0; - const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; - const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; - return if (limit == .integer and limit.integer > 0) limit.integer else 0; -} - -/// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. -fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { - const r = requested orelse return if (native > 0) @min(default, native) else default; - if (r < 0) return error.Negative; - if (r == 0) return if (native > 0) native else error.NoNative; - if (native > 0 and r > native) return error.PastNative; - return r; -} +const modelContext = @import("cuda_context.zig").modelContext; +const contextWindow = @import("cuda_context.zig").contextWindow; /// The pool now, read on the thread whose context is current; `problem` names a bad variable. fn pool(a: Allocator, io: std.Io, ctx: *const cuda.Context, problem: *[]const u8) !?budget.Pool { @@ -79,11 +65,6 @@ fn pool(a: Allocator, io: std.Io, ctx: *const cuda.Context, problem: *[]const u8 return .{ .free = free, .total = total, .reserve = reserve, .limit = limit, .unified = unified }; } -fn readMemory(_: ?*anyopaque, reset_peak: bool) ?api.Memory { - const u = cuda.usage(reset_peak); - return .{ .active = u.device, .cache = 0, .peak = u.peak, .pinned = u.host }; -} - /// A lone driver returns false at completion or true when yield hands it to the lane core. 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; @@ -111,8 +92,21 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + follow_fn: ?*const fn (*anyopaque) anyerror!void = null, + health: ?api.HealthSource = null, recording: ?*record.Recording = null, // --record: the recorder in front of the family's backend + fn followFamily(p: *anyopaque) anyerror!void { + const h: *Host = @ptrCast(@alignCast(p)); + if (h.gpu) |g| try g.ctx.makeCurrent(); + return h.follow_fn.?(h.family); + } + + /// A family's readiness reads host state only: no CUDA binding or blocking device work. + fn readiness(h: *Host, comptime F: type) void { + if (@hasDecl(F, "healthy")) h.health = .{ .ctx = h.family, .read = F.healthy }; + } + fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); h.host.stop(); @@ -155,8 +149,9 @@ const Host = struct { h.core = lanes.Engine.init(h.gpa, &h.cfg, h.backend(), h.clock.clock()); errdefer h.core.deinit(); h.host = api.LaneHost.init(h.gpa, io, &h.core, info); - h.host.memory = if (h.gpu != null) .{ .read = readMemory } else null; + h.host.memory = if (h.gpu != null) .{ .read = family_memory.read } else null; h.host.explain = explain; + h.host.health = h.health; h.host.cache = h.store; h.host.warm_quiet = if (info.warm_turns) nemotron.native.warm_quiet else 0; if (h.lone != null) h.host.lone = .{ .ctx = h, .run = loneRun, .sampled = true, .logprobs = h.inner.vtable.first_row != null, .structures = h.inner.vtable.first_masked != null, .min_budget = nemotron.native.lone_min_budget, .resumes = true }; @@ -302,7 +297,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const before = try pool(a, io, &g.ctx, problem) orelse return null; // a family split over hosts holds only its rank's share (GLM-5.3-Flash's TP) const whole = budget.weightBytes(io, o.dir); - const weights = if (@hasDecl(F, "weightShare")) F.weightShare(a, io, o.dir, whole) else whole; + const weights = if (@hasDecl(F, "placementWeightBytes")) try F.placementWeightBytes(a, io, o.dir, o.tp, o.rank, o.drafts) else if (@hasDecl(F, "weightShare")) F.weightShare(a, io, o.dir, whole) else whole; const held0 = cuda.usage(false).device; // a family over hosts reads every rank's room after its load, by this budget, and admits by the smallest const Probe = struct { @@ -319,6 +314,12 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; var probe: Probe = .{ .a = a, .io = io, .ctx = &g.ctx, .held = held0, .problem = problem }; var opts: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }; + if (@hasField(F.Options, "rank")) { + opts.tp = @intCast(o.tp); + opts.rank = @intCast(o.rank); + opts.master = if (o.master.len > 0) o.master else null; + opts.master_port = o.master_port; + } if (o.vision) |vo| { if (!@hasField(F.Options, "vision")) { problem.* = "--vision: this CUDA family reads no images; serve it without --vision"; @@ -372,7 +373,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - var room = after.room(model) -| loaded.learn_bytes; + var room = try family_memory.room(F, loaded.ctx, after.room(model) -| loaded.learn_bytes); var low: []const u8 = ""; // names the rank whose room bounds the streams, when it is not this one if (@hasField(@TypeOf(loaded), "peer_room")) if (loaded.peer_room) |pr| if (pr.bytes < room) { room = pr.bytes; @@ -403,6 +404,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; + h.readiness(F); const held = loaded.stream_bytes * (streams -| free); const cache = if (@hasDecl(F, "snap_save")) budget.cacheBytes(o.prompt_cache_gib, room, held) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ @@ -466,14 +468,17 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O .vision = if (@hasField(@TypeOf(loaded), "vision")) vision_offer.offer(loaded.vision) else null, }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + if (@hasField(@TypeOf(loaded), "follow")) if (loaded.follow) |f| { + h.follow_fn = f; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h, .follow = Host.followFamily }; + }; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } -fn toGib(bytes: u64) f64 { - return @as(f64, @floatFromInt(bytes)) / gib; -} +const toGib = family_memory.toGib; test { + _ = @import("cuda_context.zig"); _ = budget; // cuda_memory.zig's own tests, which no other test binary reaches _ = refusals; _ = lane_backend; @@ -495,16 +500,6 @@ test "every registered family and d1's decision engine are listed for capabiliti try std.testing.expectEqualStrings("lfm2_vl", families[registry.len].model_type); } -test "context windows: the family default inside the model's, 0 for the model's, never past it" { - try std.testing.expectEqual(@as(i64, 16384), try contextWindow(null, 262144, 16384)); - try std.testing.expectEqual(@as(i64, 4096), try contextWindow(null, 4096, 16384)); - try std.testing.expectEqual(@as(i64, 262144), try contextWindow(0, 262144, 16384)); - try std.testing.expectEqual(@as(i64, 65536), try contextWindow(65536, 262144, 16384)); - try std.testing.expectError(error.PastNative, contextWindow(300000, 262144, 16384)); - try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); - try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); -} - test "a cancel during a long prompt stops it, and the next request's reply is unchanged (#291)" { const gpa = std.testing.allocator; var target: lanes.fake.Fake = .{ .gpa = gpa, .prefill_chunks = 10 }; @@ -581,3 +576,7 @@ test "a cancel during a long prompt stops it, and the next request's reply is un } try std.testing.expectEqual(@as(usize, 40), restored.tokens.items.len); } + +test "family readiness reaches production lane admission without binding CUDA" { + try @import("cuda_health_test.zig").check(Host, Host.readiness, Host.serve); +} diff --git a/zig/src/native/cuda_context.zig b/zig/src/native/cuda_context.zig new file mode 100644 index 0000000000..baa74e3df9 --- /dev/null +++ b/zig/src/native/cuda_context.zig @@ -0,0 +1,33 @@ +//! Native CUDA context-window configuration and checkpoint limits. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. +pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { + const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; + const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; + if (doc != .object) return 0; + const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; + const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; + return if (limit == .integer and limit.integer > 0) limit.integer else 0; +} + +/// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. +pub fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { + const r = requested orelse return if (native > 0) @min(default, native) else default; + if (r < 0) return error.Negative; + if (r == 0) return if (native > 0) native else error.NoNative; + if (native > 0 and r > native) return error.PastNative; + return r; +} + +test "context windows: the family default inside the model's, 0 for the model's, never past it" { + try std.testing.expectEqual(@as(i64, 16384), try contextWindow(null, 262144, 16384)); + try std.testing.expectEqual(@as(i64, 4096), try contextWindow(null, 4096, 16384)); + try std.testing.expectEqual(@as(i64, 262144), try contextWindow(0, 262144, 16384)); + try std.testing.expectEqual(@as(i64, 65536), try contextWindow(65536, 262144, 16384)); + try std.testing.expectError(error.PastNative, contextWindow(300000, 262144, 16384)); + try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); + try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); +} diff --git a/zig/src/native/cuda_family_memory.zig b/zig/src/native/cuda_family_memory.zig new file mode 100644 index 0000000000..3f7cec80ab --- /dev/null +++ b/zig/src/native/cuda_family_memory.zig @@ -0,0 +1,41 @@ +//! Family-aware checkpoint and stream admission; rank coordination runs before the host admits any stream. +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("engine_api"); + +pub fn weights(comptime F: type, a: std.mem.Allocator, io: std.Io, o: api.Open, comptime fallback: anytype) !u64 { + if (@hasDecl(F, "placementWeightBytes")) return F.placementWeightBytes(a, io, o.dir, o.tp, o.rank, o.drafts); + if (@hasDecl(F, "deviceWeightBytes")) return F.deviceWeightBytes(a, io, o.dir); + return fallback(io, o.dir); +} + +pub fn room(comptime F: type, ptr: *anyopaque, local: u64) !u64 { + if (@hasDecl(F, "rankRoom")) return F.rankRoom(ptr, local); + return local; +} + +pub fn read(_: ?*anyopaque, reset_peak: bool) ?api.Memory { + const u = cuda.usage(reset_peak); + return .{ .active = u.device, .cache = 0, .peak = u.peak, .pinned = u.host }; +} + +pub fn learnBytes(loaded: anytype) u64 { + return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; +} + +pub fn toGib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / @as(f64, 1 << 30); +} + +test "rank room is agreed before stream admission and preserves a single-rank family's pool" { + const Single = struct {}; + const Pair = struct { + pub fn rankRoom(_: *anyopaque, local: u64) !u64 { + return @min(local, 19); + } + }; + var context: u8 = 0; + try std.testing.expectEqual(@as(u64, 30), try room(Single, &context, 30)); + try std.testing.expectEqual(@as(u64, 19), try room(Pair, &context, 30)); + try std.testing.expectEqual(@as(u64, 7), try room(Pair, &context, 7)); +} diff --git a/zig/src/native/cuda_follow.zig b/zig/src/native/cuda_follow.zig new file mode 100644 index 0000000000..5944e239df --- /dev/null +++ b/zig/src/native/cuda_follow.zig @@ -0,0 +1,18 @@ +//! Follower rank ownership for native families that expose a replay loop. +const std = @import("std"); +const api = @import("engine_api"); +const native = @import("cuda.zig"); + +pub fn follow(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, options: api.Open, problem: *[]const u8) !bool { + const engine = (try native.open(a, gpa, io, options, problem)) orelse return false; + defer engine.close(engine.ctx); + const replay = engine.follow orelse { + problem.* = "this family has no tensor-parallel follower loop"; + return false; + }; + replay(engine.ctx) catch |err| { + problem.* = try std.fmt.allocPrint(a, "rank {d} stopped: {s}", .{ options.rank, @errorName(err) }); + return false; + }; + return true; +} diff --git a/zig/src/native/cuda_health_test.zig b/zig/src/native/cuda_health_test.zig new file mode 100644 index 0000000000..8dbf458534 --- /dev/null +++ b/zig/src/native/cuda_health_test.zig @@ -0,0 +1,49 @@ +//! The native family's readiness source reaches the real lane host without a device. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); + +pub fn check(comptime Host: type, comptime readiness: anytype, comptime serve: anytype) !void { + const gpa = std.testing.allocator; + const Family = struct { + available: std.atomic.Value(bool) = .init(true), + pub fn healthy(ctx: *anyopaque) bool { + const f: *@This() = @ptrCast(@alignCast(ctx)); + return f.available.load(.acquire); + } + fn release(_: *anyopaque) void {} + fn event(_: *anyopaque, _: api.Id, _: *const api.Event) void { + @panic("an unavailable family must refuse admission"); + } + }; + var family: Family = .{}; + var backend: lanes.fake.Fake = .{ .gpa = gpa }; + defer backend.deinit(); + const h = try gpa.create(Host); + defer gpa.destroy(h); + h.* = .{ .gpa = gpa, .gpu = null, .family = &family, .release = Family.release, .inner = backend.backend(), .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined }; + readiness(h, struct {}); // existing families need no readiness callback + try std.testing.expect(h.health == null); + readiness(h, Family); + try serve(h, std.testing.io, .{ .exact_width = 8 }, 8, .{ .lanes = 2 }, null); + defer { + h.host.stop(); + h.core.deinit(); + h.cfg.deinit(gpa); + } + const engine = h.host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + family.available.store(false, .release); + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + const request: api.Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &family, .event = Family.event })); + family.available.store(true, .release); + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + engine.drain(); + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); +} diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index 7cdf93f18f..f8526af0cd 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -21,11 +21,24 @@ const lfm2_vision = "--vision: LFM2 reads the checkpoint's images through its ow /// A family takes --drafter only if it says so, and --learn-conversations only as Nemotron; null: `o` may open `F`. pub fn family(comptime F: type, o: api.Open) ?[]const u8 { + const ranked = @hasDecl(F, "Options") and @hasField(F.Options, "rank"); + if (o.tp != 1 and !ranked) return "this CUDA family has no tensor-parallel ranks: serve without --tp"; if (o.drafter != null and !(@hasDecl(F, "takes_drafter") and F.takes_drafter)) return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; if (o.learn_conversations and F != nemotron.native) return "--learn-conversations keeps conversation states on disk for Nemotron only on CUDA: serve this model without it"; return null; } +test "a family must consume rank options before tensor parallelism is admitted" { + const Single = struct { + pub const Options = struct { context: usize }; + }; + const Ranked = struct { + pub const Options = struct { context: usize, rank: usize }; + }; + try std.testing.expectStringStartsWith(family(Single, .{ .dir = "d", .model_type = "m", .tp = 2 }) orelse "", "this CUDA family has no tensor-parallel ranks"); + try std.testing.expect(family(Ranked, .{ .dir = "d", .model_type = "m", .tp = 2 }) == null); +} + test "--learn is refused whatever the family, and --slide or a learned-disk size alone is not" { try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x" }) orelse "", "--learn learns shared prompt prefixes on Macs only"); try std.testing.expect(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x", .learn_conversations = true }) == null); diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index 8281b19df4..aac4b59136 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -164,6 +164,7 @@ pub const Host = struct { n = 1; } out.* = .{ + .healthy = !h.closing, .running = @intFromBool(h.running != null), .waiting = @intCast(h.queued.items.len), .decode_tokens_per_second = @as(f64, @floatFromInt(tokens)) / 2.0, @@ -534,3 +535,25 @@ test "status after a reply finishes reads no freed job" { try std.testing.expectEqual(@as(u32, 0), st.running); try std.testing.expectEqual(@as(usize, 0), st.streams); } + +test "a drained Flash Next host is unhealthy before another submit" { + const gpa = std.testing.allocator; + var host: Host = .{ .gpa = gpa, .io = std.testing.io, .eng = undefined, .info_ = .{} }; + defer host.decoded.deinit(gpa); + const engine = host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + engine.drain(); + engine.status(&status, &.{}); + const Sink = struct { + fn event(_: *anyopaque, _: api.Id, _: *const api.Event) void { + @panic("a drained host must refuse admission"); + } + }; + var ctx: u8 = 0; + const request: api.Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &ctx, .event = Sink.event })); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len); + try std.testing.expect(!status.healthy); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 35b131e2cd..bd6e588c0b 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -41,6 +41,8 @@ pub const Input = struct { pub const Sink = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, delta: Value) error{Closed}!void, + /// Called only after admission, with cancellation and mailbox draining already guarded. + begin: ?*const fn (ctx: *anyopaque, prompt_len: usize) error{Closed}!void = null, }; pub const Reply = struct { @@ -194,7 +196,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer if (srv.keepalive) |k| k.end(); srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.Closed => cx.fail(.capacity, "the engine is unavailable; retry after it restarts", .{}), error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), error.Stale => cx.fail(.server, "{s}", .{api.stale_message}), }; @@ -204,6 +206,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer srv.noteRequest(prepared.prompt_len, box.cached orelse 0, gen.collected.items.len, box.stats.drafted, box.stats.accepted, box.stats.rounds, received, gen.first_ns, gen.last_ns, box.stats.prefill_seconds); errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { + if (sink) |s| if (s.begin) |begin| begin(s.ctx, prepared.prompt_len) catch break :blk error.Cancelled; if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; gen.loop(gone) catch |e| break :blk e; break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); diff --git a/zig/src/server/compact_test.zig b/zig/src/server/compact_test.zig index 986fd10922..aaa3a688eb 100644 --- a/zig/src/server/compact_test.zig +++ b/zig/src/server/compact_test.zig @@ -441,6 +441,7 @@ test "an incremental request summarizes only the new middle and a user edit is t const Watch = struct { eng: *Eng, opened: bool = false, + admitted_at_open: usize = 0, saw: bool = false, fn out(w: *Watch) openai.Out { @@ -449,6 +450,7 @@ const Watch = struct { fn open(ctx: *anyopaque, _: usize) error{Closed}!void { const w: *Watch = @ptrCast(@alignCast(ctx)); + w.admitted_at_open = w.eng.prompts.items.len; w.opened = true; w.eng.opened = true; } @@ -469,7 +471,7 @@ const Watch = struct { } }; -test "a stream compacts before it opens" { +test "a stream compacts and admits its final request before it opens" { var eng: Eng = .{ .window = 2000 }; var text: Text = .{}; defer freePrompts(&eng); @@ -486,6 +488,7 @@ test "a stream compacts before it opens" { openai.run(srv, a, watch.out(), .{ .conn = &conn }, true, raw); try std.testing.expect(watch.opened); try std.testing.expect(watch.saw); - try std.testing.expect(eng.early >= 1); - try std.testing.expect(eng.late >= 1); + try std.testing.expect(watch.admitted_at_open >= 2); // compaction and the final reply both admitted + try std.testing.expectEqual(watch.admitted_at_open, eng.early); + try std.testing.expectEqual(@as(usize, 0), eng.late); } diff --git a/zig/src/server/cortheon/cortheon_test.zig b/zig/src/server/cortheon/cortheon_test.zig index 2287eaac82..15413efef1 100644 --- a/zig/src/server/cortheon/cortheon_test.zig +++ b/zig/src/server/cortheon/cortheon_test.zig @@ -63,6 +63,8 @@ const Text = struct { const Eng = struct { script: []const []const u8 = &.{}, next: usize = 0, + refuse_at: ?usize = null, + refusal: api.SubmitError = error.Closed, prompts: std.ArrayList([]u8) = .empty, fn engine(e: *Eng) api.Engine { @@ -73,6 +75,7 @@ const Eng = struct { } fn submit(ctx: *anyopaque, id: api.Id, request: *const api.Request, sink: api.Sink) api.SubmitError!void { const e: *Eng = @ptrCast(@alignCast(ctx)); + if (e.refuse_at == e.next) return e.refusal; const buf = gpa.alloc(u8, request.prompt.len) catch return error.Busy; for (request.prompt, buf) |tok, *byte| byte.* = @intCast(tok); e.prompts.append(gpa, buf) catch {}; @@ -121,17 +124,22 @@ const Web = struct { const Got = struct { events: std.ArrayList(Value) = .empty, body: ?Value = null, + opened: bool = false, + status: ?u16 = null, a: Allocator, fn out(g: *Got) openai.Out { return .{ .ctx = g, .vt = &.{ .open = open, .event = event, .reply = reply } }; } - fn open(_: *anyopaque, _: usize) error{Closed}!void {} + fn open(ctx: *anyopaque, _: usize) error{Closed}!void { + @as(*Got, @ptrCast(@alignCast(ctx))).opened = true; + } fn event(ctx: *anyopaque, payload: ?Value) error{Closed}!void { const g: *Got = @ptrCast(@alignCast(ctx)); if (payload) |p| g.events.append(g.a, json.deepCopy(g.a, p) catch return error.Closed) catch return error.Closed; } - fn reply(ctx: *anyopaque, _: u16, payload: Value) void { + fn reply(ctx: *anyopaque, status: u16, payload: Value) void { const g: *Got = @ptrCast(@alignCast(ctx)); + g.status = status; g.body = json.deepCopy(g.a, payload) catch null; } /// Every streamed delta's text under `key`, joined. @@ -245,6 +253,22 @@ test "a fresh question on current facts is researched before the model starts" { try std.testing.expectEqualStrings("The latest is 0.17.\n\nOK", message(got.body.?).strField("content").?); } +test "research progress cannot open a stream before its final request is admitted" { + for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { + var rig: Rig = .{ .eng = .{ .script = &.{"Zig latest release notes"}, .refuse_at = 1, .refusal = reason } }; + try rig.start(true); + defer rig.stop(); + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const got = try rig.ask(arena.allocator(), "{\"stream\":true,\"messages\":[{\"role\":\"user\",\"content\":\"What is the latest release of Zig?\"}]}"); + try std.testing.expect(rig.web.searches > 0); + try std.testing.expectEqual(@as(?u16, 503), got.status); + try std.testing.expect(!got.opened); + try std.testing.expectEqual(@as(usize, 0), got.events.items.len); + try std.testing.expectEqual(@as(i64, 0), rig.srv.preparing.load(.acquire)); + } +} + test "a stream hides Cortheon's calls, shows progress as reasoning, and passes the client's calls" { var rig: Rig = .{ .eng = .{ .script = &.{ "Let me check.{\"name\": \"cortheon_fetch\", \"arguments\": {\"url\": \"https://docs.example/x\"}}", "Listing files.{\"name\": \"bash\", \"arguments\": {\"command\": \"ls\"}}" } } }; try rig.start(true); diff --git a/zig/src/server/cortheon/loop.zig b/zig/src/server/cortheon/loop.zig index 52eb5db2cc..a4fa577678 100644 --- a/zig/src/server/cortheon/loop.zig +++ b/zig/src/server/cortheon/loop.zig @@ -131,7 +131,13 @@ const Filter = struct { const Call = struct { index: i64, server: bool, client: u32 }; fn sink(f: *Filter) chat.Sink { - return .{ .ctx = f, .call = call }; + return .{ .ctx = f, .call = call, .begin = begin }; + } + + fn begin(ctx: *anyopaque, prompt_len: usize) error{Closed}!void { + const f: *Filter = @ptrCast(@alignCast(ctx)); + const inner = f.inner orelse return; + if (inner.begin) |start| try start(inner.ctx, prompt_len); } fn startRound(f: *Filter) void { diff --git a/zig/src/server/health_test.zig b/zig/src/server/health_test.zig index 272933be59..be31991c49 100644 --- a/zig/src/server/health_test.zig +++ b/zig/src/server/health_test.zig @@ -17,6 +17,7 @@ const earlier = [_][]const u8{ "status", "model", "model_ids", "max_batch_size", /// Four lanes, some memory, nothing running. const Idle = struct { + healthy: bool = true, fn engine(e: *Idle) api.Engine { return .{ .ctx = e, .vtable = &.{ .info = info, @@ -37,8 +38,8 @@ const Idle = struct { fn cancel(_: *anyopaque, _: api.Id) void {} - fn status(_: *anyopaque, out: *api.Status, _: []u32) void { - out.* = .{}; + fn status(ctx: *anyopaque, out: *api.Status, _: []u32) void { + out.* = .{ .healthy = @as(*Idle, @ptrCast(@alignCast(ctx))).healthy }; } fn memory(_: *anyopaque, _: bool) ?api.Memory { @@ -265,3 +266,32 @@ test "with keys configured, /health sends only its status, with and without a va try std.testing.expect(std.mem.startsWith(u8, try exchange(a, port, models_open), "HTTP/1.1 401 ")); try std.testing.expect(std.mem.startsWith(u8, try exchange(a, port, models_keyed), "HTTP/1.1 200 ")); } + +test "unavailable health is minimal HTTP 503 with keys enabled or disabled and with or without a valid key" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + var keys = try auth.Store.init(std.testing.allocator, &.{key}, "", null, false, &problem); + defer keys.deinit(); + for ([_]bool{ false, true }) |keyed| { + var engine: Idle = .{ .healthy = false }; + var text: Bytes = .{}; + const srv = try server.Server.init(std.testing.allocator, std.testing.io, engine.engine(), text.text(), .{ + .served_name = "private-model", + .model_ids = &.{"private-model"}, + .identity = .{ .version = "9.8.7", .backend = "test", .device = "private-device" }, + }, if (keyed) &keys else null); + defer srv.deinit(); + const lis = try listener.Listener.open(.{ .ip4 = .loopback(0) }); + defer lis.close(); + var stop: std.atomic.Value(bool) = .init(false); + const accept = try std.Thread.spawn(.{}, server.Server.serve, .{ srv, lis, &stop }); + defer quiesce(srv, &stop, accept); + for ([_][]const u8{ health_open, health_keyed }) |request| { + const reply = try exchange(a, lis.port(), request); + try std.testing.expect(std.mem.startsWith(u8, reply, "HTTP/1.1 503 ")); + try std.testing.expectEqualStrings("{\"status\": \"unavailable\"}", try bodyOf(reply)); + } + } +} diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..b508c71854 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -53,8 +53,15 @@ pub const Listener = struct { const e = posix.errno(rc); switch (e) { .SUCCESS => { + const fd: posix.socket_t = @intCast(rc); + const one: c_int = 1; + posix.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, std.mem.asBytes(&one)) catch |err| { + _ = posix.system.close(fd); + log.line("could not configure accepted TCP connection: {s}", .{@errorName(err)}); + continue; + }; nofile_logged.store(false, .release); - return .{ .fd = @intCast(rc), .peer = Threaded.addressFromPosix(&storage) }; + return .{ .fd = fd, .peer = Threaded.addressFromPosix(&storage) }; }, .INTR, .AGAIN, .CONNABORTED, .MFILE, .NFILE, .NOBUFS, .NOMEM => { if (e != .INTR) { @@ -132,3 +139,52 @@ test "peer text" { try std.testing.expectEqualStrings("127.0.0.1", peerText(&buf, .{ .ip4 = .loopback(0) })); try std.testing.expectEqualStrings("::1", peerText(&buf, .{ .ip6 = .loopback(0) })); } + +test "accepted IPv4 and IPv6 TCP streams disable Nagle for incremental responses" { + for ([_]net.IpAddress{ .{ .ip4 = .loopback(0) }, .{ .ip6 = .loopback(0) } }) |address| { + const listener = try Listener.open(address); + defer listener.close(); + const rc = posix.system.socket(if (address == .ip4) posix.AF.INET else posix.AF.INET6, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const client: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(client); + const peer: net.IpAddress = if (address == .ip4) .{ .ip4 = .loopback(listener.port()) } else .{ .ip6 = .loopback(listener.port()) }; + var storage: Threaded.PosixAddress = undefined; + const len = Threaded.addressToPosix(&peer, &storage); + if (posix.errno(posix.system.connect(client, &storage.any, len)) != .SUCCESS) return error.ConnectFailed; + const accepted = acceptWithin(listener, 2000) orelse return error.AcceptFailed; + defer _ = posix.system.close(accepted.fd); + var value: c_int = 0; + var size: posix.socklen_t = @sizeOf(c_int); + try std.testing.expect(std.c.getsockopt(accepted.fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &value, &size) == 0); + try std.testing.expect(value != 0); + } +} + +fn acceptWithin(listener: Listener, timeout_ms: u32) ?Listener.Accepted { + const Watchdog = struct { + stop: std.atomic.Value(bool) = .init(false), + done: std.atomic.Value(bool) = .init(false), + + fn run(w: *@This(), milliseconds: u32) void { + var elapsed: u32 = 0; + while (elapsed < milliseconds and !w.done.load(.acquire)) : (elapsed += 10) { + std.Io.sleep(log.io(), .fromMilliseconds(10), .awake) catch break; + } + if (!w.done.load(.acquire)) w.stop.store(true, .release); + } + }; + var watchdog: Watchdog = .{}; + const thread = std.Thread.spawn(.{}, Watchdog.run, .{ &watchdog, timeout_ms }) catch return null; + defer { + watchdog.done.store(true, .release); + thread.join(); + } + return listener.accept(&watchdog.stop); +} + +test "the accepted-stream test stops when no usable connection arrives" { + const listener = try Listener.open(.{ .ip4 = .loopback(0) }); + defer listener.close(); + try std.testing.expect(acceptWithin(listener, 20) == null); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 8c8780023c..45cf437364 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -344,10 +344,10 @@ const Run = struct { const a = r.a; const kind_ok = cx.kind != .other and cx.kind != .server; const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + r.out.vt.reply(r.out.ctx, if (kind_ok) cx.status() else 500, wrapError(a, body) catch return); } - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + /// An unsent reply: request 400, unavailable 503, failure 500, nothing when the client left. fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { switch (e) { error.Cancelled => {}, @@ -429,10 +429,25 @@ const Run = struct { const StreamSink = struct { run: *Run, tools: bool, + prompt_len: usize, + opened: bool = false, + pending: std.ArrayList(Value) = .empty, // research progress waits for actual engine admission + + fn begin(ctx: *anyopaque, _: usize) error{Closed}!void { + const s: *StreamSink = @ptrCast(@alignCast(ctx)); + if (s.opened) return; + const r = s.run; + try r.out.vt.open(r.out.ctx, s.prompt_len); + s.opened = true; + if (!s.tools and r.is_chat) try r.emit(r.chunk(roleDelta(r.a) catch return error.Closed, null) catch return error.Closed); + for (s.pending.items) |delta| try call(ctx, delta); + s.pending.clearRetainingCapacity(); + } fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { const s: *StreamSink = @ptrCast(@alignCast(ctx)); const r = s.run; + if (!s.opened) return s.pending.append(r.a, delta) catch error.Closed; if (!s.tools) { // completions stream text only, with its logprob rows when asked if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; @@ -482,7 +497,7 @@ const Run = struct { const a = r.a; const tools = r.plan.input.tools.len > 0; var cx: Cx = .{ .a = a }; - // a context-window refusal is reported inside the stream; every other refusal is still a 400 before it opens + // Context-window refusals stay streamed; other pre-admission refusals retain their HTTP status. var compaction: ?compact.Stamp = null; const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); @@ -497,12 +512,11 @@ const Run = struct { }; var handed = false; // generate gives the preparing count back from here on defer if (!handed) chat.release(r.srv, prepared.preparing); - r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - var sink_state: StreamSink = .{ .run = r, .tools = tools }; + var sink_state: StreamSink = .{ .run = r, .tools = tools, .prompt_len = prepared.prompt_len }; handed = true; - const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; + const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call, .begin = StreamSink.begin }; var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { + if (!sink_state.opened) return r.unsent(&cx, e, field); switch (e) { error.Cancelled => return, error.Refused => if (cx.kind != .other and cx.kind != .server) { diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 833f76579e..d147d96933 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -11,6 +11,9 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + var status: api.Status = .{}; + srv.engine.status(&status, &.{}); + if (!status.healthy) return conn.sendJson(503, "{\"status\": \"unavailable\"}"); if (srv.keys) |k| if (k.enabled()) return conn.sendJson(200, "{\"status\": \"ok\"}"); routes.sendValue(conn, a, 200, try openHealth(srv, a, std.mem.indexOf(u8, conn.path, "reset_peak=1") != null)); } diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index 601573da5f..b38f81ecc5 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -1,4 +1,4 @@ -//! A streamed context-window refusal is an error event, and every other refusal is a 400 before the stream opens. +//! Context-window refusals stay streamed; other refusals preserve their HTTP status before admission. const std = @import("std"); const api = @import("engine_api"); const server = @import("root.zig"); @@ -295,6 +295,36 @@ test "a stream whose client is gone when it opens gives the preparing count back try std.testing.expectEqual(@as(i64, 0), seen.preparing); } +test "disconnect at the admitted stream boundary cancels and drains its mailbox" { + const Waiting = struct { + sink: ?api.Sink = null, + submissions: usize = 0, + cancellations: usize = 0, + fn engine(e: *@This()) api.Engine { + return .{ .ctx = e, .vtable = &.{ .info = HiEngine.info, .submit = submit, .cancel = cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; + } + fn submit(ctx: *anyopaque, _: api.Id, _: *const api.Request, sink: api.Sink) api.SubmitError!void { + const e: *@This() = @ptrCast(@alignCast(ctx)); + e.submissions += 1; + e.sink = sink; + } + fn cancel(ctx: *anyopaque, id: api.Id) void { + const e: *@This() = @ptrCast(@alignCast(ctx)); + e.cancellations += 1; + const sink = e.sink.?; + e.sink = null; + sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); + } + }; + var waiting: Waiting = .{}; + const seen = try sendWith(waiting.engine(), .chat, "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"stream\":true,\"max_tokens\":2}", .{ .gone_at_open = true }); + try std.testing.expectEqual(@as(usize, 1), waiting.submissions); + try std.testing.expectEqual(@as(usize, 1), waiting.cancellations); + try std.testing.expect(waiting.sink == null); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); +} + test "a reply labels a greedy request greedy and a sampled one exact" { const greedy = try send(.chat, "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":2,\"temperature\":0}"); try std.testing.expectEqual(@as(?u16, 200), greedy.status); @@ -457,3 +487,36 @@ test "unsupported budgets leave zero unset and thinking-disabled requests valid" } try std.testing.expectEqual(@as(usize, 3), backend.submissions); } + +test "closed or busy admission is service unavailable across chat, messages and responses" { + const ClosedEngine = struct { + reason: api.SubmitError = error.Closed, + fn engine(e: *@This()) api.Engine { + return .{ .ctx = e, .vtable = &.{ .info = HiEngine.info, .submit = submit, .cancel = HiEngine.cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; + } + fn submit(ctx: *anyopaque, _: api.Id, _: *const api.Request, _: api.Sink) api.SubmitError!void { + return @as(*@This(), @ptrCast(@alignCast(ctx))).reason; + } + }; + const cases = [_]struct { route: Route, body: []const u8 }{ + .{ .route = .chat, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .messages, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .responses, .body = "{\"model\":\"test-model\",\"input\":\"x\",\"max_output_tokens\":1}" }, + }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + var wrong: usize = 0; + var engine: ClosedEngine = .{}; + for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { + engine.reason = reason; + for (cases) |c| for ([_]bool{ false, true }) |streaming| { + const body = if (streaming) try std.mem.concat(arena.allocator(), u8, &.{ c.body[0 .. c.body.len - 1], ",\"stream\":true}" }) else c.body; + const seen = try sendWith(engine.engine(), c.route, body, .{}); + if (seen.status != 503 or seen.opened or seen.preparing != 0) { + wrong += 1; + std.debug.print("ADMISSION reason={s} route={s} stream={any} status={any} opened={any} events={d} preparing={d}\n", .{ @errorName(reason), @tagName(c.route), streaming, seen.status, seen.opened, seen.events, seen.preparing }); + } + }; + } + try std.testing.expectEqual(@as(usize, 0), wrong); +} diff --git a/zig/tests/cuda/contract_flashnext_writeback.zig b/zig/tests/cuda/contract_flashnext_writeback.zig new file mode 100644 index 0000000000..3700d915db --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_writeback.zig @@ -0,0 +1,276 @@ +//! Hyper-connection writeback: guarded hidden rows, exact fused FP32 rounding and independent FP64 square sums. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const c = core.contract.spec; +const g = core.contract.gates; +const streams = 4; +const slots = 11; + +pub const Mode = enum { sums, moe, slices }; + +fn depth(_: c.Shape) usize { + return 13; +} + +fn spec(mode: Mode) c.Spec { + return .{ + .name = switch (mode) { + .sums => "cuda Flash Next fn_hc_writeback0", + .moe => "cuda Flash Next fn_hc_writeback_moe", + .slices => "cuda Flash Next fn_hc_writeback_slices", + }, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 16 }, .{ .size = 256, .step = 256, .min = 256 }, .{} }, + .served = &.{ .{ .rows = 17, .n = 2560, .k = 2 }, .{ .rows = 3, .n = 512, .k = 8 } }, + }; +} + +/// PSS is first in the output byte stream; the hidden output follows and remains part of the repeat and guard gates. +pub const Writeback = struct { + rig: rig_.Rig, + f: cuda.Function, + mode: Mode, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = undefined, + h: []u16 = &.{}, + inj: []u16 = &.{}, + y: []f32 = &.{}, + wts: []f32 = &.{}, + br: []f32 = &.{}, + dh: u64 = 0, + dinj: u64 = 0, + dy: u64 = 0, + dwts: u64 = 0, + dbr: u64 = 0, + pss: u64 = 0, + hout: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Writeback) void { + if (self.arena) |*arena| arena.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Writeback); + + pub fn kernel(self: *Writeback) g.Kernel { + return .{ .spec = spec(self.mode), .ctx = self, .vt = &vt }; + } + + pub fn prepare(self: *Writeback, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 33 or shape.n == 0 or shape.n > 2560 or shape.n % 256 != 0 or shape.k == 0 or shape.k > 8) return false; + self.rig.reset(); + if (self.arena) |*arena| arena.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + const wide = shape.n * streams; + self.h = try a.alloc(u16, shape.rows * wide); + self.inj = try a.alloc(u16, shape.rows * streams); + self.y = try a.alloc(f32, shape.rows * slots * shape.n); + self.wts = try a.alloc(f32, shape.rows * slots); + self.br = try a.alloc(f32, shape.k * shape.rows * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.h) |*b| b.* = c.toBf16(if (fill == .stress) @floatFromInt(1 + rng.uintLessThan(u32, 64)) else rng.float(f32) * 2 - 1); + for (self.inj) |*b| b.* = c.toBf16(if (fill == .stress) 1 else rng.float(f32) * 4 - 2); + for (self.y) |*v| v.* = if (fill == .stress) 0 else rng.float(f32) * 2 - 1; + for (self.wts) |*v| v.* = rng.float(f32) * 0.2; + for (self.br) |*v| v.* = if (fill == .stress) 0 else rng.float(f32) - 0.5; + if (fill != .stress and self.mode != .sums) { + self.h[0] = 0x3b34; + self.inj[0] = 0x3c9a; + for (0..slots) |k| self.y[k * shape.n] = 0; + self.y[0] = c.fromBf16(0xbe16); + self.wts[0] = 1; + for (0..shape.k) |k| self.br[k * shape.rows * shape.n] = 0; + self.br[0] = c.fromBf16(0xbe16); + } + if (fill == .nonfinite) for (0..shape.rows) |row| { + self.h[row * wide] = switch (row % 5) { + 0 => 0x7fc1, + 1 => 0x7f80, + 2 => 0xff80, + 3 => 0x8000, + else => 0x0001, + }; + }; + self.dh = try self.rig.input(std.mem.sliceAsBytes(self.h)); + self.dinj = try self.rig.input(std.mem.sliceAsBytes(self.inj)); + self.dy = try self.rig.input(std.mem.sliceAsBytes(self.y)); + self.dwts = try self.rig.input(std.mem.sliceAsBytes(self.wts)); + self.dbr = try self.rig.input(std.mem.sliceAsBytes(self.br)); + self.pss = try self.rig.output(self.count() * 4); + self.hout = if (self.mode == .sums) 0 else try self.rig.output(self.h.len * 2); + return true; + } + + pub fn run(self: *Writeback) !void { + const sh = self.shape; + var args: cuda.Args = .{}; + args.add(self.dh); + if (self.mode != .sums) args.add(self.hout); + args.add(self.pss); + switch (self.mode) { + .sums => {}, + .moe => { + args.add(self.dinj); + args.add(self.dy); + args.add(self.dwts); + }, + .slices => { + args.add(self.dbr); + args.add(self.dinj); + args.add(@as(c_longlong, @intCast(sh.rows * sh.n))); + }, + } + args.add(@as(c_int, @intCast(sh.n))); + args.add(@as(c_int, streams)); + if (self.mode == .moe) { + args.add(@as(c_int, slots - 1)); + args.add(@as(c_int, slots)); + } + if (self.mode == .slices) args.add(@as(c_int, @intCast(sh.k))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(sh.rows), .y = @intCast(sh.n / 256) }, .block = .{ .x = 256 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Writeback) usize { + return self.shape.rows * (self.shape.n / 256) * streams; + } + + fn hidden(self: *Writeback, row: usize, stream: usize, col: usize) f32 { + const sh = self.shape; + const h = c.fromBf16(self.h[(row * streams + stream) * sh.n + col]); + if (self.mode == .sums) return h; + var branch: f32 = 0; + switch (self.mode) { + .moe => for (0..slots) |k| { + const at = row * slots + k; + branch = @mulAdd(f32, self.y[at * sh.n + col], self.wts[at], branch); + }, + .slices => { + branch = self.br[row * sh.n + col]; + for (1..sh.k) |k| branch += self.br[(k * sh.rows + row) * sh.n + col]; + }, + .sums => unreachable, + } + branch = c.fromBf16(c.toBf16(branch)); + return c.fromBf16(c.toBf16(@mulAdd(f32, branch, c.fromBf16(self.inj[row * streams + stream]), h))); + } + + pub fn reference(self: *Writeback, want: []f64, mag: []f64) !void { + for (0..self.shape.rows) |row| for (0..self.shape.n / 256) |chunk| for (0..streams) |stream| { + var sum: f64 = 0; + for (0..256) |i| { + const v: f64 = self.hidden(row, stream, chunk * 256 + i); + sum += v * v; + } + const at = (row * (self.shape.n / 256) + chunk) * streams + stream; + want[at] = sum; + mag[at] = sum; + }; + } + + pub fn simulate(self: *Writeback, acc: c.Accum, order: g.Order, out: []f64) !void { + for (0..self.shape.rows) |row| for (0..self.shape.n / 256) |chunk| for (0..streams) |stream| { + var terms: [256]f64 = undefined; + for (&terms, 0..) |*term, i| { + const v: f64 = self.hidden(row, stream, chunk * 256 + i); + term.* = v * v; + } + out[(row * (self.shape.n / 256) + chunk) * streams + stream] = core.contract.rows.sumIn(acc, order, &terms); + }; + } + + pub fn check(self: *Writeback, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const inputs = try self.rig.inputBytes(); + var offset: usize = 0; + for ([_][]const u8{ std.mem.sliceAsBytes(self.h), std.mem.sliceAsBytes(self.inj), std.mem.sliceAsBytes(self.y), std.mem.sliceAsBytes(self.wts), std.mem.sliceAsBytes(self.br) }) |expected| { + if (!std.mem.eql(u8, inputs[offset..][0..expected.len], expected)) return "an immutable input changed"; + offset += expected.len; + } + const a = self.arena.?.allocator(); + const pss = try a.dupe(u8, try self.rig.bytes(self.pss, self.count() * 4)); + if (self.mode != .sums) { + const hout = try self.rig.bytes(self.hout, self.h.len * 2); + for (0..self.shape.rows) |row| for (0..streams) |stream| for (0..self.shape.n) |col| { + const at = (row * streams + stream) * self.shape.n + col; + const got = std.mem.readInt(u16, hout[at * 2 ..][0..2], .little); + const want = c.toBf16(self.hidden(row, stream, col)); + if (got != want and !(std.math.isNan(c.fromBf16(got)) and std.math.isNan(c.fromBf16(want)))) + return try std.fmt.bufPrint(&self.msg, "hidden row {d} stream {d} column {d}: {x} != fused FP32 reference {x}", .{ row, stream, col, got, want }); + }; + } + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + for (want, mag, 0..) |w, m, i| { + const got = c.read(.f32, pss, i); + const ok = if (std.math.isFinite(w)) @abs(got - w) <= c.bound(spec(self.mode), self.shape, w, m) else c.classOf(got) == c.classOf(w); + if (!ok) return try std.fmt.bufPrint(&self.msg, "PSS {d}: {d} differs from independent FP64 sum {d}", .{ i, got, w }); + } + return null; + } +}; + +test "integer square stress requires FP32 sums in every simulated order" { + var prng = std.Random.DefaultPrng.init(0x46e3); + const rng = prng.random(); + var terms: [4][256]f64 = undefined; + var want: [4]f64 = @splat(0); + for (&terms, &want) |*row, *sum| for (row) |*term| { + const v: f64 = @floatFromInt(1 + rng.uintLessThan(u32, 64)); + term.* = v * v; + sum.* += term.*; + }; + for (std.enums.values(g.Order)) |order| { + for (terms, want) |row, w| try std.testing.expectEqual(w, core.contract.rows.sumIn(.f32, order, &row)); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + var miss: f64 = 0; + for (terms, want) |row, w| { + const got = core.contract.rows.sumIn(acc, order, &row); + const bound = c.boundExact(spec(.sums), w); + miss = @max(miss, if (std.math.isFinite(got)) @abs(got - w) / bound else std.math.inf(f64)); + } + try std.testing.expect(miss >= 100); + } + } +} + +test "the cancellation fixture distinguishes fused update from rounded product" { + const h = c.fromBf16(0x3b34); + const branch = c.fromBf16(0xbe16); + const inject = c.fromBf16(0x3c9a); + const fused = c.toBf16(@mulAdd(f32, branch, inject, h)); + const old = c.toBf16(h + c.fromBf16(c.toBf16(branch * inject))); + const exact: f64 = @as(f64, h) + @as(f64, branch) * @as(f64, inject); + try std.testing.expect(fused != old); + try std.testing.expectEqual(c.toBf16(@floatCast(exact)), fused); +} + +noinline fn reuseRegistrationStack() usize { + var scratch: [4096]usize = undefined; + const words: *volatile [4096]usize = &scratch; + for (0..scratch.len) |i| words[i] = i *% 73; + return words[42]; +} + +test "registered writeback hooks survive later stack frames" { + var first = Writeback{ .rig = undefined, .f = undefined, .mode = .sums, .shape = .{ .rows = 17, .n = 2560, .k = 2 } }; + var second = Writeback{ .rig = undefined, .f = undefined, .mode = .moe, .shape = .{ .rows = 3, .n = 512, .k = 8 } }; + const Best = @import("contract_row_best.zig").Best; + var best = Best{ .rig = undefined, .f = undefined, .shape = .{ .rows = 4, .n = 4, .k = 31 } }; + const registered = [_]g.Kernel{ first.kernel(), second.kernel(), best.kernel() }; + try std.testing.expectEqual(@as(usize, 3066), reuseRegistrationStack()); + for (registered, [_]usize{ 680, 24, 16 }) |k, expected| { + try std.testing.expectEqual(expected, k.vt.count(k.ctx)); + try std.testing.expect(!try k.vt.prepare(k.ctx, .{ .rows = 0, .n = 256, .k = 1 }, .random, 51)); + try std.testing.expectEqual(expected, k.vt.count(k.ctx)); + } +} diff --git a/zig/tests/cuda/contract_mxfp8.zig b/zig/tests/cuda/contract_mxfp8.zig new file mode 100644 index 0000000000..278cb78c4c --- /dev/null +++ b/zig/tests/cuda/contract_mxfp8.zig @@ -0,0 +1,115 @@ +//! MXFP8 projections through the common non-finite, repeat, shape and accumulator gates. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const Rig = @import("contract_rig.zig").Rig; + +fn depth(shape: c.Shape) usize { + return shape.k / 16 + 24; +} + +const served = [_]c.Shape{ + .{ .rows = 1, .n = 70, .k = 1024 }, + .{ .rows = 17, .n = 70, .k = 1024 }, + .{ .rows = 33, .n = 70, .k = 1024 }, + .{ .rows = 256, .n = 70, .k = 1024 }, + .{ .rows = 16, .n = 1280, .k = 2560 }, + .{ .rows = 16, .n = 2560, .k = 640 }, + .{ .rows = 7, .n = 13952, .k = 2560 }, +}; + +pub const Product = struct { + rig: Rig, + lane: *const cuda.qmmf.Lane, + clusters: bool, + d: ?Dense = null, + weight: cuda.qmmf.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + part: u64 = 0, + + pub fn deinit(self: *Product) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Product) g.Kernel { + return .{ .spec = .{ + .name = if (self.clusters) "cuda MXFP8 lane clustered slices" else "cuda MXFP8 lane explicit slice reduction", + .accum = .f32, + .depth = depth, + .out = .bf16, + .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{ .size = 512, .step = 64, .min = 64 } }, + .served = &served, + }, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Product); + + pub fn prepare(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + self.part = 0; + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 64 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + var prng = std.Random.DefaultPrng.init(seed ^ 0x8f); + const random = prng.random(); + const n = shape.n; + const k = shape.k; + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + const scales = try a.alloc(u8, n * k / 32); + defer a.free(scales); + for (scales) |*s| s.* = if (fill == .stress) 127 else 123 + random.uintLessThan(u8, 8); + for (0..n) |j| for (0..k) |i| { + codes[j * k + i] = if (fill == .stress) 0x38 + random.uintLessThan(u8, 24) else random.uintLessThan(u8, 0x7f); + if (fill != .stress and random.boolean()) codes[j * k + i] |= 0x80; + }; + if (fill == .stress) for (0..n) |j| for (k / 2..k) |i| { + codes[j * k + i] = codes[j * k + i - k / 2] | 0x80; + }; + for (0..n) |j| for (0..k) |i| { + const scale = std.math.ldexp(@as(f64, 1), @as(i32, scales[j * (k / 32) + i / 32]) - 127); + d.w[j * k + i] = core.quant.fp8.e4m3(codes[j * k + i]) * scale; + }; + const npad = cuda.mxfp8.padded(n); + const tiled_codes = try a.alloc(u8, npad * k); + defer a.free(tiled_codes); + const packed_scales = try a.alloc(u8, npad * k / 32); + defer a.free(packed_scales); + cuda.mxfp8.packCodes(tiled_codes, codes, n, k); + cuda.mxfp8.packScales(packed_scales, scales, n, k); + self.weight = cuda.mxfp8.weight(try self.rig.input(tiled_codes), try self.rig.input(packed_scales), n, k); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const bytes = self.lane.partBytes(shape.rows, self.weight); + if (bytes > 0) self.part = try self.rig.scratch(bytes); + self.out = try self.rig.output(shape.rows * n * 2); + return true; + } + + pub fn run(self: *Product) !void { + const d = self.d.?; + try self.lane.matmul(self.rig.s, self.x, d.shape.k, d.shape.rows, self.weight, self.out, self.part); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Product) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Product, want: []f64, magnitude: []f64) !void { + try self.d.?.reference(want, magnitude); + } + + pub fn simulate(self: *Product, accumulator: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(accumulator, order, out); + } +}; diff --git a/zig/tests/cuda/contract_projection_f32.zig b/zig/tests/cuda/contract_projection_f32.zig new file mode 100644 index 0000000000..1e172dec30 --- /dev/null +++ b/zig/tests/cuda/contract_projection_f32.zig @@ -0,0 +1,159 @@ +//! Packed projections writing FP32, with independent FP64 products and mirrored accumulator stress. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const Rig = @import("contract_rig.zig").Rig; + +pub const modes = [_]cuda.qmmf.Mode{ .bf16, .fp8, .fp8g, .mxfp8 }; +const served = [_]c.Shape{ + .{ .rows = 1, .n = 70, .k = 1024 }, + .{ .rows = 17, .n = 70, .k = 1024 }, + .{ .rows = 33, .n = 70, .k = 1024 }, + .{ .rows = 256, .n = 70, .k = 1024 }, + .{ .rows = 16, .n = 1280, .k = 2560 }, + .{ .rows = 16, .n = 2560, .k = 640 }, + .{ .rows = 7, .n = 2560, .k = 10240 }, + .{ .rows = 7, .n = 13952, .k = 2560 }, +}; + +fn depth(shape: c.Shape) usize { + return shape.k / 16 + 24; +} + +pub const Product = struct { + rig: Rig, + lane: *const cuda.qmmf.Lane, + mode: cuda.qmmf.Mode, + clusters: bool, + d: ?Dense = null, + weight: cuda.qmmf.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + part: u64 = 0, + + pub fn deinit(self: *Product) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Product) g.Kernel { + const name = switch (self.mode) { + .bf16 => if (self.clusters) "cuda BF16 fp32 clustered" else "cuda BF16 fp32 reduce", + .fp8 => if (self.clusters) "cuda FP8 fp32 clustered" else "cuda FP8 fp32 reduce", + .fp8g => if (self.clusters) "cuda block FP8 fp32 clustered" else "cuda block FP8 fp32 reduce", + .mxfp8 => if (self.clusters) "cuda MXFP8 fp32 clustered" else "cuda MXFP8 fp32 reduce", + else => unreachable, + }; + return .{ .spec = .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{ .size = 512, .step = 128, .min = 128 } }, + .served = &served, + }, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Product); + + pub fn prepare(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + self.part = 0; + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 128 != 0) return false; + const d = try Dense.init(self.rig.gpu.gpa, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + try self.prepareWeight(shape, fill, seed); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const bytes = self.lane.partBytes(shape.rows, self.weight); + if (bytes > 0) self.part = try self.rig.scratch(bytes); + self.out = try self.rig.output(shape.rows * shape.n * 4); + return true; + } + + fn prepareWeight(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !void { + var arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const d = self.d.?; + const n = shape.n; + const k = shape.k; + const npad = cuda.fp8.padded(n); + var prng = std.Random.DefaultPrng.init(seed ^ 0xf32); + const random = prng.random(); + const codes = try a.alloc(u8, n * k); + const words = try a.alloc(u16, n * k); + const ue8 = try a.alloc(u8, n * k / 32); + for (ue8) |*v| v.* = if (fill == .stress) 127 else 123 + random.uintLessThan(u8, 8); + const inv = try a.alloc(f32, ((n + 127) / 128) * (k / 128)); + for (inv) |*v| v.* = if (fill == .stress) 1 else std.math.ldexp(@as(f32, 1), @as(i32, random.uintLessThan(u8, 8)) - 4); + for (0..n) |j| for (0..k) |i| { + const at = j * k + i; + codes[at] = if (fill == .stress) 0x38 + random.uintLessThan(u8, 24) else random.uintLessThan(u8, 0x7f); + if (fill != .stress and random.boolean()) codes[at] |= 0x80; + words[at] = c.toBf16(if (fill == .stress) @floatFromInt(1 + random.uintLessThan(u8, 7)) else random.float(f32) * 4 - 2); + }; + if (fill == .stress) for (0..n) |j| for (k / 2..k) |i| { + codes[j * k + i] = codes[j * k + i - k / 2] | 0x80; + words[j * k + i] = words[j * k + i - k / 2] | 0x8000; + }; + const scale: f32 = if (fill == .stress) 1 else 0.125; + for (0..n) |j| for (0..k) |i| { + const value = core.quant.fp8.e4m3(codes[j * k + i]); + d.w[j * k + i] = switch (self.mode) { + .bf16 => c.fromBf16(words[j * k + i]), + .fp8 => value * scale, + .fp8g => value * inv[(j / 128) * (k / 128) + i / 128], + .mxfp8 => value * std.math.ldexp(@as(f64, 1), @as(i32, ue8[j * (k / 32) + i / 32]) - 127), + else => unreachable, + }; + }; + if (self.mode == .bf16) { + const tiles = try a.alloc(u16, npad * k); + cuda.qmmf.packBf16(tiles, words, n, k); + self.weight = cuda.modelopt.bf16Weight(try self.rig.input(std.mem.sliceAsBytes(tiles)), n, k); + return; + } + const tiles = try a.alloc(u8, npad * k); + cuda.fp8.packCodes(tiles, codes, n, k); + const dw = try self.rig.input(tiles); + self.weight = switch (self.mode) { + .fp8 => cuda.modelopt.fp8Weight(dw, scale, n, k), + .mxfp8 => blk: { + const scales = try a.alloc(u8, npad * k / 32); + cuda.mxfp8.packScales(scales, ue8, n, k); + break :blk cuda.mxfp8.weight(dw, try self.rig.input(scales), n, k); + }, + .fp8g => blk: { + const scales = try a.alloc(f32, npad * k / 64); + cuda.fp8.packScales(scales, inv, n, k); + break :blk cuda.fp8.weight(dw, try self.rig.input(std.mem.sliceAsBytes(scales)), n, k); + }, + else => unreachable, + }; + } + + pub fn run(self: *Product) !void { + const shape = self.d.?.shape; + try self.lane.matmulOut(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out, .f32, self.part); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Product) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Product, want: []f64, magnitude: []f64) !void { + try self.d.?.reference(want, magnitude); + } + + pub fn simulate(self: *Product, accumulator: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(accumulator, order, out); + } +}; diff --git a/zig/tests/cuda/contract_qsa_ref.zig b/zig/tests/cuda/contract_qsa_ref.zig new file mode 100644 index 0000000000..b0e83b5a5a --- /dev/null +++ b/zig/tests/cuda/contract_qsa_ref.zig @@ -0,0 +1,329 @@ +//! Independent sort-based QSA membership, block-order emission and dense-row oracle. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +fn depth(_: c.Shape) usize { + return 0; +} + +pub fn spec(padded: bool) c.Spec { + return .{ + .name = if (padded) "cuda fn_qsa_select padded scores" else "cuda fn_qsa_select packed scores", + .accum = .none, + .depth = depth, + .out = .raw, + .nan = .{ .stated = "float total order includes signed NaN payloads and signed zero; ties choose earlier blocks" }, + // rows = query rows, n = first row's complete blocks, k = its tail length plus one. + .tiles = .{ .{ .size = 32 }, .{ .size = 1024, .min = 0 }, .{ .size = 4 } }, + .served = &.{ + .{ .rows = 8, .n = 1025, .k = 1 }, + .{ .rows = 1, .n = 0, .k = 2 }, + .{ .rows = 4, .n = 0, .k = 4 }, + .{ .rows = 8, .n = 512, .k = 1 }, + .{ .rows = 1, .n = 511, .k = 4 }, + .{ .rows = 1, .n = 512, .k = 4 }, + .{ .rows = 1, .n = 513, .k = 1 }, + .{ .rows = 256, .n = 513, .k = 4 }, + .{ .rows = 3, .n = 1023, .k = 2 }, + .{ .rows = 5, .n = 1024, .k = 3 }, + .{ .rows = 7, .n = 1025, .k = 4 }, + .{ .rows = 5, .n = 4097, .k = 2 }, + .{ .rows = 1, .n = 8192, .k = 4 }, + }, + }; +} + +pub const Layout = struct { + rows: usize, + nb: usize, + pos0: i32, + top: usize = 512, + ratio: usize = 4, + idw: usize = 2052, + + pub fn valid(l: Layout) bool { + if (l.rows == 0 or l.rows > 256 or l.nb == 0 or l.nb > 8263 or l.pos0 < 0) return false; + if (l.top == 0 or l.top > 512 or l.ratio == 0 or l.ratio > 32) return false; + if (l.idw < l.top * l.ratio + l.ratio - 1 or l.idw > 16384) return false; + const last = @as(usize, @intCast(l.pos0)) + l.rows; + return last <= std.math.maxInt(i32) and last / l.ratio <= l.nb; + } +}; + +/// Fixture bounds only: the production launcher is unchanged. +pub fn layout(s: c.Shape, padded: bool) ?Layout { + if (s.rows == 0 or s.rows > 256 or s.n > 8192 or s.k == 0 or s.k > 4) return null; + const end = s.n * 4 + s.k - 1; + if (end == 0) return null; + const l: Layout = .{ + .rows = s.rows, + .nb = @max(1, (end + s.rows - 1) / 4) + (if (padded) @as(usize, 7) else 0), + .pos0 = @intCast(end - 1), + }; + return if (l.valid()) l else null; +} + +const Entry = struct { score: f32, index: usize }; + +/// Numeric comparisons plus explicit IEEE sign/NaN payload policy; no radix key transform or scan is reused. +fn before(_: void, a: Entry, b: Entry) bool { + const an = std.math.isNan(a.score); + const bn = std.math.isNan(b.score); + const ab: u32 = @bitCast(a.score); + const bb: u32 = @bitCast(b.score); + const negative_a = std.math.signbit(a.score); + const negative_b = std.math.signbit(b.score); + if (an or bn) { + if (an and bn) { + if (negative_a != negative_b) return !negative_a; + const ap = ab & 0x7fffff; + const bp = bb & 0x7fffff; + if (ap != bp) return if (negative_a) ap < bp else ap > bp; + } else return if (an) !negative_a else negative_b; + } else { + if (a.score != b.score) return a.score > b.score; + if (negative_a != negative_b) return !negative_a; + } + return a.index < b.index; +} + +pub const Result = struct { + ids: []i32, + nk: []i32, + sparse: []i32, + + pub fn deinit(r: Result, a: std.mem.Allocator) void { + a.free(r.ids); + a.free(r.nk); + a.free(r.sparse); + } + + pub fn bytes(r: Result, a: std.mem.Allocator) ![]u8 { + const out = try a.alloc(u8, (r.ids.len + r.nk.len + r.sparse.len) * 4); + var i: usize = 0; + for ([_][]const i32{ r.ids, r.nk, r.sparse }) |part| for (part) |v| { + std.mem.writeInt(i32, out[i..][0..4], v, .little); + i += 4; + }; + return out; + } +}; + +/// Choose membership by a full sort, then enumerate selected blocks in original order and append live tail tokens. +pub fn select(a: std.mem.Allocator, scores: []const f32, l: Layout) !Result { + if (!l.valid() or scores.len != l.rows * l.nb) return error.InvalidFixture; + const ids = try a.alloc(i32, l.rows * l.idw); + errdefer a.free(ids); + const nk = try a.alloc(i32, l.rows); + errdefer a.free(nk); + const sparse = try a.alloc(i32, l.rows); + errdefer a.free(sparse); + @memset(ids, -1); + const entries = try a.alloc(Entry, l.nb); + defer a.free(entries); + const chosen = try a.alloc(bool, l.nb); + defer a.free(chosen); + for (0..l.rows) |r| { + const end = @as(usize, @intCast(l.pos0)) + r + 1; + const complete = end / l.ratio; + if (complete <= l.top) { + nk[r] = @intCast(end); + sparse[r] = 0; + continue; + } + for (entries[0..complete], scores[r * l.nb ..][0..complete], 0..) |*entry, value, b| + entry.* = .{ .score = value, .index = b }; + std.mem.sort(Entry, entries[0..complete], {}, before); + @memset(chosen, false); + for (entries[0..l.top]) |entry| chosen[entry.index] = true; + var at = r * l.idw; + for (chosen[0..complete], 0..) |yes, b| if (yes) { + for (0..l.ratio) |j| { + ids[at] = @intCast(b * l.ratio + j); + at += 1; + } + }; + for (complete * l.ratio..end) |token| { + ids[at] = @intCast(token); + at += 1; + } + nk[r] = @intCast(at - r * l.idw); + sparse[r] = 1; + } + return .{ .ids = ids, .nk = nk, .sparse = sparse }; +} + +pub const special_bits = [_]u32{ + 0xffc00007, 0xff800001, 0xff800000, 0x80000000, 0x00000000, + 0x7f800000, 0x7f800001, 0x7fc00001, 0x7fc00007, +}; + +pub fn fill(scores: []f32, l: Layout, kind: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // Discrete random scores deliberately create ties on both sides of the selection boundary. + for (0..l.rows) |r| { + const complete = (@as(usize, @intCast(l.pos0)) + r + 1) / l.ratio; + for (scores[r * l.nb ..][0..l.nb], 0..) |*v, b| { + v.* = if (b < complete) @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 33))) - 16) else @bitCast(@as(u32, 0x7fffffff)); + if (kind == .nonfinite and b < complete) v.* = switch (r % 4) { + 0 => 0, + 1 => @bitCast(@as(u32, 0x7fc00000) + @as(u32, @intCast(b % 1024 + 1))), + 2 => @bitCast(@as(u32, 0xffc00000) + @as(u32, @intCast(b % 1024 + 1))), + else => if (b % 2 == 0) 0.0 else -0.0, + }; + } + if (kind == .nonfinite) for (special_bits, 0..) |bits, i| { + if (i < complete) scores[r * l.nb + i] = @bitCast(bits); + }; + } +} + +pub fn compare(want: []const u8, got: []const u8) ?[]const u8 { + return if (std.mem.eql(u8, want, got)) null else "selected ids, row counts, sparse flags or untouched padding differ"; +} + +test "independent score order specifies infinities, signed zeros, signed NaNs and payloads" { + var entries: [special_bits.len]Entry = undefined; + for (&entries, special_bits, 0..) |*entry, bits, i| entry.* = .{ .score = @bitCast(bits), .index = i }; + std.mem.sort(Entry, &entries, {}, before); + for (entries, 0..) |entry, i| try std.testing.expectEqual(special_bits.len - 1 - i, entry.index); +} + +test "membership ties choose earliest blocks and output follows block order rather than rank order" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 5, .pos0 = 6, .top = 2, .ratio = 2, .idw = 6 }; + const result = try select(a, &.{ 1, 9, 9, 999, 999 }, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 2, 3, 4, 5, 6, -1 }, result.ids); + try std.testing.expectEqualSlices(i32, &.{5}, result.nk); + try std.testing.expectEqualSlices(i32, &.{1}, result.sparse); + const ranked = try select(a, &.{ 8, 0, 9, 999, 999 }, l); + defer ranked.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 4, 5, 6, -1 }, ranked.ids); +} + +test "dense-to-sparse transition leaves dense ids untouched and excludes future poisoned scores" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 3, .nb = 6, .pos0 = 4, .top = 2, .ratio = 2, .idw = 6 }; + const scores = [_]f32{ 9, 8, 999, 999, 999, 999, 7, 8, 9, 999, 999, 999, 9, 8, 7, 999, 999, 999 }; + const result = try select(a, &scores, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, -1, -1, -1, -1 }, result.ids[0..6]); + try std.testing.expectEqualSlices(i32, &.{ 2, 3, 4, 5, -1, -1 }, result.ids[6..12]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 6, -1 }, result.ids[12..]); + try std.testing.expectEqualSlices(i32, &.{ 5, 4, 5 }, result.nk); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 1 }, result.sparse); +} + +test "raw comparison catches omitted tails, wrong tie ids, unused padding and flag corruption" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 3, .pos0 = 6, .top = 2, .ratio = 2, .idw = 6 }; + const result = try select(a, &.{ 1, 1, 1 }, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 6, -1 }, result.ids); + const want = try result.bytes(a); + defer a.free(want); + const broken = try a.dupe(u8, want); + defer a.free(broken); + for ([_]usize{ 0, 16, 20, 24, 28 }) |byte| { + @memcpy(broken, want); + broken[byte] ^= 1; + try std.testing.expect(compare(want, broken) != null); + } + try std.testing.expect(compare(want, want) == null); +} + +test "nonfinite membership places signed NaNs and signed zero on the correct side of a cutoff" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 9, .pos0 = 8, .top = 2, .ratio = 1, .idw = 4 }; + var scores: [9]f32 = undefined; + for (&scores, special_bits) |*value, bits| value.* = @bitCast(bits); + const result = try select(a, &scores, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 7, 8, -1, -1 }, result.ids); + const zeros = try select(a, &.{ -0.0, 0.0 }, .{ .rows = 1, .nb = 2, .pos0 = 1, .top = 1, .ratio = 1, .idw = 2 }); + defer zeros.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 1, -1 }, zeros.ids); +} + +test "finite random sort membership agrees with an independent quadratic outranker count" { + const a = std.testing.allocator; + for ([_]usize{ 1, 3, 4, 17, 33, 65 }) |blocks| { + const l: Layout = .{ .rows = 3, .nb = blocks + 8, .pos0 = @intCast(blocks * 2 - 1), .top = 3, .ratio = 2, .idw = 9 }; + const scores = try a.alloc(f32, l.rows * l.nb); + defer a.free(scores); + fill(scores, l, .random, blocks); + const result = try select(a, scores, l); + defer result.deinit(a); + for (0..l.rows) |r| { + const end = @as(usize, @intCast(l.pos0)) + r + 1; + const complete = end / l.ratio; + const ids = result.ids[r * l.idw ..][0..l.idw]; + if (complete <= l.top) { + try std.testing.expectEqual(@as(i32, 0), result.sparse[r]); + for (ids) |id| try std.testing.expectEqual(@as(i32, -1), id); + continue; + } + var at: usize = 0; + const row = scores[r * l.nb ..][0..complete]; + for (row, 0..) |value, b| { + var better: usize = 0; + for (row, 0..) |other, j| if (other > value or (other == value and j < b)) { + better += 1; + }; + if (better < l.top) for (0..l.ratio) |j| { + try std.testing.expectEqual(@as(i32, @intCast(b * l.ratio + j)), ids[at]); + at += 1; + }; + } + for (complete * l.ratio..end) |token| { + try std.testing.expectEqual(@as(i32, @intCast(token)), ids[at]); + at += 1; + } + try std.testing.expectEqual(@as(i32, @intCast(at)), result.nk[r]); + for (ids[at..]) |id| try std.testing.expectEqual(@as(i32, -1), id); + } + } +} + +test "fixture domains retain all served and valid tile edges and explicitly refuse unsafe layouts" { + const a = std.testing.allocator; + for ([_]bool{ false, true }) |padded| { + const s = spec(padded); + for (s.served) |shape| try std.testing.expect(layout(shape, padded) != null); + const edges = try c.shapes(a, s); + defer a.free(edges); + for (edges) |shape| if (shape.k <= 4 and shape.n * 4 + shape.k > 1) + try std.testing.expect(layout(shape, padded) != null); + } + try std.testing.expect(layout(.{ .rows = 1, .n = 1, .k = 5 }, false) == null); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = -1 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 16 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 6, .ratio = 0 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 6, .idw = 2048 })); +} + +test "directed positive NaN payload cutoff spans two blocks and seeded inputs repeat exactly" { + const a = std.testing.allocator; + const l = layout(.{ .rows = 8, .n = 1025, .k = 1 }, true).?; + const scores = try a.alloc(f32, l.rows * l.nb); + defer a.free(scores); + const repeated = try a.alloc(f32, scores.len); + defer a.free(repeated); + fill(scores, l, .nonfinite, 91); + for (0..5) |_| { + fill(repeated, l, .nonfinite, 91); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(scores), std.mem.sliceAsBytes(repeated)); + } + const result = try select(a, scores, l); + defer result.deinit(a); + // Row 1's 512 largest positive NaN payloads are blocks 512 through 1023. + for (0..512) |b| for (0..4) |j| + try std.testing.expectEqual(@as(i32, @intCast((512 + b) * 4 + j)), result.ids[l.idw + b * 4 + j]); + try std.testing.expectEqual(@as(i32, 4100), result.ids[l.idw + 2048]); + try std.testing.expectEqual(@as(i32, 2049), result.nk[1]); + for (result.ids[l.idw + 2049 .. 2 * l.idw]) |id| try std.testing.expectEqual(@as(i32, -1), id); +} diff --git a/zig/tests/cuda/contract_qsa_select.zig b/zig/tests/cuda/contract_qsa_select.zig new file mode 100644 index 0000000000..e253a8d1a8 --- /dev/null +++ b/zig/tests/cuda/contract_qsa_select.zig @@ -0,0 +1,102 @@ +//! Guarded fn_qsa_select contracts against an independent sort-based selected-list oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_qsa_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Select = struct { + rig: rig_.Rig, + f: cuda.Function, + padded: bool, + l: ref.Layout = .{ .rows = 0, .nb = 0, .pos0 = 0 }, + scores: []f32 = &.{}, + want: []u8 = &.{}, + score_ptr: u64 = 0, + pos_ptr: u64 = 0, + ids_ptr: u64 = 0, + nk_ptr: u64 = 0, + sparse_ptr: u64 = 0, + + pub fn deinit(self: *Select) void { + self.rig.gpu.gpa.free(self.scores); + self.rig.gpu.gpa.free(self.want); + self.rig.deinit(); + } + + pub fn kernel(self: *Select) g.Kernel { + return .{ .spec = ref.spec(self.padded), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Select); + + pub fn prepare(self: *Select, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const l = ref.layout(shape, self.padded) orelse return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.scores); + a.free(self.want); + self.scores = &.{}; + self.want = &.{}; + self.l = l; + self.scores = try a.alloc(f32, l.rows * l.nb); + ref.fill(self.scores, l, fill, seed); + const result = try ref.select(a, self.scores, l); + defer result.deinit(a); + self.want = try result.bytes(a); + self.score_ptr = try self.rig.input(std.mem.sliceAsBytes(self.scores)); + self.pos_ptr = try self.rig.input(std.mem.asBytes(&self.l.pos0)); + self.ids_ptr = try self.rig.output(l.rows * l.idw * 4); + self.nk_ptr = try self.rig.output(l.rows * 4); + self.sparse_ptr = try self.rig.output(l.rows * 4); + return true; + } + + pub fn run(self: *Select) !void { + var args: cuda.Args = .{}; + args.add(self.score_ptr); + args.add(self.pos_ptr); + args.add(self.ids_ptr); + args.add(self.nk_ptr); + args.add(self.sparse_ptr); + args.add(@as(c_int, @intCast(self.l.nb))); + args.add(@as(c_int, @intCast(self.l.idw))); + args.add(@as(c_int, @intCast(self.l.top))); + args.add(@as(c_int, @intCast(self.l.ratio))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.l.rows) }, .block = .{ .x = 1024 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Select) usize { + return self.want.len; + } + + pub fn reference(self: *Select, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, byte| { + w.* = @floatFromInt(byte); + m.* = 0; + } + } + + pub fn check(self: *Select, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "selection changed a buffer guard"; + const inputs = try self.rig.inputBytes(); + const score_bytes = std.mem.sliceAsBytes(self.scores); + if (inputs.len != score_bytes.len + 4 or !std.mem.eql(u8, score_bytes, inputs[0..score_bytes.len]) or + !std.mem.eql(u8, std.mem.asBytes(&self.l.pos0), inputs[score_bytes.len..])) + return "selection changed scores or the position input"; + if (at == .nonfinite) { + for (ref.special_bits) |bits| { + var found = false; + for (self.scores) |value| if (@as(u32, @bitCast(value)) == bits) { + found = true; + break; + }; + if (!found) return "nonfinite fixture lacks a required signed NaN, infinity or zero"; + } + } + return ref.compare(self.want, try self.rig.outputBytes()); + } +}; diff --git a/zig/tests/cuda/contract_row_best.zig b/zig/tests/cuda/contract_row_best.zig new file mode 100644 index 0000000000..e63acda285 --- /dev/null +++ b/zig/tests/cuda/contract_row_best.zig @@ -0,0 +1,85 @@ +//! The baseline proposal row reduction through guards, first-ID checks and the independent float64 LSE oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_row_best_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Best = struct { + rig: rig_.Rig, + f: cuda.Function, + shape: c.Shape = undefined, + bits: []u16 = &.{}, + input: u64 = 0, + output: u64 = 0, + probe: u64 = 0, + + pub fn deinit(self: *Best) void { + self.rig.gpu.gpa.free(self.bits); + self.rig.deinit(); + } + + const vt = g.hooks(Best); + + pub fn kernel(self: *Best) g.Kernel { + return .{ .spec = ref.spec, .ctx = self, .vt = &vt }; + } + + pub fn prepare(self: *Best, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.admits(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.bits); + self.bits = &.{}; + self.shape = shape; + self.bits = try a.alloc(u16, shape.rows * (shape.k + 17)); + ref.fill(self.bits, shape, fill, seed); + self.input = try self.rig.input(std.mem.sliceAsBytes(self.bits)); + self.output = try self.rig.output(self.count() * 4); + self.probe = try self.rig.scratch(16); + return true; + } + + fn launch(self: *Best, input: u64, output: u64, rows: usize) !void { + var args: cuda.Args = .{}; + args.add(input); + args.add(@as(c_int, @intCast(self.shape.k))); + args.add(@as(c_longlong, @intCast(self.shape.k + 17))); + args.add(output); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn run(self: *Best) !void { + try self.launch(self.input, self.output, self.shape.rows); + } + + pub fn count(self: *Best) usize { + return self.shape.rows * 4; + } + + pub fn reference(self: *Best, want: []f64, mag: []f64) !void { + ref.reference(self.bits, self.shape, want, mag); + } + + pub fn simulate(self: *Best, acc: c.Accum, order: g.Order, out: []f64) !void { + try ref.simulate(self.rig.gpu.gpa, self.bits, self.shape, acc, order, out); + } + + pub fn check(self: *Best, gate: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const actual = try a.dupe(u8, try self.rig.outputBytes()); + defer a.free(actual); + if (ref.check(self.bits, self.shape, actual)) |why| return why; + if (gate == .shapes) for (0..self.shape.rows) |row| { + try self.launch(self.input + row * (self.shape.k + 17) * 2, self.probe, 1); + if (!std.mem.eql(u8, actual[row * 16 ..][0..16], try self.rig.bytes(self.probe, 16))) + return "a batched reduction row differs from the same input reduced alone"; + }; + if (!try self.rig.intact()) return "row reduction changed a payload guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.bits), try self.rig.inputBytes())) return "row reduction changed immutable logits"; + return null; + } +}; diff --git a/zig/tests/cuda/contract_row_best_ref.zig b/zig/tests/cuda/contract_row_best_ref.zig new file mode 100644 index 0000000000..c9739c2669 --- /dev/null +++ b/zig/tests/cuda/contract_row_best_ref.zig @@ -0,0 +1,166 @@ +//! Independent first-maximum and stable float64 log-sum-exp oracles for the proposal row reduction. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +fn depth(shape: c.Shape) usize { + return (shape.k + 1023) / 1024 + 10; +} + +const served = blk: { + const widths = [_]usize{ 31, 1, 32, 33, 1023, 1024, 1025, 79_591, 248_320 }; + var shapes: [widths.len + 1]c.Shape = undefined; + for (widths, 0..) |k, i| shapes[i] = .{ .rows = 12, .n = 4, .k = k }; + shapes[widths.len] = .{ .rows = 33, .n = 4, .k = 1025 }; + break :blk shapes; +}; + +pub const spec: c.Spec = .{ + .name = "cuda Flash Next fn_row_best", + .accum = .f32, + .depth = depth, + .ops = 2, + .out = .f32, + .nan = .{ .stated = "NaNs do not win; all NaN/-Inf retain -Inf and INT_MAX; a NaN or infinite maximum makes LSE NaN" }, + .tiles = .{ .{ .size = 32 }, .{}, .{ .size = 1024 } }, + .served = &served, +}; + +pub fn admits(shape: c.Shape) bool { + return shape.rows > 0 and shape.rows <= 64 and shape.n == 4 and shape.k > 0 and shape.k <= 248_320; +} + +pub fn fill(bits: []u16, shape: c.Shape, mode: g.Fill, seed: u64) void { + @memset(bits, 0x7fc1); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..shape.rows) |row| for (0..shape.k) |i| { + bits[row * (shape.k + 17) + i] = if (mode == .stress) 0 else switch (row % 12) { + 0 => if (i == 0 or i + 1 == shape.k) 0x4080 else 0xbf80, + 1 => if (i % 2 == 0) 0x8000 else 0, + 2 => 0x7fc1, + 3 => 0xff80, + 4 => if (i == shape.k / 3 or i + 1 == shape.k) 0x7f80 else 0xff80, + 5 => if (i % 3 == 0) 0x7fc1 else if (i % 3 == 1) 0x3f80 else 0xbf80, + 6 => if (i == shape.k / 2) 0x7f7f else 0xff7f, + 7 => if (i == shape.k / 2) 0x007f else 0x8001, + 8 => c.toBf16(rng.float(f32) * 16 - 8), + 9 => c.toBf16(rng.float(f32) * 2000 - 1000), + 10 => rng.int(u16), + else => @truncate(i *% 40503 +% 317), + }; + }; +} + +pub const Result = struct { value: f32, id: u32, lse: f64 }; + +pub fn expected(bits: []const u16) Result { + var result = Result{ .value = -std.math.inf(f32), .id = std.math.maxInt(i32), .lse = std.math.nan(f64) }; + var nan = false; + for (bits, 0..) |b, i| { + const v = c.fromBf16(b); + nan = nan or std.math.isNan(v); + if (v > result.value) { + result.value = v; + result.id = @intCast(i); + } + } + if (nan or !std.math.isFinite(result.value)) return result; + var sum: f64 = 0; + for (bits) |b| sum += @exp(@as(f64, c.fromBf16(b)) - @as(f64, result.value)); + result.lse = @as(f64, result.value) + @log(sum); + return result; +} + +pub fn reference(bits: []const u16, shape: c.Shape, want: []f64, mag: []f64) void { + for (0..shape.rows) |row| { + const e = expected(bits[row * (shape.k + 17) ..][0..shape.k]); + const at = row * 4; + want[at..][0..4].* = .{ e.value, e.lse, @as(f32, @bitCast(e.id)), 0 }; + mag[at..][0..4].* = .{ 0, 1, 0, 0 }; + } +} + +pub fn simulate(a: std.mem.Allocator, bits: []const u16, shape: c.Shape, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try a.alloc(f64, shape.k); + defer a.free(terms); + for (0..shape.rows) |row| { + const input = bits[row * (shape.k + 17) ..][0..shape.k]; + const e = expected(input); + for (input, terms) |b, *term| term.* = @exp(@as(f64, c.fromBf16(b)) - @as(f64, e.value)); + const sum = core.contract.rows.sumIn(acc, order, terms); + const at = row * 4; + out[at..][0..4].* = .{ e.value, @as(f32, @floatCast(@as(f64, e.value) + @log(sum))), @as(f32, @bitCast(e.id)), 0 }; + } +} + +pub fn check(bits: []const u16, shape: c.Shape, output: []const u8) ?[]const u8 { + for (0..shape.rows) |row| { + const e = expected(bits[row * (shape.k + 17) ..][0..shape.k]); + const actual = output[row * 16 ..][0..16]; + if (std.mem.readInt(u32, actual[0..4], .little) != @as(u32, @bitCast(e.value))) return "maximum bits differ from first-maximum oracle"; + if (std.mem.readInt(u32, actual[8..12], .little) != e.id) return "maximum ID differs from independent first-maximum oracle"; + if (std.mem.readInt(u32, actual[12..16], .little) != 0) return "reserved output field was not zeroed"; + const got = c.read(.f32, actual, 1); + if (std.math.isFinite(e.lse)) { + if (!std.math.isFinite(got) or @abs(got - e.lse) > c.bound(spec, shape, e.lse, 1)) return "LSE differs from stable independent float64 oracle"; + } else if (c.classOf(got) != c.classOf(e.lse)) return "LSE nonfinite class differs from the declared formula"; + } + return null; +} + +test "first-maximum oracle preserves first zero bits and excludes all non-winners" { + const zero = expected(&.{ 0x8000, 0 }); + try std.testing.expectEqual(@as(u32, 0x80000000), @as(u32, @bitCast(zero.value))); + try std.testing.expectEqual(@as(u32, 0), zero.id); + const tied = expected(&.{ 0xbf80, 0x4080, 0x4080 }); + try std.testing.expectEqual(@as(u32, 1), tied.id); + for ([_][]const u16{ &.{0x7fc1}, &.{0xff80}, &.{ 0x7fc1, 0xff80 } }) |input| { + const none = expected(input); + try std.testing.expectEqual(@as(u32, std.math.maxInt(i32)), none.id); + try std.testing.expectEqual(c.Class.neg_inf, c.classOf(none.value)); + try std.testing.expect(std.math.isNan(none.lse)); + } + try std.testing.expect(std.math.isNan(expected(&.{ 0x7fc1, 0x3f80 }).lse)); + try std.testing.expect(std.math.isNan(expected(&.{0x7f80}).lse)); +} + +test "equal logits prove FP32 exponential sums and expose both lower precision accumulators" { + const a = std.testing.allocator; + const shape: c.Shape = .{ .rows = 1, .n = 4, .k = 248_320 }; + const bits = try a.alloc(u16, shape.k + 17); + defer a.free(bits); + fill(bits, shape, .stress, 0); + var want: [4]f64 = undefined; + var mag: [4]f64 = undefined; + reference(bits, shape, &want, &mag); + const allow = c.boundExact(spec, want[1]); + for (std.enums.values(g.Order)) |order| { + var got: [4]f64 = undefined; + try simulate(a, bits, shape, .f32, order, &got); + try std.testing.expect(@abs(got[1] - want[1]) <= allow); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + try simulate(a, bits, shape, acc, order, &got); + try std.testing.expect(!std.math.isFinite(got[1]) or @abs(got[1] - want[1]) / allow >= 100); + } + } +} + +test "host checker rejects a last-tie result and a corrupted log-sum-exp" { + const shape: c.Shape = .{ .rows = 1, .n = 4, .k = 2 }; + var bits: [19]u16 = @splat(0x7fc1); + bits[0..2].* = .{ 0x4080, 0x4080 }; + const e = expected(bits[0..2]); + var out: [16]u8 = @splat(0); + std.mem.writeInt(u32, out[0..4], @bitCast(e.value), .little); + std.mem.writeInt(u32, out[4..8], @bitCast(@as(f32, @floatCast(e.lse))), .little); + try std.testing.expect(check(&bits, shape, &out) == null); + std.mem.writeInt(u32, out[8..12], 1, .little); + try std.testing.expect(check(&bits, shape, &out) != null); + std.mem.writeInt(u32, out[8..12], 0, .little); + std.mem.writeInt(u32, out[4..8], @bitCast(@as(f32, 4)), .little); + try std.testing.expect(check(&bits, shape, &out) != null); + try std.testing.expect(!admits(.{ .rows = 1, .n = 4, .k = 0 })); + try std.testing.expect(!admits(.{ .rows = 1, .n = 4, .k = 248_321 })); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e87a2e90e5..425a73a37a 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -7,7 +7,13 @@ const check = @import("check.zig"); const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); +const mx_k = @import("contract_mxfp8.zig"); +const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); +const writeback_k = @import("contract_flashnext_writeback.zig"); +const utility_k = @import("contract_utility.zig"); +const best_k = @import("contract_row_best.zig"); +const qsa_k = @import("contract_qsa_select.zig"); const contract = core.contract; const g = contract.gates; @@ -21,7 +27,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -86,9 +92,57 @@ pub fn run(gpu: check.Gpu) !void { } defer for (&mm) |*m| m.deinit(); for (&mm) |*m| try kernels.append(a, m.kernel()); + var mx = [_]mx_k.Product{ + .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .clusters = true }, + .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane_reduce, .clusters = false }, + }; + defer for (&mx) |*m| m.deinit(); + for (&mx) |*m| try kernels.append(a, m.kernel()); + var wide: [wide_k.modes.len * 2]wide_k.Product = undefined; + for (wide_k.modes, 0..) |mode, i| { + wide[i * 2] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = mode, .clusters = true }; + wide[i * 2 + 1] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane_reduce, .mode = mode, .clusters = false }; + } + defer for (&wide) |*m| m.deinit(); + for (&wide) |*m| try kernels.append(a, m.kernel()); var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var hc = try cuda.Module.load(gpu.d, kk.flashnext_hc); + defer hc.unload(); + var writebacks: [3]writeback_k.Writeback = undefined; + for (std.enums.values(writeback_k.Mode), &writebacks) |mode, *w| { + const symbol: [:0]const u8 = switch (mode) { + .sums => "fn_hc_writeback0", + .moe => "fn_hc_writeback_moe", + .slices => "fn_hc_writeback_slices", + }; + w.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try hc.function(symbol), .mode = mode }; + } + defer for (&writebacks) |*w| w.deinit(); + for (&writebacks) |*w| try kernels.append(a, w.kernel()); + var batch = try cuda.Module.load(gpu.d, kk.flashnext_batch); + defer batch.unload(); + var sums: [2]utility_k.Sum = undefined; + var shifts: [2]utility_k.Shift = undefined; + for ([_]bool{ false, true }, &sums, &shifts) |variant, *sum, *shift| { + sum.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_sum_slices"), .bf16 = variant }; + shift.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_shift_window"), .padded = variant }; + } + defer for (&sums) |*sum| sum.deinit(); + defer for (&shifts) |*shift| shift.deinit(); + for (&sums) |*sum| try kernels.append(a, sum.kernel()); + for (&shifts) |*shift| try kernels.append(a, shift.kernel()); + var mtp = try cuda.Module.load(gpu.d, kk.flashnext_mtp); + defer mtp.unload(); + var best = best_k.Best{ .rig = .{ .gpu = gpu, .s = s }, .f = try mtp.function("fn_row_best") }; + defer best.deinit(); + try kernels.append(a, best.kernel()); + var selections: [2]qsa_k.Select = undefined; + for ([_]bool{ false, true }, &selections) |padded, *selection| + selection.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_qsa_select"), .padded = padded }; + defer for (&selections) |*selection| selection.deinit(); + for (&selections) |*selection| try kernels.append(a, selection.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); diff --git a/zig/tests/cuda/contract_utility.zig b/zig/tests/cuda/contract_utility.zig new file mode 100644 index 0000000000..dfe44d065e --- /dev/null +++ b/zig/tests/cuda/contract_utility.zig @@ -0,0 +1,173 @@ +//! Flash Next's sum and history utilities through the guarded numerical contract. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_utility_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Sum = struct { + rig: rig_.Rig, + f: cuda.Function, + bf16: bool, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + terms: []f32 = &.{}, + part: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Sum) void { + self.rig.gpu.gpa.free(self.terms); + self.rig.deinit(); + } + + pub fn kernel(self: *Sum) g.Kernel { + return .{ .spec = ref.sumSpec(self.bf16), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Sum); + + pub fn prepare(self: *Sum, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.sumShape(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.terms); + self.terms = &.{}; + self.shape = shape; + self.terms = try a.alloc(f32, self.count() * shape.k); + ref.sumFill(self.terms, shape, fill, seed); + self.part = try self.rig.input(std.mem.sliceAsBytes(self.terms)); + self.out = try self.rig.output(self.count() * c.size(ref.sumSpec(self.bf16).out)); + return true; + } + + pub fn run(self: *Sum) !void { + var args: cuda.Args = .{}; + args.add(self.part); + args.add(self.out); + args.add(@as(c_int, if (self.bf16) 1 else 0)); + args.add(@as(c_longlong, @intCast(self.count()))); + args.add(@as(c_int, @intCast(self.shape.k))); + try cuda.launch.launch(self.f, .{ + .grid = .{ .x = @intCast((self.count() + 255) / 256) }, + .block = .{ .x = 256 }, + }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Sum) usize { + return self.shape.rows * self.shape.n; + } + + pub fn reference(self: *Sum, want: []f64, mag: []f64) !void { + ref.sumReference(self.terms, self.shape, want, mag); + } + + pub fn simulate(self: *Sum, acc: c.Accum, order: g.Order, out: []f64) !void { + ref.sumSimulate(self.terms, self.shape, acc, order, out); + } + + /// Retain the float64 error bound while additionally checking ordered f32 and single final RNE. + pub fn check(self: *Sum, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a sum launch changed a guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.terms), try self.rig.inputBytes())) + return "a sum launch changed its immutable slice inputs"; + return ref.sumCheck(self.rig.gpu.gpa, self.terms, self.shape, self.bf16, try self.rig.outputBytes(), at == .nonfinite); + } +}; + +pub const Shift = struct { + rig: rig_.Rig, + f: cuda.Function, + padded: bool, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + stride: usize = 0, + initial: []u16 = &.{}, + source: []u16 = &.{}, + want: []u16 = &.{}, + win: u64 = 0, + rows: u64 = 0, + + pub fn deinit(self: *Shift) void { + const a = self.rig.gpu.gpa; + a.free(self.initial); + a.free(self.source); + a.free(self.want); + self.rig.deinit(); + } + + pub fn kernel(self: *Shift) g.Kernel { + return .{ .spec = ref.shiftSpec(self.padded), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Shift); + + pub fn prepare(self: *Shift, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.shiftShape(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.initial); + a.free(self.source); + a.free(self.want); + self.initial = &.{}; + self.source = &.{}; + self.want = &.{}; + self.shape = shape; + self.stride = shape.n + (if (self.padded) @as(usize, 6240) else 0); + self.initial = try a.alloc(u16, self.count()); + self.source = try a.alloc(u16, shape.rows * self.stride); + ref.copyFill(self.initial, self.source, shape, self.stride, fill, seed); + self.want = try ref.copyReference(a, self.initial, self.source, shape, self.stride); + // Mutable history is an output state, never classified as immutable input. + self.win = try self.rig.output(self.count() * 2); + self.rows = try self.rig.input(std.mem.sliceAsBytes(self.source)); + return true; + } + + pub fn run(self: *Shift) !void { + try self.rig.s.synchronize(); + const history: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = self.win, .len = self.count() * 2 }; + // Every repeat starts from the same history; five successive shifts would test another operation. + try history.upload(0, std.mem.sliceAsBytes(self.initial)); + var args: cuda.Args = .{}; + args.add(self.win); + args.add(self.rows); + args.add(@as(c_longlong, @intCast(self.stride))); + args.add(@as(c_int, @intCast(self.shape.rows))); + args.add(@as(c_int, @intCast(self.shape.k))); + args.add(@as(c_int, @intCast(self.shape.n))); + try cuda.launch.launch(self.f, .{ + .grid = .{ .x = @intCast((self.shape.n + 255) / 256) }, + .block = .{ .x = 256 }, + }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Shift) usize { + return self.shape.k * self.shape.n; + } + + pub fn reference(self: *Shift, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, bits| { + w.* = c.fromBf16(bits); + m.* = 0; + } + } + + pub fn check(self: *Shift, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a history launch changed a guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.source), try self.rig.inputBytes())) + return "a history launch changed its immutable source rows"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.want), try self.rig.outputBytes())) + return "history bytes differ from append-and-retain oracle"; + if (at == .nonfinite) { + var nonfinite: usize = 0; + for (self.want) |bits| if (!std.math.isFinite(c.fromBf16(bits))) { + nonfinite += 1; + }; + if (nonfinite == 0 or nonfinite == self.want.len) + return "copy fixture did not retain both finite and nonfinite values"; + } + return null; + } +}; diff --git a/zig/tests/cuda/contract_utility_ref.zig b/zig/tests/cuda/contract_utility_ref.zig new file mode 100644 index 0000000000..7dfa1c46b0 --- /dev/null +++ b/zig/tests/cuda/contract_utility_ref.zig @@ -0,0 +1,303 @@ +//! Independent host data and oracles for ordered slice sums and in-place history copies. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub fn depth(s: c.Shape) usize { + return s.k - 1; +} + +pub fn sumSpec(bf16: bool) c.Spec { + return .{ + .name = if (bf16) "cuda fn_sum_slices bf16" else "cuda fn_sum_slices f32", + .accum = .f32, + .depth = depth, + .out = if (bf16) .bf16 else .f32, + .nan = .{ .stated = "propagates NaN and infinities; finite outputs use ordered f32 and one final RNE" }, + .tiles = .{ .{ .size = 16 }, .{ .size = 256 }, .{ .size = 4 } }, + .served = &.{ .{ .rows = 17, .n = 10240, .k = 4 }, .{ .rows = 4, .n = 2560, .k = 8 } }, + }; +} + +pub fn shiftSpec(padded: bool) c.Spec { + return .{ + .name = if (padded) "cuda fn_shift_window strided" else "cuda fn_shift_window contiguous", + .accum = .none, + .depth = depth, + .out = .bf16, + .in_place = true, + .nan = .{ .stated = "copies all bf16 bits, including NaN payloads, infinities and signed zero" }, + // rows = appended rows, n = channels, k = retained taps; stride is independently padded. + .tiles = .{ .{ .size = 16 }, .{ .size = 256 }, .{ .size = 9 } }, + .served = &.{ + .{ .rows = 5, .n = 10240, .k = 9 }, + .{ .rows = 1, .n = 10240, .k = 3 }, + .{ .rows = 4, .n = 10240, .k = 9 }, + .{ .rows = 9, .n = 257, .k = 9 }, + .{ .rows = 16, .n = 255, .k = 16 }, + .{ .rows = 32, .n = 10240, .k = 9 }, + }, + }; +} + +/// Bounded fixture domains, not validation added to the production launchers. +pub fn sumShape(s: c.Shape) bool { + return s.rows > 0 and s.rows <= 64 and s.n > 0 and s.n <= 16384 and s.k > 0 and s.k <= 8; +} + +pub fn shiftShape(s: c.Shape) bool { + return s.rows > 0 and s.rows <= 64 and s.n > 0 and s.n <= 16384 and s.k > 0 and s.k <= 16; +} + +pub fn sumFill(t: []f32, s: c.Shape, fill: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const e = s.rows * s.n; + for (t) |*v| v.* = (rng.float(f32) * 2 - 1) * + std.math.ldexp(@as(f32, 1), @as(i32, @intCast(rng.uintLessThan(u32, 17))) - 8); + switch (fill) { + .random => {}, + .nonfinite => { + t[0] = std.math.nan(f32); + t[1] = std.math.inf(f32); + t[2] = -std.math.inf(f32); + t[3] = std.math.inf(f32); + if (s.k > 1) t[e + 3] = -std.math.inf(f32); + // The initial term is loaded, not added to a synthetic positive zero. + for (0..s.k) |j| t[j * e + 4] = -0.0; + }, + .stress => for (0..s.k) |j| { + const pattern = [_]f32{ 32768, 1, -32768, 1 }; + for (0..e) |i| t[j * e + i] = pattern[j % 4]; + }, + } +} + +pub fn sumReference(t: []const f32, s: c.Shape, want: []f64, mag: []f64) void { + const e = s.rows * s.n; + for (want, mag, 0..) |*w, *m, i| { + var total: f64 = 0; + var absolute: f64 = 0; + for (0..s.k) |j| { + const term: f64 = t[j * e + i]; + total += term; + absolute += @abs(term); + } + w.* = total; + m.* = absolute; + } +} + +pub fn sumSimulate(t: []const f32, s: c.Shape, acc: c.Accum, order: g.Order, out: []f64) void { + var terms: [8]f64 = undefined; + for (out, 0..) |*v, i| { + for (terms[0..s.k], 0..) |*term, j| term.* = t[j * out.len + i]; + v.* = core.contract.rows.sumIn(acc, order, terms[0..s.k]); + } +} + +/// A separate serial f32 oracle checks the declared slice order and single final conversion. +pub fn ordered(t: []const f32, e: usize, k: usize, i: usize) f32 { + var value = t[i]; + for (1..k) |j| value += t[j * e + i]; + return value; +} + +/// Both the independent real-arithmetic bound and the documented serial f32 order must hold. +pub fn sumCheck(a: std.mem.Allocator, t: []const f32, s: c.Shape, bf16: bool, out: []const u8, nonfinite: bool) !?[]const u8 { + const e = s.rows * s.n; + const want = try a.alloc(f64, e); + defer a.free(want); + const mag = try a.alloc(f64, e); + defer a.free(mag); + sumReference(t, s, want, mag); + const spec = sumSpec(bf16); + if (out.len != e * c.size(spec.out)) return "output byte count differs from the declared shape"; + var planted: usize = 0; + for (want, mag, 0..) |w, m, i| { + const value = ordered(t, e, s.k, i); + const rounded: f32 = if (bf16) c.fromBf16(c.toBf16(value)) else value; + const got = c.read(spec.out, out, i); + if (c.classOf(w) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(w)) return "an output's class differs from the float64 oracle"; + if (std.math.isNan(rounded)) continue; + if (@as(u32, @bitCast(@as(f32, @floatCast(got)))) != @as(u32, @bitCast(rounded))) + return "an output differs from ordered f32 with one final conversion"; + if (std.math.isFinite(w) and !(@abs(got - w) <= c.bound(spec, s, w, m))) + return "an output exceeds the independent float64 error bound"; + } + if (nonfinite and (planted == 0 or planted == e)) + return "nonfinite planting did not retain both finite and nonfinite outputs"; + return null; +} + +pub fn copyFill(win: []u16, rows: []u16, s: c.Shape, stride: usize, fill: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (win) |*v| v.* = c.toBf16((rng.float(f32) * 2 - 1) * 32); + // Padding has unique nonfinite poison; it must never become an output channel. + @memset(rows, 0x7fff); + for (0..s.rows) |r| for (rows[r * stride ..][0..s.n]) |*v| { + v.* = c.toBf16((rng.float(f32) * 2 - 1) * 32); + }; + if (fill == .nonfinite) { + const bits = [_]u16{ 0x7fc1, 0x7f80, 0xff80, 0x8000, 0x0000, 0xffc7 }; + // Plant both the retained old history and appended rows. At least one finite channel remains. + for (win, 0..) |*v, i| if (i % 13 < bits.len) { + v.* = bits[i % 13]; + }; + for (0..s.rows) |r| for (rows[r * stride ..][0..s.n], 0..) |*v, ch| { + if (ch % 13 < bits.len) v.* = bits[ch % 13]; + }; + } +} + +/// Append channel rows to a host sequence, then retain its tail; no in-place update algorithm is reused. +pub fn copyReference(a: std.mem.Allocator, win: []const u16, rows: []const u16, s: c.Shape, stride: usize) ![]u16 { + const sequence = try a.alloc(u16, (s.k + s.rows) * s.n); + defer a.free(sequence); + @memcpy(sequence[0..win.len], win); + for (0..s.rows) |r| @memcpy(sequence[win.len + r * s.n ..][0..s.n], rows[r * stride ..][0..s.n]); + return a.dupe(u16, sequence[sequence.len - win.len ..]); +} + +test "ordered sum stress proves fp32 and rejects reduced accumulators in all three orders" { + const a = std.testing.allocator; + inline for (.{ false, true }) |bf16| for ([_]usize{ 4, 8 }) |k| { + const s: c.Shape = .{ .rows = 2, .n = 3, .k = k }; + var terms: [48]f32 = undefined; + const t = terms[0 .. s.rows * s.n * k]; + sumFill(t, s, .stress, 42); + var want: [6]f64 = undefined; + var mag: [6]f64 = undefined; + sumReference(t, s, &want, &mag); + const spec = sumSpec(bf16); + for ([_]g.Order{ .sequential, .pairwise, .strided }) |order| { + for ([_]c.Accum{ .f32, .f16, .bf16 }) |acc| { + var got: [6]f64 = undefined; + sumSimulate(t, s, acc, order, &got); + for (want, got) |w, v| { + const ratio = @abs(v - w) / c.boundExact(spec, w); + if (acc == .f32) try std.testing.expect(ratio <= 1) else try std.testing.expect(ratio >= 100); + } + } + } + const shapes = try c.shapes(a, spec); + defer a.free(shapes); + for (spec.served) |served| try std.testing.expect(sumShape(served)); + }; +} + +test "sum oracle preserves initial negative zero and propagates nonfinite classes" { + const s: c.Shape = .{ .rows = 1, .n = 8, .k = 4 }; + var terms: [32]f32 = undefined; + sumFill(&terms, s, .nonfinite, 71); + try std.testing.expectEqual(c.Class.nan, c.classOf(ordered(&terms, 8, 4, 0))); + try std.testing.expectEqual(c.Class.pos_inf, c.classOf(ordered(&terms, 8, 4, 1))); + try std.testing.expectEqual(c.Class.neg_inf, c.classOf(ordered(&terms, 8, 4, 2))); + try std.testing.expectEqual(c.Class.nan, c.classOf(ordered(&terms, 8, 4, 3))); + try std.testing.expectEqual(@as(u32, 0x80000000), @as(u32, @bitCast(ordered(&terms, 8, 4, 4)))); +} + +test "history oracle covers partial replacement, exact replacement and overflow with padding" { + const a = std.testing.allocator; + const win = [_]u16{ 1, 2, 3, 4, 5, 6 }; + const rows = [_]u16{ 7, 8, 99, 9, 10, 99, 11, 12, 99, 13, 14, 99 }; + for ([_]usize{ 1, 3, 4 }) |n| { + const want = try copyReference(a, &win, &rows, .{ .rows = n, .n = 2, .k = 3 }, 3); + defer a.free(want); + const expected: []const u16 = switch (n) { + 1 => &.{ 3, 4, 5, 6, 7, 8 }, + 3 => &.{ 7, 8, 9, 10, 11, 12 }, + else => &.{ 9, 10, 11, 12, 13, 14 }, + }; + try std.testing.expectEqualSlices(u16, expected, want); + } +} + +test "fixture geometry refuses zero, oversized slice counts and unsafe local history lengths" { + try std.testing.expect(!sumShape(.{ .rows = 1, .n = 1, .k = 0 })); + try std.testing.expect(!sumShape(.{ .rows = 1, .n = 1, .k = 9 })); + try std.testing.expect(!shiftShape(.{ .rows = 0, .n = 1, .k = 1 })); + try std.testing.expect(!shiftShape(.{ .rows = 1, .n = 1, .k = 17 })); + for (shiftSpec(true).served) |s| try std.testing.expect(shiftShape(s)); +} + +test "sum checker rejects wrong order, missing terms and repeated bf16 conversion" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 1, .n = 1, .k = 4 }; + const t = [_]f32{ 16777216, 1, -16777216, 1 }; + const correct = [_]f32{1}; + const reordered = [_]f32{2}; + const lost_tail = [_]f32{0}; + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&correct), false) == null); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&reordered), false) != null); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&lost_tail), false) != null); + const small = [_]f32{ 1, 0.00390625, 0.00390625, 0 }; + const once = [_]u16{0x3f81}; + const repeated = [_]u16{0x3f80}; + try std.testing.expect(try sumCheck(a, &small, s, true, std.mem.sliceAsBytes(&once), false) == null); + try std.testing.expect(try sumCheck(a, &small, s, true, std.mem.sliceAsBytes(&repeated), false) != null); +} + +test "history copy retains raw nonfinite payloads and distinguishes padding and repeated mutation" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 1, .n = 3, .k = 3 }; + const win = [_]u16{ 1, 2, 3, 0x7fc1, 0xffc7, 0x8000, 0x7f80, 0xff80, 0 }; + const source = [_]u16{ 4, 5, 6, 0x7fff }; + const once = try copyReference(a, &win, &source, s, 4); + defer a.free(once); + try std.testing.expectEqualSlices(u16, &.{ 0x7fc1, 0xffc7, 0x8000, 0x7f80, 0xff80, 0, 4, 5, 6 }, once); + const twice = try copyReference(a, once, &source, s, 4); + defer a.free(twice); + try std.testing.expect(!std.mem.eql(u16, once, twice)); + try std.testing.expect(std.mem.indexOfScalar(u16, once, 0x7fff) == null); + for ([_]bool{ false, true }) |padded| { + const spec = shiftSpec(padded); + const shapes = try c.shapes(a, spec); + defer a.free(shapes); + for (shapes) |shape| try std.testing.expect(shiftShape(shape)); + } +} + +test "seeded finite and nonfinite slice outputs satisfy both oracles and reject untouched output" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 5, .n = 17, .k = 8 }; + var t: [680]f32 = undefined; + var f: [85]f32 = undefined; + var b: [85]u16 = undefined; + for ([_]g.Fill{ .random, .nonfinite }) |fill| { + sumFill(&t, s, fill, 0x51ce); + for (&f, &b, 0..) |*fv, *bv, i| { + fv.* = ordered(&t, f.len, s.k, i); + bv.* = c.toBf16(fv.*); + } + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&f), fill == .nonfinite) == null); + try std.testing.expect(try sumCheck(a, &t, s, true, std.mem.sliceAsBytes(&b), fill == .nonfinite) == null); + f[f.len - 1] = std.math.nan(f32); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&f), fill == .nonfinite) != null); + } +} + +test "history oracle agrees with independent row-at-a-time queue including all sixteen taps" { + const a = std.testing.allocator; + for ([_]usize{ 1, 3, 9, 16 }) |taps| for ([_]usize{ 1, 5, 16, 32 }) |n| { + const s: c.Shape = .{ .rows = n, .n = 17, .k = taps }; + const stride = s.n + 7; + const win = try a.alloc(u16, taps * s.n); + defer a.free(win); + const source = try a.alloc(u16, n * stride); + defer a.free(source); + copyFill(win, source, s, stride, .nonfinite, 0x7157); + const want = try copyReference(a, win, source, s, stride); + defer a.free(want); + const queue = try a.dupe(u16, win); + defer a.free(queue); + for (0..n) |r| { + std.mem.copyForwards(u16, queue[0 .. queue.len - s.n], queue[s.n..]); + @memcpy(queue[queue.len - s.n ..], source[r * stride ..][0..s.n]); + } + try std.testing.expectEqualSlices(u16, queue, want); + }; +} diff --git a/zig/tests/cuda/experts_rows_tests.zig b/zig/tests/cuda/experts_rows_tests.zig new file mode 100644 index 0000000000..7bf59fae5d --- /dev/null +++ b/zig/tests/cuda/experts_rows_tests.zig @@ -0,0 +1,178 @@ +//! experts_rows.cu against experts.cu's decode kernel: every pair row of a window or chunk, byte for byte. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const grouped = cuda.grouped; +const Rows = cuda.experts.RowsExperts; + +/// experts_pack.cu's groups of 32 (one block: 160 uint32 words of codes, scales and biases). +const pack_symbol = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const block_bytes = 160 * 4; +const dims = 256; +const width = 128; +const routed = 16; +const experts = routed + 1; +const slots = 4; + +const Rig = struct { + gpu: Gpu, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn zeros(r: *Rig, n: usize) !u64 { + var b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(n, 16)); + // on the kernels' stream: it is non-blocking, so a memset on the default stream may land after them + try b.fill8(0, r.s.handle); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn get(r: *Rig, ptr: u64, n: usize) ![]u8 { + try r.s.synchronize(); + const out = try r.gpu.gpa.alloc(u8, n); + try r.gpu.d.check(r.gpu.d.api.cuMemcpyDtoH_v2(out.ptr, ptr, n), "cuMemcpyDtoH"); + return out; + } + + fn bf16(v: f32) u16 { + return @truncate(@as(u32, @bitCast(v)) >> 16); + } + + /// Random MLX 4-bit words with small scales and biases, packed for `count` experts of n outputs from k inputs. + fn table(r: *Rig, pack: cuda.Function, count: usize, n: usize, k: usize) !u64 { + const a = r.gpu.gpa; + const words = try a.alloc(u32, count * n * k / 8); + defer a.free(words); + for (words) |*w| w.* = r.rng.int(u32); + const groups = count * n * k / 32; + const scales = try a.alloc(u16, groups); + defer a.free(scales); + const biases = try a.alloc(u16, groups); + defer a.free(biases); + for (scales, biases) |*s, *b| { + s.* = bf16(0.002 + 0.004 * r.rng.float(f32)); + b.* = bf16(-0.015 + 0.01 * r.rng.float(f32)); + } + const out = try r.zeros(count * (n / 32) * (k / 32) * block_bytes); + var args: cuda.Args = .{}; + for ([_]u64{ try r.dev(std.mem.sliceAsBytes(words)), try r.dev(std.mem.sliceAsBytes(scales)), try r.dev(std.mem.sliceAsBytes(biases)), out }) |v| args.add(v); + for ([_]usize{ n, k / 8, k / 32, n / 32 }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(pack, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 160 } }, r.s, &args); + return out; + } + + /// Gate and up blocks side by side, the SwiGLU table the kernels read. + fn stack(r: *Rig, gate: u64, up: u64, blocks: usize) !u64 { + const g = try r.get(gate, blocks * block_bytes); + defer r.gpu.gpa.free(g); + const u = try r.get(up, blocks * block_bytes); + defer r.gpu.gpa.free(u); + const both = try r.gpu.gpa.alloc(u8, 2 * blocks * block_bytes); + defer r.gpu.gpa.free(both); + for (0..blocks) |b| { + @memcpy(both[2 * b * block_bytes ..][0..block_bytes], g[b * block_bytes ..][0..block_bytes]); + @memcpy(both[(2 * b + 1) * block_bytes ..][0..block_bytes], u[b * block_bytes ..][0..block_bytes]); + } + return r.dev(both); + } +}; + +/// The plan scratch for `pairs` pairs at `tile` pairs an item. +fn planFor(r: *Rig, pairs: usize, tile: usize) !grouped.Plan { + const wide = pairs > grouped.small; + return .{ + .members = try r.zeros(pairs * 4), + .items = try r.zeros(grouped.maxItems(pairs, experts, tile) * 12), + .counts = try r.zeros(8), + .rank = try r.zeros((if (wide) pairs else 1) * 4), + .hist = try r.zeros((if (wide) (pairs + 1023) / 1024 * experts else 1) * 4), + }; +} + +fn launch(r: *Rig, f: cuda.Function, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, p: grouped.Plan, out: u64, n: usize, units: usize) !void { + var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, p, out, n); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, r.s, &a); +} + +pub fn run(gpu: Gpu) !void { + const a = gpu.gpa; + var experts_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer experts_mod.unload(); + var rows_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_rows); + defer rows_mod.unload(); + var pack_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const router = try grouped.Router.resolve(experts_mod); + const kernels = try Rows.resolve(rows_mod); + const one_up = try experts_mod.function(Rows.symbols.decode_up); + const one_down = try experts_mod.function(Rows.symbols.decode_down); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .s = try cuda.Stream.init(gpu.d, true), .rng = prng.random() }; + defer r.s.deinit(); + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(a); + } + const pack = try pack_mod.function(pack_symbol); + const up_w = try r.stack(try r.table(pack, experts, width, dims), try r.table(pack, experts, width, dims), experts * (width / 32) * (dims / 32)); + const down_w = try r.table(pack, experts, dims, width); + for ([_]usize{ 1, 2, 5, 16, 33, 100, 300 }) |n_rows| { + const pairs = n_rows * slots; + const x = try a.alloc(u16, n_rows * dims); + defer a.free(x); + for (x) |*v| v.* = Rig.bf16(r.rng.float(f32) * 2 - 1); + // each row's routed picks distinct, then the shared expert in the last slot, as Flash Next routes + const picks = try a.alloc(i32, pairs); + defer a.free(picks); + for (0..n_rows) |row| { + var used: [routed]bool = @splat(false); + for (0..slots - 1) |j| { + var e = r.rng.uintLessThan(usize, routed); + while (used[e]) e = (e + 1) % routed; + used[e] = true; + picks[row * slots + j] = @intCast(e); + } + picks[row * slots + slots - 1] = routed; + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dpicks = try r.dev(std.mem.sliceAsBytes(picks)); + // the rows kernel: every row in one plan of 32-pair items + const p = try planFor(&r, pairs, Rows.tile); + try router.route(r.s, dpicks, pairs, experts, Rows.tile, p); + const items = pairs / Rows.tile + experts; + const act = try r.zeros(pairs * width * 2); + const y = try r.zeros(pairs * dims * 4); + try kernels.run(r.s, true, dx, dims, slots, up_w, dims / 32, width / 32, p, act, width, items); + try kernels.run(r.s, false, act, width, 0, down_w, width / 32, dims / 32, p, y, dims, items); + // the decode kernel: each row alone, its pairs written where the rows kernel writes them + const act1 = try r.zeros(pairs * width * 2); + const y1 = try r.zeros(pairs * dims * 4); + const p1 = try planFor(&r, slots, 16); + for (0..n_rows) |row| { + try router.route(r.s, dpicks + row * slots * 4, slots, experts, 16, p1); + try launch(&r, one_up, dx + row * dims * 2, dims, slots, up_w, dims / 32, width / 32, p1, act1 + row * slots * width * 2, width, slots * (width / 32)); + try launch(&r, one_down, act1 + row * slots * width * 2, width, 0, down_w, width / 32, dims / 32, p1, y1 + row * slots * dims * 4, dims, slots * (dims / 32)); + } + const got_act = try r.get(act, pairs * width * 2); + defer a.free(got_act); + const want_act = try r.get(act1, pairs * width * 2); + defer a.free(want_act); + try check.sameBytes("experts_rows SwiGLU rows", got_act, want_act); + const got_y = try r.get(y, pairs * dims * 4); + defer a.free(got_y); + const want_y = try r.get(y1, pairs * dims * 4); + defer a.free(want_y); + try check.sameBytes("experts_rows down rows", got_y, want_y); + check.pass("experts-rows: {d} rows ({d} pairs) equal the decode kernel's one-row launches byte for byte", .{ n_rows, pairs }); + } +} diff --git a/zig/tests/cuda/flashnext/capture.py b/zig/tests/cuda/flashnext/capture.py new file mode 100644 index 0000000000..ad8fa443c4 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +"""Flash Next's Python CUDA engine (drafts off, greedy) as the Zig oracle: Triton launches, prompt ids, references.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +PROMPTS = { + "short": "What is the capital of France?", + "mid": ("Explain in plain words why the sky looks blue during the day but turns red and orange at sunset. " + "Mention the role of the atmosphere and the wavelength of light, and keep it under five sentences."), +} +TOP_STEPS = 4 # sampled steps whose five best logits the teacher paths keep +LONG_TOKENS = 600 # the long prompt: module source cut to make the whole prompt about this many tokens +# prompts past the block indexer's 2,048-key budget, cut from pydoc_data.topics to exactly this many tokens +SPARSE_TOKENS = {"long3k": 3001, "long12k": 12001} +SPARSE_ASK = ("\n\nQuestion: according to the excerpt above, what does an assert statement do when Python runs with " + "optimization on, and what does an augmented assignment evaluate only once? Answer briefly.") + + +def chat(user: str) -> str: + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +def prompts(model: Path) -> dict[str, list[int]]: + from tokenizers import Tokenizer + import textwrap + + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + enc = lambda s: tok.encode(s, add_special_tokens=False).ids # noqa: E731 + out = {name: enc(chat(text)) for name, text in PROMPTS.items()} + head = enc("<|im_start|>user\nSummarize what this Python module does, then list its public functions:\n\n") + tail = enc("<|im_end|>\n<|im_start|>assistant\n\n\n\n\n") + body = enc(Path(textwrap.__file__).read_text())[:LONG_TOKENS - len(head) - len(tail)] + out["long"] = head + body + tail + from pydoc_data.topics import topics + + reference = "\n\n".join(topics[k] for k in sorted(topics)) + head = enc("<|im_start|>user\nBelow is an excerpt of the Python language reference.\n\n") + tail = enc(SPARSE_ASK + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n") + text = enc(reference) + for name, total in SPARSE_TOKENS.items(): + out[name] = head + text[:total - len(head) - len(tail)] + tail + return out + + +def specialization() -> dict: + """Each JIT function's parameter names and which ones Triton never specializes (aot_pack.py's --jit).""" + + from triton.runtime.jit import JITFunction + import gc + + out = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + out[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True, help="the checkpoint's snapshot dir") + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="", help="folder of triton_aot_manifest.py: record every Triton launch") + ap.add_argument("--tokens", type=int, default=32, help="greedy tokens per prompt") + ap.add_argument("--context", type=int, default=1536, + help="the cache window; 1025..1536 gives the port's NCH=3, QSA=0 attention specialization") + ap.add_argument("--only", default="", help="comma list of prompt names (default: every prompt the context fits)") + ap.add_argument("--teachers", default="teacher,teacher32", + help="teacher paths to run after the served one (comma list of teacher, teacher32; empty: none)") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + + import torch + from contextlib import nullcontext + + def phase(name: str, detail: bool = False): + return rec.scope(name, detail) if rec is not None else nullcontext() + + model = Path(a.model) + ids = prompts(model) + if a.only: + ids = {k: v for k, v in ids.items() if k in a.only.split(",")} + else: + ids = {k: v for k, v in ids.items() if len(v) + a.tokens < a.context} + (out / "prompts.json").write_text(json.dumps(ids) + "\n") + for name, p in ids.items(): + (out / f"{name}.ids").write_text(" ".join(map(str, p)) + "\n") + + from tensorfold.families.qwen4_exp.cuda.engine import FlashNextEngine + from tensorfold.families.qwen4_exp.cuda.forward import commit + + t0 = time.perf_counter() + with phase("startup"): + eng = FlashNextEngine(model, depth=0, max_len=a.context, context_explicit=True, streams=1, + ple_on_ssd=True, kv_dtype="bf16", graphs=False) + e = eng.e + info = {"startup_s": round(time.perf_counter() - t0, 2), "max_len": eng.max_len, "prefill_rows": eng.prefill_rows, + "nch": e.buf.attn.nch, "qsa": bool(e.buf.attn.qsa), "idw": e.buf.attn.idw, "eos": list(eng.eos), + "torch": torch.__version__, "gpu": torch.cuda.get_device_name(0), + "prompt_lengths": {k: len(v) for k, v in ids.items()}} + print(json.dumps(info), flush=True) + + results: dict[str, dict] = {} + from functools import partial + from tensorfold.families.qwen4_exp.cuda.state import Buffers + + teach = partial(_teach, phase, results, a, commit, nullcontext) + fp32 = None + for name, prompt in ids.items(): + # the served path: chunked prefill, then one token a round (draft=False is FlashNextEngine._serial) + toks: list[int] = [] + with phase(f"serial-{name}"): + torch.cuda.synchronize() + start = time.perf_counter() + stats = eng.generate(prompt, a.tokens, None, lambda new: toks.extend(map(int, new)) and False, + draft=False, stop_eos=False) + torch.cuda.synchronize() + wall = time.perf_counter() - start + results[f"serial/{name}"] = {"tokens": toks, "wall_s": round(wall, 4), **stats} + print(f"serial {name} ({len(prompt)} prompt): {toks} {stats}", flush=True) + if fp32 is None: + # teacher32: the decode-form experts (fp32 slot outputs), the arithmetic zig's flashnext-prompt runs + fp32 = eng.serial.twin() + fp32.buf = Buffers(eng.e.w, eng.e.rows, eng.e.capacity) + + for label, t in (("teacher", eng.serial), ("teacher32", fp32)): + if label in a.teachers.split(","): + teach(label, name, prompt, t, toks) + + (out / "results.json").write_text(json.dumps({"info": info, "results": results}, indent=1) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + (out / "jit.json").write_text(json.dumps(specialization(), indent=1, sort_keys=True) + "\n") + eng.close() + print("done", flush=True) + return 0 + + +def _teach(phase, results, a, commit, nullcontext, label, name, prompt, t, served): + """Every prompt token one decode row, as zig's flashnext-prompt feeds them, then greedy rows.""" + + import torch + + t.reset() + got: list[int] = [] + top: list = [] + steps = 0 + with phase(f"{label}-{name}"): + torch.cuda.synchronize() + start = time.perf_counter() + feed = list(prompt) + dec_start = None + while len(got) < a.tokens: + tok = feed.pop(0) if feed else got[-1] + with (phase(f"{label}-{name}-detail", detail=True) if name == "short" and steps < 2 + else nullcontext()): + logits = t.forward([tok]) + if not feed: + got.append(t.sample(logits[:1], [t.st.pos + 1], None)[0]) + if len(got) <= TOP_STEPS: + v, i = logits[0].float().topk(5) + top.append([[int(x), round(float(y), 4)] for x, y in zip(i.tolist(), v.tolist())]) + if dec_start is None: + torch.cuda.synchronize() + dec_start = time.perf_counter() + commit(t.w, t.st, t.buf, 1, 1) + steps += 1 + torch.cuda.synchronize() + end = time.perf_counter() + ms = (end - dec_start) / max(1, a.tokens - 1) * 1e3 + results[f"{label}/{name}"] = {"tokens": got, "wall_s": round(end - start, 4), + "decode_ms_per_token": round(ms, 3), "top5": top} + print(f"{label} {name}: {got} ({ms:.2f} ms/token) same as serial: {got == served}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture.sh b/zig/tests/cuda/flashnext/capture.sh new file mode 100755 index 0000000000..8aaf8b9ccb --- /dev/null +++ b/zig/tests/cuda/flashnext/capture.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Usage: capture.sh WORK SNAPSHOT_DIR [IMAGE]; captures Flash Next's launches, references and a kernel set under WORK. +set -euo pipefail +WORK=$(readlink -f "$1"); SNAP=$(readlink -f "$2"); IMAGE=${3:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +HF=$(dirname "$(dirname "$(dirname "$SNAP")")") +REL=${SNAP#"$HF"/} +# CONTEXT=16383 ONLY=long3k,long12k TEACHERS=teacher32 OUT=sparse: the prompts past the indexer's 2,048-key budget +CONTEXT=${CONTEXT:-1535}; ONLY=${ONLY:-}; TEACHERS=${TEACHERS-teacher,teacher32}; OUT=${OUT:-cap} +mkdir -p "$WORK/cache" "$WORK/out/$OUT" +KSET=$WORK/kernels; [ "$OUT" = cap ] || KSET=$WORK/kernels-$OUT +# SNAPSHOT_DIR is the HF snapshot (blobs at ../../blobs); the default context 1535 keeps attention dense (QSA 0) +docker run --rm --name fn-capture --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ + --user "$(id -u):$(id -g)" -e HOME=/tmp -e HF_HUB_OFFLINE=1 -e TENSORFOLD_NO_UPDATE_CHECK=1 \ + -e TENSORFOLD_MEMORY_RESERVE_GIB=2 \ + -v "$HF":/hf:ro -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture.py --model "/hf/$REL" --out "/out/$OUT" \ + --tools /src/tools/zig --tokens 32 --context "$CONTEXT" --only "$ONLY" --teachers "$TEACHERS" +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/$OUT/launches.json" --mount /cache/triton --out "$WORK/out/$OUT/manifest.json" +"$(dirname "$0")/pack.sh" "$KSET" "$WORK/out/$OUT" +echo "prompt ids (*.ids) and references (results.json) in $WORK/out/$OUT" diff --git a/zig/tests/cuda/flashnext/capture_mtp.py b/zig/tests/cuda/flashnext/capture_mtp.py new file mode 100644 index 0000000000..6cf9323f39 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_mtp.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +"""The Triton kernels Flash Next's MTP head adds to the main model's (its RMSNorms and _add_streams), no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + from tensorfold.families.qwen4_exp.cuda import glue + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + streams, dims = 4, 2560 + wide = streams * dims + with rec.scope("mtp"): + for m in (1, 7): + emb4 = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + norm_e = torch.rand((dims,), device=dev, generator=g) + mixed = torch.empty((m, dims), dtype=torch.bfloat16, device=dev) + xs = torch.empty((m, dims // 32), dtype=torch.float32, device=dev) + glue.rmsnorm(emb4[:, :dims], norm_e, 1e-6, out=mixed, xs=xs) + h = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + norm_h = torch.rand((wide,), device=dev, generator=g) + hn = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + xh = torch.empty((m, wide // 32), dtype=torch.float32, device=dev) + glue.rmsnorm(h, norm_h, 1e-6, out=hn, xs=xh) + eo = torch.randn((m, dims), device=dev, generator=g).to(torch.bfloat16) + hs = torch.randn((m * streams, dims), device=dev, generator=g).to(torch.bfloat16) + hout = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + glue.add_streams(eo, hs, hout, streams) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_mtp.sh b/zig/tests/cuda/flashnext/capture_mtp.sh new file mode 100755 index 0000000000..15ea4d36aa --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_mtp.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_mtp.sh WORK [IMAGE]: the MTP head's own Triton kernels (capture_mtp.py) into WORK/out, no checkpoint. +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" +docker run --rm --name fn-capture-mtp --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_mtp.py --out /out --tools /src/tools/zig +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" diff --git a/zig/tests/cuda/flashnext/capture_rows.py b/zig/tests/cuda/flashnext/capture_rows.py new file mode 100644 index 0000000000..85930056c3 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_rows.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""The decode kernels Flash Next's batched prompt pass runs at M > 1 (M 17 and 32), no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + from tensorfold.cuda import moe + from tensorfold.families.qwen4_exp.cuda import glue, qmm + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + streams, dims, low = 4, 2560, 320 + wide = streams * dims + + def q4(n: int, k: int): + w = torch.randint(-2**31, 2**31 - 1, (n, k // 8), dtype=torch.int32, device=dev, generator=g) + s = (torch.rand((n, k // 32), device=dev, generator=g) * 0.01).to(torch.bfloat16) + b = (torch.rand((n, k // 32), device=dev, generator=g) * 0.01 - 0.005).to(torch.bfloat16) + return qmm.make_q4(w, s, b, layout="tiled") + + downs = {n: q4(n, wide) for n in (low + streams, low)} + up = q4(wide, low) + router_rows = (torch.randn((513, dims), device=dev, generator=g) * 0.02).to(torch.bfloat16) + with rec.scope("rows"): + for m in (17, 32): + h = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + pss = torch.rand((m, dims // 256, streams), device=dev, generator=g) * 100 + scale = torch.rand((wide,), device=dev, generator=g) + normed = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + act = torch.empty((m, low), dtype=torch.bfloat16, device=dev) + xs_act = torch.empty((m, low // 32), dtype=torch.float32, device=dev) + inj = torch.empty((m, streams), dtype=torch.bfloat16, device=dev) + for n, q in downs.items(): + dn = torch.empty((m, n), dtype=torch.bfloat16, device=dev) + part = torch.empty((32 * m * n,), dtype=torch.float32, device=dev) + got = qmm.hc_down(h, pss, scale, normed, q, 1e-6, streams, out=dn, part=part) + glue.hc_reduce_act(got, act, xs_act, inj if n == low + streams else None, streams, low) + mixed = torch.empty((m, dims), dtype=torch.bfloat16, device=dev) + xs_mixed = torch.empty((m, dims // 32), dtype=torch.float32, device=dev) + qmm.hc_upmix(act, xs_act, up, normed, mixed, xs_mixed, streams) + # the up projection alone (a training pass keeps it): qmm.matmul's tiled kernel at SHAPES16 + up_rows = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + qmm._qmm[(-(-m // 16), wide // 64, 1)](act, xs_act, up.weight, up.scales, up.biases, up_rows, up_rows, m, + act.stride(0), N=wide, K=low, SK=1, BM=16, BLOCK_N=64, GPI=2, + F32=False, SBN=64, num_warps=4, num_stages=3) + x = torch.randn((m, dims), device=dev, generator=g).to(torch.bfloat16) + logits = torch.empty((m, 513), dtype=torch.float32, device=dev) + grid = (-(-m // 16), -(-513 // 32)) + moe._router[grid](x, router_rows, logits, m, x.stride(0), D=dims, NE=513, BM=16, BLOCK_E=32, BK=256, + num_warps=4, num_stages=4) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_rows.sh b/zig/tests/cuda/flashnext/capture_rows.sh new file mode 100755 index 0000000000..fcbd158b11 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_rows.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_rows.sh WORK [IMAGE]: the batched pass's M > 1 kernel forms (capture_rows.py) into WORK/out. +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" +docker run --rm --name fn-capture-rows --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_rows.py --out /out --tools /src/tools/zig +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" diff --git a/zig/tests/cuda/flashnext/capture_sparse.py b/zig/tests/cuda/flashnext/capture_sparse.py new file mode 100644 index 0000000000..dc5540df82 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_sparse.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Flash Next's sparse attention kernels and _select fixtures on synthetic tensors, no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + ap.add_argument("--fixtures", default="", help="folder for the selection fixtures (empty: none)") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + import triton + from tensorfold.families.qwen4_exp.cuda import attention as A + from tensorfold.families.qwen4_exp.cuda import glue, kvcache + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + heads, kv, hd, ih, idim, ratio = 24, 2, 256, 4, 128, 4 + cap = 8192 # nb = 2048: NB % 16 == 0, as the Zig scratch keeps it + inv = (10_000_000.0 ** (-torch.arange(32, dtype=torch.float64) / 32)).to(torch.float32).to(dev) + with rec.scope("sparse"): + for rows in (1, 17, 32): + sc = A.AttnScratch(rows, heads, hd, cap, dev) + assert sc.qsa and sc.nch == 5 and sc.nb % 16 == 0 + cache = kvcache.KVCache(cap, kv, hd, dev, "bf16") + cache.k.copy_(torch.randn(cache.k.shape, device=dev, generator=g).to(torch.bfloat16)) + cache.v.copy_(torch.randn(cache.v.shape, device=dev, generator=g).to(torch.bfloat16)) + ikc = torch.randn((cap, idim), device=dev, generator=g).to(torch.bfloat16) + pooled = torch.zeros((sc.nb, idim), dtype=torch.bfloat16, device=dev) + w = torch.rand((idim,), device=dev, generator=g) + 0.5 + # every block before the window pooled first (one row at a time would do it the same way) + pos0 = torch.zeros((1,), dtype=torch.int32, device=dev) + p0 = 3001 + A.qsa_pool(ikc, pooled, pos0, w, inv, 1e-6, sc, p0 - p0 % 32) if rows != 1 else None + pos0.fill_(p0) + A.qsa_pool(ikc, pooled, pos0, w, inv, 1e-6, sc, rows) + iq = torch.randn((rows, ih, idim), device=dev, generator=g).to(torch.bfloat16) + q = torch.randn((rows, heads, hd), device=dev, generator=g).to(torch.bfloat16) + ends = p0 + rows + blocks = min(sc.nb, max(1, triton.cdiv(ends, ratio))) + A._scores[(rows, triton.cdiv(blocks, 64))](iq, pooled, pos0, sc.scores, sc.nb, HI=ih, DI=idim, + RATIO=ratio, TOP=sc.budget // ratio, BB=64, num_warps=4) + A._launch_select(sc, pos0, rows, blocks) + A.attention(q, cache.k, cache.v, pos0, sc, rows, hd ** -0.5, context=ends, ks=cache.ks, vs=cache.vs) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + if a.fixtures: + fixtures(Path(a.fixtures), A, torch, triton, g) + print("done", flush=True) + return 0 + + +def fixtures(fx: Path, A, torch, triton, g) -> None: + """Score rows (few distinct values: ties everywhere) and the lists _select / _select_tiles make of them.""" + + fx.mkdir(parents=True, exist_ok=True) + dev = "cuda" + cases = [] + # (capacity, pos0, rows, distinct score values): register select, then rows past SELECT_REGS blocks (tiles) + for i, (cap, p0, rows, levels) in enumerate([(8192, 2040, 24, 7), (8192, 5000, 9, 0), (65536, 40000, 5, 3), + (262144, 150000, 3, 11), (262144, 262000, 2, 0)]): + sc = A.AttnScratch(rows, 24, 256, cap, dev) + if levels: + s = torch.randint(0, levels, (rows, sc.nb), device=dev, generator=g).float() / 7.0 + else: + s = torch.rand((rows, sc.nb), device=dev, generator=g) * 3.0 + s[:, ::97] = 0.0 # zeros among them, as relu sums give + sc.scores.copy_(s) + sc.ids.fill_(-1) + pos0 = torch.full((1,), p0, dtype=torch.int32, device=dev) + blocks = min(sc.nb, max(1, triton.cdiv(p0 + rows, 4))) + A._launch_select(sc, pos0, rows, blocks) + torch.cuda.synchronize() + name = f"case{i}" + sc.scores.cpu().numpy().tofile(fx / f"{name}.scores") + sc.ids.cpu().numpy().tofile(fx / f"{name}.ids") + sc.nk.cpu().numpy().tofile(fx / f"{name}.nk") + sc.sparse.cpu().numpy().tofile(fx / f"{name}.sparse") + cases.append({"name": name, "rows": rows, "pos0": p0, "nb": sc.nb, "idw": sc.idw}) + (fx / "cases.json").write_text(json.dumps(cases, indent=1) + "\n") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_sparse.sh b/zig/tests/cuda/flashnext/capture_sparse.sh new file mode 100755 index 0000000000..1598002374 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_sparse.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_sparse.sh WORK [IMAGE]: the sparse attention's kernels and selection fixtures (capture_sparse.py). +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" "$WORK/fixtures" +docker run --rm --name fn-capture-sparse --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -v "$WORK/fixtures":/fixtures -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_sparse.py --out /out --tools /src/tools/zig --fixtures /fixtures +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" diff --git a/zig/tests/cuda/flashnext/lr7/long.ids b/zig/tests/cuda/flashnext/lr7/long.ids new file mode 100644 index 0000000000..81ee38d549 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long.ids @@ -0,0 +1 @@ +248045 846 198 8917 5437 537 1092 411 12654 4536 1503 11 1179 1103 1141 569 5568 25 271 12237 1138 39834 321 20621 13 198 12237 271 2 2861 318 34 8 220 16 24 24 24 12 17 15 15 16 41550 387 13 25788 13 198 2 2861 318 34 8 220 17 15 15 17 11 220 17 15 15 18 12654 4236 4846 13 198 2 32770 539 15573 25788 361 70 1550 31 12305 4909 29 271 464 311 271 548 528 548 283 2429 1138 11210 504 359 9797 504 359 7320 504 359 9497 305 504 359 31752 504 359 8412 268 642 271 2 10906 1787 279 14541 35169 5587 310 279 2201 12 54556 198 2 35169 5587 13 220 561 1822 2781 364 3604 411 369 421 198 2 1010 34707 12258 318 4650 1088 84 15 15 64 15 8 513 2397 52483 21201 26453 13 198 83201 283 5031 83 1639 3351 15 65 3351 15 66 11674 359 271 1005 2824 11210 25 198 262 4071 198 262 2907 364 39834 6454 9121 1414 13 220 561 569 3622 16661 314 198 262 279 14623 363 321 4990 363 5245 26 279 975 5245 513 1066 1017 364 198 262 65878 310 2723 303 1906 310 49324 279 1581 16501 13 198 262 1368 488 1318 310 6379 8032 279 1822 39834 11767 11 198 262 488 3172 4507 599 310 2723 696 9797 63261 973 271 262 24876 2774 7944 2444 5090 13173 314 39834 25 198 406 2297 318 2186 25 220 22 15 8 198 285 279 6966 2297 314 18880 4965 318 37488 1387 17268 18408 198 285 369 867 8 198 406 2763 66708 318 2186 25 11363 198 285 886 421 668 381 20671 2789 310 279 1118 1500 314 18880 198 285 2468 13 220 82699 6757 279 1500 579 2297 13 198 406 16905 66708 318 2186 25 11363 198 285 886 421 668 381 20671 2789 310 660 4965 3462 279 1118 198 285 314 18880 2468 26 1048 14159 6757 1754 1500 579 2297 13 198 406 8952 56008 318 2186 25 804 8 198 285 48455 21700 303 1879 1414 310 12258 1518 4473 8427 13 198 285 8618 5474 668 3512 220 15 5074 359 5999 2073 6 12258 11 11313 383 1141 2234 198 285 303 1141 1500 13 220 1368 867 11 1754 5474 369 11421 430 264 3074 3542 13 198 406 5474 2073 318 2186 25 220 23 8 198 285 48455 21700 303 1879 1414 310 220 15 5074 359 5999 2073 6 12258 11 7014 198 285 359 31251 56008 6 369 867 13 198 406 8032 83201 318 2186 25 804 8 198 285 28592 660 35169 5587 303 279 1879 1414 539 12258 198 285 1238 5474 14046 13 220 6817 421 413 8952 56008 369 867 321 198 285 8032 83201 369 804 11 1396 5474 668 381 15625 310 264 198 285 3074 3433 0 198 406 4884 46720 6018 793 318 2186 25 867 8 198 285 28320 421 11316 64541 59429 369 2590 7854 198 285 539 1330 12258 13 220 3978 539 1581 1521 279 11767 369 198 285 318 354 46261 2752 8 65669 13 198 406 1387 17268 18408 318 2186 25 804 8 198 285 15164 4105 4860 1056 359 2998 4282 220 1368 867 11 1782 4105 668 524 198 285 381 10548 11 321 1010 4965 2493 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/long12k.ids b/zig/tests/cuda/flashnext/lr7/long12k.ids new file mode 100644 index 0000000000..e922e33cbc --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long12k.ids @@ -0,0 +1 @@ +248045 846 198 36948 369 449 47801 314 279 12654 3992 5604 13 271 760 328 2138 1 4952 198 865 6610 271 8272 11889 513 264 16219 1560 310 5479 26805 52992 198 17583 264 1957 25 271 256 1992 21152 61987 328 2138 1 7258 498 2129 7258 60 271 760 4145 1304 11 328 2138 7258 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 25 4673 63990 271 760 11244 1304 11 328 2138 7258 16 11 7258 17 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 16 25 4673 63990 46718 17 8 271 9205 42018 2303 9370 421 12880 8090 548 1 321 328 66317 1401 1 8042 198 951 279 5617 3309 7101 440 1782 4874 13 220 733 279 1428 198 13957 11 279 5617 3309 3759 12880 8090 548 1 369 328 2434 1 1172 198 7994 12779 11 328 3913 1 948 24460 369 10897 318 5454 198 1021 2904 6320 46 1764 220 561 1428 1970 13427 70307 874 1970 364 449 198 2138 4952 948 24460 369 10897 506 18612 854 13 220 6817 198 8761 424 369 24366 310 2830 279 2450 1970 364 279 7258 198 8761 4490 303 279 1412 1876 26 424 668 381 12234 430 919 314 279 198 7450 11319 13 271 27991 1320 310 12880 8090 548 1 513 11477 13 220 561 869 364 279 5617 3309 198 9521 369 10521 948 279 37999 8211 13 1358 40238 11889 198 865 17731 271 40238 11889 513 1429 310 318 265 8 7423 4874 310 2663 321 310 198 32925 7944 466 3470 314 24983 5978 25 271 256 15840 21152 61987 318 5480 1951 35036 7030 318 11531 1114 27160 735 7302 27160 8 198 256 2100 1951 257 61987 2100 318 2129 2100 4653 498 487 1293 198 256 2100 670 61987 12448 198 1025 735 32467 498 5480 1951 60 56439 198 1025 735 48975 498 5480 1951 60 84576 198 1025 735 16819 330 42514 198 1025 735 14701 198 1025 735 81987 198 1025 735 63669 2100 271 7 9538 3644 34019 5262 364 279 18890 16954 364 348 42648 330 42514 11961 198 9 33261 11961 321 348 82 86047 9 5946 271 2014 15840 4952 64013 279 7258 1103 318 28321 421 198 559 628 381 264 3074 7258 466 264 30644 70221 1103 11 279 14827 198 28726 286 264 13992 8 321 48233 279 3074 12571 1576 310 1754 314 198 1719 2100 11140 11 494 2047 310 1245 13 271 40238 369 4364 51070 11313 383 279 1304 314 279 2100 198 7080 553 3095 264 2100 369 919 314 264 24983 1576 318 276 6816 198 15812 11 14701 466 81987 681 279 24983 1576 1902 198 484 6887 2649 279 15840 321 9973 883 1141 30794 11 321 198 17812 4673 449 4531 413 279 15840 369 41565 13 220 561 5426 198 5308 2683 539 5090 4346 321 279 18987 8953 513 2574 440 279 198 17829 314 279 1576 4346 318 3924 3644 561 5129 913 198 71 32759 553 271 40238 314 449 1576 310 264 2100 1103 11 44007 42360 303 198 39708 37052 466 9057 37641 11 369 51070 4364 430 10695 13 271 9 1368 279 2100 1103 369 264 3074 2100 440 874 26849 30644 11 198 220 44007 303 71458 11 279 1576 369 12244 310 421 2100 13 271 9 18236 25 271 220 348 1368 279 2100 1103 5435 799 2100 78685 440 449 32815 3086 11 198 262 2512 264 1001 11531 1114 828 2100 25 561 1576 1902 381 449 49124 440 506 198 262 3140 430 1599 3470 430 1017 513 11218 303 279 2100 1103 11 26401 198 262 799 13 220 561 1118 3470 314 279 49124 513 12244 11 494 2047 310 198 262 1245 11 310 279 11218 1518 279 56276 2100 13 220 561 1534 3470 198 262 314 279 49124 513 12244 310 279 11218 1238 279 56276 198 262 2100 13 220 357 1103 314 279 9376 3470 303 279 49124 369 1179 198 262 12244 310 279 56276 2100 318 1719 1103 628 381 4147 553 271 220 348 18236 25 561 1576 1902 381 449 49124 440 279 1788 1324 314 3470 198 262 430 1017 513 11218 303 279 2100 1103 11 321 279 3470 513 198 262 12244 11 494 2047 310 1245 11 310 279 11811 11218 13 271 40238 314 449 1576 310 264 3074 2100 369 51070 4364 430 198 18000 82 13 271 9 1368 279 2100 369 449 12448 318 591 1590 271 220 348 1368 279 803 1503 524 11822 303 264 328 9461 1 466 328 6082 2359 1 4952 198 262 303 279 1428 1970 2424 25 279 803 369 6608 310 279 1576 303 279 198 262 1428 2136 4325 13 271 220 348 17673 25 279 803 369 6608 310 279 1576 303 279 3521 4325 198 262 466 279 15484 4325 10521 539 328 6082 2359 487 15119 13 271 220 561 803 369 40158 413 424 557 2582 6608 13 220 1061 1189 5073 279 198 220 5604 1698 364 279 1576 8334 6608 310 279 803 310 5372 198 220 6942 11 13969 279 1576 310 381 3372 38157 321 1141 44462 318 331 424 198 220 682 799 8 310 381 2512 13 271 9 1368 279 2100 369 449 6816 5604 25 561 5839 7258 303 198 220 279 5604 369 24276 13 220 1049 1220 7302 449 1576 440 198 220 86370 7944 26 413 411 369 524 279 1105 11 328 77318 1 369 198 220 8953 13 220 2844 1576 369 1179 4439 310 9501 279 12244 1576 310 198 220 279 2574 6816 26 413 424 4021 2649 279 15840 11 424 24096 198 220 449 4531 318 41547 694 524 13902 328 3776 1401 1764 271 220 6817 25 1368 279 1576 369 264 523 2774 321 279 6816 5604 198 220 13273 383 2107 10744 314 279 15840 5497 11 279 1245 23641 3008 198 220 7258 11 328 64 1926 1 628 2528 2892 449 2774 6816 466 318 331 874 198 220 2774 6816 6513 8 264 523 6816 13 220 561 2047 23641 3008 198 220 2100 328 64 1926 1 369 2590 716 430 449 2774 6816 11 6611 424 413 198 220 5689 13 220 13891 11 279 1330 54133 314 328 64 1926 1 635 524 13902 198 220 8042 310 279 1788 6816 25 413 279 1245 23641 3008 7258 198 220 18675 310 264 523 6816 11 279 2047 23641 3008 11121 264 491 198 220 2774 6816 430 279 2100 314 279 15840 25 271 257 523 2356 82 25 198 260 830 283 220 18 1734 653 523 3759 198 257 1699 283 2356 82 363 198 257 1699 1926 283 1699 1926 478 220 16 256 653 13514 1699 1926 430 220 19 9101 2356 82 1926 430 220 18 271 220 1061 3874 1503 524 13902 3668 310 15933 198 220 7944 11 1680 430 5706 3354 440 328 3573 363 3158 271 9 1368 279 2100 369 264 14701 25 561 5839 7258 303 279 198 220 5604 369 24276 13 220 1049 1220 7302 2892 264 24983 8240 198 220 1576 318 20171 430 264 1103 8 466 264 12366 1576 318 20171 430 264 10675 553 198 220 9019 11 279 94860 7258 369 24276 13 271 220 1368 279 5839 369 264 24983 8240 1576 318 20171 430 264 1103 681 279 198 220 94860 1902 7302 449 7308 13 220 1368 424 369 7968 11 279 8240 725 198 220 2984 369 3568 310 424 13 220 561 12571 869 1902 381 264 2397 40823 198 220 7308 2599 1056 279 8240 725 2984 11 321 279 8240 369 4439 198 220 310 9501 279 12244 1576 310 1141 1455 440 421 1857 13 220 1368 279 198 220 1857 369 680 314 2020 11 328 1498 1401 1 369 8953 318 28973 310 264 198 220 94860 290 8240 4021 884 491 3470 310 264 1103 553 271 220 1368 279 5839 369 264 12366 1576 318 20171 430 264 10675 681 279 198 220 94860 1902 599 264 913 17606 440 279 12366 725 1328 913 11 198 220 321 279 12366 369 1179 4439 310 1791 264 1328 55616 6505 864 13640 198 220 279 94860 310 279 12244 1576 13 220 1061 628 2892 8032 449 198 220 6149 1328 55616 6505 440 279 1788 1328 869 11 466 5479 264 491 198 220 1328 55616 6505 318 331 874 1328 440 279 1788 869 23533 553 271 220 1690 1156 37435 5978 11 279 12880 723 1163 548 33905 1654 369 2512 440 198 220 8053 5789 13 271 9 1368 279 2100 369 264 81987 25 561 5839 7258 303 279 5604 198 220 369 24276 13 220 1049 1220 7302 264 24983 8240 1576 318 20171 430 264 198 220 1103 553 220 561 12244 1576 1220 381 264 8240 1576 314 279 1788 198 220 913 13 220 9019 11 279 4570 321 8156 6608 22666 513 24276 11 198 220 303 90349 781 513 2945 26 16184 513 6942 321 279 8240 725 198 220 2984 13 220 561 13854 1220 14989 310 24959 13 1368 2892 6608 369 198 220 7968 11 279 8240 725 2984 369 3568 310 424 13 220 561 12571 198 220 13854 513 70575 310 9940 1881 6942 321 279 8240 725 2984 11 198 220 27389 13 220 16864 11 279 8240 1576 369 4439 310 8032 279 198 220 15511 440 279 3470 314 279 12244 8240 13 220 561 2984 314 279 198 220 15511 1189 381 2086 494 279 2984 314 279 12244 8240 11 198 220 8190 9722 279 2984 314 279 2100 8240 11 413 279 2100 198 220 8240 5954 424 13 271 332 6901 26033 7873 7472 64700 733 279 1428 7873 11 279 198 54304 364 11218 369 4284 310 381 279 1788 430 364 22666 11 321 198 11469 18890 369 17030 2261 279 1970 9191 9956 11 13969 198 1657 11346 1412 6397 13 271 15337 279 7044 314 15840 23197 421 73406 1881 198 1719 2047 23641 3008 321 279 1245 23641 3008 513 3260 14352 484 17140 515 318 1891 198 8422 328 64 11 292 283 292 11 264 1 70664 1330 7101 681 73406 348 53977 9 279 198 13034 314 12244 4534 7101 11822 2047 4534 6491 11 6807 198 1334 286 303 20684 13 220 1690 2774 11 279 2614 1957 22741 198 35643 15 11 220 17 60 763 271 256 830 283 498 15 11 220 16 60 198 256 585 283 220 15 198 256 585 11 830 957 60 283 220 16 11 220 17 260 653 585 369 5860 11 1179 830 957 60 369 5860 198 256 1133 2007 8 271 9538 1048 25 271 220 2972 1679 47 220 18 16 18 17 332 471 39226 36760 1188 91841 198 257 561 24332 364 279 15171 5480 1 4416 13 1358 21199 26115 15840 11889 198 1464 43414 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 2014 84027 15840 11889 198 1464 43414 271 9 19109 9 15840 369 279 10286 11 303 264 3074 4952 11 314 198 64 3759 466 6816 20572 321 449 9801 15840 198 23425 25 271 256 60736 50145 21152 61987 18720 5480 22698 7258 198 761 4241 420 318 11531 1114 27160 735 7302 27160 7025 271 760 6463 494 4472 33284 11889 369 421 1132 264 3074 198 5480 369 5252 13 271 2381 4145 4874 430 15840 11218 11 413 303 523 466 4536 6676 11 198 1719 31144 513 24276 321 9476 303 264 3175 523 466 4536 198 8843 12880 38314 548 1 421 369 264 10675 12366 494 3759 198 11072 318 76 37652 413 843 8 310 24276 31144 13 1061 6816 369 198 86 21833 321 369 9184 3354 506 279 1151 314 523 466 4536 198 2513 10993 11 413 31144 513 1669 88034 13 271 2381 22666 430 15840 11218 11 279 31144 513 24276 198 331 303 523 466 4536 6676 11 694 524 9476 13 271 2592 264 803 369 60736 303 264 709 6676 11 1179 411 803 369 2136 198 1891 421 6676 13 90117 513 2496 24276 321 9476 303 709 198 78134 13 271 2592 279 1245 1375 3008 369 2945 11 449 60736 15840 25822 198 1719 4880 15840 1518 36319 31144 318 2777 198 658 44575 553 1368 279 1245 1375 3008 369 524 2945 364 449 7258 198 5480 11 1179 279 37999 64013 279 2100 3525 364 279 1483 198 1 548 723 1163 548 33905 466 12880 84163 548 33905 1562 13 271 9538 1048 25 271 220 2972 1679 47 220 20 17 21 332 471 31058 364 12051 90117 198 257 561 13340 421 3568 18890 364 34848 1060 279 4346 314 198 257 7101 318 15693 523 7101 321 2774 7101 681 198 257 4370 314 35493 1070 1472 5853 13 271 220 2972 1679 47 220 19 23 19 332 471 3856 29642 198 257 561 13340 421 3568 279 328 86711 1 4536 310 3300 264 5129 198 257 18890 364 913 31144 421 628 381 1429 303 1064 6157 198 257 7141 321 27069 82 13 271 5221 303 2243 220 18 13 23 25 4543 60736 30152 2070 279 1788 198 13392 880 303 279 1245 1375 3008 430 5611 30152 13 56703 11 198 14265 22666 318 4650 632 2195 8878 80019 13992 22666 8 8613 264 198 54304 1412 13 1358 10266 27693 198 584 332 271 3446 303 2243 220 18 13 20 13 1358 31099 709 7044 198 1464 34195 271 256 3203 9310 727 61987 498 91430 60 328 7449 1 328 727 1 2821 591 32467 498 15704 1951 60 56439 198 3696 4241 397 1 7258 60 22698 15705 271 19678 314 12654 1766 27693 628 381 20950 321 55483 506 1599 198 7462 318 3924 348 5816 13754 187525 328 11093 1 22666 11 328 7449 364 1 321 328 7449 198 4058 1 628 1132 381 1429 303 279 2407 314 264 75140 709 13 271 24407 4364 440 328 7449 687 1 18890 513 2590 75140 198 20752 11 1442 413 781 635 524 6435 328 11093 1 466 328 7449 1 20209 13 271 2064 369 264 328 32769 1401 1 310 958 264 328 28726 494 1 7258 4613 279 2407 198 1020 264 75140 709 13 271 2014 3010 314 264 75140 709 25 271 256 3203 687 2821 9453 16 11 1626 17 1590 198 950 635 92099 363 198 950 2454 1010 14308 13754 363 271 5221 303 2243 220 18 13 22 25 328 11093 1 321 328 7449 1 513 1381 20209 26 198 1666 17737 781 998 1132 11421 430 1680 4613 279 2407 314 264 198 5816 13754 709 13 1358 760 328 7449 364 1 4952 198 1464 41615 271 256 3203 5305 21152 61987 328 7449 1 364 21152 271 2014 348 299 26930 49124 9 5529 449 12880 1269 261 548 1 1654 421 198 19194 391 4523 449 348 299 26930 14650 11961 864 628 1562 198 299 26930 1970 303 1141 12880 276 419 548 1 1654 13 271 760 328 7449 364 1 4952 5954 16219 18971 888 198 299 26930 5199 4635 13 271 760 2614 1970 25 271 256 3203 364 39678 303 84788 25 198 950 15035 5803 198 256 745 25 198 950 15035 5803 17 271 3742 5067 78157 13185 310 25 271 256 5199 283 318 23954 8 198 256 5199 283 913 26202 19768 1269 261 3676 1947 8 198 256 4162 283 2912 271 256 1345 4162 25 198 950 1380 25 198 1779 39678 283 2454 913 26202 19768 276 419 3676 1947 8 198 950 3525 13809 6322 51317 25 198 1779 4162 283 3439 198 950 745 25 198 1779 15035 5803 198 256 745 25 198 950 15035 5803 17 271 9538 1048 12880 1269 261 548 33905 321 12880 276 419 548 33905 364 3447 13 271 2064 369 264 328 32769 1401 1 310 958 449 328 7449 364 1 4952 4732 279 2407 198 1020 264 75140 709 13 1358 760 328 7449 440 1 4952 198 1464 41830 271 256 3203 6406 21152 61987 328 7449 1 440 21152 271 2014 348 299 26930 2193 6436 9 369 264 348 2078 6436 9 421 369 2858 198 951 25410 10993 303 1141 348 1883 9 321 348 13259 9 5245 13 271 760 2614 1970 25 271 256 3203 440 13073 78847 430 39678 25 198 950 15035 5803 271 284 5067 78157 13185 310 25 271 256 6436 283 318 3152 61670 8 198 256 264 1883 283 913 53314 19768 64 1883 548 198 256 264 13259 283 913 53314 19768 64 13259 548 198 256 869 283 2454 264 1883 53314 8 198 256 4062 52797 283 3439 271 256 1380 25 198 950 39678 283 869 198 950 15035 5803 198 256 3525 25 198 950 4062 52797 283 2912 198 950 413 524 2454 264 13259 53314 11 348 7545 62538 3009 75458 198 1779 4673 198 256 5326 25 198 950 413 524 4062 52797 25 198 1779 2454 264 13259 53314 11 2168 11 2168 11 2168 8 271 9538 1048 12880 64 1883 548 33905 321 12880 64 13259 548 33905 364 3447 13 271 2064 369 264 328 32769 1401 1 310 958 449 328 7449 440 1 4952 4732 279 198 2513 314 264 75140 709 13 271 9538 1048 25 271 220 2972 1679 47 220 19 24 17 332 471 4315 27693 440 3203 321 2454 18890 198 257 561 13340 421 1801 1766 27693 264 5976 41954 7059 303 198 257 12654 11 321 3568 12250 18890 13 1358 27382 11497 318 7729 8 198 865 11862 271 2014 12448 29853 430 449 18546 369 264 803 13 220 3382 3644 21808 11497 198 429 20209 364 74931 7044 321 3644 93141 321 10649 364 198 73693 314 33791 321 10649 13 271 4350 279 803 369 6608 310 449 1576 11 15977 314 279 18546 34237 198 8761 1576 13 3095 264 803 369 524 6608 11 449 4621 310 14989 424 198 31207 264 328 657 1401 1 4531 13 271 332 16295 803 48509 2631 64700 3095 449 12448 421 1414 1769 13273 303 198 64 523 7044 11690 440 1330 466 777 50359 5587 321 198 26189 524 809 303 1330 466 777 69351 11 424 369 6306 264 348 1929 198 591 9 314 421 523 13 9391 4874 513 22773 310 264 4860 1304 198 14372 1970 369 7658 364 1070 13 220 561 17454 47640 279 198 1005 803 11 440 6187 69351 6747 321 264 3074 50359 198 4069 290 11 303 3929 314 279 803 13 220 1690 3010 11 279 12448 12880 72967 1 198 499 44216 303 264 523 6725 328 38617 1 668 381 22773 310 8728 38617 548 72967 3158 198 1919 17454 369 9209 314 279 80316 34826 2193 303 864 198 1719 12448 369 1429 13 220 1368 279 22773 803 369 8744 1248 198 12265 261 1056 220 17 20 20 5587 681 7873 4364 60735 362 1189 198 71 61608 13 1368 279 523 803 16661 1132 314 69351 11 874 198 1405 1570 369 2725 13 1358 43 36552 198 584 271 29292 11222 886 321 5639 72710 321 5090 23311 198 75 36552 25 271 256 22419 61987 886 34600 735 5639 34600 198 2215 735 7308 735 2153 3946 735 712 3141 862 271 79520 314 264 22419 34237 449 1576 314 279 2574 913 318 889 11 198 9363 11 7308 11 18484 1406 1324 11 6150 1324 8 440 279 2574 198 927 13 220 561 869 1189 381 9556 6921 303 279 1105 314 18484 1406 198 429 48271 318 22523 8 72710 13 220 3382 3644 24661 1091 364 3447 13 271 2324 72710 7618 310 49835 795 4346 11 321 15612 279 198 1640 725 9284 369 2599 2894 1056 1141 869 13 220 27869 198 13764 35840 314 72710 440 279 1788 869 318 48219 279 1788 198 13581 19934 303 279 1957 1414 466 264 2086 30520 8 1189 6637 198 1719 1788 1576 466 264 2086 1576 440 279 1788 869 13 1358 9962 4694 6816 2528 198 34581 271 760 2614 5245 628 381 4364 310 30069 279 7063 314 198 8843 2528 318 784 314 11 15840 310 11 466 34872 314 328 87 2557 871 364 198 1005 12740 13 271 1640 4692 78913 3676 701 11 803 8 271 256 20505 948 279 1581 6816 2528 14107 440 449 198 256 328 3776 1401 1 318 48219 12880 447 8843 548 33905 24096 449 198 256 328 3776 1401 1 1521 348 591 9 369 524 449 2774 6816 466 449 198 256 6816 303 279 523 4757 364 328 701 4961 466 12880 447 548 33905 314 264 348 591 9 198 256 3234 24096 328 3776 1401 1764 220 1061 1654 1220 2892 198 256 460 279 318 68600 8 6816 869 466 4673 449 328 3776 1401 1 198 256 4531 13 271 256 6817 421 413 279 6816 369 1669 1472 279 4472 16456 11 198 256 12880 78913 548 33905 369 524 2512 13 220 318 1919 369 449 45299 51505 31179 198 256 1881 12880 78913 548 33905 321 12880 84163 91514 3158 8 1061 369 2725 2107 364 198 256 14588 7715 321 1521 5752 12880 78913 548 33905 1000 599 198 256 874 1560 310 2528 975 7944 314 279 2774 13 220 6817 421 506 198 256 3140 364 2774 7101 11 488 628 12062 2702 2444 539 524 198 256 37029 866 2663 303 279 2774 6816 10675 318 7834 198 256 4370 37029 1070 303 2361 1576 553 220 3382 279 198 256 12880 447 8843 548 33905 1654 3559 364 264 1560 310 3404 615 2702 198 256 2444 888 6816 2528 13 271 1640 4692 447 8843 3676 701 11 803 8 271 256 20505 632 92649 310 4072 6816 74992 364 198 256 12740 314 279 523 13 1368 279 523 1048 18091 12880 78913 548 48812 198 256 279 14827 668 524 381 2512 7014 12880 447 8843 548 33905 2892 198 256 6526 424 20335 466 24096 449 328 3776 1401 3158 1061 1654 198 256 1220 460 279 318 68600 8 6816 869 466 4673 449 198 256 328 3776 1401 1 4531 13 733 1906 310 5471 23064 48575 303 198 256 411 1654 11 1141 7873 1220 2590 1562 279 2255 523 198 256 1654 440 279 1788 803 310 2528 866 7944 424 3749 11 364 198 256 3010 11 328 1640 4692 447 8843 3676 701 11 803 90765 271 256 6817 25 271 257 1061 1654 1189 1990 381 29715 290 948 3221 685 3175 5245 198 257 430 279 1067 314 17457 27757 4417 3992 18890 466 198 257 5617 3309 5568 13 3382 9493 1654 18054 13 271 256 1690 3529 15739 6816 74992 11 24096 449 89304 1484 198 256 328 1640 4692 78913 548 1 440 5789 328 2220 1 321 328 591 3158 271 1640 4692 84163 3676 701 11 803 11 869 8 271 256 20505 948 449 6816 15840 369 16693 13 220 1061 369 2512 198 256 4370 314 279 4472 16456 318 72 1673 13 3436 279 869 303 279 198 256 2774 10675 553 348 591 9 369 279 6816 803 11 348 927 9 369 279 198 256 869 310 381 12244 310 424 13 271 256 1368 12880 84163 548 33905 6587 310 9501 310 449 2774 6816 11 424 198 256 1220 1562 279 2255 523 1654 440 279 1788 803 11 364 3010 11 198 256 328 1640 4692 84163 3676 701 11 803 11 869 90765 271 256 1690 3529 15739 6816 30152 11 24096 449 89304 198 256 1484 328 1640 4692 84163 548 1 440 5789 328 2220 487 328 591 487 328 927 3158 271 1640 4692 9302 2896 3676 701 11 803 8 271 256 8639 12880 84163 548 33905 694 364 6816 34872 4370 314 198 256 15840 13 220 1061 1220 1132 381 11205 413 328 9302 2746 2557 1 369 198 256 21716 364 279 1576 13 271 256 1690 3529 15739 6816 18686 880 11 24096 449 89304 1484 198 256 328 1640 4692 9302 2896 548 1 440 5789 328 2220 1 321 328 591 3158 271 1640 4692 3614 3676 701 8 271 256 20505 948 328 3614 33905 369 2512 383 279 1576 13 1473 49124 1902 381 198 256 5743 13 328 3614 33905 31637 279 5743 49124 310 264 1103 321 198 256 20218 424 13 1358 9962 4694 4536 6816 2528 198 2960 8442 271 19747 4874 12880 78913 548 1 321 12880 3614 548 1 628 381 1048 1429 310 198 91125 2528 310 4536 7944 13 561 12880 78913 548 1 709 506 198 1719 4536 2119 1220 4054 799 5515 864 369 279 803 314 449 198 8843 321 460 279 23424 869 466 4673 449 328 3776 1401 3158 198 2592 449 6816 369 524 1669 383 264 4536 1576 1472 279 4472 198 20377 11 585 1673 13 328 1640 4692 447 8843 548 48812 1179 12880 78913 548 1 369 198 1773 290 303 279 4536 12880 8217 548 1 1518 17987 449 328 3776 1401 3158 198 2592 1669 11 424 369 2512 440 279 6816 803 321 279 1067 369 198 74936 13 271 760 12880 3614 548 1 709 1220 4054 874 5789 11 321 460 449 198 1947 470 314 8796 421 10551 279 4874 15156 383 4536 13 1368 198 27803 11 411 709 34282 279 5129 328 3614 33905 2624 383 264 198 4211 13 271 2381 264 777 6699 1046 2540 46439 314 279 4536 7465 318 14874 198 11987 11 5706 11 4831 23618 799 628 716 279 12880 1005 548 1 6816 198 1020 264 4536 1576 310 264 35959 314 328 8969 26094 900 3158 1690 3010 25 271 256 1120 5530 198 256 494 4346 1120 13317 900 271 256 523 6054 8039 3223 72729 900 1590 198 950 687 1259 29828 3676 701 1590 198 1779 460 281 6 41197 81 8039 313 701 4692 591 548 11325 271 950 687 1259 84163 3676 701 11 6175 11 869 1590 198 1779 1133 925 12887 49369 313 2896 92 1076 837 198 1779 2184 25658 84163 3676 2896 11 869 8 271 256 5530 25949 88970 591 548 906 548 1005 548 283 6054 8039 3223 271 8839 25 271 220 3761 5562 4536 12880 78913 548 1 321 6047 4536 12880 1005 548 1 1132 198 220 7556 1353 8339 1801 1608 279 6816 2528 18890 1317 5774 198 220 30745 279 4536 36536 318 47132 539 1970 2785 279 4536 11 466 198 220 4417 264 5604 310 279 4536 725 36536 10675 8 369 74872 13 271 5221 303 2243 220 18 13 20 25 12880 1005 548 1 4536 6816 369 1381 44375 13 271 3446 303 2243 220 18 13 22 25 12880 78913 548 1 321 12880 3614 548 1 4536 7944 13 271 9538 1048 25 271 220 2972 1679 47 220 20 21 17 332 471 13317 1259 78913 548 321 1259 3614 548 198 257 3744 53283 279 12880 78913 548 1 321 12880 3614 548 1 5568 383 13067 13 1358 60856 286 3744 24887 198 1464 857 271 760 2614 5245 1132 3668 948 449 2774 314 279 523 198 747 1988 279 1654 318 64 748 18330 348 51347 9 523 8 7701 303 449 198 9 7862 9 523 318 1719 15933 1902 381 303 2892 279 6171 725 523 198 34491 466 303 279 523 10675 364 799 314 1141 6357 553 220 733 279 198 49927 3559 11 1001 1719 6816 828 18675 310 279 6816 6483 803 369 198 1719 1328 314 279 3234 303 279 6171 523 515 12880 8217 548 3158 271 1640 4692 447 3676 701 11 2774 11 6171 5675 8 271 256 20505 310 615 279 6816 314 279 6171 523 318 1005 6816 198 256 2528 8 466 314 449 2774 314 421 523 318 4696 6816 198 256 2528 553 561 9801 348 7862 9 5515 369 279 6171 523 11 1345 198 256 348 4696 9 369 279 2774 421 279 6816 557 23975 1472 11 198 256 466 328 3928 1 948 279 6816 369 23975 1472 279 348 7862 19315 271 256 1061 1654 1220 460 279 23424 6816 869 466 4673 449 198 256 328 3776 1401 1 4531 13 271 256 2972 1679 47 220 17 20 17 332 28151 421 12880 447 548 33905 369 39571 440 799 466 1330 198 256 5789 13 220 12654 725 1765 5617 3309 45933 1761 411 198 256 24332 26 4611 11 424 369 4222 421 1010 4688 23280 7141 198 256 599 45933 421 1325 2107 5789 13 220 12654 725 1765 198 256 12880 447 8843 548 33905 7873 2590 15734 303 2107 5789 198 256 3315 781 513 2483 466 524 13 271 1640 4692 723 3676 701 11 2774 11 869 8 271 256 20505 310 716 279 6816 383 449 2774 348 4696 9 314 279 6171 198 256 523 310 264 491 869 11 348 927 19315 271 256 6817 11 7594 12880 723 548 33905 466 12880 4394 548 33905 4203 279 2993 314 198 256 15933 310 264 1001 673 15933 10878 220 3382 18239 10430 3744 24887 364 198 256 777 3447 13 271 1640 4692 4394 3676 701 11 2774 8 271 256 20505 310 3572 279 6816 383 449 2774 348 4696 9 314 279 198 256 6171 523 13 271 41310 314 45933 1189 1048 599 279 12880 2220 1005 548 1 6816 198 27803 25 271 1640 4692 2220 1005 548 271 256 561 6816 12880 2220 1005 548 1 369 31233 539 279 328 79221 1 4536 198 256 430 36588 279 523 1332 411 1576 557 4364 318 14874 411 198 256 33744 628 7543 303 15135 24715 15603 314 8476 523 198 256 7944 553 1690 1562 4635 11 424 1189 12832 421 449 2774 314 279 198 256 2574 913 318 269 264 35959 8 369 3481 466 2483 430 279 1118 198 256 65269 5515 318 1891 3010 11 351 29292 7062 411 6816 364 198 256 632 10573 5245 421 513 11205 303 351 553 1358 14732 10430 3744 24887 198 1464 590 271 623 4437 11 264 15933 369 449 1576 6816 440 1001 30018 198 53992 9059 799 6483 6816 2528 682 978 49724 539 5245 198 258 279 15933 11176 25 220 12880 447 548 48812 12880 723 548 48812 321 198 1 548 4394 91514 3158 1368 866 314 1782 5245 513 4364 364 449 1576 11 424 198 284 1018 310 381 264 15933 13 271 760 1581 7465 364 6816 2528 369 310 615 11 716 11 466 3572 198 1719 6816 494 449 1576 725 10675 13 1690 2774 11 328 64 1926 1 682 264 198 20377 8516 5732 440 328 64 4692 8217 548 659 87 72203 1179 198 42993 2784 19768 8217 548 659 87 72203 321 13942 1472 279 2255 6632 314 198 42993 2784 9646 42328 2197 565 32748 13 271 10883 11 413 279 6749 5072 869 369 449 1576 25471 799 314 279 198 51347 5245 11 1179 12654 1189 2723 279 1581 7465 321 198 21731 279 15933 1654 4370 13 220 10646 411 13273 303 279 198 25317 738 8516 13402 383 864 15933 5245 998 4364 321 198 4993 781 998 2512 13 271 760 5732 1406 364 15933 27757 369 264 10649 11 328 64 1926 3158 2500 198 1719 5789 513 32878 13402 383 328 64 763 271 15555 6920 198 256 561 43070 321 3140 4047 1562 369 948 1156 1970 5774 198 256 63821 264 15933 1654 25 262 328 87 4692 447 3676 64 90765 271 2443 25073 198 256 1368 10649 310 449 1576 2774 11 328 64 1926 1 369 22773 1083 279 198 256 1562 25 328 1267 2784 19768 8217 548 659 87 642 4692 447 3676 64 11 913 2784 578 3158 271 1890 25073 198 256 1368 10649 310 264 523 11 328 32 1926 1 369 22773 1083 279 1562 25 198 256 328 32 4692 8217 548 659 87 642 4692 447 3676 3928 11 357 90765 271 18701 25073 198 256 357 56150 18054 1680 430 328 9241 4205 11 264 553 87 1 25628 198 256 328 64 4692 1005 47719 76 298 548 1 364 264 2255 523 328 33 1 2614 328 32 1 321 1179 198 256 4523 328 33 4692 8217 548 659 87 642 4692 447 3676 64 11 357 90765 220 1368 524 264 15933 11 328 87 1 198 256 369 5743 33702 13 271 2381 2774 34519 11 279 51274 314 15933 27757 13402 198 263 864 15933 5245 513 4364 13 220 357 15933 628 6762 866 198 72474 314 12880 447 548 48812 12880 723 548 33905 321 12880 4394 91514 3158 220 1368 424 198 26189 524 6762 12880 447 548 48812 1179 30745 279 6816 668 460 198 1719 15933 1576 4924 7014 1017 369 264 869 303 279 1576 725 198 4696 10675 13 220 1368 279 15933 18091 12880 723 548 33905 321 4979 198 1 548 4394 548 48812 424 369 264 795 15933 26 413 424 18091 13468 11 424 369 198 64 2397 13545 15933 13 220 49941 11 795 45933 6762 2107 198 1 548 447 548 33905 321 12880 723 548 48812 1345 2397 13545 45933 599 1066 279 198 1 548 447 548 33905 1654 13 220 2791 45933 440 12880 447 548 33905 321 12880 723 548 33905 198 57889 4979 12880 4394 91514 871 4364 2590 2723 264 311 17829 303 449 198 4696 10675 13 220 733 12503 11 2397 13545 45933 628 381 198 1908 40787 539 12740 13 271 29292 5245 318 15693 1782 34816 440 8168 21527 1 321 198 92943 25901 871 513 11205 430 2397 13545 45933 13 220 60694 11 198 46230 628 80403 321 2723 5245 13 220 1061 5954 3714 198 46230 310 20604 26228 421 1684 494 975 12740 314 279 198 23310 523 13 271 760 328 3573 33905 709 369 11205 430 264 795 15933 13 198 10865 391 11 12740 4021 2723 279 7465 314 264 3234 13 1358 548 48093 548 198 41615 271 9 548 48093 548 9 2070 586 310 20335 15444 795 3492 318 4650 198 12811 8 321 22380 279 9400 314 12880 8217 548 1 321 35832 12768 1062 548 9 198 17818 1657 20335 13866 303 35832 48093 548 9 466 2420 303 264 2594 5946 271 760 3433 6568 888 1608 12880 8217 548 1 628 381 4927 13 16260 198 20377 4478 628 381 11602 12573 430 1575 13 271 1640 4692 48093 548 271 256 1061 523 3759 628 381 12244 264 886 11 49124 11 466 8240 198 256 314 8796 440 3759 4874 1429 539 12740 13 220 35832 48093 548 9 198 256 28538 3433 364 279 13866 7101 321 27028 279 198 256 16465 9400 314 12880 8217 548 1 321 35832 12768 1062 548 9 364 1754 198 256 2774 13 271 21003 383 1608 35832 48093 548 77340 271 9 3095 14722 5672 494 264 523 1973 35832 48093 548 11961 279 12880 8217 548 1 321 198 220 35832 12768 1062 548 9 6816 314 279 12740 668 2590 381 15156 13 271 9 16643 264 12880 8217 548 1 3759 11 12740 4021 381 12244 491 198 220 7101 524 9711 303 279 35832 48093 548 9 7044 13 220 80427 310 198 220 9501 310 449 632 30214 3759 803 24096 328 3776 1401 3158 1368 198 220 8476 15840 314 491 7101 369 12321 11 1179 884 198 220 6952 548 8217 548 14482 310 279 8240 314 8796 303 279 35832 48093 548 9 198 220 17466 13 271 9 16643 264 35832 12768 1062 548 9 3759 364 1754 2774 11 6632 25471 198 220 35832 48093 548 9 635 524 1761 328 12768 14619 1 310 1141 12740 13 1368 198 220 7234 5604 1761 369 4221 11 1179 884 6952 548 12768 1062 548 14482 310 279 198 220 8240 314 8796 303 279 35832 48093 548 9 17466 13 271 9 35832 48093 548 9 513 11205 506 279 523 2119 539 6611 198 220 45933 364 1754 3759 803 13 220 1577 264 1067 11 523 7944 198 220 4021 381 1429 310 716 1581 2663 364 2774 7101 4364 198 220 539 35832 48093 548 243343 5752 11 279 523 6816 1000 31153 279 198 220 15933 15840 13 271 9 561 1852 314 264 35832 48093 548 9 17466 369 524 6973 310 279 523 198 220 1332 424 369 4364 13 220 35832 48093 548 9 13866 303 6357 513 2420 198 220 303 1623 6632 13 4213 11 1623 65878 668 615 264 12880 8217 548 1 198 220 321 35832 12768 1062 548 9 7014 781 1048 6762 35832 48093 548 9 318 7949 1220 198 220 1132 6435 4874 314 866 348 34325 9 15250 553 271 9 1368 264 523 18091 264 9167 1048 4364 303 264 2255 523 11 279 2774 198 220 3759 4364 539 279 2255 523 9167 369 78909 318 11347 539 198 220 46616 1141 15933 5774 494 279 2255 523 553 1061 198 220 42812 279 7063 314 279 1957 5438 13 220 733 279 3725 11 264 198 220 1716 1189 381 3568 310 5190 411 13 271 9 328 77318 1 668 381 8953 413 2397 3092 35832 48093 548 9 513 4364 364 264 198 220 523 13849 494 264 1535 9521 28365 1 5617 3309 913 1 1680 430 198 220 328 390 487 328 9363 487 321 328 23814 3158 271 9 5586 2397 29868 348 1947 470 9 1189 381 12244 310 35832 48093 548 19315 271 9 1368 264 328 34491 1 369 1429 310 9501 35832 48093 548 11961 279 10675 6679 198 220 668 381 1429 430 279 9167 4874 13 561 2663 314 279 10675 628 381 198 220 1429 310 3300 791 12 8843 4479 18034 421 668 381 37382 539 198 220 328 79221 652 5069 33905 321 12234 303 279 2468 314 328 8389 363 3158 271 9 12880 1005 548 1 15840 4138 1132 413 2107 6632 599 279 1788 198 220 35832 48093 548 19315 271 9 27869 40443 440 5081 1678 15065 2594 6632 628 381 198 220 1429 11 694 1132 799 2594 369 5252 310 599 7944 3354 539 198 220 15250 318 1719 975 22372 1902 599 4147 9167 47410 8 471 25709 198 220 4673 328 77318 3158 271 9 1368 449 348 6640 9 369 1429 364 35832 48093 548 9 1179 264 348 51347 9 369 198 220 3354 364 1754 314 279 14650 725 2663 13 4213 11 279 35832 48093 548 9 198 220 6816 668 381 449 4147 14650 13 1358 3776 14619 198 30896 271 2014 6816 5604 369 264 5839 7854 539 264 4031 321 264 803 25 271 256 16819 330 42514 61987 5839 22052 12448 271 760 5839 1902 14989 310 449 1576 314 264 913 421 11222 198 8843 14619 11 864 1379 5978 635 13 220 1061 1576 369 1179 198 1038 290 310 7936 279 6816 6483 803 369 279 12448 13 561 913 198 429 869 8677 369 10521 539 279 1576 13 220 27869 53231 198 1020 279 1788 6816 5604 1189 7302 2086 5978 13 271 1919 5492 628 381 30650 539 65581 279 198 1 548 447 8843 548 33905 1654 466 279 12880 78913 548 33905 1654 13 220 561 198 1 548 447 8843 548 33905 1654 369 2512 1118 321 2892 4523 264 869 198 269 24096 328 3776 1401 1 413 279 6816 369 524 2420 13 271 2592 449 328 3776 1401 1 369 8953 321 279 1576 682 264 12880 78913 548 33905 198 4249 11 421 1654 369 2512 430 264 31686 13 1358 21199 26115 15840 11889 198 13052 37745 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 36117 7258 198 865 271 50 12393 279 10993 314 348 5816 13754 9 383 449 348 11093 470 9 1576 13 2885 198 3138 381 1429 4613 264 348 5816 13754 709 19315 271 256 2454 21239 61987 328 11093 1 5839 271 3446 303 2243 220 18 13 20 13 1358 20682 33633 7289 198 34581 271 760 7620 33633 7289 599 279 20066 10304 198 41150 13 220 6817 421 1010 314 1439 7289 1048 3668 310 3529 2397 12 198 19000 4346 13 220 33554 494 279 2279 5497 11 1017 513 1132 1330 198 41150 11 799 364 12197 63062 19024 321 799 364 61103 198 31750 25 271 256 295 21239 61987 560 21239 735 295 21239 63669 560 21239 735 295 21239 93504 295 21239 735 198 1025 295 21239 34397 1 560 21239 735 295 21239 16397 560 21239 735 198 1025 295 21239 76104 560 21239 198 256 264 21239 61987 295 21239 735 264 21239 57488 295 21239 735 264 21239 25675 295 21239 271 760 63669 318 75305 1633 8 5497 34237 279 1918 314 1141 5789 13 198 760 5789 1902 2892 2107 381 4947 11 466 799 5515 1902 381 449 198 11326 321 279 975 1902 381 264 8240 13 733 279 4588 1105 11 279 198 36142 513 15625 310 264 4047 913 321 1179 53070 3658 13 198 623 279 14827 1105 11 8240 51623 369 10345 26 264 7968 198 265 53897 7911 34237 449 4147 8240 13 271 1919 5488 628 381 30650 1608 279 3175 12880 23402 548 33905 321 198 1 548 8454 355 548 33905 5245 13 271 760 93504 318 266 8 5497 369 10287 310 381 1429 364 5979 198 75305 1633 13 220 2233 57610 12654 4346 4072 411 5497 13 271 3446 303 2243 220 18 13 20 13 271 760 16397 318 27822 8 321 34397 1 318 29457 12436 8 19024 7302 279 198 13104 1128 314 836 5789 13 220 561 23311 5789 513 1118 198 73596 310 264 4047 913 13 14074 314 24959 34237 264 2153 11 1345 198 29457 12436 314 24959 2961 303 449 7308 26 279 1067 369 421 198 1020 34780 12436 440 279 3260 29457 515 709 8977 310 279 198 1334 13 220 14074 539 6942 24096 279 328 17461 49513 1401 1 4531 13 271 1919 5488 628 381 30650 1608 279 3175 12880 371 3155 340 548 33905 321 198 1 548 86715 526 340 548 33905 5245 13 271 760 76104 318 84934 8 5497 34237 279 25472 494 279 12436 314 198 1719 1118 5515 539 279 2018 13 220 561 23311 5789 513 1118 198 73596 310 264 4047 913 13 220 357 6942 1245 5515 24096 279 198 1 17461 49513 1401 1 4531 13 220 561 5789 1189 381 18484 1406 198 36142 11 378 1257 2487 328 18 13 16 19 4 15 13 22 1 16327 328 15 13 18 19 1 318 10957 328 18 13 16 19 1 16327 328 19 9 15 13 22 478 198 15 13 18 19 3158 8 220 561 57844 5497 2590 34237 264 1067 440 279 1788 1777 198 299 1141 2018 26336 318 269 6942 1174 279 10424 869 314 279 1067 369 198 6418 391 8882 1056 279 10424 869 314 279 2018 26336 498 16 906 271 760 6220 12436 321 57844 19024 513 8173 539 279 2614 198 16417 25 328 87 606 318 87 320 88 4653 88 478 318 87 4 88 90765 220 16093 12436 321 57844 513 1048 198 14844 440 279 5617 3309 709 328 596 2508 363 763 328 596 2508 2007 11 374 8 606 198 2007 320 88 11 830 4 88 90765 498 17 906 271 623 5089 310 15899 279 57844 5488 383 4947 11 279 76104 198 7636 369 1048 72052 539 886 5978 310 2649 2235 10970 198 889 35358 318 18092 3750 430 35282 553 220 561 18890 364 198 889 35358 369 7246 303 279 12654 11300 16699 11 198 2720 3964 10970 894 86508 13 271 760 348 84934 9 5488 628 381 30650 1608 279 3175 12880 2508 548 33905 198 4249 13 271 760 6220 12436 5497 11 279 57844 5497 11 321 279 328 596 2508 33905 198 1628 513 524 4364 364 6150 4947 13 220 11746 11 5335 310 264 198 79665 1406 1324 1608 279 328 3325 33905 709 413 8053 13 271 760 57488 318 698 662 8 5497 34237 279 2542 314 1141 5789 13 220 561 198 15889 1902 2892 2107 381 4947 466 2107 381 22958 314 279 1788 198 1267 13 220 733 279 4588 1105 11 279 4947 513 15625 310 264 4047 913 198 429 1179 3568 3658 13 733 279 14827 1105 11 279 22958 513 198 19551 268 639 13 271 1919 5488 628 381 30650 1608 279 3175 12880 698 548 33905 321 198 1 548 81 698 548 33905 5245 13 271 760 25675 318 1899 25580 8 5497 34237 279 6463 314 1141 5789 13 198 760 23311 5789 513 1118 15625 310 264 4047 913 13 271 1919 5488 628 381 30650 1608 279 3175 12880 1899 548 33905 1654 13 1358 20682 94589 7289 198 13052 9 271 4699 314 279 2250 94589 7289 682 264 2086 10304 2119 25 271 256 321 21239 61987 6206 21239 735 321 21239 13012 1 6206 21239 198 256 52131 21239 61987 321 21239 735 52131 21239 38591 1 321 21239 198 256 466 21239 220 61987 52131 21239 735 466 21239 58383 52131 21239 271 760 13012 1 5497 34237 279 94589 3449 314 1141 5789 11 864 1902 198 1323 24959 466 799 314 1070 1902 381 264 2446 1576 65581 198 1 548 429 548 33905 466 12880 11008 548 33905 3175 5245 13 271 760 38591 1 5497 34237 279 94589 67521 318 86055 2639 8 314 1141 198 15889 11 864 1902 381 24959 466 799 314 1070 1902 381 264 2446 198 1640 65581 12880 68890 548 33905 466 12880 12288 269 548 33905 3175 5245 13 271 760 58383 5497 34237 279 94589 318 94964 8 2639 314 1141 5789 11 198 7949 1902 381 24959 466 799 314 1070 1902 381 264 2446 1576 198 1908 48008 12880 269 548 33905 466 12880 7389 548 33905 3175 5245 13 1358 2010 23359 198 584 422 271 2010 5978 513 1429 539 279 7873 310 3875 1001 66561 12 198 49134 828 31097 12654 1970 1680 430 264 709 2407 13 2302 1684 198 1445 709 5978 1521 781 1459 1357 6435 264 5604 310 836 198 9461 10993 4424 13 220 5927 5978 513 5743 539 279 5617 12 198 258 328 19810 33905 709 321 628 381 26543 494 709 5978 198 19139 836 12880 1787 548 1 6816 13 3382 1048 279 328 1787 1 4536 13 271 5865 286 12880 1787 548 1 24096 449 89304 1484 328 1640 4692 78913 548 1 198 4058 5789 328 2220 1 321 1535 548 1787 548 67008 271 32 1970 1576 628 381 15234 466 24276 539 11946 424 318 62643 314 264 198 2348 886 8 310 279 328 11411 33905 466 328 13764 33905 220 5617 3309 5568 13 271 9538 561 5129 913 27980 364 777 1928 13 1358 760 13150 45823 2907 198 865 11862 271 1919 1576 369 16138 1429 539 81987 318 3924 326 407 793 553 220 1049 11222 198 2083 3175 7289 13 220 2532 369 6681 799 24792 45823 1576 11 6725 198 1 41295 45823 1 318 64 5617 3309 803 553 220 328 1267 10406 636 45823 8 33905 18082 279 198 1 41295 45823 1 35680 13 271 2064 369 5158 430 328 41295 45823 1 466 37431 3158 1358 760 17545 2907 198 584 37745 271 1919 1576 369 5743 539 5568 421 1459 1357 20335 460 264 198 927 13 220 1049 11222 874 3175 7289 13 220 2532 369 6681 799 819 198 1640 11 6725 328 3928 1 318 64 5617 3309 803 553 220 328 1267 25864 8 33905 18082 279 198 23310 35680 13 271 2064 369 5158 430 328 3928 3158 1358 900 23359 198 584 422 271 900 5978 3875 279 5090 1576 4346 13 220 1473 1576 725 913 369 198 5045 290 539 279 5617 3309 709 328 1267 363 3158 220 2532 513 874 3175 198 36901 383 4346 13 220 561 5129 4536 328 8969 1 18091 4874 364 198 528 5129 5617 3309 4346 13 271 4036 513 5158 1040 411 25 3919 1005 359 390 6073 3158 1358 6675 7289 198 865 332 271 256 466 4306 220 61987 321 4306 735 466 4306 328 269 1 321 4306 198 256 321 4306 61987 524 4306 735 321 4306 328 429 1 524 4306 198 256 524 4306 61987 11960 735 328 1856 1 524 4306 271 623 279 2193 314 6774 7289 11 321 1048 948 22666 513 198 2506 539 2444 6195 11889 11 279 2614 2663 513 31233 198 299 867 25 328 3913 487 328 3928 487 23311 6942 314 660 4346 11 321 4147 198 18034 321 23107 318 15693 8796 11 43716 11 11140 11 198 67 42723 11 7062 321 12431 17991 1367 553 220 1942 975 2663 513 31233 198 299 804 13 220 2570 37435 5978 628 30069 836 7793 869 539 198 751 1256 286 264 12880 2554 548 33905 1654 13 271 760 5497 328 1856 1 34237 328 2434 1 413 1141 5515 369 867 11 328 3913 1 198 57913 13 271 760 7258 328 87 321 374 1 1118 64013 348 87 243343 413 348 87 9 369 867 11 1141 198 927 369 5743 26 5752 11 348 88 9 369 24276 321 279 12571 869 198 284 5743 13 271 760 7258 328 87 466 374 1 1118 64013 348 87 243343 413 348 87 9 369 804 11 1141 869 198 284 5743 26 5752 11 348 88 9 369 24276 321 279 12571 869 369 198 74936 13 271 8839 421 13468 328 429 1 6130 328 269 1 8623 279 869 321 913 781 198 671 310 328 3913 1 321 328 2434 487 694 4598 460 279 1483 24276 198 14064 13 220 1061 369 6807 5222 11 378 1257 2487 413 328 82 1 369 264 886 421 198 5274 381 12215 539 264 1581 869 413 424 369 4147 11 279 7258 198 39441 466 359 7724 14482 34237 279 12321 869 13 220 8938 328 1856 1 682 310 1791 264 198 902 869 11 424 4523 264 2623 869 15029 314 279 913 314 1141 198 14064 318 1891 3010 11 328 1856 359 7724 14482 18082 328 3913 1 4598 1056 328 4456 3158 8 1358 760 328 8689 1 4952 198 865 17731 271 256 1387 21152 61987 328 8689 1 271 1 8689 1 1189 1132 11822 80316 519 2586 23283 303 264 328 1891 1 466 328 3361 1 198 10185 11 694 524 23283 303 264 709 466 523 7044 2785 421 198 10185 13 271 2064 84591 279 22746 82000 6136 11 41244 279 9801 328 1449 1 198 79248 413 279 6136 682 799 13 271 2592 264 328 1891 1 6136 369 30245 539 328 8689 487 279 6136 2444 2100 198 432 6902 1141 1428 869 13 271 4350 328 8689 1 15734 2444 680 314 264 328 1485 1 4952 440 264 328 37879 1 198 79248 11 421 328 37879 1 20984 369 15234 1518 2098 9101 279 198 10185 13 1358 2196 15031 39571 5978 198 13052 332 271 1640 4692 6450 3676 701 36038 2736 1076 2387 271 256 20505 948 279 2774 369 1001 42431 828 430 264 709 26 413 411 1654 198 256 369 4364 11 328 87 9125 16 11 1344 17 11 2423 9646 16757 45351 310 198 256 328 1267 2007 19768 6450 3676 87 11 1344 16 11 45347 3158 1358 53435 198 17731 271 32 1562 6526 264 39571 1576 318 68 1257 2487 264 348 1628 3636 440 264 10451 198 3092 3878 314 348 15889 77340 271 256 1562 336 61987 5839 32467 498 14064 1951 498 487 1293 735 59036 60 56439 198 256 5515 1951 285 61987 65269 41997 498 2129 56276 8119 49627 60 198 5015 498 2129 20209 41997 60 198 3696 735 56276 8119 49627 498 2129 20209 41997 60 198 3696 735 20209 41997 198 256 65269 41997 61987 65269 5457 318 2129 65269 5457 4653 198 256 65269 5457 406 61987 15840 27160 735 63669 7258 198 256 56276 8119 49627 61987 3376 61946 7258 735 15698 5457 8 198 967 318 2129 63669 7258 735 15096 15698 5457 4653 198 256 20209 41997 256 61987 318 19258 5457 735 65667 1 7258 8 198 3388 318 2129 15698 5457 735 15096 65667 1 7258 4653 198 256 15698 5457 260 61987 12448 35036 7258 271 2014 9801 26849 30644 1189 381 2945 1238 279 65269 321 198 19258 5789 694 1503 524 7556 279 50921 13 271 760 5839 1902 14989 310 264 39571 1576 318 846 37435 198 20752 11 5617 3309 5568 11 5245 314 5617 3309 5978 11 523 198 18630 11 5245 314 523 12740 11 321 660 5978 3322 264 198 1 548 6450 548 33905 1654 513 39571 553 220 1942 5515 22666 513 198 13764 12484 1518 279 1562 369 16693 13 220 5044 8042 310 3644 198 4987 16954 364 279 18890 314 15440 348 15704 9 11140 13 271 2592 15698 5789 513 2945 11 781 513 1118 15625 310 198 936 2910 5789 11 430 10695 13 220 5339 11 264 1103 314 9194 4233 15250 369 198 6898 364 279 15440 4868 13 220 1368 1017 513 443 65269 198 15889 11 781 513 8826 303 279 1118 443 15250 13 220 9019 11 364 1754 198 19258 5515 11 279 12448 369 1429 310 7995 279 198 5816 6604 286 9167 318 331 279 12448 369 279 1788 430 279 1118 15440 198 15704 803 11 279 1118 9167 369 1429 11 321 748 383 553 220 1368 279 9167 369 198 43967 9896 11 264 328 77318 1 4531 369 8953 13 17673 11 279 198 14064 369 8826 303 279 9167 11 20621 424 318 16296 413 279 7258 369 198 1 3928 487 424 39248 279 9167 553 220 3095 660 5789 599 978 14789 11 198 1719 15250 421 513 1990 9194 4233 513 9896 440 279 11811 198 2186 869 494 279 709 7044 13 220 318 3549 2663 513 198 5248 48129 11 2957 11 948 279 709 369 4364 26 8190 11 264 24983 1576 198 20171 430 264 1103 466 10675 1429 430 1581 869 668 381 5904 539 198 528 6526 421 1459 1357 13441 449 5515 869 364 279 11811 198 20877 26 411 1220 5802 381 29784 5946 220 1368 1017 513 866 9194 4233 198 48093 364 864 874 1581 869 369 5024 11 264 328 77318 1 4531 198 284 8953 13 220 17673 11 279 1103 314 9896 15250 369 1429 430 279 198 14064 1103 364 279 1562 13 271 332 6901 26033 7873 7472 64700 1473 7873 1189 3300 198 44697 3309 5568 6483 65269 4868 635 524 599 4874 11 1442 198 331 781 513 3260 29259 515 364 279 7193 314 9417 11 321 864 198 17977 1278 4021 381 16713 539 15698 13 220 733 351 29292 11 411 369 279 1105 198 1891 5568 11205 303 351 421 958 328 13432 2648 74691 27742 33905 310 4563 198 48055 5789 13 271 2592 1017 513 777 65269 5789 1056 1017 513 15440 5555 198 48093 11 264 328 77318 1 4531 369 8953 11 7014 264 15440 5555 198 940 279 18890 15171 15441 1 369 2945 26 303 411 1105 11 421 15440 198 15704 20539 264 13992 8222 279 13230 65269 5789 198 77167 449 4147 13992 413 1017 998 874 13230 65269 5789 553 271 2592 866 15698 5515 1503 524 7618 310 264 15440 5555 198 591 11 264 328 77318 1 4531 369 8953 11 7014 264 15440 5555 198 940 279 18890 65667 15441 1 369 2945 26 303 411 1105 11 421 15440 198 15704 20539 264 10675 8222 279 13230 15698 198 15889 318 940 279 20209 430 6679 321 279 5515 2663 430 198 5816 6604 286 2663 681 466 264 318 902 8 4147 10675 413 1017 998 874 198 325 1084 15698 5789 13 271 2592 279 18890 15171 27188 1 7701 303 279 709 1562 11 328 27188 1 198 24028 14989 310 449 348 1947 470 19315 220 33027 494 1439 5199 4635 513 198 83 2757 430 413 781 998 4945 65269 5789 13 220 1690 279 1562 198 92845 2007 16 11 830 17 11 348 88 11 830 18 11 830 19 11250 413 348 88 9 64013 310 264 8240 348 88 16 11961 4444 11 348 88 44 11961 198 559 369 13185 310 264 1562 440 380 10 19 65269 5789 348 87 16 11961 348 87 17 11961 198 32144 16 11961 4444 11 348 88 44 11961 348 87 18 11961 348 87 19 19315 271 32 27735 314 411 369 421 7643 279 15171 27188 1 18890 1189 198 39420 348 10378 9 11135 15698 5789 11 424 369 14789 348 14372 9 198 1719 15698 5789 318 429 866 65667 27188 1 5789 1317 1436 3559 553 198 4272 25 271 256 11765 687 281 2784 11 292 1590 198 256 2423 257 1133 2784 11 292 8 198 256 2423 198 256 11765 281 1818 28 16 11 12492 17 11 578 198 256 220 17 220 16 198 256 11765 281 2784 28 16 11 12492 17 11 578 198 256 26289 1371 318 3444 3110 1562 1483 1590 198 257 2793 3919 49863 20507 1500 220 16 11 303 361 4211 29 198 256 24237 25 281 363 2597 5081 2663 364 15698 5515 359 64 6 198 256 11765 281 7 16 11 12492 17 11 578 198 256 220 16 220 17 271 2064 369 17957 364 2107 15698 5789 321 279 15171 27188 1 18890 198 951 381 1429 303 279 1788 1562 11 748 303 6380 411 20684 1503 524 198 58768 29783 13 271 2592 279 18890 65667 27188 1 7701 303 279 709 1562 11 198 1 27188 1 1902 14989 310 264 348 39446 11961 279 8529 314 864 513 198 83 2757 430 4945 15698 5789 13 1368 264 5555 12219 264 1328 198 4497 2582 978 2574 264 869 318 1634 449 11135 15698 5515 11 466 198 1445 2361 30047 286 681 264 328 77318 1 4531 369 8953 13 271 4350 65667 27188 1 369 1429 11 1754 1328 303 411 12366 1902 381 264 198 889 13 8618 869 494 279 12366 369 12244 310 279 1118 15440 198 15704 16515 364 15698 15840 6483 803 369 5951 310 279 198 767 13 357 1328 1144 524 381 264 12654 12448 318 68 1257 13 1535 2721 65884 35087 37 2918 369 198 35788 11 7643 424 668 524 2353 866 15440 5555 421 1362 198 1323 13866 553 1368 1017 369 874 2353 310 264 15440 5555 279 1328 18509 198 12306 369 14129 539 279 65667 1 5555 11 413 1017 369 799 11 466 413 1017 198 284 524 11 264 328 77318 1 4531 369 8953 13 271 1774 278 4868 1608 279 18890 15171 15441 1 466 65667 15441 1 198 32157 381 1429 430 65269 5515 15250 466 430 15698 5515 198 11072 13 271 5221 303 2243 220 18 13 20 25 5534 6526 4054 866 1324 314 63669 321 198 175267 1 30047 793 11 65269 5789 1189 1732 49124 30047 793 198 28625 3842 321 15698 5789 1189 1732 10675 30047 793 3376 332 1764 198 36030 10660 539 2972 1679 47 220 19 19 23 159034 271 32 1562 2590 4523 1010 869 11 10451 328 3928 487 7014 424 24096 449 198 7557 13 220 2500 411 869 369 23424 13402 383 279 913 314 279 198 91907 1576 13 271 2592 424 369 2218 271 64 1156 37435 709 25 198 256 561 1970 2424 364 279 709 369 15234 11 11946 424 279 198 256 5515 1103 13 220 561 1118 3065 279 1970 2424 668 635 369 10403 279 198 256 15440 4868 310 279 5789 26 411 369 7246 303 3644 198 256 5534 16954 13 220 3095 279 1970 2424 49703 264 328 671 1 198 256 4952 11 411 28151 279 460 869 314 279 709 1562 13 271 64 5617 3309 709 466 1654 25 198 256 561 1067 369 685 310 279 37999 26 1436 31959 3309 22815 364 279 198 256 26884 314 5617 3309 5568 321 5245 13 271 64 523 1576 25 198 256 357 491 2774 314 421 523 369 5743 13 271 64 523 2774 1654 25 198 256 561 11811 1156 37435 709 369 2512 11 440 449 5515 198 256 1103 421 369 799 4860 1056 279 5515 1103 314 279 1562 25 279 198 256 2774 8772 279 1118 5515 13 271 64 523 2774 25 198 256 561 523 1902 6762 264 12880 6450 548 33905 1654 26 279 2376 369 1179 279 198 256 1788 430 413 421 1654 557 2512 13 1358 1890 16954 198 865 9 271 32 523 7044 18091 264 523 1576 318 3924 3644 561 5129 198 1267 27980 1590 271 256 523 727 262 61987 498 91430 60 328 1005 1 68259 498 1267 6533 60 498 12457 663 60 22698 15705 198 256 40443 61987 32467 498 14064 1951 60 56439 198 256 68259 256 61987 12448 271 32 523 7044 369 449 31097 4952 13 220 561 40443 1103 198 41547 6486 264 1103 314 2255 6632 318 3924 6017 565 32748 364 777 198 40400 5533 681 748 1754 1455 303 279 1103 1220 14989 310 264 523 198 1640 864 5954 35959 286 13 220 32749 1973 449 40443 1103 198 12457 11 539 1581 11 494 279 2255 523 328 1640 4961 15612 11 271 256 523 32320 25 198 950 1440 271 284 13185 310 271 256 523 32320 4253 1590 198 950 1440 271 760 523 725 15705 369 1179 15234 303 264 491 10993 3899 318 3924 198 82459 321 10649 681 1608 264 13238 3354 2136 4325 321 279 198 9596 3521 4325 13 318 68253 11 279 15705 5435 9712 198 1628 16954 5946 220 3095 279 523 725 15705 31941 10993 11 1141 198 50947 3899 369 41784 694 1141 2136 4325 369 6568 13 498 20 60 357 198 1005 1576 369 1179 3354 1608 279 40443 1103 364 279 2255 198 8588 321 279 6568 2136 4325 364 279 6816 10675 13 198 760 523 803 369 6608 310 411 523 1576 303 279 3889 2136 198 2159 13 271 760 1906 303 864 7944 513 4364 303 279 523 2407 369 198 22885 2683 303 279 491 523 725 12880 8217 548 3158 220 6817 421 411 369 14294 198 3138 1245 1238 279 523 369 3354 321 1132 364 6632 421 998 198 9423 1608 279 7044 18890 13 271 1890 9400 628 381 30650 16421 1608 2197 565 32748 13 271 19465 628 1048 381 34816 25 1066 1040 948 46963 5568 11 271 256 554 69 16 9125 8 198 256 554 69 17 198 256 523 32320 25 1440 271 284 16757 13185 310 271 256 523 32320 25 1440 198 256 32320 283 281 16 9125 2299 69 17 7584 2537 578 271 760 15977 5426 364 279 48973 22666 513 279 1788 430 364 198 1628 60565 13 220 561 1067 369 1179 6608 310 279 523 803 13 271 5221 303 2243 220 18 13 24 25 32749 1189 381 34816 440 866 2610 198 1 28973 27160 3158 56703 11 279 30394 557 1691 777 198 26993 520 26 1436 2972 1679 47 220 21 16 19 332 364 3447 13 271 32 1103 314 913 4868 1189 381 2574 303 9057 37641 6849 198 10378 279 523 725 803 13 1061 14378 310 1064 913 1716 382 421 198 1719 523 369 13553 13 2332 15135 11 279 913 4868 628 381 29413 198 1445 279 523 725 12880 1267 6533 548 1 6816 13 3382 20627 6632 364 198 6183 13 271 5221 303 2243 271 14162 25 3956 310 279 47801 3294 11 1092 1503 449 1992 4952 635 948 12654 8213 440 24460 383 11 321 1092 1503 449 54109 15840 14989 1132 2957 30 21134 25899 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/long3k.ids b/zig/tests/cuda/flashnext/lr7/long3k.ids new file mode 100644 index 0000000000..04e5bd5759 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long3k.ids @@ -0,0 +1 @@ +248045 846 198 36948 369 449 47801 314 279 12654 3992 5604 13 271 760 328 2138 1 4952 198 865 6610 271 8272 11889 513 264 16219 1560 310 5479 26805 52992 198 17583 264 1957 25 271 256 1992 21152 61987 328 2138 1 7258 498 2129 7258 60 271 760 4145 1304 11 328 2138 7258 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 25 4673 63990 271 760 11244 1304 11 328 2138 7258 16 11 7258 17 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 16 25 4673 63990 46718 17 8 271 9205 42018 2303 9370 421 12880 8090 548 1 321 328 66317 1401 1 8042 198 951 279 5617 3309 7101 440 1782 4874 13 220 733 279 1428 198 13957 11 279 5617 3309 3759 12880 8090 548 1 369 328 2434 1 1172 198 7994 12779 11 328 3913 1 948 24460 369 10897 318 5454 198 1021 2904 6320 46 1764 220 561 1428 1970 13427 70307 874 1970 364 449 198 2138 4952 948 24460 369 10897 506 18612 854 13 220 6817 198 8761 424 369 24366 310 2830 279 2450 1970 364 279 7258 198 8761 4490 303 279 1412 1876 26 424 668 381 12234 430 919 314 279 198 7450 11319 13 271 27991 1320 310 12880 8090 548 1 513 11477 13 220 561 869 364 279 5617 3309 198 9521 369 10521 948 279 37999 8211 13 1358 40238 11889 198 865 17731 271 40238 11889 513 1429 310 318 265 8 7423 4874 310 2663 321 310 198 32925 7944 466 3470 314 24983 5978 25 271 256 15840 21152 61987 318 5480 1951 35036 7030 318 11531 1114 27160 735 7302 27160 8 198 256 2100 1951 257 61987 2100 318 2129 2100 4653 498 487 1293 198 256 2100 670 61987 12448 198 1025 735 32467 498 5480 1951 60 56439 198 1025 735 48975 498 5480 1951 60 84576 198 1025 735 16819 330 42514 198 1025 735 14701 198 1025 735 81987 198 1025 735 63669 2100 271 7 9538 3644 34019 5262 364 279 18890 16954 364 348 42648 330 42514 11961 198 9 33261 11961 321 348 82 86047 9 5946 271 2014 15840 4952 64013 279 7258 1103 318 28321 421 198 559 628 381 264 3074 7258 466 264 30644 70221 1103 11 279 14827 198 28726 286 264 13992 8 321 48233 279 3074 12571 1576 310 1754 314 198 1719 2100 11140 11 494 2047 310 1245 13 271 40238 369 4364 51070 11313 383 279 1304 314 279 2100 198 7080 553 3095 264 2100 369 919 314 264 24983 1576 318 276 6816 198 15812 11 14701 466 81987 681 279 24983 1576 1902 198 484 6887 2649 279 15840 321 9973 883 1141 30794 11 321 198 17812 4673 449 4531 413 279 15840 369 41565 13 220 561 5426 198 5308 2683 539 5090 4346 321 279 18987 8953 513 2574 440 279 198 17829 314 279 1576 4346 318 3924 3644 561 5129 913 198 71 32759 553 271 40238 314 449 1576 310 264 2100 1103 11 44007 42360 303 198 39708 37052 466 9057 37641 11 369 51070 4364 430 10695 13 271 9 1368 279 2100 1103 369 264 3074 2100 440 874 26849 30644 11 198 220 44007 303 71458 11 279 1576 369 12244 310 421 2100 13 271 9 18236 25 271 220 348 1368 279 2100 1103 5435 799 2100 78685 440 449 32815 3086 11 198 262 2512 264 1001 11531 1114 828 2100 25 561 1576 1902 381 449 49124 440 506 198 262 3140 430 1599 3470 430 1017 513 11218 303 279 2100 1103 11 26401 198 262 799 13 220 561 1118 3470 314 279 49124 513 12244 11 494 2047 310 198 262 1245 11 310 279 11218 1518 279 56276 2100 13 220 561 1534 3470 198 262 314 279 49124 513 12244 310 279 11218 1238 279 56276 198 262 2100 13 220 357 1103 314 279 9376 3470 303 279 49124 369 1179 198 262 12244 310 279 56276 2100 318 1719 1103 628 381 4147 553 271 220 348 18236 25 561 1576 1902 381 449 49124 440 279 1788 1324 314 3470 198 262 430 1017 513 11218 303 279 2100 1103 11 321 279 3470 513 198 262 12244 11 494 2047 310 1245 11 310 279 11811 11218 13 271 40238 314 449 1576 310 264 3074 2100 369 51070 4364 430 198 18000 82 13 271 9 1368 279 2100 369 449 12448 318 591 1590 271 220 348 1368 279 803 1503 524 11822 303 264 328 9461 1 466 328 6082 2359 1 4952 198 262 303 279 1428 1970 2424 25 279 803 369 6608 310 279 1576 303 279 198 262 1428 2136 4325 13 271 220 348 17673 25 279 803 369 6608 310 279 1576 303 279 3521 4325 198 262 466 279 15484 4325 10521 539 328 6082 2359 487 15119 13 271 220 561 803 369 40158 413 424 557 2582 6608 13 220 1061 1189 5073 279 198 220 5604 1698 364 279 1576 8334 6608 310 279 803 310 5372 198 220 6942 11 13969 279 1576 310 381 3372 38157 321 1141 44462 318 331 424 198 220 682 799 8 310 381 2512 13 271 9 1368 279 2100 369 449 6816 5604 25 561 5839 7258 303 198 220 279 5604 369 24276 13 220 1049 1220 7302 449 1576 440 198 220 86370 7944 26 413 411 369 524 279 1105 11 328 77318 1 369 198 220 8953 13 220 2844 1576 369 1179 4439 310 9501 279 12244 1576 310 198 220 279 2574 6816 26 413 424 4021 2649 279 15840 11 424 24096 198 220 449 4531 318 41547 694 524 13902 328 3776 1401 1764 271 220 6817 25 1368 279 1576 369 264 523 2774 321 279 6816 5604 198 220 13273 383 2107 10744 314 279 15840 5497 11 279 1245 23641 3008 198 220 7258 11 328 64 1926 1 628 2528 2892 449 2774 6816 466 318 331 874 198 220 2774 6816 6513 8 264 523 6816 13 220 561 2047 23641 3008 198 220 2100 328 64 1926 1 369 2590 716 430 449 2774 6816 11 6611 424 413 198 220 5689 13 220 13891 11 279 1330 54133 314 328 64 1926 1 635 524 13902 198 220 8042 310 279 1788 6816 25 413 279 1245 23641 3008 7258 198 220 18675 310 264 523 6816 11 279 2047 23641 3008 11121 264 491 198 220 2774 6816 430 279 2100 314 279 15840 25 271 257 523 2356 82 25 198 260 830 283 220 18 1734 653 523 3759 198 257 1699 283 2356 82 363 198 257 1699 1926 283 1699 1926 478 220 16 256 653 13514 1699 1926 430 220 19 9101 2356 82 1926 430 220 18 271 220 1061 3874 1503 524 13902 3668 310 15933 198 220 7944 11 1680 430 5706 3354 440 328 3573 363 3158 271 9 1368 279 2100 369 264 14701 25 561 5839 7258 303 279 198 220 5604 369 24276 13 220 1049 1220 7302 2892 264 24983 8240 198 220 1576 318 20171 430 264 1103 8 466 264 12366 1576 318 20171 430 264 10675 553 198 220 9019 11 279 94860 7258 369 24276 13 271 220 1368 279 5839 369 264 24983 8240 1576 318 20171 430 264 1103 681 279 198 220 94860 1902 7302 449 7308 13 220 1368 424 369 7968 11 279 8240 725 198 220 2984 369 3568 310 424 13 220 561 12571 869 1902 381 264 2397 40823 198 220 7308 2599 1056 279 8240 725 2984 11 321 279 8240 369 4439 198 220 310 9501 279 12244 1576 310 1141 1455 440 421 1857 13 220 1368 279 198 220 1857 369 680 314 2020 11 328 1498 1401 1 369 8953 318 28973 310 264 198 220 94860 290 8240 4021 884 491 3470 310 264 1103 553 271 220 1368 279 5839 369 264 12366 1576 318 20171 430 264 10675 681 279 198 220 94860 1902 599 264 913 17606 440 279 12366 725 1328 913 11 198 220 321 279 12366 369 1179 4439 310 1791 264 1328 55616 6505 864 13640 198 220 279 94860 310 279 12244 1576 13 220 1061 628 2892 8032 449 198 220 6149 1328 55616 6505 440 279 1788 1328 869 11 466 5479 264 491 198 220 1328 55616 6505 318 331 874 1328 440 279 1788 869 23533 553 271 220 1690 1156 37435 5978 11 279 12880 723 1163 548 33905 1654 369 2512 440 198 220 8053 5789 13 271 9 1368 279 2100 369 264 81987 25 561 5839 7258 303 279 5604 198 220 369 24276 13 220 1049 1220 7302 264 24983 8240 1576 318 20171 430 264 198 220 1103 553 220 561 12244 1576 1220 381 264 8240 1576 314 279 1788 198 220 913 13 220 9019 11 279 4570 321 8156 6608 22666 513 24276 11 198 220 303 90349 781 513 2945 26 16184 513 6942 321 279 8240 725 198 220 2984 13 220 561 13854 1220 14989 310 24959 13 1368 2892 6608 369 198 220 7968 11 279 8240 725 2984 369 3568 310 424 13 220 561 12571 198 220 13854 513 70575 310 9940 1881 6942 321 279 8240 725 2984 11 198 220 27389 13 220 16864 11 279 8240 1576 369 4439 310 8032 279 198 220 15511 440 279 3470 314 279 12244 8240 13 220 561 2984 314 279 198 220 15511 1189 381 2086 494 279 2984 314 279 12244 8240 11 198 220 8190 9722 279 2984 314 279 2100 8240 11 413 279 2100 198 220 8240 5954 424 13 271 332 6901 26033 7873 7472 64700 733 279 1428 7873 11 279 198 54304 364 11218 369 4284 310 381 279 1788 430 364 22666 11 321 198 11469 18890 369 17030 2261 279 1970 9191 9956 11 13969 198 1657 11346 1412 6397 13 271 15337 279 7044 314 15840 23197 421 73406 1881 198 1719 2047 23641 3008 321 279 1245 23641 3008 513 3260 14352 484 17140 515 318 1891 198 8422 328 64 11 292 283 292 11 264 1 70664 1330 7101 681 73406 348 53977 9 279 198 13034 314 12244 4534 7101 11822 2047 4534 6491 11 6807 198 1334 286 303 20684 13 220 1690 2774 11 279 2614 1957 22741 198 35643 15 11 220 17 60 763 271 256 830 283 498 15 11 220 16 60 198 256 585 283 220 15 198 256 585 11 830 957 60 283 220 16 11 220 17 260 653 585 369 5860 11 1179 830 957 60 369 5860 198 256 1133 2007 8 271 9538 1048 25 271 220 2972 1679 47 220 18 16 18 17 332 471 39226 36760 1188 91841 198 257 561 24332 364 279 15171 5480 1 4416 13 1358 21199 26115 15840 11889 198 1464 43414 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 2014 84027 15840 11889 198 1464 43414 271 9 19109 9 15840 369 279 10286 11 303 264 3074 4952 11 314 198 64 3759 466 6816 20572 321 449 9801 15840 198 23425 25 271 256 60736 50145 21152 61987 18720 5480 22698 7258 198 761 4241 420 318 11531 1114 27160 735 7302 27160 7025 271 760 6463 494 4472 33284 11889 369 421 1132 264 3074 198 5480 369 5252 13 271 2381 4145 4874 430 15840 11218 11 413 303 523 466 4536 6676 11 198 1719 31144 513 24276 321 9476 303 264 3175 523 466 4536 198 8843 12880 38314 548 1 421 369 264 10675 12366 494 3759 198 11072 318 76 37652 413 843 8 310 24276 31144 13 1061 6816 369 198 86 21833 321 369 9184 3354 506 279 1151 314 523 466 4536 198 2513 10993 11 413 31144 513 1669 88034 13 271 2381 22666 430 15840 11218 11 279 31144 513 24276 198 331 303 523 466 4536 6676 11 694 524 9476 13 271 2592 264 803 369 60736 303 264 709 6676 11 1179 411 803 369 2136 198 1891 421 6676 13 90117 513 2496 24276 321 9476 303 709 198 78134 13 271 2592 279 1245 1375 3008 369 2945 11 449 60736 15840 25822 198 1719 4880 15840 1518 36319 31144 318 2777 198 658 44575 553 1368 279 1245 1375 3008 369 524 2945 364 449 7258 198 5480 11 1179 279 37999 64013 279 2100 3525 364 279 1483 198 1 548 723 1163 548 33905 466 12880 84163 548 33905 1562 13 271 9538 1048 25 271 220 2972 1679 47 220 20 17 21 332 471 31058 364 12051 90117 198 257 561 13340 421 3568 18890 364 34848 1060 279 4346 314 198 257 7101 318 15693 523 7101 321 2774 7101 681 198 257 4370 314 35493 1070 1472 5853 13 271 220 2972 1679 47 220 19 23 19 332 471 3856 29642 198 257 561 13340 421 3568 279 328 86711 1 4536 310 3300 264 5129 198 257 18890 364 913 31144 421 628 381 1429 303 1064 6157 198 257 7141 321 27069 82 13 271 5221 303 2243 220 18 13 23 25 4543 60736 30152 2070 279 1788 198 13392 880 303 279 1245 1375 3008 430 5611 30152 13 56703 11 198 14265 22666 318 4650 632 2195 8878 80019 13992 22666 8 8613 264 198 54304 1412 13 1358 10266 27693 198 584 332 271 3446 303 2243 220 18 13 20 13 1358 31099 709 7044 198 1464 34195 271 256 3203 9310 727 61987 498 91430 60 328 7449 1 328 727 1 2821 591 32467 498 15704 1951 60 56439 198 3696 4241 397 1 7258 60 22698 15705 271 19678 314 12654 1766 27693 628 381 20950 321 55483 506 1599 198 7462 318 3924 348 5816 13754 187525 328 11093 1 22666 11 328 7449 364 1 321 328 7449 198 4058 1 628 1132 381 1429 303 279 2407 314 264 75140 709 13 271 24407 4364 440 328 7449 687 1 18890 513 2590 75140 198 20752 11 1442 413 781 635 524 6435 328 11093 1 466 328 7449 1 20209 13 271 2064 369 264 328 32769 1401 1 310 958 264 328 28726 494 1 7258 4613 279 2407 198 1020 264 75140 709 13 271 2014 3010 314 264 75140 709 25 271 256 3203 687 2821 9453 16 11 1626 17 1590 198 950 635 92099 363 198 950 2454 1010 14308 13754 363 271 5221 303 2243 220 271 14162 25 3956 310 279 47801 3294 11 1092 1503 449 1992 4952 635 948 12654 8213 440 24460 383 11 321 1092 1503 449 54109 15840 14989 1132 2957 30 21134 25899 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/mid.ids b/zig/tests/cuda/flashnext/lr7/mid.ids new file mode 100644 index 0000000000..e2a6571190 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/mid.ids @@ -0,0 +1 @@ +248045 846 198 814 20139 303 13983 4105 3069 279 12515 5686 6105 2261 279 1834 694 10263 2438 321 18016 506 41564 13 82210 279 3364 314 279 16078 321 279 43796 314 3000 11 321 2426 424 1172 4097 22157 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/prompts.json b/zig/tests/cuda/flashnext/lr7/prompts.json new file mode 100644 index 0000000000..c067510088 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/prompts.json @@ -0,0 +1 @@ +{"short": [248045, 846, 198, 3710, 369, 279, 6511, 314, 9338, 30, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], "mid": [248045, 846, 198, 814, 20139, 303, 13983, 4105, 3069, 279, 12515, 5686, 6105, 2261, 279, 1834, 694, 10263, 2438, 321, 18016, 506, 41564, 13, 82210, 279, 3364, 314, 279, 16078, 321, 279, 43796, 314, 3000, 11, 321, 2426, 424, 1172, 4097, 22157, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], "long": [248045, 846, 198, 8917, 5437, 537, 1092, 411, 12654, 4536, 1503, 11, 1179, 1103, 1141, 569, 5568, 25, 271, 12237, 1138, 39834, 321, 20621, 13, 198, 12237, 271, 2, 2861, 318, 34, 8, 220, 16, 24, 24, 24, 12, 17, 15, 15, 16, 41550, 387, 13, 25788, 13, 198, 2, 2861, 318, 34, 8, 220, 17, 15, 15, 17, 11, 220, 17, 15, 15, 18, 12654, 4236, 4846, 13, 198, 2, 32770, 539, 15573, 25788, 361, 70, 1550, 31, 12305, 4909, 29, 271, 464, 311, 271, 548, 528, 548, 283, 2429, 1138, 11210, 504, 359, 9797, 504, 359, 7320, 504, 359, 9497, 305, 504, 359, 31752, 504, 359, 8412, 268, 642, 271, 2, 10906, 1787, 279, 14541, 35169, 5587, 310, 279, 2201, 12, 54556, 198, 2, 35169, 5587, 13, 220, 561, 1822, 2781, 364, 3604, 411, 369, 421, 198, 2, 1010, 34707, 12258, 318, 4650, 1088, 84, 15, 15, 64, 15, 8, 513, 2397, 52483, 21201, 26453, 13, 198, 83201, 283, 5031, 83, 1639, 3351, 15, 65, 3351, 15, 66, 11674, 359, 271, 1005, 2824, 11210, 25, 198, 262, 4071, 198, 262, 2907, 364, 39834, 6454, 9121, 1414, 13, 220, 561, 569, 3622, 16661, 314, 198, 262, 279, 14623, 363, 321, 4990, 363, 5245, 26, 279, 975, 5245, 513, 1066, 1017, 364, 198, 262, 65878, 310, 2723, 303, 1906, 310, 49324, 279, 1581, 16501, 13, 198, 262, 1368, 488, 1318, 310, 6379, 8032, 279, 1822, 39834, 11767, 11, 198, 262, 488, 3172, 4507, 599, 310, 2723, 696, 9797, 63261, 973, 271, 262, 24876, 2774, 7944, 2444, 5090, 13173, 314, 39834, 25, 198, 406, 2297, 318, 2186, 25, 220, 22, 15, 8, 198, 285, 279, 6966, 2297, 314, 18880, 4965, 318, 37488, 1387, 17268, 18408, 198, 285, 369, 867, 8, 198, 406, 2763, 66708, 318, 2186, 25, 11363, 198, 285, 886, 421, 668, 381, 20671, 2789, 310, 279, 1118, 1500, 314, 18880, 198, 285, 2468, 13, 220, 82699, 6757, 279, 1500, 579, 2297, 13, 198, 406, 16905, 66708, 318, 2186, 25, 11363, 198, 285, 886, 421, 668, 381, 20671, 2789, 310, 660, 4965, 3462, 279, 1118, 198, 285, 314, 18880, 2468, 26, 1048, 14159, 6757, 1754, 1500, 579, 2297, 13, 198, 406, 8952, 56008, 318, 2186, 25, 804, 8, 198, 285, 48455, 21700, 303, 1879, 1414, 310, 12258, 1518, 4473, 8427, 13, 198, 285, 8618, 5474, 668, 3512, 220, 15, 5074, 359, 5999, 2073, 6, 12258, 11, 11313, 383, 1141, 2234, 198, 285, 303, 1141, 1500, 13, 220, 1368, 867, 11, 1754, 5474, 369, 11421, 430, 264, 3074, 3542, 13, 198, 406, 5474, 2073, 318, 2186, 25, 220, 23, 8, 198, 285, 48455, 21700, 303, 1879, 1414, 310, 220, 15, 5074, 359, 5999, 2073, 6, 12258, 11, 7014, 198, 285, 359, 31251, 56008, 6, 369, 867, 13, 198, 406, 8032, 83201, 318, 2186, 25, 804, 8, 198, 285, 28592, 660, 35169, 5587, 303, 279, 1879, 1414, 539, 12258, 198, 285, 1238, 5474, 14046, 13, 220, 6817, 421, 413, 8952, 56008, 369, 867, 321, 198, 285, 8032, 83201, 369, 804, 11, 1396, 5474, 668, 381, 15625, 310, 264, 198, 285, 3074, 3433, 0, 198, 406, 4884, 46720, 6018, 793, 318, 2186, 25, 867, 8, 198, 285, 28320, 421, 11316, 64541, 59429, 369, 2590, 7854, 198, 285, 539, 1330, 12258, 13, 220, 3978, 539, 1581, 1521, 279, 11767, 369, 198, 285, 318, 354, 46261, 2752, 8, 65669, 13, 198, 406, 1387, 17268, 18408, 318, 2186, 25, 804, 8, 198, 285, 15164, 4105, 4860, 1056, 359, 2998, 4282, 220, 1368, 867, 11, 1782, 4105, 668, 524, 198, 285, 381, 10548, 11, 321, 1010, 4965, 2493, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271]} diff --git a/zig/tests/cuda/flashnext/lr7/results-sparse.json b/zig/tests/cuda/flashnext/lr7/results-sparse.json new file mode 100644 index 0000000000..1a4198200a --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/results-sparse.json @@ -0,0 +1,364 @@ +{ + "info": { + "startup_s": 95.44, + "max_len": 16384, + "prefill_rows": 4096, + "nch": 5, + "qsa": true, + "idw": 2052, + "eos": [ + 248046, + 248044 + ], + "torch": "2.13.0a0+9186a08b2c.nv26.07", + "gpu": "NVIDIA GB10", + "prompt_lengths": { + "long3k": 3001, + "long12k": 12001 + } + }, + "results": { + "serial/long3k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1510, + 2138, + 63, + 4952, + 70307, + 874, + 1970, + 318, + 275, + 369, + 13057, + 11693, + 553, + 733, + 449, + 54109, + 15840 + ], + "wall_s": 16.7331, + "prefill_s": 2.4387, + "cached": 0, + "drafts": false, + "decode_s": 14.2676, + "rounds": 31, + "decode_tps": 2.17 + }, + "teacher32/long3k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1970, + 13427, + 70307, + 2972, + 2083, + 1970, + 332, + 364, + 449, + 1510, + 2138, + 63, + 4952, + 13, + 271, + 623, + 449 + ], + "wall_s": 86.062, + "decode_ms_per_token": 30.764, + "top5": [ + [ + [ + 4350, + 21.25 + ], + [ + 2278, + 18.25 + ], + [ + 16, + 18.0 + ], + [ + 27775, + 17.625 + ], + [ + 2014, + 17.375 + ] + ], + [ + [ + 12654, + 22.375 + ], + [ + 24460, + 20.125 + ], + [ + 4162, + 19.125 + ], + [ + 1542, + 16.125 + ], + [ + 32237, + 15.75 + ] + ], + [ + [ + 8213, + 24.75 + ], + [ + 369, + 19.5 + ], + [ + 1542, + 13.875 + ], + [ + 24460, + 13.25 + ], + [ + 4162, + 13.1875 + ] + ], + [ + [ + 440, + 23.875 + ], + [ + 303, + 13.9375 + ], + [ + 383, + 12.75 + ], + [ + 4058, + 12.5625 + ], + [ + 24460, + 12.375 + ] + ] + ] + }, + "serial/long12k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1510, + 2138, + 63, + 4952, + 70307, + 2972, + 2083, + 1970, + 332, + 318, + 275, + 369, + 13057, + 33993, + 14, + 56497, + 553 + ], + "wall_s": 6.012, + "prefill_s": 5.117, + "cached": 0, + "drafts": false, + "decode_s": 0.8949, + "rounds": 31, + "decode_tps": 34.64 + }, + "teacher32/long12k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1970, + 13427, + 70307, + 2972, + 2083, + 1970, + 332, + 364, + 449, + 1510, + 2138, + 63, + 4952, + 13, + 271, + 623, + 449 + ], + "wall_s": 338.3455, + "decode_ms_per_token": 30.42, + "top5": [ + [ + [ + 4350, + 21.375 + ], + [ + 27775, + 19.375 + ], + [ + 10865, + 18.125 + ], + [ + 2278, + 18.125 + ], + [ + 16, + 18.0 + ] + ], + [ + [ + 12654, + 22.5 + ], + [ + 24460, + 19.75 + ], + [ + 4162, + 19.25 + ], + [ + 1542, + 15.9375 + ], + [ + 32237, + 15.5 + ] + ], + [ + [ + 8213, + 24.375 + ], + [ + 369, + 19.25 + ], + [ + 1542, + 13.5625 + ], + [ + 4162, + 12.9375 + ], + [ + 725, + 12.9375 + ] + ], + [ + [ + 440, + 23.125 + ], + [ + 303, + 13.3125 + ], + [ + 383, + 12.5625 + ], + [ + 24460, + 12.5 + ], + [ + 4058, + 12.4375 + ] + ] + ] + } + } +} diff --git a/zig/tests/cuda/flashnext/lr7/results.json b/zig/tests/cuda/flashnext/lr7/results.json new file mode 100644 index 0000000000..50ff059a2d --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/results.json @@ -0,0 +1,920 @@ +{ + "info": { + "startup_s": 93.63, + "max_len": 1536, + "prefill_rows": 4096, + "nch": 3, + "qsa": false, + "idw": 2052, + "eos": [ + 248046, + 248044 + ], + "torch": "2.13.0a0+9186a08b2c.nv26.07", + "gpu": "NVIDIA GB10", + "prompt_lengths": { + "short": 19, + "mid": 52, + "long": 600 + } + }, + "results": { + "serial/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 14.6758, + "prefill_s": 0.3383, + "cached": 0, + "drafts": false, + "decode_s": 14.3201, + "rounds": 31, + "decode_tps": 2.16 + }, + "teacher/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 1.494, + "decode_ms_per_token": 28.396, + "top5": [ + [ + [ + 760, + 24.0 + ], + [ + 3710, + 15.875 + ], + [ + 3733, + 14.3125 + ], + [ + 561, + 14.3125 + ], + [ + 1206, + 13.625 + ] + ], + [ + [ + 6511, + 25.0 + ], + [ + 3177, + 14.1875 + ], + [ + 1994, + 13.875 + ], + [ + 5839, + 13.25 + ], + [ + 1822, + 13.125 + ] + ], + [ + [ + 314, + 25.5 + ], + [ + 321, + 19.125 + ], + [ + 3177, + 16.0 + ], + [ + 11, + 14.3125 + ], + [ + 318, + 14.1875 + ] + ], + [ + [ + 9338, + 28.5 + ], + [ + 279, + 17.75 + ], + [ + 2972, + 17.25 + ], + [ + 47358, + 15.1875 + ], + [ + 11751, + 14.75 + ] + ] + ] + }, + "teacher32/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 1.4951, + "decode_ms_per_token": 27.503, + "top5": [ + [ + [ + 760, + 23.875 + ], + [ + 3710, + 15.625 + ], + [ + 561, + 14.375 + ], + [ + 3733, + 14.1875 + ], + [ + 1206, + 13.5625 + ] + ], + [ + [ + 6511, + 24.75 + ], + [ + 3177, + 14.0 + ], + [ + 1994, + 13.625 + ], + [ + 1822, + 13.3125 + ], + [ + 5839, + 13.25 + ] + ], + [ + [ + 314, + 25.125 + ], + [ + 321, + 19.375 + ], + [ + 3177, + 16.0 + ], + [ + 318, + 14.4375 + ], + [ + 11, + 14.375 + ] + ], + [ + [ + 9338, + 28.125 + ], + [ + 2972, + 17.625 + ], + [ + 279, + 17.625 + ], + [ + 47358, + 14.6875 + ], + [ + 11751, + 14.625 + ] + ] + ] + }, + "serial/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 1.3186, + "prefill_s": 0.3477, + "cached": 0, + "drafts": false, + "decode_s": 0.9708, + "rounds": 31, + "decode_tps": 31.93 + }, + "teacher/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 2.2782, + "decode_ms_per_token": 28.083, + "top5": [ + [ + [ + 29108, + 23.25 + ], + [ + 13672, + 19.0 + ], + [ + 15542, + 18.25 + ], + [ + 760, + 16.25 + ], + [ + 10025, + 14.75 + ] + ], + [ + [ + 4009, + 25.625 + ], + [ + 5603, + 15.625 + ], + [ + 3000, + 14.375 + ], + [ + 37728, + 13.9375 + ], + [ + 4022, + 13.5625 + ] + ], + [ + [ + 5435, + 22.75 + ], + [ + 369, + 20.0 + ], + [ + 7701, + 19.875 + ], + [ + 16661, + 19.0 + ], + [ + 27891, + 18.75 + ] + ], + [ + [ + 660, + 26.75 + ], + [ + 264, + 21.0 + ], + [ + 1599, + 19.0 + ], + [ + 1396, + 17.75 + ], + [ + 7736, + 16.75 + ] + ] + ] + }, + "teacher32/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 2.2279, + "decode_ms_per_token": 27.005, + "top5": [ + [ + [ + 29108, + 23.125 + ], + [ + 13672, + 18.875 + ], + [ + 15542, + 18.125 + ], + [ + 760, + 16.25 + ], + [ + 37728, + 14.75 + ] + ], + [ + [ + 4009, + 25.5 + ], + [ + 5603, + 15.625 + ], + [ + 3000, + 14.3125 + ], + [ + 37728, + 13.9375 + ], + [ + 4022, + 13.5625 + ] + ], + [ + [ + 5435, + 22.75 + ], + [ + 369, + 19.875 + ], + [ + 7701, + 19.75 + ], + [ + 16661, + 19.125 + ], + [ + 27891, + 18.75 + ] + ], + [ + [ + 660, + 26.75 + ], + [ + 264, + 21.0 + ], + [ + 1599, + 19.0 + ], + [ + 1396, + 17.625 + ], + [ + 7736, + 16.875 + ] + ] + ] + }, + "serial/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 1049, + 369, + 5995, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981, + 421 + ], + "wall_s": 1.8294, + "prefill_s": 0.852, + "cached": 0, + "drafts": false, + "decode_s": 0.9773, + "rounds": 31, + "decode_tps": 31.72 + }, + "teacher/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 11116, + 5839, + 7193, + 369, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981 + ], + "wall_s": 17.2408, + "decode_ms_per_token": 29.059, + "top5": [ + [ + [ + 13962, + 22.375 + ], + [ + 1919, + 20.0 + ], + [ + 760, + 19.125 + ], + [ + 27775, + 18.75 + ], + [ + 550, + 18.375 + ] + ], + [ + [ + 20858, + 22.5 + ], + [ + 13317, + 19.875 + ], + [ + 3437, + 16.125 + ], + [ + 33655, + 14.5 + ], + [ + 2972, + 14.4375 + ] + ], + [ + [ + 314, + 22.625 + ], + [ + 198, + 21.75 + ], + [ + 271, + 21.0 + ], + [ + 25, + 15.25 + ], + [ + 13317, + 13.0 + ] + ], + [ + [ + 279, + 22.75 + ], + [ + 13317, + 20.125 + ], + [ + 5534, + 18.875 + ], + [ + 12654, + 17.5 + ], + [ + 1510, + 17.5 + ] + ] + ] + }, + "teacher32/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 11116, + 5839, + 7193, + 369, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981 + ], + "wall_s": 17.1927, + "decode_ms_per_token": 28.848, + "top5": [ + [ + [ + 13962, + 22.375 + ], + [ + 1919, + 20.0 + ], + [ + 760, + 19.125 + ], + [ + 27775, + 18.625 + ], + [ + 550, + 18.375 + ] + ], + [ + [ + 20858, + 22.5 + ], + [ + 13317, + 20.0 + ], + [ + 3437, + 16.125 + ], + [ + 33655, + 14.5625 + ], + [ + 2972, + 14.4375 + ] + ], + [ + [ + 314, + 22.625 + ], + [ + 198, + 21.875 + ], + [ + 271, + 21.0 + ], + [ + 25, + 15.25 + ], + [ + 13317, + 13.0 + ] + ], + [ + [ + 279, + 22.75 + ], + [ + 13317, + 20.125 + ], + [ + 5534, + 18.625 + ], + [ + 1510, + 17.625 + ], + [ + 12654, + 17.5 + ] + ] + ] + } + } +} diff --git a/zig/tests/cuda/flashnext/lr7/short.ids b/zig/tests/cuda/flashnext/lr7/short.ids new file mode 100644 index 0000000000..23c34f9c55 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/short.ids @@ -0,0 +1 @@ +248045 846 198 3710 369 279 6511 314 9338 30 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/pack.sh b/zig/tests/cuda/flashnext/pack.sh new file mode 100755 index 0000000000..8e3c27d874 --- /dev/null +++ b/zig/tests/cuda/flashnext/pack.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Usage: pack.sh OUT_KERNELS CAPTURE_OUT...: one kernel set from the captures' manifest.json and jit.json. +set -euo pipefail +OUT=$1; shift +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +ARGS=() +for d in "$@"; do ARGS+=(--manifest "$d/manifest.json" --jit "$d/jit.json"); done +python3 -B "$REPO/tools/zig/triton_aot_pack.py" "${ARGS[@]}" --out "$OUT" diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig new file mode 100644 index 0000000000..806a7207bf --- /dev/null +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -0,0 +1,282 @@ +//! Flash Next's own attention glue and its indexer's pooling and scores against float64 references. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Rig = rig.Rig; +const bf = rig.bf; +const rbf = rig.rbf; +const ulps = rig.ulps; +const sumTol = rig.sumTol; + +const q_heads = 24; +const kv_heads = 2; +const head_dim = 256; +const index_heads = 4; +const index_dim = 128; +const half = 32; +const proj_n = 13952; +const eps: f32 = 1e-6; +const cap = 3072; +const idw = 2052; + +/// RoPE's inverse frequencies as the pass uploads them (base 10,000,000 over 32 pairs). +fn inverse(r: *Rig) ![]f32 { + const inv = try r.a.alloc(f32, half); + for (inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + return inv; +} + +/// The kernels' rotate-half RoPE of normed values xn, xpn (feature d of 2 * half rotated) at `pos`. +fn rope(xn: f64, xpn: f64, d: usize, pos: usize, inv: []const f32) f64 { + if (d >= 2 * half) return xn; + const i = if (d < half) d else d - half; + const ang: f64 = @as(f32, @floatFromInt(pos)) * inv[i]; + return if (d < half) xn * @cos(ang) - xpn * @sin(ang) else xpn * @sin(ang) + xn * @cos(ang); +} + +fn partner(d: usize) usize { + return if (d < half) d + half else if (d < 2 * half) d - half else d; +} + +/// A head of `width` features at src normed by w, rounded, roped at pos and rounded: the float64 chain. +fn prepped(p: []const u16, src: usize, width: usize, w: []const f32, pos: usize, inv: []const f32, d: usize) f64 { + var ss: f64 = 0; + for (0..width) |i| ss += bf(p[src + i]) * bf(p[src + i]); + const rinv = 1 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + const xn = rbf(bf(p[src + d]) * rinv * w[d]); + const xpn = rbf(bf(p[src + partner(d)]) * rinv * w[partner(d)]); + return rbf(rope(xn, xpn, d, pos, inv)); +} + +fn prep(r: *Rig, g: *const flash.Glue) !void { + const counts = [_]usize{ 1, 3, 17 }; + const most = counts[counts.len - 1]; + const pos0 = 700; + const p = r.bfs(most * proj_n, 2); + const qw = r.f32s(head_dim, 0.5, 1.5); + const kw = r.f32s(head_dim, 0.5, 1.5); + const iw = r.f32s(index_dim, 0.5, 1.5); + const inv = try inverse(r); + const dp = try r.dev(std.mem.sliceAsBytes(p)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const dqw = try r.dev(std.mem.sliceAsBytes(qw)); + const dkw = try r.dev(std.mem.sliceAsBytes(kw)); + const diw = try r.dev(std.mem.sliceAsBytes(iw)); + const dinv = try r.dev(std.mem.sliceAsBytes(inv)); + var qs: [counts.len]u64 = undefined; + var iqs: [counts.len]u64 = undefined; + var kcs: [counts.len]u64 = undefined; + var vcs: [counts.len]u64 = undefined; + var ikcs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + qs[c] = try r.out(m * q_heads * head_dim * 2); + iqs[c] = try r.out(m * index_heads * index_dim * 2); + kcs[c] = try r.out(cap * kv_heads * head_dim * 2); + vcs[c] = try r.out(cap * kv_heads * head_dim * 2); + ikcs[c] = try r.out(cap * index_dim * 2); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m); + } + const last = counts.len - 1; + const q = try r.back(u16, qs[last], most * q_heads * head_dim); + const iq = try r.back(u16, iqs[last], most * index_heads * index_dim); + const kc = try r.back(u16, kcs[last], cap * kv_heads * head_dim); + const vc = try r.back(u16, vcs[last], cap * kv_heads * head_dim); + const ikc = try r.back(u16, ikcs[last], cap * index_dim); + for (0..most) |row| { + const pos = pos0 + row; + const base = row * proj_n; + for (0..q_heads) |h| for (0..head_dim) |d| { + const want = prepped(p, base + h * 2 * head_dim, head_dim, qw, pos, inv, d); + try r.within("prep queries", bf(q[(row * q_heads + h) * head_dim + d]), want, ulps(want, 1) + 1e-3); + }; + for (0..kv_heads) |h| for (0..head_dim) |d| { + const want = prepped(p, base + q_heads * 2 * head_dim + h * head_dim, head_dim, kw, pos, inv, d); + const at = (pos * kv_heads + h) * head_dim + d; + try r.within("prep keys", bf(kc[at]), want, ulps(want, 1) + 1e-3); + try check.expect(vc[at] == p[base + q_heads * 2 * head_dim + kv_heads * head_dim + h * head_dim + d], "prep values: row {d} head {d} copied", .{ row, h }); + }; + for (0..index_heads) |h| for (0..index_dim) |d| { + const want = prepped(p, base + q_heads * 2 * head_dim + 2 * kv_heads * head_dim + h * index_dim, index_dim, iw, pos, inv, d); + try r.within("prep indexer queries", bf(iq[(row * index_heads + h) * index_dim + d]), want, ulps(want, 1) + 1e-3); + }; + for (0..index_dim) |d| try check.expect(ikc[pos * index_dim + d] == p[base + q_heads * 2 * head_dim + 2 * kv_heads * head_dim + index_heads * index_dim + d], "prep indexer keys: row {d} copied", .{row}); + } + for (counts[0..last], 0..) |m, c| { + try r.sameRows("prep queries", u16, qs[c], q[0 .. m * q_heads * head_dim]); + try r.sameRows("prep keys", u16, kcs[c] + pos0 * kv_heads * head_dim * 2, kc[pos0 * kv_heads * head_dim ..][0 .. m * kv_heads * head_dim]); + } + // the output gate: bf16(o / (1 + e^-gate)) and its group sums + const o = r.bfs(most * q_heads * head_dim, 2); + const out = try r.out(most * q_heads * head_dim * 2); + const xs = try r.out(most * q_heads * head_dim / 32 * 4); + try g.attnGate(r.s, try r.dev(std.mem.sliceAsBytes(o)), dp, out, xs, most, proj_n, q_heads, head_dim); + const go = try r.back(u16, out, most * q_heads * head_dim); + const gx = try r.back(f32, xs, most * q_heads * head_dim / 32); + for (0..most) |row| for (0..q_heads) |h| { + for (0..head_dim / 32) |gi| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..32) |k| { + const d = gi * 32 + k; + const at = (row * q_heads + h) * head_dim + d; + const want = bf(o[at]) * rig.sigmoid(bf(p[row * proj_n + h * 2 * head_dim + head_dim + d])); + try r.within("output gate", bf(go[at]), want, ulps(want, 0.5) + @abs(want) * 0x1p-18); + sum += bf(go[at]); + abs += @abs(bf(go[at])); + } + try r.within("output gate sums", gx[(row * q_heads + h) * (head_dim / 32) + gi], sum, sumTol(abs, 32)); + } + }; +} + +/// The indexer's pooled blocks and each row's block scores. +fn indexer(r: *Rig, g: *const flash.Glue) !void { + const pos0 = 2196; + const rows = 5; + const ikc = r.bfs(cap * index_dim, 2); + const w = r.f32s(index_dim, 0.5, 1.5); + const inv = try inverse(r); + const dikc = try r.dev(std.mem.sliceAsBytes(ikc)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const pooled = try r.out((cap / 4) * index_dim * 2); + try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows); + const got = try r.back(u16, pooled, (cap / 4) * index_dim); + var b: usize = pos0 / 4; + while (4 * b + 4 <= pos0 + rows) : (b += 1) { + var x: [index_dim]f64 = undefined; + var ss: f64 = 0; + for (&x, 0..) |*v, d| { + var acc: f64 = 0; + for (0..4) |k| acc += bf(ikc[(4 * b + k) * index_dim + d]); + v.* = rbf(acc / 4); + ss += v.* * v.*; + } + const rinv = 1 / @sqrt(ss / index_dim + eps); + for (0..index_dim) |d| { + const xn = rbf(x[d] * rinv * w[d]); + const xpn = rbf(x[partner(d)] * rinv * w[partner(d)]); + const want = rbf(rope(xn, xpn, d, 4 * b, inv)); + try r.within("indexer pool", bf(got[b * index_dim + d]), want, ulps(want, 1) + 1e-3); + } + } + // scores of every complete block for rows past 512 of them + const keys = r.bfs((cap / 4) * index_dim, 1); + const iq = r.bfs(rows * index_heads * index_dim, 1); + const nb = 784; + const dkeys = try r.dev(std.mem.sliceAsBytes(keys)); + const diq = try r.dev(std.mem.sliceAsBytes(iq)); + var scs: [2]u64 = undefined; + for ([_]usize{ 1, rows }, 0..) |m, c| { + scs[c] = try r.out(m * nb * 4); + try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4); + } + const sc = try r.back(f32, scs[1], rows * nb); + for (0..rows) |row| { + const complete = (pos0 + row + 1) / 4; + for (0..complete) |blk| { + var total: f64 = 0; + var abs: f64 = 0; + for (0..index_heads) |h| { + var dot: f64 = 0; + for (0..index_dim) |d| { + const t = bf(iq[(row * index_heads + h) * index_dim + d]) * bf(keys[blk * index_dim + d]); + dot += t; + abs += @abs(t); + } + total += @max(dot, 0); + } + const want = total / @sqrt(@as(f64, index_dim)); + try r.within("indexer scores", sc[row * nb + blk], want, sumTol(abs, 2 * index_dim) / @sqrt(@as(f64, index_dim))); + } + } + try r.sameRows("indexer scores", f32, scs[0], sc[0 .. (pos0 + 1) / 4]); +} + +/// Each row's attention over its keys (dense from position 0, or a sparse row's list) in 512-key chunks, merged. +fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { + const counts = [_]usize{ 1, 3, 20 }; // up to 16 rows read the next tile ahead; 20 take one tile at a time + const rows = counts[counts.len - 1]; + const pos0: usize = if (sparse) 2600 else 1100; + const q = r.bfs(rows * q_heads * head_dim, 1); + const kc = r.bfs(cap * kv_heads * head_dim, 1); + const vc = r.bfs(cap * kv_heads * head_dim, 1); + const ids = try r.a.alloc(i32, rows * idw); + const nk = try r.a.alloc(i32, rows); + const spr = try r.a.alloc(i32, rows); + for (0..rows) |row| { + nk[row] = @intCast(if (sparse) idw - 1 - row % 3 else pos0 + row + 1); + spr[row] = @intFromBool(sparse); + // a sparse row's list: distinct keys before its position, in a scrambled order + for (0..idw) |i| ids[row * idw + i] = @intCast((i * 1237 + row * 31) % (pos0 + row + 1)); + } + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dkc = try r.dev(std.mem.sliceAsBytes(kc)); + const dvc = try r.dev(std.mem.sliceAsBytes(vc)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ @intCast(pos0), 0, 0, 0 })); + const dids = try r.dev(std.mem.sliceAsBytes(ids)); + const dnk = try r.dev(std.mem.sliceAsBytes(nk)); + const dspr = try r.dev(std.mem.sliceAsBytes(spr)); + const nch = 5; + const po = try r.out(rows * nch * q_heads * head_dim * 4); + const pm = try r.out(rows * nch * q_heads * 4); + const pl = try r.out(rows * nch * q_heads * 4); + var outs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * q_heads * head_dim * 2); + const chunks = if (sparse) (idw + 511) / 512 else (pos0 + m + 511) / 512; + try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true); + try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, kv_heads, true); + } + const got = try r.back(u16, outs[counts.len - 1], rows * q_heads * head_dim); + const name = if (sparse) "attention sparse" else "attention dense"; + const s = try r.a.alloc(f64, idw + pos0 + rows); + for (0..rows) |row| for (0..q_heads) |h| { + const kvh = h / (q_heads / kv_heads); + const n: usize = @intCast(nk[row]); + var top = -std.math.inf(f64); + for (0..n) |j| { + const key: usize = if (sparse) @intCast(ids[row * idw + j]) else j; + var dot: f64 = 0; + for (0..head_dim) |d| dot += bf(q[(row * q_heads + h) * head_dim + d]) * bf(kc[(key * kv_heads + kvh) * head_dim + d]); + s[j] = dot * 0.0625; + top = @max(top, s[j]); + } + var l: f64 = 0; + for (s[0..n]) |*v| { + v.* = @exp(v.* - top); + l += v.*; + } + for (0..head_dim) |d| { + var o: f64 = 0; + var abs: f64 = 0; + for (0..n) |j| { + const key: usize = if (sparse) @intCast(ids[row * idw + j]) else j; + const v = bf(vc[(key * kv_heads + kvh) * head_dim + d]); + o += s[j] * v; + abs += s[j] * @abs(v); + } + const want = o / l; + // p goes into p . v as bf16 (2^-9 of each term), its sum l stays fp32 + try r.within(name, bf(got[(row * q_heads + h) * head_dim + d]), want, abs / l * 0x1p-8 + ulps(want, 0.5)); + } + }; + for (counts[0 .. counts.len - 1], 0..) |m, c| try r.sameRows(name, u16, outs[c], got[0 .. m * q_heads * head_dim]); +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(59); + var r = try Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + try prep(&r, &g); + try indexer(&r, &g); + try attention(&r, &g, false); + try attention(&r, &g, true); + r.report("flashnext-attn: q, k and indexer heads, the output gate, pooled blocks and scores, dense and sparse attention"); +} diff --git a/zig/tests/cuda/flashnext_dense_hc_tests.zig b/zig/tests/cuda/flashnext_dense_hc_tests.zig new file mode 100644 index 0000000000..7afa56cc9d --- /dev/null +++ b/zig/tests/cuda/flashnext_dense_hc_tests.zig @@ -0,0 +1,146 @@ +//! Dense hyper-connection normalization, projections and mixing against an independent FP64 chain. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const bf = rig.bf; +const rbf = rig.rbf; +const d = 2560; +const streams = 4; +const wide = d * streams; +const low = 320; +const ndn = low + streams; +const most = 17; + +fn tensor(values: []const u16, n: usize, k: usize) core.safetensors.Tensor { + return .{ .dtype = .bf16, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(values) }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(0xde05e); + var r = try rig.Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + const h = r.bfs(most * wide, 1); + const gamma = r.f32s(wide, 0.5, 1.5); + const dw = r.bfs(ndn * wide, 0.01); + const uw = r.bfs(wide * low, 0.03); + var down = try flash.projection.dense(r.a, gpu.d, tensor(dw, ndn, wide)); + defer down.free(); + var up = try flash.projection.dense(r.a, gpu.d, tensor(uw, wide, low)); + defer up.free(); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dg = try r.dev(std.mem.sliceAsBytes(gamma)); + const pss = try r.out(most * (d / 256) * streams * 4); + try g.hcWriteback(r.s, dh, dh, pss, 0, 0, 0, 0, most * d, most, d, streams, 0, 1, 1, 1); + const sums = try r.back(f32, pss, most * (d / 256) * streams); + const norm = try r.out(most * wide * 2); + const projected = try r.out(most * ndn * 4); + const act = try r.out(most * low * 2); + const xs = try r.out(most * (low / 32) * 4); + const inj = try r.out(most * streams * 2); + const up_out = try r.out(most * wide * 2); + const mixed = try r.out(most * d * 2); + const xsm = try r.out(most * (d / 32) * 4); + const part = try r.out(8 * most * ndn * 4); + var norms: []u16 = undefined; + var dots: []f32 = undefined; + var acts: []u16 = undefined; + var injs: []u16 = undefined; + var ups: []u16 = undefined; + var mixes: []u16 = undefined; + var mix_sums: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try g.hcNormalize(r.s, dh, pss, dg, norm, 1e-6, rows, d, streams); + try lane.matmulF32(r.s, norm, wide, rows, down.native.?, projected, part); + try g.hcReduceAct(r.s, projected, act, xs, inj, 1, rows, streams, low, ndn, 1); + try lane.matmul(r.s, act, low, rows, up.native.?, up_out, part); + try g.hcDenseMix(r.s, up_out, norm, mixed, xsm, rows, d, streams); + if (rows == most) { + norms = try r.back(u16, norm, most * wide); + dots = try r.back(f32, projected, most * ndn); + acts = try r.back(u16, act, most * low); + injs = try r.back(u16, inj, most * streams); + ups = try r.back(u16, up_out, most * wide); + mixes = try r.back(u16, mixed, most * d); + mix_sums = try r.back(f32, xsm, most * (d / 32)); + } else { + try r.sameRows("dense HC normalize", u16, norm, norms[0 .. rows * wide]); + try r.sameRows("dense HC down FP32", f32, projected, dots[0 .. rows * ndn]); + try r.sameRows("dense HC activation", u16, act, acts[0 .. rows * low]); + try r.sameRows("dense HC inject", u16, inj, injs[0 .. rows * streams]); + try r.sameRows("dense HC up BF16", u16, up_out, ups[0 .. rows * wide]); + try r.sameRows("dense HC mix", u16, mixed, mixes[0 .. rows * d]); + try r.sameRows("dense HC group sums", f32, xsm, mix_sums[0 .. rows * (d / 32)]); + } + } + for (0..most) |row| { + for (0..streams) |st| { + var total: f64 = 0; + for (0..d / 256) |c| total += sums[(row * (d / 256) + c) * streams + st]; + const inv = 1 / @sqrt(total / d + 1e-6); + for (0..d) |col| { + const at = row * wide + st * d + col; + const want = bf(h[at]) * inv * gamma[st * d + col]; + try r.within("dense norm FP32 gamma", bf(norms[at]), want, rig.ulps(want, 0.5) + @abs(want) * 0x1p-20); + } + } + for (0..ndn) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..wide) |i| { + const v = bf(norms[row * wide + i]) * bf(dw[col * wide + i]); + sum += v; + magnitude += @abs(v); + } + try r.within("dense down FP32", dots[row * ndn + col], sum, magnitude * 2e-6 + @abs(sum) * 1e-6 + 1e-30); + const v = rbf(rbf(dots[row * ndn + col]) / streams); + if (col < low) { + const want = rbf(v * rig.sigmoid(v)); + try r.within("dense activation", bf(acts[row * low + col]), want, rig.ulps(want, 1)); + } else { + const want = rbf(2 * rbf(rig.sigmoid(v))); + try r.within("dense inject", bf(injs[row * streams + col - low]), want, rig.ulps(want, 1)); + } + } + for (0..wide) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..low) |i| { + const v = bf(acts[row * low + i]) * bf(uw[col * low + i]); + sum += v; + magnitude += @abs(v); + } + try r.within("dense up BF16", bf(ups[row * wide + col]), sum, rig.ulps(sum, 0.5) + magnitude * 2e-6); + } + for (0..d / 32) |group| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..32) |i| { + const col = group * 32 + i; + var total: f64 = 0; + var terms: f64 = 0; + for (0..streams) |st| { + const at = row * wide + st * d + col; + const term = rbf(rbf(rig.sigmoid(bf(ups[at]))) * bf(norms[at])); + total += term; + terms += @abs(term); + } + const want = rbf(total / streams); + const got = bf(mixes[row * d + col]); + try r.within("dense stream mix", got, want, rig.ulps(terms / streams, 1) + rig.ulps(want, 0.5)); + sum += got; + magnitude += @abs(got); + } + try r.within("dense group sums", mix_sums[row * (d / 32) + group], sum, rig.sumTol(magnitude, 32)); + } + } + r.report("dense HC: FP32 gamma, BF16 weights, FP32 down and one-round up; rows 1/7/17"); +} diff --git a/zig/tests/cuda/flashnext_draft_affine_tests.zig b/zig/tests/cuda/flashnext_draft_affine_tests.zig new file mode 100644 index 0000000000..a160eb8e41 --- /dev/null +++ b/zig/tests/cuda/flashnext_draft_affine_tests.zig @@ -0,0 +1,61 @@ +//! The proposal projection's quantizer, packing and GPU arithmetic against an independent FP64 reference. +const std = @import("std"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const k = 2560; +const n = 137; +const most = 7; + +fn half(bytes: []const u8, i: usize) f64 { + return rig.bf(std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)); +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0xaff14e); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + const source = r.bfs(n * k, 0.125); + const words = try a.alloc(u8, n * k / 2); + const scales = try a.alloc(u8, n * k / 16); + const biases = try a.alloc(u8, scales.len); + for (0..n) |i| try flash.draft_affine.row(std.mem.sliceAsBytes(source[i * k ..][0..k]), words[i * k / 2 ..][0 .. k / 2], scales[i * k / 16 ..][0 .. k / 16], biases[i * k / 16 ..][0 .. k / 16]); + var layout = try flash.qmm.pack(a, words, scales, biases, n, k / 8); + defer layout.deinit(a); + const dw = try r.dev(layout.weight); + const ds = try r.dev(layout.scales); + const db = try r.dev(layout.biases); + const x = r.bfs(most * k, 0.25); + const sums = try a.alloc(f32, most * k / 32); + for (sums, 0..) |*v, g| { + v.* = 0; + for (x[g * 32 ..][0..32]) |h| v.* += @floatCast(rig.bf(h)); + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dsum = try r.dev(std.mem.sliceAsBytes(sums)); + const out = try r.out(most * n * 2); + var widest: []u16 = undefined; + for ([_]usize{ most, 1, 3 }) |rows| { + try flash.qmm.matmul(gpu.d, r.s, dx, dsum, dw, ds, db, out, rows, n, k); + if (rows == most) widest = try r.back(u16, out, most * n) else try r.sameRows("affine proposal rows", u16, out, widest[0 .. rows * n]); + } + for (0..most) |row| for (0..n) |col| { + var total: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const at = col * k + i; + const code = (words[at / 2] >> @as(u3, @intCast((i % 2) * 4))) & 15; + const g = col * (k / 32) + i / 32; + const weight = @as(f64, @floatFromInt(code)) * half(scales, g) + half(biases, g); + const term = rig.bf(x[row * k + i]) * weight; + total += term; + magnitude += @abs(term); + } + const want = rig.rbf(total); + try r.within("affine proposal FP64", rig.bf(widest[row * n + col]), want, rig.ulps(want, 1) + rig.sumTol(magnitude, k)); + }; + r.report("group-32 proposal quantization and affine projection, every column; rows 1/3/7 equal"); +} diff --git a/zig/tests/cuda/flashnext_fidelity_tests.zig b/zig/tests/cuda/flashnext_fidelity_tests.zig new file mode 100644 index 0000000000..3f3c793ce1 --- /dev/null +++ b/zig/tests/cuda/flashnext_fidelity_tests.zig @@ -0,0 +1,127 @@ +//! Emit full BF16 score rows at fixed histories for comparison with an independent model implementation. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const context = 8192; +const chunk = 32; + +const Case = struct { id: []const u8, prompt: []const u32, continuation: []const u32 }; + +fn validate(cases: []const Case) !void { + if (cases.len == 0 or cases.len > 64) return error.BadFixture; + for (cases, 0..) |case, i| { + if (case.id.len == 0 or case.id.len > 64 or case.prompt.len == 0 or case.continuation.len > 64 or case.prompt.len + case.continuation.len > context) return error.BadFixture; + for (case.id) |c| if (!std.ascii.isAlphanumeric(c) and c != '-') return error.BadFixture; + for (cases[0..i]) |prior| if (std.mem.eql(u8, prior.id, case.id)) return error.BadFixture; + for (case.prompt) |id| if (id >= flash.prompt.vocabulary) return error.InvalidToken; + for (case.continuation) |id| if (id >= flash.prompt.vocabulary) return error.InvalidToken; + } +} + +pub fn run(gpu: check.Gpu, model: []const u8, rank: usize, port: u16, fixture: []const u8, output: []const u8) !void { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, fixture, gpu.gpa, .limited(1 << 22)); + defer gpu.gpa.free(text); + const parsed = try std.json.parseFromSlice([]const Case, gpu.gpa, text, .{}); + defer parsed.deinit(); + try validate(parsed.value); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, model, context, chunk); + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const e = flash.prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, model, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = e; + e.collective = &session.collective; + defer { + session.deinit(); + e.deinit(); + } + if (e.chunk != chunk) return error.RankIdentity; + var logits = try cuda.DeviceBuffer.alloc(gpu.d, flash.prompt.vocabulary * 4); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("fidelity follower completed every fixed-history case", .{}); + return; + } + for (parsed.value) |case| { + var seq = try session.create(context); + defer session.drop(&seq); + const row_bytes = flash.prompt.vocabulary * 2; + const rows = try gpu.gpa.alloc(u8, (case.continuation.len + 1) * row_bytes); + defer gpu.gpa.free(rows); + try session.forward(flash.Tri, tri, &seq, case.prompt, .{}); + @memcpy(rows[0..row_bytes], seq.logits); + for (case.continuation, 0..) |token, i| { + try session.forward(flash.Tri, tri, &seq, &.{token}, .{}); + @memcpy(rows[(i + 1) * row_bytes ..][0..row_bytes], seq.logits); + } + const name = try std.fmt.allocPrint(gpu.gpa, "{s}.bf16", .{case.id}); + defer gpu.gpa.free(name); + const path = try std.fs.path.join(gpu.gpa, &.{ output, name }); + defer gpu.gpa.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = rows }); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(rows, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + check.pass("fixed history {s}: prompt={d} score-rows={d} sha256={s}", .{ case.id, case.prompt.len, case.continuation.len + 1, &hex }); + } +} + +pub fn runSingle(gpu: check.Gpu, model: []const u8, fixture: []const u8, output: []const u8, serial: bool) !void { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, fixture, gpu.gpa, .limited(1 << 22)); + defer gpu.gpa.free(text); + const parsed = try std.json.parseFromSlice([]const Case, gpu.gpa, text, .{}); + defer parsed.deinit(); + try validate(parsed.value); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const e = try flash.prompt.Engine.init(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, model); + defer e.deinit(); + for (parsed.value) |case| { + var seq = try e.newSeq(context); + defer seq.deinit(gpu.gpa); + const row_bytes = flash.prompt.vocabulary * 2; + const rows = try gpu.gpa.alloc(u8, (case.continuation.len + 1) * row_bytes); + defer gpu.gpa.free(rows); + if (serial) { + for (case.prompt) |token| try e.forward(flash.Tri, tri, &seq, &.{token}, .{}); + } else try e.forward(flash.Tri, tri, &seq, case.prompt, .{}); + @memcpy(rows[0..row_bytes], seq.logits); + for (case.continuation, 0..) |token, i| { + try e.forward(flash.Tri, tri, &seq, &.{token}, .{}); + @memcpy(rows[(i + 1) * row_bytes ..][0..row_bytes], seq.logits); + } + const name = try std.fmt.allocPrint(gpu.gpa, "{s}.bf16", .{case.id}); + defer gpu.gpa.free(name); + const path = try std.fs.path.join(gpu.gpa, &.{ output, name }); + defer gpu.gpa.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = rows }); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(rows, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + check.pass("fixed history {s}: prompt={d} score-rows={d} sha256={s}", .{ case.id, case.prompt.len, case.continuation.len + 1, &hex }); + } +} + +test "fixed histories reject unsafe output identities and invalid token or context bounds" { + const good: Case = .{ .id = "short-1", .prompt = &.{1}, .continuation = &.{2} }; + try validate(&.{good}); + var bad = good; + bad.id = "../bad"; + try std.testing.expectError(error.BadFixture, validate(&.{bad})); + try std.testing.expectError(error.BadFixture, validate(&.{ good, good })); + bad = good; + bad.prompt = &.{flash.prompt.vocabulary}; + try std.testing.expectError(error.InvalidToken, validate(&.{bad})); +} diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig new file mode 100644 index 0000000000..974ce95d5e --- /dev/null +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -0,0 +1,385 @@ +//! Flash Next's own hyper-connection and n-gram glue against float64 chains that round where the kernels round. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Rig = rig.Rig; +const bf = rig.bf; +const rbf = rig.rbf; +const ulps = rig.ulps; +const sumTol = rig.sumTol; + +const dims = 2560; +const streams = 4; +const wide = streams * dims; +const low = 320; +const n_down = low + streams; +const slices = 32; +const nc = dims / 256; +const eps: f32 = 1e-6; +const counts = [_]usize{ 1, 3, 17 }; +const most = counts[counts.len - 1]; + +/// Tiled 4-bit rows on the device: words [N/64][K/32][64][4], scales and biases [K/32][N] (cuda_hc.zig's layout). +const Tiled = struct { w: []u32, s: []u16, b: []u16, n: usize, k: usize, dw: u64, ds: u64, db: u64 }; + +fn tiled(r: *Rig, n: usize, k: usize, scale: f32) !Tiled { + const w = r.words(n * k / 8); + const s = r.bfs(n * k / 32, scale); + for (s) |*v| v.* = rig.bf16(@floatCast(@abs(bf(v.*)))); + const b = r.bfs(n * k / 32, scale * 4); + const tw = try flash.hc.tileWords(r.a, std.mem.sliceAsBytes(w), n, k / 8); + const ts = try flash.hc.transposeBf16(r.a, std.mem.sliceAsBytes(s), n, k / 32); + const tb = try flash.hc.transposeBf16(r.a, std.mem.sliceAsBytes(b), n, k / 32); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k, .dw = try r.dev(tw), .ds = try r.dev(ts), .db = try r.dev(tb) }; +} + +/// Row `col` of a tiled matrix over inputs [k0, k0 + 32 g): sum over groups of (x . q) s + (sum x) b, and its |terms|. +fn rowDot(t: Tiled, col: usize, x: []const f64, k0: usize, groups: usize, xs: ?[]const f64) [2]f64 { + var total: f64 = 0; + var abs: f64 = 0; + for (0..groups) |gi| { + const g = k0 / 32 + gi; + var p: f64 = 0; + var sum: f64 = 0; + var pa: f64 = 0; + for (0..32) |k| { + const q = rig.nibble(t.w, t.k / 8, col, g * 32 + k); + p += x[gi * 32 + k] * q; + pa += @abs(x[gi * 32 + k]) * q; + sum += x[gi * 32 + k]; + } + if (xs) |v| sum = v[gi]; + const s = bf(t.s[col * (t.k / 32) + g]); + const b = bf(t.b[col * (t.k / 32) + g]); + total += p * s + sum * b; + abs += pa * @abs(s) + @abs(sum * b); + } + return .{ total, abs }; +} + +fn writeback(r: *Rig, g: *const flash.Glue, mode: usize) !void { + const h = r.bfs(most * wide, 1); + const inj = r.bfs(most * streams, 2); + const y = r.f32s(most * 11 * dims, -1, 1); + const wts = r.f32s(most * 11, 0, 0.2); + const br = r.f32s(8 * most * dims, -0.5, 0.5); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dinj = try r.dev(std.mem.sliceAsBytes(inj)); + const dy = try r.dev(std.mem.sliceAsBytes(y)); + const dwts = try r.dev(std.mem.sliceAsBytes(wts)); + const dbr = try r.dev(std.mem.sliceAsBytes(br)); + var outs: [counts.len]u64 = undefined; + var psss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * wide * 2); + psss[c] = try r.out(m * nc * streams * 4); + try g.hcWriteback(r.s, dh, outs[c], psss[c], dbr, dinj, dy, dwts, most * dims, m, dims, streams, mode, 10, 11, 8); + } + const name = switch (mode) { + 0 => "writeback sums", + 2 => "writeback moe", + else => "writeback slices", + }; + const hout = if (mode == 0) h else try r.back(u16, outs[counts.len - 1], most * wide); + const pss = try r.back(f32, psss[counts.len - 1], most * nc * streams); + for (0..most) |row| { + for (0..dims) |d| { + if (mode == 0) break; + var acc: f64 = 0; + if (mode == 2) { + for (0..11) |k| acc += y[(row * 11 + k) * dims + d] * wts[row * 11 + k]; + } else for (0..8) |k| { + acc += br[k * most * dims + row * dims + d]; + } + const branch = rbf(acc); + for (0..streams) |s| { + const at = row * wide + s * dims + d; + const want = rbf(bf(h[at]) + branch * bf(inj[row * streams + s])); + try r.within(name, bf(hout[at]), want, ulps(want, 1)); + } + } + for (0..nc) |c| for (0..streams) |s| { + var t: f64 = 0; + for (0..256) |i| t += bf(hout[row * wide + s * dims + c * 256 + i]) * bf(hout[row * wide + s * dims + c * 256 + i]); + try r.within(name, pss[(row * nc + c) * streams + s], t, sumTol(t, 256)); + }; + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + if (mode != 0) try r.sameRows(name, u16, outs[c], hout[0 .. m * wide]); + try r.sameRows(name, f32, psss[c], pss[0 .. m * nc * streams]); + } +} + +/// The read-out chain as a pass runs it: down (normed rows, K slices), the slices' sum and activation, the up mix. +fn readout(r: *Rig, g: *const flash.Glue) !void { + const h = r.bfs(most * wide, 1); + const pss = try r.a.alloc(f32, most * nc * streams); + for (0..most) |row| for (0..nc) |c| for (0..streams) |s| { + var t: f32 = 0; + for (0..256) |i| t += @floatCast(bf(h[row * wide + s * dims + c * 256 + i]) * bf(h[row * wide + s * dims + c * 256 + i])); + pss[(row * nc + c) * streams + s] = t; + }; + const scale = r.f32s(wide, 0.5, 1.5); + const down = try tiled(r, n_down, wide, 0.004); + const up = try tiled(r, wide, low, 0.05); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dpss = try r.dev(std.mem.sliceAsBytes(pss)); + const dscale = try r.dev(std.mem.sliceAsBytes(scale)); + var norms: [counts.len]u64 = undefined; + var parts: [counts.len]u64 = undefined; + var acts: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + var injs: [counts.len]u64 = undefined; + var mixes: [counts.len]u64 = undefined; + var xsms: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + norms[c] = try r.out(m * wide * 2); + parts[c] = try r.out(slices * m * n_down * 4); + acts[c] = try r.out(m * low * 2); + xss[c] = try r.out(m * low / 32 * 4); + injs[c] = try r.out(m * streams * 2); + mixes[c] = try r.out(m * dims * 2); + xsms[c] = try r.out(m * dims / 32 * 4); + try g.hcDown(r.s, dh, dpss, dscale, norms[c], down.dw, down.ds, down.db, 0, parts[c], eps, m, n_down, wide, dims, nc, streams, slices); + try g.hcReduceAct(r.s, parts[c], acts[c], xss[c], injs[c], slices, m, streams, low, n_down, 1); + try g.hcUpmix(r.s, acts[c], xss[c], up.dw, up.ds, up.db, norms[c], mixes[c], xsms[c], m, wide, low, dims, streams); + } + const last = counts.len - 1; + const normed = try r.back(u16, norms[last], most * wide); + const part = try r.back(f32, parts[last], slices * most * n_down); + const act = try r.back(u16, acts[last], most * low); + const xs = try r.back(f32, xss[last], most * low / 32); + const inj = try r.back(u16, injs[last], most * streams); + const mixed = try r.back(u16, mixes[last], most * dims); + const xsm = try r.back(f32, xsms[last], most * dims / 32); + const xv = try r.a.alloc(f64, wide); + for (0..most) |row| { + for (0..streams) |s| { + var t: f64 = 0; + for (0..nc) |c| t += pss[(row * nc + c) * streams + s]; + const rinv = 1 / @sqrt(t / dims + eps); + for (0..dims) |d| { + const at = s * dims + d; + const want = bf(h[row * wide + at]) * rinv * scale[at]; + try r.within("down normed", bf(normed[row * wide + at]), want, ulps(want, 0.5) + @abs(want) * 0x1p-20); + xv[at] = bf(normed[row * wide + at]); + } + } + // each K slice's sums from the stored normed values + const per = wide / slices; + for (0..slices) |z| for (0..n_down) |col| { + const d = rowDot(down, col, xv[z * per ..][0..per], z * per, per / 32, null); + try r.within("down slices", part[(z * most + row) * n_down + col], d[0], d[1] * 0x1p-17 + 1e-9); + }; + var asum: [low / 32]f64 = @splat(0); + var aabs: [low / 32]f64 = @splat(0); + for (0..n_down) |col| { + var v: f64 = 0; + for (0..slices) |z| v += part[(z * most + row) * n_down + col]; + const vv = rbf(rbf(v) / streams); + if (col < low) { + const want = rbf(vv * rig.sigmoid(vv)); + try r.within("reduce act", bf(act[row * low + col]), want, ulps(want, 1)); + asum[col / 32] += bf(act[row * low + col]); + aabs[col / 32] += @abs(bf(act[row * low + col])); + } else { + const want = rbf(2 * rbf(rig.sigmoid(vv))); + try r.within("reduce inject", bf(inj[row * streams + col - low]), want, ulps(want, 1)); + } + } + for (0..low / 32) |gi| try r.within("reduce act sums", xs[row * (low / 32) + gi], asum[gi], sumTol(aabs[gi], 32)); + // the up mix from the stored activations, their sums and the normed streams + const av = try r.a.alloc(f64, low); + for (av, 0..) |*o, i| o.* = bf(act[row * low + i]); + const xsv = try r.a.alloc(f64, low / 32); + for (xsv, 0..) |*o, i| o.* = xs[row * (low / 32) + i]; + for (0..dims / 32) |j| { + var gsum: f64 = 0; + var gabs: f64 = 0; + for (0..32) |dd| { + const d = j * 32 + dd; + var total: f64 = 0; + var tabs: f64 = 0; + for (0..streams) |s| { + const u = rowDot(up, s * dims + d, av, 0, low / 32, xsv)[0]; + const term = rbf(rbf(rig.sigmoid(rbf(u))) * bf(normed[row * wide + s * dims + d])); + total += term; + tabs += @abs(term); + } + const want = rbf(total / streams); + try r.within("up mix", bf(mixed[row * dims + d]), want, ulps(tabs / streams, 1) + ulps(want, 0.5)); + gsum += bf(mixed[row * dims + d]); + gabs += @abs(bf(mixed[row * dims + d])); + } + try r.within("up mix sums", xsm[row * (dims / 32) + j], gsum, sumTol(gabs, 32)); + } + } + for (counts[0..last], 0..) |m, c| { + try r.sameRows("down normed", u16, norms[c], normed[0 .. m * wide]); + for (0..slices) |z| try r.sameRows("down slices", f32, parts[c] + z * m * n_down * 4, part[z * most * n_down ..][0 .. m * n_down]); + try r.sameRows("reduce act", u16, acts[c], act[0 .. m * low]); + try r.sameRows("reduce act sums", f32, xss[c], xs[0 .. m * low / 32]); + try r.sameRows("reduce inject", u16, injs[c], inj[0 .. m * streams]); + try r.sameRows("up mix", u16, mixes[c], mixed[0 .. m * dims]); + try r.sameRows("up mix sums", f32, xsms[c], xsm[0 .. m * dims / 32]); + } +} + +/// The n-gram gate then its conv over [tail; rows], and the MTP head's stream adds. +fn ngram(r: *Rig, g: *const flash.Glue) !void { + const keys = r.bfs(most * wide, 1); + const vals = r.bfs(most * dims, 1); + const h = r.bfs(most * wide, 1); + const nk = r.f32s(wide, 0.5, 1.5); + const nq = r.f32s(wide, 0.5, 1.5); + const ncw = r.f32s(wide, 0.5, 1.5); + const tail = r.bfs(9 * wide, 1); + const cw = r.bfs(wide * 4, 0.5); + const dk = try r.dev(std.mem.sliceAsBytes(keys)); + const dv = try r.dev(std.mem.sliceAsBytes(vals)); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dnk = try r.dev(std.mem.sliceAsBytes(nk)); + const dnq = try r.dev(std.mem.sliceAsBytes(nq)); + const dnc = try r.dev(std.mem.sliceAsBytes(ncw)); + const dtail = try r.dev(std.mem.sliceAsBytes(tail)); + const dcw = try r.dev(std.mem.sliceAsBytes(cw)); + var gated: [counts.len]u64 = undefined; + var gpss: [counts.len]u64 = undefined; + var houts: [counts.len]u64 = undefined; + var nrows: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + gated[c] = try r.out(m * wide * 2); + gpss[c] = try r.out(m * streams * 4); + houts[c] = try r.out(m * wide * 2); + nrows[c] = try r.out(m * wide * 2); + try g.pleGate(r.s, dk, dv, dh, dnk, dnq, gated[c], gpss[c], eps, m, dims, streams); + try g.pleConv(r.s, gated[c], gpss[c], dnc, dtail, dcw, dh, houts[c], nrows[c], eps, m, dims, streams, 4, 3); + } + const last = counts.len - 1; + const gt = try r.back(u16, gated[last], most * wide); + const ps = try r.back(f32, gpss[last], most * streams); + const hout = try r.back(u16, houts[last], most * wide); + const nrow = try r.back(u16, nrows[last], most * wide); + for (0..most) |row| for (0..streams) |s| { + const at = row * wide + s * dims; + var kss: f64 = 0; + var qss: f64 = 0; + for (0..dims) |d| { + kss += bf(keys[at + d]) * bf(keys[at + d]); + qss += bf(h[at + d]) * bf(h[at + d]); + } + const kr = 1 / @sqrt(kss / dims + eps); + const qr = 1 / @sqrt(qss / dims + eps); + var dot: f64 = 0; + for (0..dims) |d| dot += rbf(rbf(bf(keys[at + d]) * kr * nk[s * dims + d]) * rbf(bf(h[at + d]) * qr * nq[s * dims + d])); + var gate = rbf(dot / @sqrt(@as(f64, dims))); + const mag = @sqrt(@max(@abs(gate), 1e-6)); + gate = rbf(if (gate > 0) mag else if (gate < 0) -mag else 0); + const sg = rbf(rig.sigmoid(gate)); + var gss: f64 = 0; + for (0..dims) |d| { + const want = rbf(sg * bf(vals[row * dims + d])); + try r.within("ngram gate", bf(gt[at + d]), want, ulps(want, 1)); + gss += bf(gt[at + d]) * bf(gt[at + d]); + } + try r.within("ngram gate sums", ps[row * streams + s], gss, sumTol(gss, dims)); + }; + // the conv of row r reads rows r - 9 + 3t, taken from the tail before the chunk + for (0..most) |row| for (0..wide) |e| { + const s = e / dims; + var acc: f64 = 0; + for (0..4) |t| { + const at = row + 3 * t; + const xv = if (at < 9) bf(tail[at * wide + e]) else blk: { + const rr = at - 9; + const rinv = 1 / @sqrt(@as(f64, ps[rr * streams + s]) / dims + eps); + break :blk rbf(bf(gt[rr * wide + e]) * rinv * ncw[e]); + }; + acc += bf(cw[e * 4 + t]) * xv; + } + const conv = rbf(acc); + const ple = rbf(bf(gt[row * wide + e]) + rbf(conv * rig.sigmoid(conv))); + const want = rbf(bf(h[row * wide + e]) + ple); + try r.within("ngram conv", bf(hout[row * wide + e]), want, ulps(ple, 1) + ulps(want, 0.5)); + const rinv = 1 / @sqrt(@as(f64, ps[row * streams + s]) / dims + eps); + const own = rbf(bf(gt[row * wide + e]) * rinv * ncw[e]); + try r.within("ngram conv rows", bf(nrow[row * wide + e]), own, ulps(own, 1)); + }; + for (counts[0..last], 0..) |m, c| { + try r.sameRows("ngram gate", u16, gated[c], gt[0 .. m * wide]); + try r.sameRows("ngram gate sums", f32, gpss[c], ps[0 .. m * streams]); + try r.sameRows("ngram conv rows", u16, nrows[c], nrow[0 .. m * wide]); + } + // the MTP head's input streams: bf16(e + h), each sum exact before its one rounding + const e = r.bfs(most * dims, 1); + const out = try r.out(most * wide * 2); + try g.addStreams(r.s, try r.dev(std.mem.sliceAsBytes(e)), dh, out, most, dims, streams); + const got = try r.back(u16, out, most * wide); + for (0..most) |row| for (0..wide) |i| { + const want = rbf(bf(e[row * dims + i % dims]) + bf(h[row * wide + i])); + try check.expect(bf(got[row * wide + i]) == want, "add streams: row {d} value {d} is the rounded exact sum", .{ row, i }); + }; +} + +fn writebackRounding(r: *Rig, g: *const flash.Glue) !void { + const n: usize = 256; + const m: usize = 17; + const ss: usize = 4; + const h = r.bfs(m * ss * n, 1); + const inject = r.bfs(m * ss, 0.125); + const values = r.bfs(m * n, 1); + h[0] = 0x3b34; + inject[0] = 0x3c9a; + values[0] = 0xbe16; + const branch = try r.a.alloc(f32, m * n); + for (values, branch) |v, *b| b.* = @floatCast(bf(v)); + const weights = try r.a.alloc(f32, m); + @memset(weights, 1); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const di = try r.dev(std.mem.sliceAsBytes(inject)); + const db = try r.dev(std.mem.sliceAsBytes(branch)); + const dw = try r.dev(std.mem.sliceAsBytes(weights)); + for ([_]usize{ 2, 4 }) |mode| { + const output = try r.out(m * ss * n * 2); + const sums = try r.out(m * ss * 4); + try g.hcWriteback(r.s, dh, output, sums, db, di, db, dw, m * n, m, n, ss, mode, 0, 1, 1); + const got = try r.back(u16, output, h.len); + var worst: f64 = 0; + var improved: usize = 0; + for (got, 0..) |v, i| { + const row = i / (ss * n); + const stream = i / n % ss; + const d = i % n; + const x = bf(values[row * n + d]); + const gate = bf(inject[row * ss + stream]); + const exact = bf(h[i]) + x * gate; + const old = rbf(bf(h[i]) + rbf(x * gate)); + const error_now = @abs(bf(v) - exact); + const error_old = @abs(old - exact); + try check.expect(error_now <= error_old, "writeback mode {d} row {d} stream {d} column {d}: FP64 error {e} exceeds prior {e}", .{ mode, row, stream, d, error_now, error_old }); + try check.expect(bf(v) == rbf(exact), "writeback mode {d} column {d}: one rounding of the full update", .{ mode, i }); + if (error_now < error_old) improved += 1; + worst = @max(worst, error_now); + } + try check.expect(improved > 0, "writeback mode {d}: cancellation fixture improves over double rounding", .{mode}); + std.debug.print("PASS writeback mode {d}: {d} FP64 columns, 0 worse, {d} improved, worst absolute {e}\n", .{ mode, got.len, improved, worst }); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(53); + var r = try Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + try writebackRounding(&r, &g); + for ([_]usize{ 0, 2, 4 }) |mode| try writeback(&r, &g, mode); + try readout(&r, &g); + try ngram(&r, &g); + r.report("flashnext-hc: write-backs, the read-out chain, the n-gram gate and conv, stream adds"); +} diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig new file mode 100644 index 0000000000..ebb937ae71 --- /dev/null +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -0,0 +1,77 @@ +//! Full-checkpoint experiment: two-rank prompt, snapshot replay and concurrent decode agree byte for byte. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; +const ranks = flash.rank_engine; + +pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: []const u8) !void { + const context = 256; + const chunk = 32; + const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); + defer gpu.gpa.free(source); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpu.gpa); + var pieces = std.mem.tokenizeAny(u8, source, " \r\n\t"); + while (pieces.next()) |word| try ids.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10)); + if (ids.items.len == 0 or ids.items.len > context / 2) return error.BadPrompt; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const identity = try ranks.identity(gpu.gpa, gpu.io, dir, context, chunk); + const session = try ranks.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const engine = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = engine; + defer { + session.deinit(); + engine.deinit(); + } + engine.collective = &session.collective; + if (engine.chunk != chunk) return error.RankIdentity; + const room = try session.room((try gpu.ctx.memInfo()).free); + try check.expect(room > prompt.seqBytes(context) * 3, "both ranks have room for experimental caches", .{}); + var logits = try cuda.DeviceBuffer.alloc(gpu.d, 2 * prompt.vocabulary * 2); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("rank 1 replay completed", .{}); + return; + } + var first = try session.create(context); + defer session.drop(&first); + try session.forward(flash.Tri, tri, &first, ids.items, .{}); + check.pass("two-rank prompt completed ({d} tokens)", .{ids.items.len}); + for (0..prompt.vocabulary) |i| { + const value: f32 = @bitCast(@as(u32, std.mem.readInt(u16, first.logits[i * 2 ..][0..2], .little)) << 16); + try check.expect(std.math.isFinite(value), "finite initial logits at {d}", .{i}); + } + const kept = try session.save(&first); + defer session.dropSaved(kept); + const token = prompt.argmax(first.logits); + try session.forward(flash.Tri, tri, &first, &.{token}, .{}); + const expected = try gpu.gpa.dupe(u8, first.logits); + defer gpu.gpa.free(expected); + try session.restore(&first, kept); + try session.forward(flash.Tri, tri, &first, &.{token}, .{}); + try check.sameBytes("two-rank snapshot restore and replay", first.logits, expected); + try session.restore(&first, kept); + var second = try session.create(context); + defer session.drop(&second); + try session.forward(flash.Tri, tri, &second, ids.items, .{}); + try session.round(flash.Tri, tri, &.{ &first, &second }, &.{ token, token }, logits.ptr); + try check.sameBytes("two-rank concurrent first stream", first.logits, expected); + try check.sameBytes("two-rank concurrent second stream", second.logits, expected); + for (0..32) |_| { + const next = prompt.argmax(first.logits); + std.debug.print("TOKEN {d}\n", .{next}); + try session.round(flash.Tri, tri, &.{&first}, &.{next}, logits.ptr); + } + check.pass("two-rank full-model prompt, snapshot replay, concurrent byte equality and 32 decode rows", .{}); +} diff --git a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig new file mode 100644 index 0000000000..7f9b81f1b4 --- /dev/null +++ b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig @@ -0,0 +1,165 @@ +//! Intermediate-channel NVFP4 partitions: exact source cuts, owned upload and FP32 expert partials against FP64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const d = 64; +const width = 64; +const experts = 2; +const slots = 3; +const most = 17; +const rank_width = width / 2; +const fx = cuda.experts.Nvfp4Experts; + +const Source = struct { + codes: [experts][3][d * width / 2]u8 = undefined, + scales: [experts][3][d * width / 16]u8 = undefined, + globals: [experts][3][4]u8 = undefined, + input: [4]u8 = undefined, + releases: usize = 0, + + fn init() Source { + var s: Source = .{}; + for (0..experts) |e| for (0..3) |p| { + for (&s.codes[e][p], 0..) |*v, i| v.* = @truncate(i * 13 + e * 37 + p * 19); + for (&s.scales[e][p], 0..) |*v, i| v.* = @intCast(0x20 + (i % 3) * 8); + std.mem.writeInt(u32, &s.globals[e][p], @bitCast(@as(f32, if (p == 2) 0.25 else 0.125)), .little); + }; + std.mem.writeInt(u32, &s.input, @bitCast(@as(f32, 1)), .little); + return s; + } + + pub fn has(_: *Source, name: []const u8) bool { + return std.mem.endsWith(u8, name, ".weight") or std.mem.endsWith(u8, name, ".weight_scale") or + std.mem.endsWith(u8, name, ".weight_scale_2") or std.mem.endsWith(u8, name, ".input_scale"); + } + + pub fn get(s: *Source, name: []const u8) !core.safetensors.Tensor { + const marker = std.mem.indexOf(u8, name, ".experts.") orelse return error.MissingTensor; + const tail = name[marker + ".experts.".len ..]; + const end = std.mem.indexOfScalar(u8, tail, '.') orelse return error.MissingTensor; + const e = try std.fmt.parseInt(usize, tail[0..end], 10); + if (e >= experts) return error.MissingTensor; + const p: usize = if (std.mem.indexOf(u8, tail, ".gate_proj.") != null) 0 else if (std.mem.indexOf(u8, tail, ".up_proj.") != null) 1 else if (std.mem.indexOf(u8, tail, ".down_proj.") != null) 2 else return error.MissingTensor; + if (std.mem.endsWith(u8, name, ".weight")) return tensor(.u8, &.{ width, d / 2 }, &s.codes[e][p]); + if (std.mem.endsWith(u8, name, ".weight_scale")) return tensor(.f8_e4m3, &.{ width, d / 16 }, &s.scales[e][p]); + if (std.mem.endsWith(u8, name, ".weight_scale_2")) return tensor(.f32, &.{}, &s.globals[e][p]); + if (std.mem.endsWith(u8, name, ".input_scale")) return tensor(.f32, &.{}, &s.input); + return error.MissingTensor; + } + + pub fn release(s: *Source) void { + s.releases += 1; + } +}; + +fn tensor(dtype: core.safetensors.DType, shape: []const usize, bytes: []const u8) core.safetensors.Tensor { + var t: core.safetensors.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +fn weight(s: *const Source, e: usize, p: usize, n: usize, k: usize) f64 { + const byte = s.codes[e][p][n * (d / 2) + k / 2]; + const code: u4 = if (k % 2 == 0) @truncate(byte) else @truncate(byte >> 4); + const scale = s.scales[e][p][n * (d / 16) + k / 16]; + const global: f32 = @bitCast(std.mem.readInt(u32, &s.globals[e][p], .little)); + return core.quant.nvfp4.e2m1(code) * core.quant.fp8.e4m3(scale) * global; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0x2f4); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var source = Source.init(); + var left = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 0, d, width, experts); + defer left.free(); + var right = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 1, d, width, experts); + defer right.free(); + try check.expect(source.releases == 2 * experts * 3, "every source projection released after its owned rank cut", .{}); + const x = r.bfs(most * d, 0.125); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const picks = try a.alloc(i32, most * slots); + for (0..most) |row| { + picks[row * slots] = @intCast(row % experts); + picks[row * slots + 1] = @intCast((row + 1) % experts); + picks[row * slots + 2] = experts; // shared expert is a separate format, skipped here + } + const dp = try r.dev(std.mem.sliceAsBytes(picks)); + const plan: cuda.grouped.Plan = .{ + .members = try r.out(most * slots * 4), + .items = try r.out((most * slots + experts + 1) * 12), + .counts = try r.out(8), + .rank = try r.out(most * slots * 4), + .hist = try r.out((experts + 1) * 4), + }; + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try cuda.grouped.Router.resolve(mod); + var prior: [2]?[]f32 = .{ null, null }; + for ([_]usize{ 4, 1 }) |warps| { + var kernels = try fx.loadWarps(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count)), warps); + defer kernels.unload(); + const act = try r.out(most * slots * rank_width * 2); + const out = try r.out(most * slots * d * 4); + const combined = try a.alloc(f32, most * slots * d); + @memset(combined, 0); + const reference_full = try a.alloc(f64, combined.len); + @memset(reference_full, 0); + const reference_abs = try a.alloc(f64, combined.len); + @memset(reference_abs, 0); + for ([_]*const flash.modelopt_experts.Owned{ &left, &right }, 0..) |owned, rank| { + var all: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, most * slots * d * 4, r.s.handle), "cuMemsetD8Async"); + const items = cuda.grouped.maxItems(rows * slots, experts + 1, 16); + try router.route(r.s, dp, rows * slots, experts + 1, 16, plan); + try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, owned.view(), plan, act, items, experts); + try kernels.run(r.s, .down_f32, .{ .x = act, .stride = rank_width }, owned.view(), plan, out, items, experts); + if (rows == most) all = try r.back(f32, out, most * slots * d) else try r.sameRows("NVFP4 rank output rows", f32, out, all[0 .. rows * slots * d]); + } + if (prior[rank]) |ref| try check.sameBytes("NVFP4 CTA grouping leaves every rank output byte-exact", std.mem.sliceAsBytes(all), std.mem.sliceAsBytes(ref)) else prior[rank] = all; + for (0..most) |row| for (0..slots) |slot| { + const e: usize = @intCast(picks[row * slots + slot]); + if (e == experts) { + for (all[(row * slots + slot) * d ..][0..d]) |v| try check.expect(v == 0, "shared slot untouched by NVFP4", .{}); + continue; + } + var activation: [rank_width]f64 = undefined; + for (&activation, 0..) |*v, j| { + const channel = rank * rank_width + j; + var gate: f64 = 0; + var up: f64 = 0; + for (0..d) |k| { + gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, channel, k); + up += rig.bf(x[row * d + k]) * weight(&source, e, 1, channel, k); + } + gate = rig.rbf(gate); + up = rig.rbf(up); + v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); + } + for (0..d) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (activation, 0..) |v, k| { + const term = v * weight(&source, e, 2, j, rank * rank_width + k); + want += term; + mag += @abs(term); + } + const at = (row * slots + slot) * d + j; + try r.within("NVFP4 rank FP32 down", all[at], want, rig.sumTol(mag, rank_width) + 1e-6); + combined[at] += all[at]; + reference_full[at] += want; + reference_abs[at] += mag; + } + }; + } + for (combined, reference_full, reference_abs) |got, want, mag| try r.within("NVFP4 summed rank partials", got, want, rig.sumTol(mag, width) + 2e-6); + r.report("NVFP4 TP cuts: both intermediate halves, exact source reader and packing, independent FP64 reference, FP32 partials and untouched shared slot; rows 1/7/17 equal"); + } +} diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig new file mode 100644 index 0000000000..6bbe109dfb --- /dev/null +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -0,0 +1,121 @@ +//! Real two-process/two-device NCCL sums and bounded disconnect/timeout behavior over the framed control link. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const wire = flash.rank_wire; + +pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { + const missing = std.mem.eql(u8, mode, "missing"); + const started = wire.now(); + if (missing) { + try std.testing.expectError(error.RankTimeout, wire.Link.open(gpu.io, "127.0.0.1", port, rank, 200)); + try check.expect(wire.now() - started < 3 * std.time.ns_per_s, "missing peer startup is bounded", .{}); + check.pass("rank {d}: missing peer startup timed out within bound", .{rank}); + return; + } + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 30_000); + defer link.close(); + var no_delay: c_int = 0; + var option_len: std.posix.socklen_t = @sizeOf(c_int); + try check.expect(std.c.getsockopt(link.fd, std.posix.IPPROTO.TCP, std.posix.TCP.NODELAY, &no_delay, &option_len) == 0 and no_delay == 1, "both connected rank sockets disable Nagle", .{}); + const until = wire.deadline(30_000); + const hello: wire.Header = .{ .kind = .hello, .sequence = 1, .object = 77, .arg = 123, .bytes = 7 }; + if (rank == 0) { + const raw = try hello.encode(); + try link.send(raw[0..3], until); + try link.send(raw[3..], until); + try link.send("payload", until); + const ack = try link.receiveHeader(until); + try check.expect(ack.kind == .ack and ack.sequence == 1 and ack.bytes == 0, "exact hello acknowledgment", .{}); + } else { + try std.testing.expectEqualDeep(hello, try link.receiveHeader(until)); + var payload: [7]u8 = undefined; + try link.receive(&payload, until); + try check.sameBytes("fragmented control frame payload", &payload, "payload"); + try link.sendHeader(.{ .kind = .ack, .sequence = 1 }, until); + } + var collective = try flash.collective.Collective.init(link, rank, if (std.mem.eql(u8, mode, "timeout")) 200 else 30_000); + defer collective.deinit(); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + if (std.mem.eql(u8, mode, "eof") or std.mem.eql(u8, mode, "timeout")) { + if (rank == 1) { + // A deliberate peer omission: the owner either disconnects or stays alive without entering the collective. + if (std.mem.eql(u8, mode, "eof")) { + _ = std.c.shutdown(link.fd, std.posix.SHUT.RDWR); + } else { + const delay: std.c.timespec = .{ .sec = 0, .nsec = 600_000_000 }; + _ = std.c.nanosleep(&delay, null); + } + check.pass("rank 1: injected {s}", .{mode}); + return; + } + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&@as(f32, 1))); + defer values.free(); + const start = wire.now(); + completeSum(&collective, stream, values.ptr, 1) catch |err| { + try check.expect(err == error.RankDisconnected or err == error.RankTimeout or err == error.NcclFailed, "failure is surfaced ({s})", .{@errorName(err)}); + try check.expect(collective.failed and collective.comm == null, "failed communicator is aborted and cannot be reused", .{}); + try check.expect(wire.now() - start < 3 * std.time.ns_per_s, "failed collective terminates within bound", .{}); + try std.testing.expectError(error.RankDisconnected, collective.sum(stream, values.ptr, 1)); + check.pass("rank 0: {s} surfaced as {s}, aborted, and refused reuse within bound", .{ mode, @errorName(err) }); + return; + }; + return error.ExpectedRankFailure; + } + for ([_]usize{ 0, 1, 17, 4099, 1_000_003 }, 0..) |count, seq| { + const host = try gpu.gpa.alloc(f32, @max(1, count)); + defer gpu.gpa.free(host); + for (host, 0..) |*v, i| v.* = @floatFromInt(i % 127 + rank); + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer values.free(); + try collective.sum(stream, values.ptr, count); + try collective.sum(stream, values.ptr, count); + try collective.finish(stream); + try values.download(0, std.mem.sliceAsBytes(host)); + for (host[0..count], 0..) |actual, i| try check.expect(actual == @as(f32, @floatFromInt(4 * (i % 127) + 2)), "rank {d}: exact sum at {d}", .{ rank, i }); + // Acknowledge completion before either peer can close or start its next reduction. + const h: wire.Header = .{ .kind = .ready, .sequence = seq + 2, .arg = count }; + if (rank == 0) { + try link.sendHeader(h, wire.deadline(30_000)); + try std.testing.expectEqualDeep(h, try link.receiveHeader(wire.deadline(30_000))); + } else { + try std.testing.expectEqualDeep(h, try link.receiveHeader(wire.deadline(30_000))); + try link.sendHeader(h, wire.deadline(30_000)); + } + } + for ([_]usize{ 0, 1, 17, 124_160 }) |count| { + const host = try gpu.gpa.alloc(u16, @max(1, count)); + defer gpu.gpa.free(host); + for (host, 0..) |*v, i| v.* = @truncate(i * 31 + rank * 0x8001); + var local = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer local.free(); + var output = try cuda.DeviceBuffer.alloc(gpu.d, @max(2, count * 4)); + defer output.free(); + try collective.gather(stream, local.ptr, output.ptr, count); + try collective.finish(stream); + const all = try gpu.gpa.alloc(u16, count * 2); + defer gpu.gpa.free(all); + if (count > 0) try output.download(0, std.mem.sliceAsBytes(all)); + for (0..2) |owner| for (0..count) |i| { + const expected: u16 = @truncate(i * 31 + owner * 0x8001); + try check.expect(all[owner * count + i] == expected, "rank {d}: gathered BF16 bits at {d}/{d}", .{ rank, owner, i }); + }; + } + const done: wire.Header = .{ .kind = .ready, .sequence = 99 }; + if (rank == 0) { + try link.sendHeader(done, wire.deadline(30_000)); + try std.testing.expectEqualDeep(done, try link.receiveHeader(wire.deadline(30_000))); + } else { + try std.testing.expectEqualDeep(done, try link.receiveHeader(wire.deadline(30_000))); + try link.sendHeader(done, wire.deadline(30_000)); + } + check.pass("rank {d}: BF16 all-gather preserves every bit and rank column order for empty, scalar, tail and vocabulary-half rows", .{rank}); + check.pass("rank {d}: two-process chained NCCL FP32 sums exact for empty, scalar, tail and 1,000,003 values; fragmented framed control preserved", .{rank}); +} + +fn completeSum(c: *flash.collective.Collective, stream: cuda.Stream, values: u64, count: usize) !void { + try c.sum(stream, values, count); + try c.finish(stream); +} diff --git a/zig/tests/cuda/flashnext_rig.zig b/zig/tests/cuda/flashnext_rig.zig new file mode 100644 index 0000000000..9d917f5a1c --- /dev/null +++ b/zig/tests/cuda/flashnext_rig.zig @@ -0,0 +1,119 @@ +//! The Flash Next glue tests' shared rig: device buffers, bf16 values, float64 rounding points and tolerances. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); + +pub fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +pub fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +/// A float64 value at a kernel's bf16 rounding point (through fp32, as the kernel rounds). +pub fn rbf(x: f64) f64 { + return bf(bf16(@floatCast(x))); +} + +pub fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +/// `n` bf16 ulps of |want| (one ulp: 2^-7 of it), with a floor for zeros. +pub fn ulps(want: f64, n: f64) f64 { + return n * @abs(want) * std.math.pow(f64, 2, -7) + 1e-6; +} + +/// An fp32 sum's tolerance: `terms` roundings of at most 2^-24 of the sum of |terms|. +pub fn sumTol(abs_sum: f64, terms: usize) f64 { + return abs_sum * @as(f64, @floatFromInt(terms)) * std.math.pow(f64, 2, -24) + 1e-30; +} + +pub const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// The largest |got - want| / tolerance seen, by check name. + worst: std.StringArrayHashMapUnmanaged(f64) = .empty, + + pub fn init(gpu: check.Gpu, a: std.mem.Allocator, rng: std.Random) !Rig { + return .{ .gpu = gpu, .a = a, .s = try cuda.Stream.init(gpu.d, false), .rng = rng }; + } + + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + r.s.deinit(); + } + + pub fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + /// A device buffer of `bytes`, every byte 0xff (a value the kernel misses reads NaN). + pub fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(bytes, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); + return b.ptr; + } + + pub fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try r.a.alloc(T, n); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } + + /// bf16 values within [-scale, scale). + pub fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + const v = r.a.alloc(u16, n) catch unreachable; + for (v) |*x| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale); + return v; + } + + /// fp32 values within [lo, hi). + pub fn f32s(r: *Rig, n: usize, lo: f32, hi: f32) []f32 { + const v = r.a.alloc(f32, n) catch unreachable; + for (v) |*x| x.* = lo + r.rng.float(f32) * (hi - lo); + return v; + } + + pub fn words(r: *Rig, n: usize) []u32 { + const v = r.a.alloc(u32, n) catch unreachable; + for (v) |*x| x.* = r.rng.int(u32); + return v; + } + + pub fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); + const gop = try r.worst.getOrPut(r.a, what); + if (!gop.found_existing) gop.value_ptr.* = 0; + gop.value_ptr.* = @max(gop.value_ptr.*, q); + } + + /// Bytes of a launch over fewer rows against the same rows of the widest launch. + pub fn sameRows(r: *Rig, what: []const u8, comptime T: type, ptr: u64, want: []const T) !void { + try check.sameBytes(what, std.mem.sliceAsBytes(try r.back(T, ptr, want.len)), std.mem.sliceAsBytes(want)); + } + + pub fn report(r: *Rig, what: []const u8) void { + var it = r.worst.iterator(); + var line: std.ArrayList(u8) = .empty; + while (it.next()) |e| line.print(r.a, " {s} {d:.3};", .{ e.key_ptr.*, e.value_ptr.* }) catch {}; + check.pass("{s}: every row count's bits equal; worst tolerances used:{s}", .{ what, line.items }); + } +}; + +/// q (4 bits) of an MLX row's value i: word i / 8, nibble i % 8. +pub fn nibble(w: []const u32, row_words: usize, row: usize, i: usize) f64 { + return @floatFromInt((w[row * row_words + i / 8] >> @intCast(4 * (i % 8))) & 0xF); +} diff --git a/zig/tests/cuda/flashnext_shared_tests.zig b/zig/tests/cuda/flashnext_shared_tests.zig new file mode 100644 index 0000000000..1077d0e4dc --- /dev/null +++ b/zig/tests/cuda/flashnext_shared_tests.zig @@ -0,0 +1,125 @@ +//! Independent stage-wise FP64 references for the separate MXFP8 shared expert and its reserved slot. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const d = 2560; +const w = 640; +const most = 17; +const slots = 11; + +const Weight = struct { + codes: []u8, + scales: []u8, + n: usize, + k: usize, + + fn init(a: std.mem.Allocator, n: usize, k: usize, salt: usize) !Weight { + const codes = try a.alloc(u8, n * k); + const scales = try a.alloc(u8, n * k / 32); + for (codes, 0..) |*v, i| v.* = @intCast(0x20 + ((i * 13 + salt * 7) % 32) + (if ((i + salt) % 3 == 0) @as(usize, 128) else 0)); + for (scales, 0..) |*v, i| v.* = @intCast(121 + (i + salt) % 3); + return .{ .codes = codes, .scales = scales, .n = n, .k = k }; + } + + fn value(self: Weight, n: usize, k: usize) f64 { + const exponent: i32 = @as(i32, self.scales[n * (self.k / 32) + k / 32]) - 127; + return core.quant.fp8.e4m3(self.codes[n * self.k + k]) * std.math.pow(f64, 2, @floatFromInt(exponent)); + } + + fn upload(self: Weight, a: std.mem.Allocator, driver: *const cuda.Driver) !flash.projection.Projection { + return flash.projection.mxfp8(a, driver, tensor(.f8_e4m3, self.codes, self.n, self.k), tensor(.u8, self.scales, self.n, self.k / 32)); + } +}; + +fn tensor(dtype: core.safetensors.DType, bytes: []const u8, n: usize, k: usize) core.safetensors.Tensor { + return .{ .dtype = dtype, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = bytes }; +} + +fn projectionReference(r: *rig.Rig, name: []const u8, x: []const u16, weight: Weight, out: anytype, rows: usize, comptime bf16: bool) !void { + for (0..rows) |row| for (0..weight.n) |n| { + var total: f64 = 0; + var magnitude: f64 = 0; + for (0..weight.k) |k| { + const term = rig.bf(x[row * weight.k + k]) * weight.value(n, k); + total += term; + magnitude += @abs(term); + } + const at = row * weight.n + n; + if (bf16) { + try r.within(name, rig.bf(out[at]), rig.rbf(total), rig.ulps(rig.rbf(total), 1)); + } else try r.within(name, out[at], total, rig.sumTol(magnitude, weight.k) + 1e-6); + }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0x5a4ed); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + const gw = try Weight.init(a, w, d, 1); + const uw = try Weight.init(a, w, d, 2); + const dw = try Weight.init(a, d, w, 3); + var gate = try gw.upload(a, gpu.d); + defer gate.free(); + var up = try uw.upload(a, gpu.d); + defer up.free(); + var down = try dw.upload(a, gpu.d); + defer down.free(); + const x = r.bfs(most * d, 0.125); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const g = try r.out(most * w * 2); + const u = try r.out(most * w * 2); + const act = try r.out(most * w * 2); + const out = try r.out(most * d * 4); + const temp = try r.out(8 * most * d * 4); + const sentinel = try a.alloc(f32, most * slots * d); + @memset(sentinel, 0.375); + const y = try r.dev(std.mem.sliceAsBytes(sentinel)); + var all_g: []u16 = undefined; + var all_u: []u16 = undefined; + var all_act: []u16 = undefined; + var all_out: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try lane.matmul(r.s, dx, d, rows, gate.native.?, g, temp); + try lane.matmul(r.s, dx, d, rows, up.native.?, u, temp); + try glue.sharedSwiGLU(r.s, g, u, act, rows, w); + try lane.matmulF32(r.s, act, w, rows, down.native.?, out, temp); + try glue.sharedScatter(r.s, out, y, rows, d, slots); + if (rows == most) { + all_g = try r.back(u16, g, most * w); + all_u = try r.back(u16, u, most * w); + all_act = try r.back(u16, act, most * w); + all_out = try r.back(f32, out, most * d); + } else { + try r.sameRows("shared gate rows", u16, g, all_g[0 .. rows * w]); + try r.sameRows("shared up rows", u16, u, all_u[0 .. rows * w]); + try r.sameRows("shared activation rows", u16, act, all_act[0 .. rows * w]); + try r.sameRows("shared FP32 output rows", f32, out, all_out[0 .. rows * d]); + } + } + try projectionReference(&r, "shared MXFP8 gate", x, gw, all_g, most, true); + try projectionReference(&r, "shared MXFP8 up", x, uw, all_u, most, true); + for (all_g, all_u, all_act) |gate_value, up_value, actual| { + const f = rig.bf(gate_value); + const want = rig.rbf(rig.rbf(f / (1 + @exp(-f))) * rig.bf(up_value)); + try r.within("shared SwiGLU rounding", rig.bf(actual), want, rig.ulps(want, 1)); + } + try projectionReference(&r, "shared MXFP8 FP32 down", all_act, dw, all_out, most, false); + const scattered = try r.back(f32, y, most * slots * d); + for (scattered, 0..) |actual, i| { + const slot = (i / d) % slots; + const want = if (slot == slots - 1) all_out[(i / (slots * d)) * d + i % d] else @as(f32, 0.375); + try check.expect(actual == want, "shared scatter preserves routed slots and copies exact FP32 values", .{}); + } + try std.testing.expectError(error.RowShape, glue.sharedScatter(r.s, out, y, most, d, 0)); + r.report("MXFP8 shared expert, BF16 projection/activation seams, FP32 down, reserved-slot scatter; rows 1/7/17 equal"); +} diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig new file mode 100644 index 0000000000..999d773cd0 --- /dev/null +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -0,0 +1,80 @@ +//! A two-rank verification window agrees with serial rows and rolls back at partial and full acceptance seams. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; +const mtp = flash.mtp; + +pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: []const u8) !void { + const context = 256; + const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); + defer gpu.gpa.free(source); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpu.gpa); + var words = std.mem.tokenizeAny(u8, source, " \r\n\t"); + while (words.next()) |word| try ids.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10)); + if (ids.items.len == 0 or ids.items.len > 128) return error.BadPrompt; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + var identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, dir, context, 32); + identity.mtp = 1; + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const e = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = e; + e.collective = &session.collective; + defer { + session.deinit(); + e.deinit(); + } + if (e.chunk != 32) return error.RankIdentity; + var logits = try cuda.DeviceBuffer.alloc(gpu.d, 2 * prompt.vocabulary * 2); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("rank 1 verification windows and prefix commits completed", .{}); + return; + } + var serial = try session.create(context); + defer session.drop(&serial); + try session.forward(flash.Tri, tri, &serial, ids.items, .{}); + const base = try session.save(&serial); + defer session.dropSaved(base); + const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71, 99, 4, 80, 122, 7, 25, 61, 94 }; + const row_bytes = prompt.vocabulary * 2; + const expected = try gpu.gpa.alloc(u8, forced.len * row_bytes); + defer gpu.gpa.free(expected); + for (forced, 0..) |token, i| { + try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); + @memcpy(expected[i * row_bytes ..][0..row_bytes], serial.logits); + } + var window = try session.create(context); + defer session.drop(&window); + var draft = try mtp.Draft.init(gpu.gpa, e.driver, stream, context, 1, 32); + defer draft.deinit(); + const rows = try gpu.gpa.alloc(u8, expected.len); + defer gpu.gpa.free(rows); + for ([_]usize{ 1, 3, 7, 15, 16 }) |keep| { + try session.restore(&window, base); + try session.verify(flash.Tri, tri, &window, &draft, &forced); + try stream.synchronize(); + try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(rows.ptr, draft.logits.ptr, rows.len), "cuMemcpyDtoH"); + try check.sameBytes("each two-rank window row equals its serial teacher-forced row", rows, expected); + try session.keep(flash.Tri, tri, &window, &draft, keep); + try check.expect(window.fed == ids.items.len + keep, "kept prefix position is exact", .{}); + try session.restore(&serial, base); + for (forced[0..keep]) |token| try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); + try check.expect(std.mem.eql(i64, &window.history, &serial.history), "kept host n-gram history matches serial", .{}); + try session.forward(flash.Tri, tri, &serial, &.{13}, .{}); + try session.forward(flash.Tri, tri, &window, &.{13}, .{}); + try check.sameBytes("rejected suffix cannot change the next row after prefix commit", window.logits, serial.logits); + check.pass("two-rank forced window: all rows and next row byte-exact at keep={d}/{d}", .{ keep, forced.len }); + } +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..6633ebed77 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -17,6 +17,8 @@ const prompt_cuts = @import("prompt_cuts.zig"); const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); +const mxfp8_tests = @import("mxfp8_tests.zig"); +const projection_f32_tests = @import("projection_f32_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); @@ -39,8 +41,34 @@ const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const experts_rows_tests = @import("experts_rows_tests.zig"); +const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); +const flashnext_fidelity_tests = @import("flashnext_fidelity_tests.zig"); +const flashnext_window_tests = @import("flashnext_window_tests.zig"); +const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); +const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); +const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); +const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); +const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); +const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); +const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); + const usage = \\usage: tf-cuda-test + \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches + \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-fidelity-single MODEL FIXTURE OUTPUT single-rank reference isolation + \\ flashnext-fidelity-rank MODEL RANK PORT FIXTURE OUTPUT fixed-history full score rows + \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity + \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) + \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes + \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance + \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 + \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) + \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 + \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) + \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ info driver, device and library versions \\ smoke copies, fills, launches, argument packing, module globals \\ graph stream capture, explicit graphs, node and whole-exec updates @@ -80,6 +108,8 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance + \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes @@ -128,6 +158,20 @@ fn arg(rest: []const [:0]const u8, i: usize) ![]const u8 { } fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { + if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), false); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-single-serial")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), true); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); + if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); + if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); + if (std.mem.eql(u8, cmd, "flashnext-draft-affine")) return flashnext_draft_affine_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); if (std.mem.eql(u8, cmd, "graph")) return runtime_tests.graphs(gpu); @@ -155,6 +199,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); + if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); 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, "modelopt")) return modelopt_tests.run(gpu); diff --git a/zig/tests/cuda/mxfp8_tests.zig b/zig/tests/cuda/mxfp8_tests.zig new file mode 100644 index 0000000000..5417d3d1fd --- /dev/null +++ b/zig/tests/cuda/mxfp8_tests.zig @@ -0,0 +1,143 @@ +//! Native MXFP8: lossless packing, independent FP64 products, strided rows and row-count invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +fn tobf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +pub fn run(gpu: check.Gpu) !void { + try scaleRange(gpu); + const a = gpu.gpa; + const n = 70; // a partial tile and padded rows + const k = 1024; // split-K, both cluster and explicit reduction + const m = 256; // fused prompt reduction + const ldx = k + 64; + const npad = cuda.mxfp8.padded(n); + var prng = std.Random.DefaultPrng.init(0x8e80_5eed); + const rng = prng.random(); + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + for (codes) |*v| { + v.* = rng.int(u8); + if (v.* & 0x7f == 0x7f) v.* ^= 1; + } + const scales = try a.alloc(u8, n * (k / 32)); + defer a.free(scales); + for (scales) |*v| v.* = 121 + rng.uintLessThan(u8, 8); + // A dedicated row makes UE8M0 byte 0 observable in a normal BF16 result. + @memset(codes[0..k], 0x38); // E4M3 1 + @memset(scales[0 .. k / 32], 0); + const packed_codes = try a.alloc(u8, npad * k); + defer a.free(packed_codes); + cuda.mxfp8.packCodes(packed_codes, codes, n, k); + const packed_scales = try a.alloc(u8, npad * (k / 32)); + defer a.free(packed_scales); + cuda.mxfp8.packScales(packed_scales, scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, packed_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, packed_scales); + defer ds.free(); + const w = cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k); + const x = try a.alloc(u16, m * ldx); + defer a.free(x); + @memset(x, 0x7fc0); // stride padding is poison, not another activation + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = tobf(rng.float(f32) * 2); + } + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); + defer dy.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, 8 * m * n * 4); + defer part.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var baseline: ?[]u8 = null; + defer if (baseline) |b| a.free(b); + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try gpu.ctx.capability()) >= 90; + const lin: cuda.qlinear.Linear = .{ .lane = &lane }; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + for ([_]bool{ false, true }) |prompt| { + try dy.fill8(0xff, s.handle); + const in: cuda.qlinear.Rows = .{ .x = dx.ptr, .ldx = ldx, .part = part.ptr }; + if (prompt) try lin.prompt(s, .{ .mxfp8 = w }, in, dy.ptr, rows) else try lin.decode(s, .{ .mxfp8 = w }, in, dy.ptr, rows); + try s.synchronize(); + const got = try a.alloc(u8, rows * n * 2); + defer a.free(got); + try dy.download(0, got); + if (baseline) |b| { + try check.sameBytes("MXFP8 row-count/dispatch/cluster invariance", got, b[0..got.len]); + } else baseline = try a.dupe(u8, got); + for (0..rows) |row| for (0..n) |col| { + var ref: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const scale = std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127); + const v = bf(x[row * ldx + i]) * core.quant.fp8.e4m3(codes[col * k + i]) * scale; + ref += v; + magnitude += @abs(v); + } + const actual = bf(std.mem.readInt(u16, got[(row * n + col) * 2 ..][0..2], .little)); + const tolerance = @abs(ref) / 128 + magnitude / 4096 + 0x1p-133; + try check.expect(@abs(actual - ref) <= tolerance, "MXFP8 FP64 row {d}, column {d}: {e} versus {e}, tolerance {e}", .{ row, col, actual, ref, tolerance }); + }; + } + } + } + check.pass("MXFP8: strided rows 1,7,17,33,256 equal across dispatch and reduction paths; FP64 including UE8M0 byte zero", .{}); +} + +/// Every UE8M0 encoding, through the production matmul: each dot product is exactly one times its scale. +fn scaleRange(gpu: check.Gpu) !void { + const a = gpu.gpa; + const n = 256; + const k = 64; + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + @memset(codes, 0x38); + var scales: [n * 2]u8 = undefined; + for (0..n) |r| @memset(scales[r * 2 ..][0..2], @intCast(r)); + const tiled_codes = try a.alloc(u8, n * k); + defer a.free(tiled_codes); + cuda.mxfp8.packCodes(tiled_codes, codes, n, k); + var tiles: [n * 2]u8 = undefined; + cuda.mxfp8.packScales(&tiles, &scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, tiled_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, &tiles); + defer ds.free(); + var x: [k]u16 = @splat(0x3c80); // 1/64: 64 equal terms sum to one + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, n * 2); + defer dy.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + try lane.matmul(s, dx.ptr, k, 1, cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k), dy.ptr, 0); + try s.synchronize(); + var got: [n * 2]u8 = undefined; + try dy.download(0, &got); + for (0..n) |r| { + const bits = std.mem.readInt(u16, got[r * 2 ..][0..2], .little); + if (r == 255) { + try check.expect(std.math.isNan(bf(bits)), "UE8M0 reserved byte produces NaN", .{}); + } else { + const want: u16 = if (r == 0) 0x40 else @intCast(r << 7); + try check.expect(bits == want, "UE8M0 {d}: BF16 bits {x} versus {x}", .{ r, bits, want }); + } + } + check.pass("MXFP8: all 256 UE8M0 encodings preserve the specified value, including subnormal and NaN", .{}); +} diff --git a/zig/tests/cuda/nemotron/test_aot_manifest.py b/zig/tests/cuda/nemotron/test_aot_manifest.py index e87ff8551c..7a82c87028 100755 --- a/zig/tests/cuda/nemotron/test_aot_manifest.py +++ b/zig/tests/cuda/nemotron/test_aot_manifest.py @@ -37,10 +37,10 @@ def fake_cache(cache: Path) -> dict: for leaf in ("__grp__fake_kernel.json", "fake_kernel.cubin", "fake_kernel.ptx")} info = {"name": "fake_kernel", "function": "fakemod.FakeKernel", "source": {"file": "tensorfold/families/nemotron_h/cuda/kernels.py", "line": 10}, - "params": ["p0", "n", "BLOCK", "SCALE"], - "signature": {"p0": "*i8", "n": "i32", "BLOCK": "constexpr", "SCALE": "constexpr"}, + "params": ["p0", "n", "BLOCK", "SCALE", "ROPE"], + "signature": {"p0": "*i8", "n": "i32", "BLOCK": "constexpr", "SCALE": "constexpr", "ROPE": "constexpr"}, "constexprs": {"BLOCK": 64, "SCALE": {"float": 0.125, "fp32_bits": "0x3e000000", - "fp64_bits": "0x3fc0000000000000"}, "n": 3}, + "fp64_bits": "0x3fc0000000000000"}, "n": 3, "ROPE": None}, "attrs": {"p0": [["tt.divisibility", 16]], "n": [["tt.contiguity", 1]]}, "metadata": {"num_warps": 4}, "files": files, "n_regs": 24, "n_spills": 0} return {"generator": "test", "written_unix": 0, "kernels": {HASH: info}, @@ -73,7 +73,7 @@ def main() -> int: assert k["launches"] == {"startup": 3} and k["grids"] == [[128, 1, 1]] jit = out / "jit.json" - jit.write_text(json.dumps({"fakemod.FakeKernel": {"params": ["p0", "n", "BLOCK", "SCALE"], + jit.write_text(json.dumps({"fakemod.FakeKernel": {"params": ["p0", "n", "BLOCK", "SCALE", "ROPE"], "do_not_specialize": [], "no_align": []}}) + "\n") # the manifest shape, with a cache and a jit file named on the command line packed = work / "packed" @@ -112,7 +112,18 @@ def main() -> int: print(r.stdout + r.stderr) return 1 assert (default / "aot.json").read_text() == (packed / "aot.json").read_text() - print("test_aot_manifest: one kernel through the manifest, the dump and the defaults; entries checked") + + # a second capture's jit.json beside the first: merged, the same set again + jit2 = out / "jit2.json" + jit2.write_text(json.dumps({"othermod.Other": {"params": ["x"], "do_not_specialize": ["x"], "no_align": []}})) + merged = work / "merged" + r = subprocess.run([sys.executable, "-B", str(PACK), "--manifest", str(manifest), "--cache", str(cache), + "--jit", str(jit), "--jit", str(jit2), "--out", str(merged)], capture_output=True, text=True) + if r.returncode != 0: + print(r.stdout + r.stderr) + return 1 + assert (merged / "aot.json").read_text() == (packed / "aot.json").read_text() + print("test_aot_manifest: one kernel through the manifest, the dump, the defaults and two jit files; entries checked") return 0 diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig index 8f1d295ec1..9eb6005c30 100644 --- a/zig/tests/cuda/norm_rows_tests.zig +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -171,35 +171,93 @@ fn gatedRms(r: *Rig, n: *const cuda.norms.Norms, heads: usize, dv: usize) !void } } -fn embed(r: *Rig, d: usize) !void { +fn embed(r: *Rig, d: usize, group: usize, copies: usize) !void { var e = try cuda.embed_rows.EmbedRows.load(r.gpu.d); defer e.unload(); const vocab = 300; const w = try r.a.alloc(u32, vocab * d / 8); for (w) |*x| x.* = r.rng.int(u32); - const sc = r.bfs(vocab * d / 64, 0.1); - const bi = r.bfs(vocab * d / 64, 0.5); - const tab = r.bfs(vocab * d, 1); + const sc = r.bfs(vocab * d / group, 0.1); + const bi = r.bfs(vocab * d / group, 0.5); + const tab = try r.a.alloc(u16, vocab * d); + for (tab, 0..) |*x, i| x.* = @truncate(i); // every BF16 bit pattern, including NaN payloads const ids = try r.a.alloc(i32, most); for (ids) |*x| x.* = @intCast(r.rng.uintLessThan(usize, vocab)); const did = try r.dev(std.mem.sliceAsBytes(ids)); - const o4 = try r.out(most * d * 2); + const o4 = try r.out(most * copies * d * 2); const o16 = try r.out(most * d * 2); - try e.affine4(r.s, did, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(sc)), try r.dev(std.mem.sliceAsBytes(bi)), o4, most, d); - try e.bf16(r.s, did, try r.dev(std.mem.sliceAsBytes(tab)), o16, most, d); - const g4 = try r.back(u16, o4, most * d); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const ds = try r.dev(std.mem.sliceAsBytes(sc)); + const db = try r.dev(std.mem.sliceAsBytes(bi)); + if (group == 64 and copies == 1) try e.affine4(r.s, did, dw, ds, db, o4, most, d) else try e.affine4Groups(r.s, did, dw, ds, db, o4, most, d, group, copies); + const dt = try r.dev(std.mem.sliceAsBytes(tab)); + try e.bf16(r.s, did, dt, o16, most, d); + const oc = try r.out(most * copies * d * 2); + try e.bf16Copies(r.s, did, dt, oc, 0, d, copies); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc, most, d, 0)); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc + 2, most, d, copies)); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc, most, d - 1, copies)); + for (counts) |rows| { + try e.bf16Copies(r.s, did, dt, oc, rows, d, copies); + const got = try r.back(u16, oc, rows * copies * d); + for (0..rows) |row| { + const token: usize = @intCast(ids[row]); + for (0..copies) |c| try check.sameBytes("embed BF16 exact copies", std.mem.sliceAsBytes(got[(row * copies + c) * d ..][0..d]), std.mem.sliceAsBytes(tab[token * d ..][0..d])); + } + } + const g4 = try r.back(u16, o4, most * copies * d); const g16 = try r.back(u16, o16, most * d); for (0..most) |row| { const t: usize = @intCast(ids[row]); for (0..d) |i| { const q: u4 = @truncate(w[t * d / 8 + i / 8] >> @intCast(4 * (i % 8))); - const want = aref.affine4(q, bf(sc[t * d / 64 + i / 64]), bf(bi[t * d / 64 + i / 64])); - try r.within("embed affine4", bf(g4[row * d + i]), want, nref.tolerance(want, .bf16)); + const want = aref.affine4(q, bf(sc[t * d / group + i / group]), bf(bi[t * d / group + i / group])); + try r.within("embed affine4", bf(g4[row * copies * d + i]), want, nref.tolerance(want, .bf16)); } + for (1..copies) |c| try check.sameBytes("embed affine4 copies", std.mem.sliceAsBytes(g4[(row * copies + c) * d ..][0..d]), std.mem.sliceAsBytes(g4[row * copies * d ..][0..d])); try check.sameBytes("embed bf16", std.mem.sliceAsBytes(g16[row * d ..][0..d]), std.mem.sliceAsBytes(tab[t * d ..][0..d])); } } +/// RMS rows `x_ld` apart with fp32 weights, and the stored outputs' sums of each `group`. +fn rmsSums(r: *Rig, n: *const cuda.norms.Norms, width: usize, x_ld: usize, group: usize) !void { + const x = r.bfs(most * x_ld, 1); + const wf = r.weights(width); + const wvals = try r.a.alloc(f64, width); + for (wvals, wf) |*o, v| o.* = v; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dw = try r.dev(std.mem.sliceAsBytes(wf)); + var ys: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + ys[c] = try r.out(m * width * 2); + xss[c] = try r.out(m * width / group * 4); + try n.rmsSums(r.s, dx, x_ld, dw, ys[c], xss[c], m, width, group, eps); + } + const y = try r.back(u16, ys[counts.len - 1], most * width); + const xs = try r.back(f32, xss[counts.len - 1], most * width / group); + const xv = try r.a.alloc(f64, width); + const want = try r.a.alloc(f64, width); + for (0..most) |row| { + for (xv, 0..) |*v, i| v.* = bf(x[row * x_ld + i]); + nref.rms(xv, wvals, eps, want); + for (want, 0..) |w, i| try r.within("rms sums", bf(y[row * width + i]), w, nref.tolerance(w, .bf16)); + for (0..width / group) |j| { + var t: f64 = 0; + var abs: f64 = 0; + for (y[row * width + j * group ..][0..group]) |v| { + t += bf(v); + abs += @abs(bf(v)); + } + try r.within("rms sums group sums", xs[row * (width / group) + j], t, abs * @as(f64, @floatFromInt(group)) * std.math.pow(f64, 2, -24) + 1e-30); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes("rms sums", std.mem.sliceAsBytes(try r.back(u16, ys[c], m * width)), std.mem.sliceAsBytes(y[0 .. m * width])); + try check.sameBytes("rms sums group sums", std.mem.sliceAsBytes(try r.back(f32, xss[c], m * width / group)), std.mem.sliceAsBytes(xs[0 .. m * width / group])); + } +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -220,8 +278,12 @@ pub fn run(gpu: check.Gpu) !void { try gatedRms(&r, &n, 48, 128); // the 27B's GDN output try gatedRms(&r, &n, 5, 64); try gatedRms(&r, &n, 3, 256); - try embed(&r, 5120); // the 27B's hidden - try embed(&r, 2688); - try embed(&r, 192); - check.pass("norm rows: add_rms_h, gated rms, embed rows within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); + try embed(&r, 5120, 64, 1); // the 27B's hidden + try embed(&r, 2688, 64, 1); + try embed(&r, 192, 64, 1); + try embed(&r, 2560, 32, 4); // Flash Next: group 32, a copy a hyper-connection stream + try embed(&r, 96, 32, 1); + try rmsSums(&r, &n, 2560, 10240, 32); // Flash Next's MTP head: the embedding, its first stream of four + try rmsSums(&r, &n, 10240, 10240, 32); // and the four streams normed as one row + check.pass("norm rows: add_rms_h, gated rms, rms with group sums, embed rows within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); } diff --git a/zig/tests/cuda/projection_f32_host_tests.zig b/zig/tests/cuda/projection_f32_host_tests.zig new file mode 100644 index 0000000000..5d38e2e29c --- /dev/null +++ b/zig/tests/cuda/projection_f32_host_tests.zig @@ -0,0 +1,109 @@ +//! Exercise the fixture's actual poison-and-dispatch path with CPU buffers and no-op CUDA launches. +const std = @import("std"); +const cuda = @import("cuda"); +const fixture = @import("projection_f32_tests.zig"); +const Rig = @import("quant_ref_tests.zig").RigT; + +const Mock = struct { + var destination: []u8 = &.{}; + var scratch: []u8 = &.{}; + var fills: usize = 0; + var launches: usize = 0; + var fresh = true; + + fn fill(ptr: cuda.abi.DevicePtr, byte: u8, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + const bytes: [*]u8 = @ptrFromInt(ptr); + @memset(bytes[0..n], byte); + fills += 1; + return cuda.abi.success; + } + + fn observe() cuda.abi.Result { + for (destination) |byte| fresh = fresh and byte == 0xff; + for (scratch) |byte| fresh = fresh and byte == 0xff; + launches += 1; + return cuda.abi.success; + } + + fn plain(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + return observe(); + } + + fn extended(_: *const cuda.abi.LaunchConfig, _: cuda.abi.Function, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + return observe(); + } + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + fn copy(dst: ?*anyopaque, src: cuda.abi.DevicePtr, n: usize) callconv(.c) cuda.abi.Result { + const out: [*]u8 = @ptrCast(dst.?); + const in: [*]const u8 = @ptrFromInt(src); + @memcpy(out[0..n], in[0..n]); + return cuda.abi.success; + } +}; + +test "every F32 and BF16 fixture dispatch rejects a no-op after a valid stale prefix" { + const a = std.testing.allocator; + var driver: cuda.Driver = undefined; + driver.api.cuMemsetD8Async = Mock.fill; + driver.api.cuLaunchKernel = Mock.plain; + driver.api.cuLaunchKernelEx = Mock.extended; + driver.api.cuStreamSynchronize = Mock.sync; + driver.api.cuMemcpyDtoH_v2 = Mock.copy; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(1) }; + const f: cuda.Function = .{ .d = &driver, .handle = null }; + var lane: cuda.qmmf.Lane = .{ + .mod = .{ .d = &driver, .handle = null }, + .fns = @splat(@splat(f)), + .f32_fns = @splat(@splat(f)), + .reduce = @splat(f), + .clusters = false, + }; + const weight: cuda.qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 70, .k = 1024, .npad = 128 }; + const elements = 256 * weight.n; + const expected = try a.alloc(f64, elements); + defer a.free(expected); + const tolerance = try a.alloc(f64, elements); + defer a.free(tolerance); + @memset(expected, 0); + @memset(tolerance, 1); + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var rig: Rig = .{ + .gpu = .{ .d = &driver, .ctx = undefined, .gpa = a, .io = std.testing.io }, + .a = arena.allocator(), + .s = stream, + .rng = undefined, + }; + for ([_]cuda.qmmf.Out{ .f32, .bf16 }) |format| { + const width: usize = if (format == .f32) 4 else 2; + const destination = try a.alloc(u8, elements * width); + defer a.free(destination); + const scratch = try a.alloc(u8, 8 * elements * 4); + defer a.free(scratch); + const out: cuda.DeviceBuffer = .{ .d = &driver, .ptr = @intFromPtr(destination.ptr), .len = destination.len }; + const part: cuda.DeviceBuffer = .{ .d = &driver, .ptr = @intFromPtr(scratch.ptr), .len = scratch.len }; + Mock.destination = destination; + Mock.scratch = scratch; + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + @memset(destination, 0); + @memset(scratch, 0); + const count = rows * weight.n; + try rig.close("valid stale prefix", out.ptr, format == .f32, expected[0..count], tolerance[0..count]); + Mock.fills = 0; + Mock.launches = 0; + Mock.fresh = true; + try fixture.dispatch(&lane, stream, 0, 1088, rows, weight, out, part, format); + try std.testing.expectEqual(@as(usize, 2), Mock.fills); + try std.testing.expect(Mock.launches > 0); + try std.testing.expect(Mock.fresh); + try std.testing.expectError(error.TestFailed, rig.close("no-op projection", out.ptr, format == .f32, expected[0..count], tolerance[0..count])); + } + } + } +} diff --git a/zig/tests/cuda/projection_f32_tests.zig b/zig/tests/cuda/projection_f32_tests.zig new file mode 100644 index 0000000000..3edd66a243 --- /dev/null +++ b/zig/tests/cuda/projection_f32_tests.zig @@ -0,0 +1,139 @@ +//! FP32 lane output before residual/rank combination: independent FP64, rounding and row/reduction invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const reference = @import("quant_ref_tests.zig"); +const Rig = reference.RigT; +const bf = reference.helpers.bfOf; +const bf16 = reference.helpers.bf16Of; +const Mode = cuda.qmmf.Mode; + +/// One fixture dispatch starts with fresh destination and reduction poison, including a reused row prefix. +pub fn dispatch(lane: *const cuda.qmmf.Lane, s: cuda.Stream, x: u64, ldx: usize, rows: usize, w: cuda.qmmf.Weight, out: cuda.DeviceBuffer, part: cuda.DeviceBuffer, format: cuda.qmmf.Out) !void { + try out.fill8(0xff, s.handle); + try part.fill8(0xff, s.handle); + try lane.matmulOut(s, x, ldx, rows, w, out.ptr, format, part.ptr); +} + +pub fn run(gpu: check.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(0xf32_5eed); + 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 lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + for ([_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }) |mode| try projection(&r, &lane, mode); +} + +fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { + const n = 70; + const k = 1024; + const m = 256; + const ldx = k + 64; + const npad = cuda.fp8.padded(n); + const a = r.a; + const codes = try a.alloc(u8, n * k); + for (codes) |*v| v.* = reference.helpers.fp8CodeOf(r); + const scales = try a.alloc(u8, n * (k / 16)); + for (scales) |*v| v.* = if (mode == .mxfp8) 121 + r.rng.uintLessThan(u8, 8) else 0x20 + r.rng.uintLessThan(u8, 0x28); + const inv = try a.alloc(f32, ((n + 127) / 128) * (k / 128)); + for (inv) |*v| v.* = r.rng.float(f32) * 0.01 + 0.001; + const w16 = try r.bfs(n * k, 0.25); + const packed16 = try a.alloc(u16, npad * k); + const tiled_codes = std.mem.sliceAsBytes(packed16); + const global: f32 = 0.0123; + var w: cuda.qmmf.Weight = undefined; + switch (mode) { + .bf16 => { + cuda.qmmf.packBf16(packed16, w16, n, k); + w = cuda.modelopt.bf16Weight(try r.dev(tiled_codes), n, k); + }, + .fp4 => { + const words = try a.alloc(u32, npad * k / 8); + cuda.nvfp4.packWords(words, codes[0 .. n * k / 2], n, k); + const tiles = try a.alloc(u8, npad * (k / 16)); + cuda.nvfp4.packScales(tiles, scales, n, k); + w = cuda.nvfp4.weight(try r.dev(std.mem.sliceAsBytes(words)), try r.dev(tiles), global, n, k); + }, + .fp8, .fp8g, .mxfp8 => { + cuda.fp8.packCodes(tiled_codes[0 .. npad * k], codes, n, k); + const dw = try r.dev(tiled_codes[0 .. npad * k]); + w = switch (mode) { + .fp8 => cuda.modelopt.fp8Weight(dw, global, n, k), + .fp8g => blk: { + const tiles = try a.alloc(f32, npad * (k / 64)); + cuda.fp8.packScales(tiles, inv, n, k); + break :blk cuda.fp8.weight(dw, try r.dev(std.mem.sliceAsBytes(tiles)), n, k); + }, + .mxfp8 => blk: { + const tiles = try a.alloc(u8, npad * (k / 32)); + cuda.mxfp8.packScales(tiles, scales[0 .. n * k / 32], n, k); + break :blk cuda.mxfp8.weight(dw, try r.dev(tiles), n, k); + }, + else => unreachable, + }; + }, + } + const x = try a.alloc(u16, m * ldx); + @memset(x, 0x7fc0); + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = bf16((r.rng.float(f32) - 0.5) * 2); + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + var out_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, m * n * 4); + defer out_buf.free(); + var bf16_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, m * n * 2); + defer bf16_buf.free(); + var part_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, 8 * m * n * 4); + defer part_buf.free(); + const out = out_buf.ptr; + const out_bf16 = bf16_buf.ptr; + const want = try a.alloc(f64, m * n); + const tol = try a.alloc(f64, m * n); + for (0..m) |row| for (0..n) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const code = codes[col * k + i]; + const weight: f64 = switch (mode) { + .bf16 => bf(w16[col * k + i]), + .fp8 => core.quant.fp8.e4m3(code) * global, + .fp8g => core.quant.fp8.e4m3(code) * inv[(col / 128) * (k / 128) + i / 128], + .mxfp8 => core.quant.fp8.e4m3(code) * std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127), + .fp4 => blk: { + const byte = codes[col * (k / 2) + i / 2]; + break :blk core.quant.nvfp4.e2m1(if (i % 2 == 0) byte & 15 else byte >> 4) * core.quant.fp8.e4m3(scales[col * (k / 16) + i / 16]) * global; + }, + }; + const v = bf(x[row * ldx + i]) * weight; + sum += v; + magnitude += @abs(v); + } + want[row * n + col] = sum; + tol[row * n + col] = magnitude * 2e-6 + @abs(sum) * 1e-6 + 1e-30; + }; + var baseline: ?[]u8 = null; + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try r.gpu.ctx.capability()) >= 90; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + try dispatch(lane, r.s, dx, ldx, rows, w, out_buf, part_buf, .f32); + try r.close("FP32 projection", out, true, want[0 .. rows * n], tol[0 .. rows * n]); + const got = try r.back(out, rows * n * 4); + if (baseline) |b| try check.sameBytes("FP32 row/reduction invariance", got, b[0..got.len]) else baseline = got; + try dispatch(lane, r.s, dx, ldx, rows, w, bf16_buf, part_buf, .bf16); + const rounded = try r.back(out_bf16, rows * n * 2); + for (0..rows * n) |i| { + const v = std.mem.readInt(u32, got[i * 4 ..][0..4], .little); + try check.expect(bf16(@bitCast(v)) == std.mem.readInt(u16, rounded[i * 2 ..][0..2], .little), "FP32 rounds once to the existing BF16 result", .{}); + } + } + } + check.pass("{t}: FP32 outputs preserve FP64 accuracy, BF16 rounding and row/reduction invariance", .{mode}); +} diff --git a/zig/tests/cuda/quant_ref_host_tests.zig b/zig/tests/cuda/quant_ref_host_tests.zig new file mode 100644 index 0000000000..dc13f4235e --- /dev/null +++ b/zig/tests/cuda/quant_ref_host_tests.zig @@ -0,0 +1,83 @@ +//! CPU tests of Rig.close: only stream synchronization and readback are replaced; the comparison is the GPU gate's. +const std = @import("std"); +const cuda = @import("cuda"); +const Rig = @import("quant_ref_tests.zig").RigT; + +const Readback = struct { + var syncs: usize = 0; + var copies: usize = 0; + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + syncs += 1; + return cuda.abi.success; + } + + fn copy(dst: ?*anyopaque, src: cuda.abi.DevicePtr, n: usize) callconv(.c) cuda.abi.Result { + const out: [*]u8 = @ptrCast(dst.?); + const in: [*]const u8 = @ptrFromInt(src); + @memcpy(out[0..n], in[0..n]); + copies += 1; + return cuda.abi.success; + } +}; + +fn close(comptime T: type, got: []const T, want: []const f64, tol: []const f64) !void { + var d: cuda.Driver = undefined; + d.api.cuStreamSynchronize = Readback.sync; + d.api.cuMemcpyDtoH_v2 = Readback.copy; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var rig: Rig = .{ + .gpu = .{ .d = &d, .ctx = undefined, .gpa = std.testing.allocator, .io = std.testing.io }, + .a = arena.allocator(), + .s = .{ .d = &d, .handle = null }, + .rng = undefined, + }; + Readback.syncs = 0; + Readback.copies = 0; + const result = rig.close("CPU reference comparison", @intFromPtr(got.ptr), T == f32, want, tol); + try std.testing.expectEqual(@as(usize, 1), Readback.syncs); + try std.testing.expectEqual(@as(usize, 1), Readback.copies); + return result; +} + +const zero = [_]f64{ 0, 0, 0 }; +const tolerance = [_]f64{ 1, 1, 1 }; + +test "finite FP32 and BF16 outputs pass at the bound and fail beyond it" { + try close(f32, &.{ 0.5, 1, 0 }, &zero, &tolerance); + try close(u16, &.{ 0x3f00, 0x3f80, 0 }, &zero, &tolerance); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 2, 0, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x4000, 0, 0 }, &zero, &tolerance)); +} + +test "FP32 NaN cannot disappear before a later finite result" { + for ([_]u32{ 0x7fc00000, 0xffffffff }) |bits| { + for (0..3) |i| { + var got = [_]f32{ 0, 0, 0 }; + got[i] = @bitCast(bits); + try std.testing.expectError(error.TestFailed, close(f32, &got, &zero, &tolerance)); + } + } +} + +test "BF16 NaN cannot disappear before a later finite result" { + for ([_]u16{ 0x7fc0, 0xffff }) |bits| { + for (0..3) |i| { + var got = [_]u16{ 0, 0, 0 }; + got[i] = bits; + try std.testing.expectError(error.TestFailed, close(u16, &got, &zero, &tolerance)); + } + } +} + +test "NaN cannot erase an earlier out-of-bound or infinite result" { + try std.testing.expectError(error.TestFailed, close(f32, &.{ 0, std.math.inf(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 0, -std.math.inf(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0, 0x7f80, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0, 0xff80, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 2, std.math.nan(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ std.math.inf(f32), std.math.nan(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x4000, 0xffff, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x7f80, 0xffff, 0 }, &zero, &tolerance)); +} diff --git a/zig/tests/cuda/quant_ref_tests.zig b/zig/tests/cuda/quant_ref_tests.zig index 657cd388b7..a0f40b664c 100644 --- a/zig/tests/cuda/quant_ref_tests.zig +++ b/zig/tests/cuda/quant_ref_tests.zig @@ -58,6 +58,7 @@ const Rig = struct { for (want, tol, 0..) |w, t, i| { const g: f64 = if (f32_out) std.mem.bytesToValue(f32, raw[4 * i ..][0..4]) else bf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); const q = @abs(g - w) / t; + try check.expect(std.math.isFinite(q), "{s}: non-finite error at element {d} (got {e}, reference {e}, tolerance {e})", .{ name, i, g, w, t }); if (!(q <= most)) { most = q; at = i; diff --git a/zig/tests/cuda/route_rows_tests.zig b/zig/tests/cuda/route_rows_tests.zig new file mode 100644 index 0000000000..5cb1795c86 --- /dev/null +++ b/zig/tests/cuda/route_rows_tests.zig @@ -0,0 +1,145 @@ +//! route_rows.cu's router logits and softmax top-k rows against float64, bits equal at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const rref = core.route_ref; +const aref = core.act_ref; +const counts = [_]usize{ 1, 3, 16, 17, 40 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + worst: f64 = 0, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); + return b.ptr; + } + + fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try r.a.alloc(T, n); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } + + fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + const v = r.a.alloc(u16, n) catch unreachable; + for (v) |*x| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale); + return v; + } + + fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); + r.worst = @max(r.worst, q); + } +}; + +/// Router logits at every row count (the few-row and the many-row launches): fp64 dots, then the same bits a row. +fn logits(r: *Rig, g: *const cuda.route_rows.RouteRows, d: usize, e: usize) !void { + const ld = d + 64; // rows with something after them + const x = r.bfs(most * ld, 1); + const w = r.bfs(e * d, 0.05); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + var outs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * e * 4); + try g.logits(r.s, dx, ld, dw, outs[c], m, d, e); + } + const got = try r.back(f32, outs[counts.len - 1], most * e); + const xv = try r.a.alloc(f64, d); + const wv = try r.a.alloc(f64, d); + for (0..most) |row| { + for (xv, 0..) |*o, i| o.* = bf(x[row * ld + i]); + for (0..e) |ex| { + for (wv, 0..) |*o, i| o.* = bf(w[ex * d + i]); + try r.within("route logits", got[row * e + ex], rref.dot(xv, wv), rref.dotTolerance(xv, wv)); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| try check.sameBytes("route logits rows", std.mem.sliceAsBytes(try r.back(f32, outs[c], m * e)), std.mem.sliceAsBytes(got[0 .. m * e])); +} + +/// Top-k of each row (a row of ties too): the picks the fp64 rule makes, weights and the shared gate within bf16. +fn topk(r: *Rig, g: *const cuda.route_rows.RouteRows, experts: usize, k: usize, shared: bool) !void { + const ld = experts + 1; + const slots = k + @intFromBool(shared); + const l = try r.a.alloc(f32, most * ld); + for (l) |*v| v.* = (r.rng.float(f32) * 2 - 1) * 4; + for (0..experts) |i| l[i] = @floatFromInt(i % 3); // row 0: every pick a tie broken by the lower id + const dl = try r.dev(std.mem.sliceAsBytes(l)); + var picks: [counts.len]u64 = undefined; + var wts: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + picks[c] = try r.out(m * slots * 4); + wts[c] = try r.out(m * slots * 4); + try g.topk(r.s, dl, ld, experts, k, picks[c], wts[c], slots, m, shared); + } + const p = try r.back(i32, picks[counts.len - 1], most * slots); + const wt = try r.back(f32, wts[counts.len - 1], most * slots); + const row = try r.a.alloc(f64, experts); + const ids = try r.a.alloc(usize, k); + const want = try r.a.alloc(f64, k); + for (0..most) |i| { + for (row, 0..) |*o, j| o.* = l[i * ld + j]; + rref.topk(row, ids, want); + for (0..k) |j| { + try check.expect(p[i * slots + j] == @as(i32, @intCast(ids[j])), "route top-k row {d} pick {d}: expert {d}, the rule picks {d}", .{ i, j, p[i * slots + j], ids[j] }); + try r.within("route top-k weight", wt[i * slots + j], want[j], aref.tolerance(want[j])); + } + if (shared) { + try check.expect(p[i * slots + k] == @as(i32, @intCast(experts)), "route top-k row {d}: the shared slot names expert {d}", .{ i, p[i * slots + k] }); + const gate = aref.sigmoid(bf(bf16(l[i * ld + experts]))); + try r.within("route shared gate", wt[i * slots + k], gate, aref.tolerance(gate)); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes("route top-k picks", std.mem.sliceAsBytes(try r.back(i32, picks[c], m * slots)), std.mem.sliceAsBytes(p[0 .. m * slots])); + try check.sameBytes("route top-k weights", std.mem.sliceAsBytes(try r.back(f32, wts[c], m * slots)), std.mem.sliceAsBytes(wt[0 .. m * slots])); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(41); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var g = try cuda.route_rows.RouteRows.load(gpu.d); + defer g.unload(); + try logits(&r, &g, 2560, 513); // Flash Next: 512 experts and the shared expert's gate row + try logits(&r, &g, 2688, 128); + try logits(&r, &g, 1040, 37); // a K tail past the staged 256s, and experts past a block + try topk(&r, &g, 512, 10, true); + try topk(&r, &g, 128, 8, false); + check.pass("route rows: logits and top-k within float64 tolerance (worst {d:.3}), the fp64 rule's picks, every row count's bits equal", .{r.worst}); +}