diff --git a/RUNBOOK.md b/RUNBOOK.md index fe9d257802..f1b41f9851 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -277,6 +277,12 @@ Each settings file names `rank`, `library` and `links`; each link names its peer [Speed-up settings](docs/speed-up-mode.md#4-write-the-settings-files) has the schema and link setup. That guide's two-Mac Flash Next commands load the checkpoint directly, as one Mac does. +Qwen3.8-27B serves the same way: the same `$HOME/models/qwen3.8-27b` checkpoint on both Macs, one settings file per rank, rank 1 started first and rank 0 within five minutes, with `--parallel 1` because a pair serves one greedy stream. +Send requests to rank 0 only, and expect no drafter, no prompt cache and no learning: `--drafter`, `--prompt-cache-gib` and `--learn` are refused on a pair. +Each rank waits up to five minutes for the other's ready word, then exits with `the peer never signalled ready in five minutes`. +When rank 0 stops, rank 1 stops following and exits by itself within about three seconds: each rank runs a link watchdog that ends the process when the peer stops signalling for three seconds, so a dead link never leaves a rank serving alone. Restart both after a link failure: that fails the request in flight and leaves the pair down until both ranks run again. +[Qwen3.8-27B on two Macs](docs/recipes/qwen3.8-27b.md#two-macs-over-thunderbolt) has the commands and what the split covers. + On one Mac, GLM loads only when its weights and caches fit in 70% of RAM, 179.2 GB on a 256 GB Mac, and leaves the MTP head on disk when only the head is over. `GLM_LOAD_LIMIT_GB=N` sets that limit in GB, up to the GPU's recommended working set; a value that isn't a number above 0 stops the server. Past 70% of RAM the server warns at startup: a load that wires most of the RAM can stall macOS until its watchdog restarts the Mac, so run nothing else large beside it. diff --git a/build.zig b/build.zig index 60eadf478a..df21706e98 100644 --- a/build.zig +++ b/build.zig @@ -153,7 +153,7 @@ pub fn build(b: *std.Build) void { const gemma4_host_mods = metalEngineModules(b, b.graph.host, .Debug, draft_ids); gemma4_host_mods.sources.addImport("nemotron_packed_metallib", embedded(b, b.addWriteFiles().add("gemma4-compile-empty.metallib", ""), "nemotron_packed.metallib")); @import("zig/build/gemma4_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes, gemma4_host_mods.sources); - @import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes); + test_step.dependOn(@import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes)); } @import("zig/build/qwen27_host.zig").tests(b); const xgrammar = grammar_build.option(b); diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 2dd2a3a5e4..7d9da03c77 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -20,7 +20,7 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho and the magnitude of the terms each output sums, and can replay the sums in another accumulator. - Its served envelope: the shapes the engine calls it at, and for each dimension its tile, the step a valid size moves by and the least valid size. -- Its accumulator (`f32` for every kernel here) and the depth of its longest chain of roundings at a shape, plus any +- Its accumulator (`f32` where a kernel sums) and the depth of its longest chain of roundings at a shape, plus any fp32 operations after the sums (`ops`) and intermediate roundings to the output type (`mids`). - Whether it may write in place. If it replaces another kernel, the kernel whose bits it must equal. - Its NaN rule: `propagate` (as float64 arithmetic on its formula does), or a stated rule the test checks itself @@ -28,7 +28,8 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho - For gate 5, its launch in `zig/src/core/contract/registry.zig`: threads a group and threadgroup bytes (the source's, and any its encoder sets) on Metal, or symbol, image, threads a block, dynamic shared bytes and cluster width on CUDA, with the limits its compiles recorded. HIP launches are in `zig/src/core/contract/hip_records.zig`: symbol, code - object, threads a workgroup and dynamic LDS bytes. + object, threads a workgroup and dynamic LDS bytes. The seam's records are in `metal_pair_records.zig`, which + `registry.zig` concatenates. ## The gates @@ -99,7 +100,8 @@ The program exits 1 on any failed gate, any control that stayed green, or any li the inputs for a shape and a gate's fill (usually through an input family), `run` launches once and waits, `count` and `reference` give the outputs and their float64 values, `simulate` replays the sums for gate 4. 2. Add it to the runner's list, with its served shapes and tiles. -3. Add its launches to `registry.zig`, and fill the records from `--offline` (Metal) or `kernel-contract-limits` +3. Add its launches to `registry.zig`, or to a records file it concatenates as the seam's are, and fill them from + `--offline` (Metal) or `kernel-contract-limits` (CUDA: nvcc 13.0's in `cuda`, and each later release's set where that release compiles it differently); a HIP kernel's go to `hip_records.zig`, filled from `hip-kernel-contract-limits`. 4. Post the `SUMMARY` lines from each backend you ran in the pull request, and say which chips you could not run. @@ -126,11 +128,17 @@ float64 output class, without an inherited exception. ## Coverage -Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm -(4 bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit +Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm (4 +bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's -prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of -tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10. +prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles), and the row-split TP seam's +five kernels (`pair_rows.metal`, launched by `pair_rows.zig`): `pair_add_rms` on the fp32 stream, and the bf16 frame's +`pair_stage`, `pair_gather`, `pair_argmax` and `pair_pick`. `pair_stage` reads only its own rank's half of a frame and +copies it into the slot the relay sends, and `pair_gather` writes both halves of every row, its own half natively and +the peer's from the relayed words, with no arithmetic between them: no exchange rounds a sum across ranks, and +`pair_argmax` and `pair_pick` state their own NaN and tie rules. Gate 5 also holds the rest of tile_qmm's pipelines +and GLM-5.3-Flash's hc_expand_10. + CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and diff --git a/docs/known-issues.md b/docs/known-issues.md index f5b37fdde9..82a1aa080f 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -37,7 +37,12 @@ release fixes it, and the release notes say which one. since 1.0.10's faster products. `FZ_M3_LANES=1` keeps the earlier kernels and their tokens. - @Bizuayeu's regression run on shipped 1.1.0 loaded the whole GLM-5.3-Flash NVFP4 checkpoint on two GB10s at TP=2 in 302.1 s, kept one token hash per prompt across three runs and found a passphrase in 199,307 tokens; scored prompts and more than two ranks were not covered ([#1040](https://github.com/ashhart/TensorFold/issues/1040)). - DeepSeek-V4.1 CUDA pieces land in this release, but no release binary serves DeepSeek checkpoints until its whole-model gate ([#827](https://github.com/ashhart/TensorFold/pull/827)) passes. - +- Qwen3.8-27B in speed-up mode, two Macs over Thunderbolt: a pair serves one greedy stream, drafts nothing, keeps no + prompt cache and does not learn, because rank 1 has no drafter checkpoint and each Mac holds only its own halves of a + slot's state. `--drafter`, `--prompt-cache-gib`, `--learn` and `--parallel` above 1 are refused, and a request that + samples or carries a grammar fails. Each output projection splits by output rows: each rank reads the whole input K + and reduces its rows as one Mac does, and a `pair_gather` joins the two bf16 halves with no arithmetic between them: + pair logits are byte-identical to one Mac's on the recorded fixtures. ## Living Weights (`--slide`) diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index fcc4cdb43a..7a07d14976 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -72,6 +72,50 @@ least 256 tokens beyond the previous chunk start, or after 2,048 tokens if no ea Prefix reuse resumes only at these chunk starts, so a follow-up can reuse the state before its previous reply. The plan comes from rendered tokens; a template without detected markers uses 2,048-token chunks. +### Two Macs over Thunderbolt + +Speed-up mode splits one Qwen3.8-27B reply across two Macs over a Thunderbolt 5 cable: both Macs hold the same +`TensorFold/Qwen3.8-27B-MLX-4bit` checkpoint, each serves it from its own directory with its own settings file, and +rank 1 starts before rank 0. +[Speed-up mode](../speed-up-mode.md) builds the server, sets up MCDMA and writes the two settings files; this +subsection covers what the 27B pair serves and what it leaves out. +Send requests to rank 0 only; rank 1 holds the other half of the weights and replays rank 0's rounds. + +```bash +# rank 1, on the second Mac: start it first +zig-out/native/bin/tensorfold-native serve ~/models/qwen3.8-27b \ + --speed-up rank1.json --port 4101 --parallel 1 --context 32768 --no-update-check +# rank 0, on the first Mac: start it within five minutes of rank 1 +zig-out/native/bin/tensorfold-native serve ~/models/qwen3.8-27b \ + --speed-up rank0.json --port 4100 --parallel 1 --context 32768 --no-update-check +``` + +`--parallel 1` is required: a pair serves one greedy stream, the default is 8, and any other value stops with +`--speed-up serves one greedy stream: pass --parallel 1`. +`--context` defaults to 32768 for this model, and `--no-update-check` is the compatibility switch it is on one Mac, +so beyond `--parallel` and the settings file each command is what that Mac runs on its own. + +Each rank runs every layer and splits the even widths: 12 of the 24 attention heads, 2 of the 4 key/value heads, +half of the linear attention's heads and value channels, 8,704 of the 17,408-wide MLP, and 124,160 of the untied +lm_head's 248,320 rows. +Each projection's output rows are computed on the rank that owns them, with the full-K reduction a single Mac uses, so +no cross-rank partial sum exists. A seam's `pair_gather` carries the two bf16 halves, input activations and output rows, +bit-exactly, and the pair needs no tie policy. +Embeddings and norms stay whole on both Macs, so the 248,320-row table is not split and a token id means the same +thing on both ranks. +At each decided token the pair takes the argmax of the 248,320-row logits gathered from the two output halves, +exactly as one Mac does. + +The pair refuses what needs one Mac's whole state rather than quietly serving without it: `--drafter`, because rank 1 +has no drafter checkpoint, and `--prompt-cache-gib` and `--learn`, because each Mac holds only its own halves of a +slot's state. +A pair therefore keeps no prompt cache, learns nothing and drafts nothing. +A request that samples or carries a grammar fails too, since a pair serves one greedy stream. + +A pair's logits are byte-identical to one Mac's on the recorded fixtures, so its reply is a single Mac's token for +token. The pair-vs-single gate is `tp_check` (`tf-qwen27-tp-check`), which returns +`error.PairTokensDifferFromSingleMac` on any token difference, with no exemptions. + ## Weights other than 4-bit The M5 lane kernels accept MLX affine 2-, 3-, 4-, 5-, 6- and 8-bit projections in groups of 64. diff --git a/tools/qwen27/pair_accept.py b/tools/qwen27/pair_accept.py new file mode 100644 index 0000000000..e19f5971b9 --- /dev/null +++ b/tools/qwen27/pair_accept.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Accept a paired TensorFold serve against the checked-in single-host receipts. + +run send every oracle case twice; both replies must be identical, and their token_sha must equal the oracle's +warm one chat turn, then that conversation's next turn: the reply must report cached prompt tokens + +The oracle (qwen38_27b_oracle.json, next to this file) is the single-Mac, undrafted, temperature-0 record. +A pair passes when the same token ids come out, so neither Mac's rank count nor the seam can show up in the text. +Every case is compared to its oracle token_sha and must match: no case is exempt, because accepting a known +different value would hide exactly the pair-only divergence this run exists to catch. +""" +import argparse +import json +import sys +import urllib.error +import urllib.request + +DEFAULT_MODEL = "qwen3.8-27b" + + +def post(base, path, body, timeout): + request = urllib.request.Request( + base.rstrip("/") + path, + data=json.dumps(body).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read()) + + +def reply_text(reply, endpoint): + choice = reply["choices"][0] + if endpoint.endswith("/chat/completions"): + return choice.get("message", {}).get("content") or "" + return choice.get("text", "") + + +def token_sha(reply): + return reply.get("tensorfold", {}).get("token_sha") + + +def request_body(case, model): + body = { + "model": model, + "max_tokens": case["max_tokens"], + "temperature": 0.0, + "seed": 1234, + } + if case["endpoint"].endswith("/chat/completions"): + body["messages"] = [{"role": "user", "content": case["prompt"]}] + body["chat_template_kwargs"] = {"enable_thinking": bool(case.get("thinking", False))} + else: + body["prompt"] = case["prompt"] + return body + + +def load_oracle(path): + with open(path, "r", encoding="utf-8") as handle: + return json.load(handle) + + +def run(base, oracle, model, repeat, timeout, only): + cases = [c for c in oracle["cases"] if only is None or c["name"] in only] + if not cases: + print("pair_accept: no case matches --only", file=sys.stderr) + return 2 + if only is not None: + print("pair_accept: running {} of {} cases; --only is a subset, not a full acceptance".format(len(cases), len(oracle["cases"]))) + failures = 0 + for case in cases: + body = request_body(case, model) + shas = [] + texts = [] + for _ in range(repeat): + reply = post(base, case["endpoint"], body, timeout) + shas.append(token_sha(reply)) + texts.append(reply_text(reply, case["endpoint"])) + same = len(set(shas)) == 1 and len(set(texts)) == 1 + if same and shas[0] == case["token_sha"]: + print("PASS {} token_sha={} ({} identical runs)".format(case["name"], shas[0], repeat)) + else: + failures += 1 + if not same: + why = "the {} runs differ".format(repeat) + else: + why = "token_sha {} != oracle {}".format(shas[0], case["token_sha"]) + print("FAIL {} {} shas={}".format(case["name"], why, ",".join(str(s) for s in shas))) + total = len(cases) + print("pair_accept: {}/{} cases match the oracle on {}; {} did not".format(total - failures, total, base, failures)) + return 0 if failures == 0 else 1 + + +def cached_tokens(reply): + usage = reply.get("usage") or {} + details = usage.get("prompt_tokens_details") or {} + value = details.get("cached_tokens") + return value if value is not None else usage.get("cached_tokens") + + +def warm(base, oracle, model, timeout): + case = next((c for c in oracle["cases"] if c["name"].startswith("chat")), None) + if case is None: + print("pair_accept: the oracle has no chat case to build a conversation from", file=sys.stderr) + return 2 + first = post(base, case["endpoint"], request_body(case, model), timeout) + used = cached_tokens(first) + text = reply_text(first, case["endpoint"]) + if not text: + print("FAIL warm: the first turn returned no text", file=sys.stderr) + return 1 + second_body = { + "model": model, + "max_tokens": 32, + "temperature": 0.0, + "seed": 1234, + "messages": [ + {"role": "user", "content": case["prompt"]}, + {"role": "assistant", "content": text}, + {"role": "user", "content": "Now give the same example for a 3 by 3 matrix."}, + ], + "chat_template_kwargs": {"enable_thinking": False}, + } + second = post(base, case["endpoint"], second_body, timeout) + cached = cached_tokens(second) + print("warm: turn 1 prefill marked {}; turn 2 cached={}".format(used, cached)) + if cached is None or cached <= 0: + print("FAIL warm: the next turn reused no prompt state (cached={})".format(cached)) + return 1 + print("PASS warm: the next turn resumed {} cached prompt tokens".format(cached)) + return 0 + + +def main(): + parser = argparse.ArgumentParser(description="Accept a paired TensorFold serve against single-host receipts") + sub = parser.add_subparsers(dest="command", required=True) + for name in ("run", "warm"): + p = sub.add_parser(name) + p.add_argument("--base-url", default="http://127.0.0.1:4100") + p.add_argument("--oracle", default="tools/qwen27/qwen38_27b_oracle.json") + p.add_argument("--model", default=DEFAULT_MODEL) + p.add_argument("--timeout", type=float, default=600.0) + if name == "run": + p.add_argument("--repeat", type=int, default=2) + p.add_argument("--only", nargs="*", default=None) + args = parser.parse_args() + oracle = load_oracle(args.oracle) + try: + if args.command == "run": + return run(args.base_url, oracle, args.model, args.repeat, args.timeout, args.only) + return warm(args.base_url, oracle, args.model, args.timeout) + except urllib.error.URLError as err: + print("pair_accept: {} did not answer: {}".format(args.base_url, err), file=sys.stderr) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/qwen27/pair_supervise.sh b/tools/qwen27/pair_supervise.sh new file mode 100644 index 0000000000..1cb4d570ca --- /dev/null +++ b/tools/qwen27/pair_supervise.sh @@ -0,0 +1,317 @@ +#!/bin/sh +# Two-Mac paired serving, supervised: rank 1 first, rank 0 second, both restarted when either stops answering. +# Run it with no verb for the verbs and the environment variables it reads. +set -eu + +usage() { + cat <<'USAGE' +usage: pair_supervise.sh up|down|status|watch + + up start rank 1 on $PEER, then rank 0 here, and wait for both to answer /health + down stop both ranks this script started (the pids in the two pidfiles; never by command-line pattern) + status pids, ports and each rank's last log lines + watch up, then restart both whenever either rank stops answering (unattended) + +Nothing here names a machine or a user path: every location is environment overridable. +required PEER rank 1's ssh target (key auth: ssh -o BatchMode=yes must work) + MODEL the checkpoint both ranks load + RANK0_SETTINGS rank 0's --speed-up settings file + PEER_SETTINGS rank 1's --speed-up settings file + BIN rank 0's tensorfold-native + PEER_BIN rank 1's tensorfold-native (same bytes: the check refuses a mismatch) +optional PORT0=4100 PORT1=4101 CONTEXT=32768 KEEP_WARM=600 SERVE_FLAGS='' SERVE_ENVS='' + STATE=${TMPDIR:-/tmp}/tensorfold-pair local logs, pidfile and the port count + PEER_STATE=${TMPDIR:-/tmp}/tensorfold-pair the same on rank 1's Mac, in that Mac's own TMPDIR + PEER_MODEL=$MODEL SSH_OPTS='-o BatchMode=yes -o ConnectTimeout=5' + POLL=2 seconds between liveness probes and the first restart delay + STABLE=60 seconds a pair must stay up before the restart delay drops back to POLL + STOP_WAIT=12 seconds a rank may take to exit after the stop signal, before it is killed + READY_TIMEOUT=300 seconds either rank may take from process start to answering /health + ALLOW_MISMATCH=1 start even though the two binaries differ (only for a deliberate A/B test) + +The peer's host key is verified: SSH_OPTS keeps BatchMode=yes and the default asks the known-hosts file, so an +unknown or changed key fails the connection instead of being accepted. Set SSH_OPTS yourself to override that +policy deliberately; nothing here turns the check off for you. +USAGE +} + + +PORT0=${PORT0:-4100} +PORT1=${PORT1:-4101} +CONTEXT=${CONTEXT:-32768} +KEEP_WARM=${KEEP_WARM:-600} +POLL=${POLL:-2} +STABLE=${STABLE:-60} +STOP_WAIT=${STOP_WAIT:-12} +READY_TIMEOUT=${READY_TIMEOUT:-300} +STATE=${STATE:-${TMPDIR:-/tmp}/tensorfold-pair} +PEER_STATE=${PEER_STATE:-${TMPDIR:-/tmp}/tensorfold-pair} +SSH_OPTS=${SSH_OPTS:--o BatchMode=yes -o ConnectTimeout=5} +SERVE_FLAGS=${SERVE_FLAGS:-} +SERVE_ENVS=${SERVE_ENVS:-} +PEER_MODEL=${PEER_MODEL:-${MODEL:-}} +PEER=${PEER:-} +MODEL=${MODEL:-} +RANK0_SETTINGS=${RANK0_SETTINGS:-} +PEER_SETTINGS=${PEER_SETTINGS:-} +BIN=${BIN:-} +PEER_BIN=${PEER_BIN:-${BIN}} +ALLOW_MISMATCH=${ALLOW_MISMATCH:-0} + +LOG0=$STATE/rank0.log +PID0=$STATE/rank0.pid +LOG1=$PEER_STATE/rank1.log +PID1=$PEER_STATE/rank1.pid + +say() { printf '[pair] %s\n' "$*"; } + +peer() { ssh -n $SSH_OPTS "$PEER" "$@"; } + +now() { date +%s; } + +# /health is unauthenticated and answers only once the rank is serving, which both ranks reach after the handshake. +health0() { curl -fsS -m 2 -o /dev/null "http://127.0.0.1:$PORT0/health" 2>/dev/null; } +health1() { peer "curl -fsS -m 2 -o /dev/null http://127.0.0.1:$PORT1/health" 2>/dev/null; } + +# A pidfile holds the pid and the binary path this script started, so a stop works on a host where MODEL, settings +# or even BIN are no longer in the environment, and it can tell the rank apart from a later process the kernel handed +# the same pid. The pid is the first word; everything after it is the binary path. +local_pid() { [ -f "$PID0" ] && cut -d' ' -f1 "$PID0" 2>/dev/null || true; } +local_bin() { [ -f "$PID0" ] && cut -d' ' -f2- "$PID0" 2>/dev/null || true; } +peer_pid() { peer "cut -d' ' -f1 '$PID1' 2>/dev/null" 2>/dev/null || true; } +peer_bin() { peer "cut -d' ' -f2- '$PID1' 2>/dev/null" 2>/dev/null || true; } + +alive0() { p=$(local_pid); [ -n "$p" ] && kill -0 "$p" 2>/dev/null; } +alive1() { peer "kill -0 \$(cut -d' ' -f1 '$PID1' 2>/dev/null) 2>/dev/null" 2>/dev/null; } + +need_peer() { + [ -n "$PEER" ] || { say "set PEER to the rank 1 ssh target"; return 1; } +} + +# Only a start needs the checkpoint and the two settings files: stopping a pair must work on a host where the model +# has been moved or was never mounted, which is exactly when someone runs 'down'. +need_up() { + need_peer || return 1 + [ -n "$MODEL" ] || { say "set MODEL to the checkpoint directory"; return 1; } + [ -n "$RANK0_SETTINGS" ] || { say "set RANK0_SETTINGS to rank 0's settings file"; return 1; } + [ -n "$PEER_SETTINGS" ] || { say "set PEER_SETTINGS to rank 1's settings file"; return 1; } + [ -n "$BIN" ] || { say "set BIN to rank 0's tensorfold-native (PEER_BIN defaults to it)"; return 1; } + [ -n "$PEER_BIN" ] || { say "set PEER_BIN to rank 1's tensorfold-native"; return 1; } +} + +check_binaries() { + if [ "$ALLOW_MISMATCH" = 1 ]; then say "binary check skipped (ALLOW_MISMATCH=1)"; return 0; fi + local here there + here=$(shasum -a 256 "$BIN" | cut -d' ' -f1) + there=$(peer "shasum -a 256 '$PEER_BIN' | cut -d' ' -f1" 2>/dev/null || echo missing) + if [ "$here" != "$there" ]; then + say "the two ranks are not the same build: $BIN $here, $PEER_BIN $there" + say "fix: copy the same binary to both Macs, or set ALLOW_MISMATCH=1 for a deliberate A/B test" + return 1 + fi + say "binaries match: $(printf %s "$here" | cut -c1-12)" +} + +# Rank 1 waits for rank 0 inside the handshake, so it cannot answer /health first: start it, let it reach that wait. +start_rank1() { + peer "mkdir -p '$PEER_STATE'" + say "starting rank 1 on $PEER (it waits for rank 0 in the pair handshake)" + # The background subshell execs the rank, so the pid written to the pidfile is the rank itself and not a launcher + # shell that would die first and leave the rank running: 'down' signals the recorded pid and nothing else. + peer "( cd '$PEER_STATE' && exec nohup env $SERVE_ENVS '$PEER_BIN' serve '$PEER_MODEL' --speed-up '$PEER_SETTINGS'" \ + "--port $PORT1 --context $CONTEXT --keep-warm $KEEP_WARM --parallel 1 --no-update-check $SERVE_FLAGS" \ + ">$LOG1 2>&1 ) & echo \$! '$PEER_BIN' > '$PID1'" + sleep 2 + if ! alive1; then + say "rank 1 exited at once: last lines of $PEER:$LOG1" + peer "tail -n 5 '$LOG1' 2>/dev/null" || true + return 1 + fi + say "rank 1 is up (pid $(peer_pid)), waiting for the handshake" +} + +start_rank0() { + say "starting rank 0 here" + cd "$STATE" + # The launcher is a subshell so the rank's death is never this shell's job to report: a supervisor log stays one + # line per event instead of one per killed child. It execs the rank for the same reason as rank 1: the recorded pid + # must be the rank itself, so a stop signals exactly the process this script started. + pid=$( + ( exec nohup env $SERVE_ENVS "$BIN" serve "$MODEL" --speed-up "$RANK0_SETTINGS" --port "$PORT0" \ + --context "$CONTEXT" --keep-warm "$KEEP_WARM" --parallel 1 --no-update-check $SERVE_FLAGS >"$LOG0" 2>&1 ) & + echo $! + ) + echo "$pid $BIN" >"$PID0" + say "rank 0 pid $(local_pid)" +} + +# Poll until this rank answers /health; the ports come up only after the pair handshake, so the count of refusals +# is the readiness delay, not a startup race the client can lose a request to. +wait_serving() { + start=$(now) + refused=0 + while :; do + if health0 && health1; then + say "both ranks answer: rank 0 on :$PORT0 and rank 1 on :$PORT1, $(( $(now) - start ))s after rank 0 started" + [ "$refused" -gt 0 ] && say "rank 0's port refused $refused probes while the pair handshake completed" + return 0 + fi + if ! alive0; then + say "rank 0 exited before it served: last lines of $LOG0" + tail -n 8 "$LOG0" 2>/dev/null || true + return 1 + fi + if ! alive1; then + say "rank 1 exited before it served: last lines of $PEER:$LOG1" + peer "tail -n 8 '$LOG1' 2>/dev/null" || true + return 1 + fi + refused=$((refused + 1)) + if [ $(( $(now) - start )) -ge "$READY_TIMEOUT" ]; then + say "no answer after ${READY_TIMEOUT}s: last lines of each log" + tail -n 8 "$LOG0" 2>/dev/null || true + peer "tail -n 8 '$LOG1' 2>/dev/null" || true + return 1 + fi + sleep 1 + done +} + +up() { + need_up || return 1 + mkdir -p "$STATE" + command -v curl >/dev/null || { say "curl is required for the readiness probe"; return 1; } + if health0 && health1; then say "already serving"; return 0; fi + check_binaries + down + start_rank1 + start_rank0 + wait_serving +} + +# A stop is TERM, then KILL if the rank is still there: a draining rank exits in a second, and a wedged one must not +# hold a restart. kill() returns 0 for a process that is still shutting down, so each stop waits the pid out first. +# The pidfile can outlive its rank, so the recorded binary is checked first: a pid the kernel has reused is not ours +# and is never signalled. +stop_local() { + pid=$1 + bin=$2 + if [ -n "$bin" ]; then + cmd=$(ps -ww -p "$pid" -o command= 2>/dev/null) || cmd= + if [ -z "$cmd" ]; then + say "rank 0 (pid $pid) is no longer running" + return 0 + fi + case $cmd in + *"$bin"*) ;; + *) say "rank 0: pidfile names pid $pid, whose command is not $bin: not signalling it"; return 0 ;; + esac + fi + kill "$pid" 2>/dev/null || { say "rank 0 (pid $pid) is already gone"; return 0; } + i=0 + while kill -0 "$pid" 2>/dev/null && [ "$i" -lt $((STOP_WAIT * 5)) ]; do sleep 0.2; i=$((i + 1)); done + if kill -0 "$pid" 2>/dev/null; then + say "rank 0 did not stop within ${STOP_WAIT}s: killing it (pid $pid)" + kill -9 "$pid" 2>/dev/null || true + i=0 + while kill -0 "$pid" 2>/dev/null && [ "$i" -lt 25 ]; do sleep 0.2; i=$((i + 1)); done + fi + say "stopped rank 0 (pid $pid)" +} + +# The peer runs its own stop loop so the signals and the waits all happen on that Mac. Only the pid in the pidfile is +# a candidate, after the recorded binary is checked the same way as rank 0: a stale pidfile or a reused pid means no +# signal. The pidfile is removed whether or not anything was signalled, so a second stop is a no-op. +stop_peer() { + q=$(peer_pid) + if [ -z "$q" ]; then + say "rank 1: no pidfile at $PEER:$PID1 (already stopped or never started)" + peer "rm -f '$PID1'" 2>/dev/null || true + return 0 + fi + b=$(peer_bin) + [ -n "$b" ] || b=$PEER_BIN + if [ -n "$b" ]; then + cmd=$(peer "ps -ww -p '$q' -o command= 2>/dev/null" 2>/dev/null || true) + if [ -z "$cmd" ]; then + say "rank 1 (pid $q) is no longer running" + peer "rm -f '$PID1'" 2>/dev/null || true + return 0 + fi + case $cmd in + *"$b"*) ;; + *) say "rank 1: pidfile names pid $q, whose command is not $b: not signalling it" + peer "rm -f '$PID1'" 2>/dev/null || true + return 0 ;; + esac + fi + say "stopping rank 1 on $PEER (pid $q)" + peer "kill '$q' 2>/dev/null || { echo '[pair] rank 1 (pid $q) is already gone'; exit 0; };" \ + "i=0; while kill -0 '$q' 2>/dev/null && [ \$i -lt $((STOP_WAIT * 5)) ]; do sleep 0.2; i=\$((i + 1)); done;" \ + "if kill -0 '$q' 2>/dev/null; then echo '[pair] rank 1 will not stop: killing it (pid $q)';" \ + "kill -9 '$q' 2>/dev/null; sleep 1; fi;" \ + "echo '[pair] stopped rank 1 (pid $q)'" 2>/dev/null || true + peer "rm -f '$PID1'" 2>/dev/null || true +} + +# down stops only what this script recorded: the two pidfiles, each checked against the binary written beside the +# pid. A missing or stale pidfile is reported and forgotten, never followed by a search for likely-looking processes. +down() { + need_peer || return 1 + p=$(local_pid) + if [ -n "$p" ]; then + stop_local "$p" "$(local_bin)" + else + say "rank 0: no pidfile at $PID0 (already stopped or never started)" + fi + rm -f "$PID0" + stop_peer + return 0 +} + +status() { + need_peer || return 1 + if health0; then say "rank 0: serving on :$PORT0 (pid $(local_pid))"; else say "rank 0: not answering on :$PORT0"; fi + if health1; then say "rank 1: serving on :$PORT1 (pid $(peer_pid))"; else say "rank 1: not answering on :$PORT1"; fi + say "last rank 0 lines: $LOG0" + tail -n 3 "$LOG0" 2>/dev/null || true + say "last rank 1 lines: $PEER:$LOG1" + peer "tail -n 3 '$LOG1' 2>/dev/null" || true +} + +# One supervisor tick: a rank that stopped answering is stopped on both Macs and started again, so a dead peer +# never leaves a half pair up and no human has to notice. The wait between attempts doubles while the pair keeps +# failing, and drops back to POLL only after a run that stayed up for STABLE seconds: a crash loop must not spin. +watch() { + trap 'say "supervisor stopping: stopping both ranks"; down; exit 0' INT TERM + restarts=0 + delay=$POLL + while :; do + if up; then + began=$(now) + while health0 && health1; do + sleep "$POLL" + done + restarts=$((restarts + 1)) + if ! health0; then say "rank 0 stopped answering: restarting both (restart $restarts)"; fi + if ! health1; then say "rank 1 stopped answering: restarting both (restart $restarts)"; fi + if [ $(( $(now) - began )) -ge "$STABLE" ]; then delay=$POLL; fi + else + restarts=$((restarts + 1)) + say "the pair did not come up: restarting both (restart $restarts)" + fi + down + say "next attempt in ${delay}s" + sleep "$delay" + delay=$((delay * 2)) + [ "$delay" -le 60 ] || delay=60 + done +} + +case ${1:-usage} in + up) up ;; + down) down ;; + status) status ;; + watch) watch ;; + *) usage; exit 2 ;; +esac diff --git a/tools/qwen27/qwen38_27b_oracle.json b/tools/qwen27/qwen38_27b_oracle.json new file mode 100644 index 0000000000..84396e0a33 --- /dev/null +++ b/tools/qwen27/qwen38_27b_oracle.json @@ -0,0 +1,44 @@ +{ + "description": "Single-host Qwen3.8-27B 4-bit receipts: one Mac, no drafter, temperature 0, seed 1234. A paired run must reproduce every token_sha here.", + "runtime": "tensorfold-native 1.1.2", + "checkpoint_revision": "22d8d538154e0e5b6f8dcb1fcc74a60a0798104c", + "quantization": "4-bit affine, group 64", + "decoding": { + "temperature": 0.0, + "seed": 1234, + "drafter": false, + "parallel": 1 + }, + "cases": [ + { + "name": "code-128", + "endpoint": "/v1/completions", + "prompt": "Write a short Python function that computes the Fibonacci sequence and explain it.", + "max_tokens": 128, + "token_sha": "f5c9409cfe40" + }, + { + "name": "code-256", + "endpoint": "/v1/completions", + "prompt": "Write a short Python function that computes the Fibonacci sequence and explain it.", + "max_tokens": 256, + "token_sha": "2b28aa2108de" + }, + { + "name": "chat-128", + "endpoint": "/v1/chat/completions", + "thinking": false, + "prompt": "Explain how matrix multiplication uses a GPU in plain English, then give a small numerical example.", + "max_tokens": 128, + "token_sha": "b69b0269917f" + }, + { + "name": "chat-256", + "endpoint": "/v1/chat/completions", + "thinking": false, + "prompt": "Explain how matrix multiplication uses a GPU in plain English, then give a small numerical example.", + "max_tokens": 256, + "token_sha": "a06fab348fbe" + } + ] +} diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 841e4c38ff..8be682befa 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -107,6 +107,9 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const family: []const std.Build.Module.Import = &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal } }; _ = program(b, "tf-qwen27-run", "Build the native greedy, chunk and prompt-speed check on real weights", "zig/tests/qwen27/run.zig", target, true, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal }, .{ .name = "core_lane", .module = core_lane }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources }, .{ .name = "qwen_gdn_source", .module = s.gdn_source } }); + // the TP gate before any wire: ranks = 1 parity, then two rank-local shards through the loopback link + const tp_check = program(b, "tf-qwen27-tp-check", "The 27B TP gate: ranks = 1 parity, then the two-rank loopback stream (SKIP without a checkpoint)", "zig/tests/qwen27/tp_check.zig", target, true, family); + test_step.dependOn(&b.addRunArtifact(tp_check).step); _ = program(b, "tf-qwen27-attention-fused", "Build the fused-chain versus tail-kernel attention byte check on this Mac's tensor units, no models", "zig/tests/qwen27/attention_fused.zig", target, false, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources } }); _ = program(b, "tf-qwen27-attention-precision", "Build decode attention's error against a float64 reference on this Mac, no models", "zig/tests/qwen27/attention_precision.zig", target, false, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources } }); _ = program(b, "tf-qwen27-projection-cases", "Build real-input production projection precision replay", "zig/tests/qwen27/projection_cases.zig", target, true, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal } }); diff --git a/zig/build/qwen27_host.zig b/zig/build/qwen27_host.zig index 7ef1d6d56d..0eef06db25 100644 --- a/zig/build/qwen27_host.zig +++ b/zig/build/qwen27_host.zig @@ -42,9 +42,14 @@ pub fn tests(b: *std.Build) void { b.step("test-qwen27-equality-host", "Shared lanes drafted, resumed and concurrent equality on the fake target").dependOn(&equal.step); } -pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.Module, lanes: *std.Build.Module) void { +pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Step { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const host = b.createModule(.{ .root_source_file = b.path("zig/src/native/qwen27_served.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "engine_api", .module = api } } }); + // pair refusals are decided before a model loads: the host's own tests run on CPU alone + const served = b.step("test-qwen27-served-host", "Rank 1's named refusal and the pair's --vision refusal, no shaders or GPU"); + const served_run = b.addRunArtifact(b.addTest(.{ .root_module = host })); + served_run.has_side_effects = true; // a rerun must re-execute the tests, not replay a cached pass + served.dependOn(&served_run.step); const boundary = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/qwen27_bank_boundary_host.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "tensorfold", .module = engine }} }), .filters = &.{"qwen27 bank boundary:"} }); b.step("test-qwen27-bank-boundary-host", "Runner pending-keep guards with no device or shader execution").dependOn(&b.addRunArtifact(boundary).step); const object = b.addObject(.{ .name = "qwen27-speed-compile", .root_module = b.createModule(.{ @@ -56,6 +61,7 @@ pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.M }) }); _ = object.getEmittedBin(); b.step("check-qwen27-speed-host", "Compile the actual Qwen27 startup and encoding without shaders").dependOn(&object.step); + return served; } fn tryNames(b: *std.Build, names: *std.ArrayList([]const u8), text: []const u8) void { diff --git a/zig/cost_rule_cases.zig b/zig/cost_rule_cases.zig index 91cd4fb9d4..42b3f715ba 100644 --- a/zig/cost_rule_cases.zig +++ b/zig/cost_rule_cases.zig @@ -16,7 +16,9 @@ test "Nemotron initial and drift passes share interleaved head sampling; the 27B try std.testing.expect(std.mem.indexOf(u8, body, "if (b.head != null) &head[0] else null") != null); try std.testing.expect(std.mem.indexOf(u8, body, "drifted(&head, &r.head)") != null); try std.testing.expect(std.mem.indexOf(u8, body, "fastest(heads") == null); + const warm = @embedFile("src/native/qwen27_warm.zig"); // the window timing moved here from the lane backend + try std.testing.expect(std.mem.indexOf(u8, warm, "cost_rule.fastest(timer, i)") != null); + try std.testing.expect(std.mem.indexOf(u8, warm, "cost_rule.measure(") == null); const qwen = @embedFile("src/native/qwen27_lanes.zig"); - try std.testing.expect(std.mem.indexOf(u8, qwen, "cost_rule.fastest(timer, i)") != null); try std.testing.expect(std.mem.indexOf(u8, qwen, "cost_rule.measure(") == null); } diff --git a/zig/kernels/metal/core/pair_rows.metal b/zig/kernels/metal/core/pair_rows.metal new file mode 100644 index 0000000000..2170b489a3 --- /dev/null +++ b/zig/kernels/metal/core/pair_rows.metal @@ -0,0 +1,150 @@ +// Two Macs' tensor-parallel exchanges: the host relays parts through shared events; these kernels only do arithmetic. +#include +using namespace metal; + +// Sync page words (shared with pair_rows.zig): FAILED = the host gave up on a seam; OBSERVED = this side's refusals. +constant constexpr uint FAILED = 16, OBSERVED = 20; + +// True once the host gave up: FAILED is set before the wait releases, so a later kernel poisons its output rows. +inline bool pair_link_broken(device atomic_uint* sync, uint t) { + if (atomic_load_explicit(&sync[FAILED], memory_order_relaxed) == 0u) return false; + if (t == 0) atomic_fetch_add_explicit(&sync[OBSERVED], 1u, memory_order_relaxed); + return true; +} + +inline float pair_lanes(float v) { + for (ushort d = 16; d > 0; d >>= 1) v += simd_shuffle_xor(v, d); + return v; +} + +// Kernels' argument block, byte-identical to pair_rows.zig's PairRows: four u32 then four f32; no payload element. +struct PairRows { + uint width, stride, rank, x; + float eps, pad0, pad1, pad2; +}; + +// bf16 bits as the exact fp32 value they name (argmax reads a bf16 logits stream this way). +inline float pair_bf16(uint bits) { + return as_type(bits << 16); +} + +// Row r: h += rank 0's part + rank 1's part, then y = w * (h / rms(h)); payload is fp32 (streamBytes). +kernel void pair_add_rms(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device float* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device float* h [[buffer(4)]], const device bfloat* w [[buffer(5)]], + device float* y [[buffer(6)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) y[ulong(r) * a.stride + i] = NAN; + return; + } + device float* x = h + ulong(r) * a.stride; + float s = 0.0f; + for (uint i = t; i < a.width; i += 256) { + const ulong at = ulong(r) * a.width + i; + const float other = as_type(atomic_load_explicit(&peer[at], memory_order_relaxed)); + const float sum = a.rank == 0 ? mine[at] + other : other + mine[at]; + const float v = x[i] + sum; + x[i] = v; + s = fma(v, v, s); + } + s = pair_lanes(s); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + float var = 0.0f; + for (uint g = 0; g < 8; ++g) var += part[g]; + const float rs = precise::rsqrt(var / float(a.width) + a.eps); + for (uint i = t; i < a.width; i += 256) y[ulong(r) * a.stride + i] = float(w[i]) * (x[i] * rs); +} + +// Row r: this rank's own half of the frame row into the exchange slot; a bf16 copy, so it carries the frame's bits. +kernel void pair_stage(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* src [[buffer(2)]], device ushort* dst [[buffer(3)]], + uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + const ulong in = ulong(r) * a.stride; + const ulong out = ulong(r) * a.width; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) dst[out + i] = 0x7FC0u; // bf16 NaN: the link stays down + return; + } + for (uint i = t; i < a.width; i += 256) dst[out + i] = src[in + i]; +} + +// Row r: this rank's half at its own offset in the frame row and the peer's at the other; a bf16 copy, no arithmetic. +kernel void pair_gather(device atomic_uint* sync [[buffer(0)]], constant PairRows& a [[buffer(1)]], + const device ushort* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device ushort* out [[buffer(4)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + const ulong mrow = ulong(r) * a.width; + const ulong orow = ulong(r) * a.stride; + const ulong own = a.rank == 0 ? orow : orow + a.width; + const ulong got = a.rank == 0 ? orow + a.width : orow; + if (pair_link_broken(sync, t)) { + for (uint i = t; i < a.width; i += 256) { + out[own + i] = 0x7FC0u; + out[got + i] = 0x7FC0u; + } + return; + } + for (uint i = t; i < a.width; i += 256) { + // The relay wrote the peer's half: an untracked frame read, so one atomic word (two bf16) at a time. + const uint word = atomic_load_explicit(&peer[(mrow + i) >> 1], memory_order_relaxed); + out[own + i] = mine[mrow + i]; + out[got + i] = ushort(word >> ((i & 1) * 16)); + } +} + +// Row r's argmax over this rank's logits (ids lo..lo+n): (fp32 value, id), lower id on tie; dims.w 0 = fp32, 1 = bf16. +kernel void pair_argmax(const device float* logits [[buffer(0)]], constant uint4& dims [[buffer(1)]], + device uint* out [[buffer(2)]], uint r [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const ulong row = ulong(r) * dims.x; + const device float* xf = logits + row; + const device ushort* xb = reinterpret_cast(logits) + row; + const bool wide = dims.w != 0u; + float best = -INFINITY; + uint at = 0xFFFFFFFFu; + for (uint i = t; i < dims.z; i += 1024) { + const float v = wide ? pair_bf16(uint(xb[i])) : xf[i]; + if (v > best) best = v, at = dims.y + i; + } + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + } + threadgroup float vb[32]; + threadgroup uint va[32]; + if (lane == 0) vb[sg] = best, va[sg] = at; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg != 0) return; + best = vb[lane]; + at = va[lane]; + for (ushort o = 16; o > 0; o >>= 1) { + const float ob = simd_shuffle_xor(best, o); + const uint oa = simd_shuffle_xor(at, o); + if (ob > best || (ob == best && oa < at)) best = ob, at = oa; + } + if (lane == 0) out[2 * r] = as_type(best), out[2 * r + 1] = at; +} + +// Each row's pick over both Macs' argmaxes: the larger value, the lower id on a tie (one Mac's argmax). +kernel void pair_pick(device atomic_uint* sync [[buffer(0)]], constant uint2& rx [[buffer(1)]], + const device uint* mine [[buffer(2)]], device atomic_uint* peer [[buffer(3)]], + device uint* picks [[buffer(4)]], uint t [[thread_index_in_threadgroup]]) { + if (pair_link_broken(sync, t)) { + if (t == 0) for (uint r = 0; r < rx.x; ++r) picks[r] = 0xFFFFFFFFu; // no such token: visibly wrong, never a token + return; + } + if (t != 0) return; + for (uint r = 0; r < rx.x; ++r) { + const float vm = as_type(mine[2 * r]); + const uint am = mine[2 * r + 1]; + const float vp = as_type(atomic_load_explicit(&peer[2 * r], memory_order_relaxed)); + const uint ap = atomic_load_explicit(&peer[2 * r + 1], memory_order_relaxed); + picks[r] = (vp > vm || (vp == vm && ap < am)) ? ap : am; + } +} diff --git a/zig/kernels/metal/qwen27/attention_prompt.metal b/zig/kernels/metal/qwen27/attention_prompt.metal index 8ba07e0501..9a90453b56 100644 --- a/zig/kernels/metal/qwen27/attention_prompt.metal +++ b/zig/kernels/metal/qwen27/attention_prompt.metal @@ -4,7 +4,8 @@ using namespace metal; using namespace mpp::tensor_ops; constant constexpr int PG = 6, PD = 256, PH = 128, PCK = 512, PTK = 32; -struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; }; +// kv_heads is the rank's, not the model's: a split attention writes local node rows. +struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; int kv_heads; }; typedef vec q27f8; typedef vec q27b8; typedef vec q27h8; @@ -166,5 +167,6 @@ kernel void q27_attn_prompt_merge(device const float* PO [[buffer(0)]], device c m = nm; } const int node = r / PG, head = hk * PG + r % PG; - for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * 4) + head) * PD + lane * DP + i] = static_cast(o[i] / l); + const int64_t at = ((int64_t)node * (PG * args.kv_heads) + head) * PD + lane * DP; + for (int i = 0; i < DP; i++) OUT[at + i] = static_cast(o[i] / l); } diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 51b14b5a40..402d7dd2f4 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -34,6 +34,8 @@ pub const core_exlr8_dense = @embedFile("core/exlr8_dense.metal"); pub const core_exlr8_prompt = @embedFile("core/exlr8_prompt.metal"); pub const core_exlr8_dense_prompt = @embedFile("core/exlr8_dense_prompt.metal"); pub const core_pair = @embedFile("core/pair.metal"); +/// The row-split seam's kernels (see `pair_rows.zig`): `pair.metal` above stays as it is for glm_dsa. +pub const core_pair_rows = @embedFile("core/pair_rows.metal"); /// GLM-5.3 (glm_moe_dsa): MLA with rope over a bf16 latent cache, the DSA indexer and its top-k, fp32 glue. pub const glm_dsa_mla = @embedFile("glm_dsa/mla.metal"); diff --git a/zig/qwen27_core.zig b/zig/qwen27_core.zig index ebd50da149..56997e711c 100644 --- a/zig/qwen27_core.zig +++ b/zig/qwen27_core.zig @@ -33,3 +33,6 @@ pub const bf16_topk_gpu = @import("bf16_topk_gpu"); pub const tree_commit_gpu = @import("tree_commit_gpu"); pub const tree_slots = @import("src/core/tree_slots.zig"); + +/// The peer window and its seam kernel: the 27B's tensor-parallel exchange, and its one-Mac loopback. +pub const pair = @import("src/core/pair_rows.zig"); diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index b38e46eccd..ba65cf7d2e 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -19,6 +19,7 @@ test { _ = @import("contract/vit_records.zig"); _ = @import("contract/simd_pairs_records.zig"); _ = @import("contract/glm_rows_records.zig"); + _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; _ = rows; diff --git a/zig/src/core/contract/metal_pair_records.zig b/zig/src/core/contract/metal_pair_records.zig new file mode 100644 index 0000000000..66717449de --- /dev/null +++ b/zig/src/core/contract/metal_pair_records.zig @@ -0,0 +1,26 @@ +//! The pair seam's five launches (core/pair_rows.metal) with the M1/M2 limits the offline compile measures. +const std = @import("std"); +const limits = @import("limits.zig"); +const registry = @import("registry.zig"); + +pub const records = [_]limits.MetalLaunch{ + .{ .name = "metal pair add_rms pair_add_rms", .threads = 256, .tg_memory = 32, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair stage pair_stage", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair gather pair_gather", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair argmax pair_argmax", .threads = 1024, .tg_memory = 256, .m1 = 1024, .m2 = 1024 }, + .{ .name = "metal pair pick pair_pick", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, +}; + +test "the pair seam's five launches are in the live and offline collector and fit every family" { + for (records) |record| { + var found: usize = 0; + for (registry.all_metal) |active| if (std.mem.eql(u8, active.name, record.name)) { + try std.testing.expectEqualDeep(record, active); + found += 1; + }; + try std.testing.expectEqual(@as(usize, 1), found); + try std.testing.expectEqual(@as(usize, 0), try registry.checkMetal(std.testing.allocator, &.{record}, true)); + // A zero here would be a launch no M1/M2 compiler was asked about. + try std.testing.expect(record.m1 > 0 and record.m2 > 0); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 8d5e1ee270..0072127c61 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -570,7 +570,7 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal; +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/src/core/pair_rows.zig b/zig/src/core/pair_rows.zig new file mode 100644 index 0000000000..2fd8c5005c --- /dev/null +++ b/zig/src/core/pair_rows.zig @@ -0,0 +1,488 @@ +//! Two Macs' exchanges: GPU part, posted(x), host MCDMA relay, landed(x); the event wait is a real scheduler hand-off. +const std = @import("std"); +const platform = @import("sys"); +const mtl = @import("metal"); +const fabric = @import("fabric"); +const sources = @import("kernel_sources"); +const settings = @import("../families/flashnext/tp_settings.zig"); +const Fence = @import("fence.zig").Fence; +const Owner = @import("../families/flashnext/owned.zig").Owner; +const pair_layout = @import("pair_rows_layout.zig"); +const pair_loopback = @import("pair_rows_loopback.zig"); + +pub const Settings = settings.Settings; +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const command_words = pair_layout.command_words; +pub const Slot = pair_layout.Slot; +pub const seam_failed = pair_layout.seam_failed; +pub const Seam = pair_layout.Seam; +pub const Elem = pair_layout.Elem; +pub const elem_bytes = pair_layout.elem_bytes; +pub const elem_name = pair_layout.elem_name; +pub const elems_per_word = pair_layout.elems_per_word; +pub const partBytes = pair_layout.partBytes; +pub const stageBytes = pair_layout.stageBytes; +pub const gatherBytes = pair_layout.gatherBytes; +pub const pickBytes = pair_layout.pickBytes; +pub const streamBytes = pair_layout.streamBytes; +pub const frameBytes = pair_layout.frameBytes; +pub const toFrame = pair_layout.toFrame; +pub const fromElem = pair_layout.fromElem; +pub const toElem = pair_layout.toElem; +pub const seedValue = pair_layout.seedValue; +pub const Loopback = pair_loopback.Loopback; + +const PAGE = pair_layout.PAGE; +const pages = pair_layout.pages; +const CMD_BYTES = pair_layout.CMD_BYTES; +const JOBS = pair_layout.JOBS; +const slotOf = pair_layout.slotOf; +const Layout = pair_layout.Layout; +const FLAG = pair_layout.FLAG; +const HELLO = pair_layout.HELLO; +const CMD_FLAG = pair_layout.CMD_FLAG; +const CMD_ACK = pair_layout.CMD_ACK; +const READY = pair_layout.READY; +const BEAT = pair_layout.BEAT; +const POSTED = pair_layout.POSTED; +const FAILED = pair_layout.FAILED; +const OBSERVED = pair_layout.OBSERVED; +const writeElem = pair_layout.writeElem; +const PairRows = pair_layout.PairRows; +const seam = pair_layout.seam; + +/// Mach's clock: 24 MHz, so a millisecond is 24,000 ticks -- the clock the probes' seam numbers are in. +const ticks_per_ms: u64 = 24_000; + +/// How long the host waits for the GPU to reach a published seam: a hung GPU must fail the request, not wedge it. +const post_timeout_ms: u64 = 5_000; +/// The peer's part wait: 250 ms, 300x the measured tail (p99 ~700 us), so a dead peer fails fast. +const peer_timeout_ms: u64 = 250; +/// The cold-link peer wait: the first request pays one-time Metal pipelines (1.07 s measured, 0.14 s warm). +const cold_peer_timeout_ms: u64 = 5_000; +/// How long after the link's first relay the cold bound applies: 8 s covers a longer first prompt (prefill chunks). +const cold_window_ms: u64 = 8_000; + +/// How long rank 1 waits for rank 0's acknowledgement of a replay: a stuck peer fails the request, it is not wedged. +const cmd_ack_timeout_ms: u64 = 60_000; + +/// How long either rank waits for the other's ready word: a cold 27B load on the far Mac, five minutes at most. +const ready_timeout_ms: u64 = 300_000; + +/// The close's last flush: microseconds when the peer is there, and never the library's whole timeout when it is gone. +const last_flush_ns: u64 = 2 * std.time.ns_per_s; + +threadlocal var raised = false; + +/// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. +pub fn raise() void { + if (raised) return; + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + raised = true; +} + +/// Whether 32-bit counter a has reached b across the wrap. +fn reached(a: u32, b: u32) bool { + return @as(i32, @bitCast(a -% b)) >= 0; +} + +fn ticksNow() u64 { + return platform.c.mach_absolute_time(); +} + +/// A short sleep, for the host threads' waits that are not on the seam's critical path. +fn sleepUs(us: u64) void { + const ts: platform.c.timespec = .{ .sec = @intCast(us / 1_000_000), .nsec = @intCast(us % 1_000_000 * 1000) }; + _ = platform.c.nanosleep(&ts, null); +} + +pub const Pair = struct { + rank: u32, + peer: u32, + l: Layout, + ep: *fabric.mcdma.Endpoint, + rd: fabric.rdma.Rdma, + win: []u8, + wbuf: mtl.Buffer, + add_rms: mtl.Pipeline, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Orders the window's untracked writes across the encoders a seam splits. + fence: Fence, + /// GPU -> host: seam x's part is written, relay it. Host -> GPU: the peer's part for x is in place, add it. + posted: mtl.SharedEvent, + landed: mtl.SharedEvent, + /// Exchanges encoded so far: the same count on both Macs. + x: u64 = 0, + bytes: []u32, + queued: std.atomic.Value(u32) = .init(0), + thread: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + broken: std.atomic.Value(bool) = .init(false), + /// The exchange the service thread is on: the round number in a failure line. + seam_x: std.atomic.Value(u32) = .init(0), + /// The request this pair is on, for failure lines; `beginRound` bumps it -- a label only, failures are permanent. + round: std.atomic.Value(u32) = .init(0), + /// Commands sent (rank 0) or taken (rank 1) so far. + cmd: u64 = 0, + /// Beats this Mac's watchdog has signalled to the peer: its counter for the peer's silence check. + beats: u64 = 0, + /// The GPU objects this pair holds: released in `deinit`, the window's buffer before the memory behind it. + owned: Owner(mtl) = .{}, + /// The allocator `init` was given: `bytes`, the settings arena and this struct go back to it. + gpa: std.mem.Allocator, + /// The parsed settings file: the endpoint's config points into it, so it outlives `ep`. + arena: *std.heap.ArenaAllocator, + + /// Connect to the peer named in `settings_path` (speed-up mode's file), exchanges up to `bytes`, and start sending. + pub fn init(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, settings_path: []const u8, bytes: usize) !*Pair { + const l: Layout = .{ .slot = pages(bytes) }; + const text = try std.Io.Dir.cwd().readFileAlloc(io, settings_path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const arena = try gpa.create(std.heap.ArenaAllocator); + errdefer gpa.destroy(arena); + arena.* = .init(gpa); + errdefer arena.deinit(); + // the endpoint's config keeps `s.links`, so the parse outlives `text` in the arena, not on the stack + const s = try settings.read(arena.allocator(), text); + if (s.rank > 1 or s.links.len != 1) return error.PairTwoRanksOneLink; + const lib = try gpa.dupeSentinel(u8, s.library, 0); + defer gpa.free(lib); + const ep = try fabric.mcdma.Endpoint.create(gpa, lib, .{ .rank = s.rank, .ranks = 2, .window_bytes = l.window(), .staging_bytes = 4 << 20, .links = s.links, .timeout_ns = 60 * std.time.ns_per_s, .connect_timeout_ns = 300 * std.time.ns_per_s }); + errdefer ep.deinit(); + var owned: Owner(mtl) = .{}; + errdefer owned.deinit(); + const rd = ep.rdma(); + const win = rd.window(); + const lib_m = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib_m.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const p = try gpa.create(Pair); + errdefer gpa.destroy(p); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + p.* = .{ + .rank = s.rank, + .peer = 1 - s.rank, + .l = l, + .ep = ep, + .rd = rd, + .win = win, + .wbuf = try owned.keep(gpa, try device.bufferNoCopy(win.ptr, win.len, opts)), + .add_rms = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_add_rms", false)), + .stage_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_stage", false)), + .gather_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_gather", false)), + .argmax_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_argmax", false)), + .pick_pipe = try owned.keep(gpa, try mtl.Pipeline.init(device, lib_m, "pair_pick", false)), + .fence = fence, + .posted = posted, + .landed = landed, + .bytes = try gpa.alloc(u32, JOBS), + .gpa = gpa, + .arena = arena, + .owned = owned, + }; + owned = .{}; // the pair owns them now: the error path must not release them twice + errdefer gpa.free(p.bytes); + // both Macs up before the first exchange: a word each way + try rd.signal(p.peer, l.flags() + HELLO, 1); + while (@atomicLoad(u64, p.word64(l.flags() + HELLO), .acquire) < 1) std.atomic.spinLoopHint(); + p.thread = try std.Thread.spawn(.{}, service, .{p}); + return p; + } + + /// Rank 0 tells rank 1 it is done (its `next` returns null); then the sending thread and the link close. + pub fn deinit(p: *Pair) void { + if (p.rank == 0 and !p.failed()) p.send(0, &.{}) catch {}; + p.stop.store(true, .release); + if (p.thread) |t| t.join(); + if (!p.failed()) flushLast(p); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.owned.deinit(); // the window's buffer before the endpoint frees the memory it aliases + p.ep.deinit(); + const gpa = p.gpa; + gpa.free(p.bytes); + p.arena.deinit(); + gpa.destroy(p.arena); + gpa.destroy(p); + } + + /// Hand the last seam over without the library's whole timeout: a peer that is gone leaves nothing to deliver. + fn flushLast(p: *Pair) void { + if (fabric.words.checkedNowNs()) |now| p.rd.setDeadline(now + last_flush_ns) else |_| return; + p.rd.flush() catch {}; + } + + fn word64(p: *const Pair, off: usize) *u64 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + fn word32(p: *const Pair, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Seams the host gave up on: nonzero means the run is invalid. + pub fn gaveUp(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// Seams the GPU itself refused after seeing FAILED: this side's own record of a failed relay. + pub fn observed(p: *const Pair) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + OBSERVED * 4), .acquire); + } + + /// The exchange the service thread is on (or the one it gave up on, if it has). + pub fn failedSeam(p: *const Pair) u32 { + return p.seam_x.load(.acquire); + } + + /// Labels a request; never a reset: after a seam failure the pair is dead (a half-written ring corrupts). + pub fn beginRound(p: *Pair) u32 { + return p.round.fetchAdd(1, .monotonic) + 1; + } + + /// Where this rank's part for exchange x goes, and where the peer's lands: the relay's two addresses. + pub fn sendOffset(p: *const Pair, x: u32) usize { + return p.l.send(x); + } + pub fn recvOffset(p: *const Pair, x: u32) usize { + return p.l.recv(x); + } + + /// The link is down: failed relay, dead peer, GPU never at the seam; check after every seam and command buffer. + pub fn failed(p: *const Pair) bool { + return p.broken.load(.acquire) or p.gaveUp() != 0; + } + + /// The peer's watchdog is alive: its last beat, a word its watchdog wrote here; silence means the peer is gone. + pub fn heartbeats(p: *const Pair) u64 { + return @atomicLoad(u64, p.word64(p.l.flags() + pair_layout.BEAT), .acquire); + } + + /// Tell the peer this rank's watchdog is alive; a failed signal is the peer gone, not nothing to do. + pub fn beat(p: *Pair) void { + p.beats +%= 1; + p.rd.signal(p.peer, p.l.flags() + pair_layout.BEAT, p.beats) catch p.broken.store(true, .release); + } + + /// The next exchange: take it before the producer writes, then pass the same x to stage/gather/addRms/argmax. + pub fn nextX(p: *const Pair) u32 { + return @truncate(p.x + 1); + } + + /// The exchange slot for this seam's part: where this rank's rows go, and where the peer's land. + pub fn part(p: *const Pair) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// Publish x and its bytes to the service thread; x must be `part`'s, so an out-of-order one panics in release too. + fn queue(p: *Pair, x: u32, bytes: usize) void { + if (x != @as(u32, @truncate(p.x + 1))) @panic("pair: an exchange was encoded out of order (part() not called first)"); + if (bytes > p.l.slot) @panic("pair: this exchange's part does not fit the window's slot (speed-up's rows too large)"); + p.x += 1; + p.bytes[x % JOBS] = @intCast(bytes); + p.queued.store(x, .release); + } + + /// After `part()`: send, wait for the peer's, h += both ranks', y = w * h / rms(h); part is `streamBytes` (fp32). + pub fn addRms(p: *Pair, e: *mtl.ComputeEncoder, x: u32, h: Ref, w: Ref, y: Ref, rows: u32, width: u32, eps: f32) void { + p.queue(x, streamBytes(rows, width)); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.add_rms); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = width, .rank = p.rank, .x = x, .eps = eps }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(h.buf, h.off, 4); + e.setBuffer(w.buf, w.off, 5); + e.setBuffer(y.buf, y.off, 6); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// Before `seam`: this rank's `width` columns of `rows` rows from `src` into the slot, bf16 to bf16, no arithmetic. + pub fn stage(p: *Pair, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("pair: stage() wants rows > 0 and a width inside the source row"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x, .eps = 0 }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// After `stage` (or the producer): relay, then pack this rank's `width` at `rank * width` and the peer's after. + pub fn gather(p: *Pair, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("pair: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("pair: gather()'s frame row is narrower than both halves"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("pair: gather()'s part is not the region part() returned"); + p.queue(x, stageBytes(rows, width)); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x, .eps = 0 }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// Each row's argmax over this rank's n fp32 logits from id lo, row stride `stride`, then the pick over both Macs'. + pub fn argmax(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + /// The same seam over bf16 logits (Qwen3.8-27B's `frame.logits`): one kernel, `dims.w` selects the load width. + pub fn argmaxBf16(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// `dtype` 0 is fp32 logits, 1 is bf16; the exchange is (fp32 value, id) either way, so pair_pick is one path. + fn argmaxWidth(p: *Pair, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("pair: argmax() wants rows > 0 and n > 0"); + p.queue(x, pickBytes(rows)); + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Rank 0: command `kind` (never 0, which ends rank 1's loop) and its words, once rank 1 took the one before. + pub fn send(p: *Pair, kind: u32, words: []const u32) !void { + if (p.rank != 0) return error.FollowsPeer; + if (words.len > command_words) return error.CommandTooLong; + raise(); + try p.reach(p.l.flags() + CMD_ACK, p.cmd, cmd_ack_timeout_ms); + p.cmd += 1; + const head: [2]u32 = .{ kind, @intCast(words.len) }; + const out = p.l.cmdOut(); + @memcpy(p.win[out..][0..8], std.mem.asBytes(&head)); + @memcpy(p.win[out + 8 ..][0 .. 4 * words.len], std.mem.sliceAsBytes(words)); + try p.ep.writeSignalFrom(p.peer, out, p.l.cmdIn(), 8 + 4 * words.len, p.l.flags() + CMD_FLAG, p.cmd); + } + + /// Rank 1: rank 0's next command's words into `out` and its kind; null once rank 0 is done or this Mac stops. + pub fn next(p: *Pair, gpa: std.mem.Allocator, out: *std.ArrayList(u32)) !?u32 { + raise(); + p.reach(p.l.flags() + CMD_FLAG, p.cmd + 1, 0) catch |err| return if (err == error.PairStopped) null else err; + p.cmd += 1; + const in = p.l.cmdIn(); + const head: *const [2]u32 = @ptrCast(@alignCast(p.win.ptr + in)); + if (head[1] > command_words) return error.CommandOutOfStep; + out.clearRetainingCapacity(); + try out.appendSlice(gpa, @as([*]const u32, @ptrCast(@alignCast(p.win.ptr + in + 8)))[0..head[1]]); + const kind = head[0]; + try p.rd.signal(p.peer, p.l.flags() + CMD_ACK, p.cmd); + return if (kind == 0) null else kind; + } + + /// Both Macs loaded before the first command: a word each way, each waited for with the load's bound. + pub fn ready(p: *Pair) !void { + try p.rd.signal(p.peer, p.l.flags() + READY, 1); + try p.reach(p.l.flags() + READY, 1, ready_timeout_ms); + } + + /// This Mac's command wait ends (rank 1 at close). + pub fn stopFollowing(p: *Pair) void { + p.stop.store(true, .release); + } + + /// Wait until this Mac's word at `off` is `value`: spins, then short sleeps, or fails after `limit_ms` (0: none). + fn reach(p: *Pair, off: usize, value: u64, limit_ms: u64) !void { + const t0 = ticksNow(); + var spins: usize = 0; + while (@atomicLoad(u64, p.word64(off), .acquire) < value) { + if (p.stop.load(.acquire)) return error.PairStopped; + if (p.failed()) return error.PairLinkFailed; + if (limit_ms != 0 and ticksNow() -% t0 > limit_ms * ticks_per_ms) return error.PairPeerSilent; + spins += 1; + if (spins > 200_000) sleepUs(20) else std.atomic.spinLoopHint(); + } + } + + /// Give up on x: release every event wait with the sentinel, so the blocked GPU resumes and the request aborts. + fn fail(p: *Pair, x: u32, comptime what: []const u8, peer_ms: u64) void { + p.broken.store(true, .release); + _ = @atomicRmw(u32, p.word32(p.l.sync() + FAILED * 4), .Add, 1, .monotonic); + std.log.err("pair rank {d}: request {d} seam {d} failed: {s}; the link is down and stays down (posted {d}, flag {d}, queued {d}, GPU-observed {d}; deadlines {d} ms post, {d} ms peer): restart both processes. The blocked GPU is released with the failure sentinel", .{ + p.rank, p.round.load(.acquire), x, what, p.posted.value(), @atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire), p.queued.load(.acquire), p.observed(), post_timeout_ms, peer_ms, + }); + p.landed.set(seam_failed); + } + + /// One exchange at a time: part written, relay it, peer's lands, release the GPU; every wait has a deadline. + fn service(p: *Pair) void { + raise(); + var x: u32 = 1; + // Ticks after which this link is no longer cold: set at the first relay, when both ranks have just loaded. + var cold_until: u64 = 0; + while (!p.stop.load(.acquire)) { + p.seam_x.store(x, .release); + // 1. The GPU must reach the seam: sleep while idle, spin once published, fail if it never arrives. + var armed: u64 = 0; + while (p.posted.value() < @as(u64, x)) { + if (p.stop.load(.acquire)) return; + if (!reached(p.queued.load(.acquire), x)) { + sleepUs(20); + continue; + } + if (armed == 0) armed = ticksNow(); + if (ticksNow() -% armed > post_timeout_ms * ticks_per_ms) return p.fail(x, "this Mac's GPU never reached the seam", peer_timeout_ms); + std.atomic.spinLoopHint(); + } + // 2. Relay: part to the peer's recv slot, arrival flag in the same MCDMA write, data and count together. + const at = x % JOBS; + p.ep.writeSignalFrom(p.peer, p.l.send(x), p.l.recv(x), p.bytes[at], p.l.flags() + FLAG, x) catch |err| { + std.log.err("pair rank {d}: seam {d} relay write failed: {s}", .{ p.rank, x, @errorName(err) }); + return p.fail(x, "the relay write failed", peer_timeout_ms); + }; + // 3. The peer's relay puts its part in our slot; its flag is proof; the cold bound, then the steady one. + if (x == 1) cold_until = ticksNow() +% cold_window_ms * ticks_per_ms; + const cold = cold_until != 0 and ticksNow() < cold_until; + const budget_ms = if (cold) cold_peer_timeout_ms else peer_timeout_ms; + const t1 = ticksNow(); + while (@atomicLoad(u64, p.word64(p.l.flags() + FLAG), .acquire) < x) { + if (p.stop.load(.acquire)) return; + if (ticksNow() -% t1 > budget_ms * ticks_per_ms) return p.fail(x, "the peer's part never landed", budget_ms); + std.atomic.spinLoopHint(); + } + // 4. Release the GPU: both parts are in; the event's wait is the acquire that makes them visible. + @atomicStore(u32, p.word32(p.l.sync() + POSTED * 4), x, .release); + p.landed.set(x); + x +%= 1; + } + } +}; + +test "the window's regions do not overlap and its words sit on their pages" { + const l: Layout = .{ .slot = pages(streamBytes(256, 6144)) }; + try std.testing.expect(l.send(1) >= l.recv(1) + l.slot and l.send(0) >= 2 * l.slot + PAGE and l.send(0) + l.slot <= l.send(1)); + try std.testing.expect(l.cmdIn() >= l.send(1) + l.slot and l.cmdOut() >= l.cmdIn() + CMD_BYTES + 64); + try std.testing.expect(l.flags() >= l.cmdOut() + 8 + 4 * command_words and BEAT + 8 <= PAGE and l.window() % PAGE == 0); + try std.testing.expect(reached(1, 0xFFFFFFFF) and !reached(5, 6)); +} + +test "the parity ring alternates and one sentinel releases every exchange" { + try std.testing.expect(slotOf(1) == .b and slotOf(2) == .a and slotOf(3) == .b); + const l: Layout = .{ .slot = pages(stageBytes(16, 5120)) }; + try std.testing.expect(l.recv(1) == l.slot and l.recv(2) == 0 and l.send(2) == 2 * l.slot + PAGE); + try std.testing.expectEqual(pages(stageBytes(16, 5120)), l.slot); // the 27B's 16-row prefill part, page-rounded + try std.testing.expect(seam_failed > std.math.maxInt(u32)); // every waitFor(x) in a buffer is satisfied by it +} diff --git a/zig/src/core/pair_rows_layout.zig b/zig/src/core/pair_rows_layout.zig new file mode 100644 index 0000000000..a1a45d6157 --- /dev/null +++ b/zig/src/core/pair_rows_layout.zig @@ -0,0 +1,167 @@ +//! The seam window's regions and the payload ABI, shared by the link and the host-copy Loopback. +const std = @import("std"); +const mtl = @import("metal"); +const Fence = @import("fence.zig").Fence; + +pub const Ref = @import("exlr8_projection.zig").Ref; + +pub const PAGE = 16384; +pub fn pages(n: usize) usize { + return (n + PAGE - 1) / PAGE * PAGE; +} +/// A command's words at most: a 262,144-token prompt and its head. +pub const command_words = 262144 + 64; +pub const CMD_BYTES = pages(8 + 4 * command_words); +pub const JOBS = 4096; + +/// Which parity slot x uses: rank r cannot finish x without the peer's x, so the peer is at most one exchange ahead. +pub const Slot = enum(u1) { a = 0, b = 1 }; + +pub fn slotOf(x: u32) Slot { + return @fromBackingInt(@intCast(x % 2)); +} + +/// The window on both Macs: parts in/out by parity (a page before each out), commands, flags, host words. +pub const Layout = struct { + slot: usize, + + pub fn recv(l: Layout, x: u32) usize { + return @as(usize, @backingInt(slotOf(x))) * l.slot; + } + pub fn send(l: Layout, x: u32) usize { + return 2 * l.slot + @as(usize, @backingInt(slotOf(x))) * (PAGE + l.slot) + PAGE; + } + pub fn cmdIn(l: Layout) usize { + return 4 * l.slot + 2 * PAGE; + } + pub fn cmdOut(l: Layout) usize { + return l.cmdIn() + CMD_BYTES + PAGE; + } + pub fn flags(l: Layout) usize { + return l.cmdOut() + CMD_BYTES; + } + pub fn sync(l: Layout) usize { + return l.flags() + PAGE; + } + pub fn window(l: Layout) usize { + return l.sync() + PAGE; + } +}; +pub const FLAG = 0; // u64 in the flag page: the last exchange whose peer part has landed (written by the peer's relay) +pub const HELLO = 64; +pub const CMD_FLAG = 128; // u64, at rank 1: the last command rank 0 has written +pub const CMD_ACK = 192; // u64, at rank 0: the last command rank 1 has copied out +pub const READY = 256; // u64: the peer has loaded its model +pub const BEAT = 320; // u64: the peer's watchdog ran; a rank idle in no round still sees a dead peer through this + +// The sync page's words shared with pair_rows.metal: POSTED last released exchange, FAILED give-ups, OBSERVED refusals. +pub const POSTED = 0; +pub const FAILED = 16; +pub const OBSERVED = 20; + +pub const seam_failed: u64 = std.math.maxInt(u64); + +pub const Seam = struct { x: u32, ref: Ref }; + +/// `Elem` is f32: the fp32 stream (`addRms`) sums exactly; `stage`/`gather` copy bf16 halves, never arithmetic. +pub const Elem = f32; +/// One element's bytes; a payload size in this file is `rows * width * elem_bytes`, never a literal. +pub const elem_bytes: usize = @sizeOf(Elem); +/// The element's name, for a log or a receipt that must say which payload it measured. +pub const elem_name: []const u8 = switch (Elem) { + u16 => "bf16", + f32 => "f32", + else => @compileError("pair: Elem is u16 (a bf16's bits) or f32"), +}; +/// Elements in one 32-bit window word: the relay and the link move whole words, so bf16 packs two and f32 one. +pub const elems_per_word: usize = @sizeOf(u32) / elem_bytes; + +/// One seam payload's bytes: `rows` rows of `width` `Elem` each -- the wire's, the window slot's and the relay's unit. +pub fn partBytes(rows: u32, width: u32) usize { + return @as(usize, rows) * width * elem_bytes; +} + +/// `stage`'s outgoing copy: this rank's `cols` of `rows` rows, bf16 -- the frame seam's payload, not an `Elem` payload. +pub fn stageBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * cols * @sizeOf(u16); +} + +/// `gather`'s `out`: both halves of `rows` rows packed into the frame, so the frame row is `2 * cols` wide. +pub fn gatherBytes(rows: u32, cols: u32) usize { + return @as(usize, rows) * 2 * cols * @sizeOf(u16); +} + +/// The head seam's exchange: an (fp32 value, u32 id) a row under either element; pair_pick compares logits, not words. +pub fn pickBytes(rows: u32) usize { + return @as(usize, rows) * @sizeOf(f32) * 2; +} + +/// `addRms`'s payload: fp32 candidates whatever `Elem` is; that path adds into the fp32 stream and rounds once. +pub fn streamBytes(rows: u32, width: u32) usize { + return @as(usize, rows) * width * @sizeOf(f32); +} + +/// A frame tensor's bytes: the bf16 grid `stage` reads and `gather` writes its halves into. +pub fn frameBytes(rows: u32, width: u32) usize { + return @as(usize, rows) * width * @sizeOf(u16); +} + +/// One fp32 value as the frame stores it (bf16, nearest even), the host side of the kernels' frame elements. +pub fn toFrame(v: f32) u16 { + return bf16Bits(v); +} + +/// An element's fp32 value: the payload's arithmetic is fp32 on both Macs, so a payload operand enters here. +pub fn fromElem(v: Elem) f32 { + return switch (Elem) { + u16 => @bitCast(@as(u32, v) << 16), + f32 => v, + else => @compileError("pair: Elem is u16 (a bf16's bits) or f32"), + }; +} + +/// One fp32 value as the payload stores it: bf16 rounds nearest-even, f32 is exact. +pub fn toElem(v: f32) Elem { + return switch (Elem) { + u16 => bf16Bits(v), + f32 => v, + else => @compileError("pair: Elem is u16 (a bf16's bits) or f32"), + }; +} + +/// One rounding to bf16, nearest-even: the host side of a frame element (pair_rows.metal copies the bits it is given). +fn bf16Bits(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u +% 0x7FFF +% ((u >> 16) & 1)) >> 16); +} + +/// The seed pattern as a payload element: 1 + m/128, exact in bf16 and f32, so expected bytes hold either way. +pub fn seedValue(r: u32, i: u32, salt: u16) Elem { + const k = r *% 5120 +% i *% 37; + const m: u32 = (salt & 0x1F) + ((k *% 53) & 0x1F); + return toElem(1.0 + @as(f32, @floatFromInt(m)) / 128.0); +} + +pub fn writeElem(out: []u8, at: usize, v: Elem) void { + @memcpy(out[at * elem_bytes ..][0..elem_bytes], std.mem.asBytes(&v)); +} + +/// One frame element at `at` of a bf16 payload: `stage`'s and `gather`'s slot is the frame's bits, two bytes each. +pub fn writeHalf(out: []u8, at: usize, v: f32) void { + std.mem.writeInt(u16, out[at * @sizeOf(u16) ..][0..@sizeOf(u16)], toFrame(v), .little); +} + +/// The kernels' argument block: no field is `Elem`, so its layout and pair_rows.metal's struct do not move with `Elem`. +pub const PairRows = extern struct { width: u32, stride: u32, rank: u32, x: u32, eps: f32, pad0: f32 = 0, pad1: f32 = 0, pad2: f32 = 0 }; + +/// Between encoders only (core/segments.zig:59-66): fence, close, posted, wait, reopen, wait -- a real hand-off. +pub fn seam(enc: *mtl.ComputeEncoder, fence: Fence, posted: mtl.SharedEvent, landed: mtl.SharedEvent, x: u32) void { + const cb = enc.commandBuffer(); + const mode = enc.dispatchType(); + fence.update(enc.*); + enc.end(); + cb.signal(posted, @as(u64, x)); + cb.waitFor(landed, @as(u64, x)); + enc.* = cb.compute(mode); + fence.wait(enc.*); +} diff --git a/zig/src/core/pair_rows_loopback.zig b/zig/src/core/pair_rows_loopback.zig new file mode 100644 index 0000000000..223375e893 --- /dev/null +++ b/zig/src/core/pair_rows_loopback.zig @@ -0,0 +1,243 @@ +//! The same seam with a host copy instead of the link: the numeric reference the real transport is checked against. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const Fence = @import("fence.zig").Fence; +const pair_layout = @import("pair_rows_layout.zig"); + +const FAILED = pair_layout.FAILED; +const JOBS = pair_layout.JOBS; +const Layout = pair_layout.Layout; +const PAGE = pair_layout.PAGE; +const PairRows = pair_layout.PairRows; +const Seam = pair_layout.Seam; +const elems_per_word = pair_layout.elems_per_word; +const pages = pair_layout.pages; +const stageBytes = pair_layout.stageBytes; +const pickBytes = pair_layout.pickBytes; +const seam = pair_layout.seam; +const seam_failed = pair_layout.seam_failed; +const seedValue = pair_layout.seedValue; +const writeHalf = pair_layout.writeHalf; +const Ref = pair_layout.Ref; + +/// The same seam with a host copy instead of the link: the numeric reference; bytes arrive via `arrive`/`seed`. +pub const Loopback = struct { + allocator: std.mem.Allocator, + rank: u32, + l: Layout, + win: []u8, + wbuf: mtl.Buffer, + stage_pipe: mtl.Pipeline, + gather_pipe: mtl.Pipeline, + argmax_pipe: mtl.Pipeline, + pick_pipe: mtl.Pipeline, + /// Per-exchange part bytes (stage/gather and argmax queue different sizes): the relay copies exactly this much. + bytes: []u32, + fence: Fence, + /// This rank's parts: the GPU signals it per exchange, and the relay reads `bytesFor(x)` after that. + posted: mtl.SharedEvent, + /// The peer's parts: the relay sets it once the copy has landed, and the seam waits for it. + landed: mtl.SharedEvent, + x: u64 = 0, + relay: ?std.Thread = null, + stop: std.atomic.Value(bool) = .init(false), + + pub fn init(device: mtl.Device, bytes: usize, rank: u32) !Loopback { + const l: Layout = .{ .slot = pages(bytes) }; + // the window is a no-copy Metal buffer: page-aligned, page-multiple allocation, never the gpa's good luck + const win = try std.heap.page_allocator.alloc(u8, l.window()); + errdefer std.heap.page_allocator.free(win); + std.debug.assert(@intFromPtr(win.ptr) % PAGE == 0 and win.len % PAGE == 0); + @memset(win, 0); + const lib = try mtl.Library.fromSource(device, sources.core_pair_rows, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const wbuf = try device.bufferNoCopy(win.ptr, win.len, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer wbuf.deinit(); + const fence = try Fence.init(device); + errdefer fence.deinit(); + const posted = try device.sharedEvent(); + errdefer posted.deinit(); + const landed = try device.sharedEvent(); + errdefer landed.deinit(); + const stage_pipe = try mtl.Pipeline.init(device, lib, "pair_stage", false); + errdefer stage_pipe.deinit(); + const gather_pipe = try mtl.Pipeline.init(device, lib, "pair_gather", false); + errdefer gather_pipe.deinit(); + const argmax_pipe = try mtl.Pipeline.init(device, lib, "pair_argmax", false); + errdefer argmax_pipe.deinit(); + const pick_pipe = try mtl.Pipeline.init(device, lib, "pair_pick", false); + errdefer pick_pipe.deinit(); + const sizes = try std.heap.page_allocator.alloc(u32, JOBS); + errdefer std.heap.page_allocator.free(sizes); + @memset(sizes, 0); + return .{ .allocator = std.heap.page_allocator, .rank = rank, .l = l, .win = win, .wbuf = wbuf, .stage_pipe = stage_pipe, .gather_pipe = gather_pipe, .argmax_pipe = argmax_pipe, .pick_pipe = pick_pipe, .bytes = sizes, .fence = fence, .posted = posted, .landed = landed }; + } + + pub fn deinit(p: *Loopback) void { + p.stopRelay(); + p.stage_pipe.deinit(); + p.gather_pipe.deinit(); + p.argmax_pipe.deinit(); + p.pick_pipe.deinit(); + p.allocator.free(p.bytes); + p.wbuf.deinit(); + p.posted.deinit(); + p.landed.deinit(); + p.fence.deinit(); + p.allocator.free(p.win); + } + + fn word32(p: *const Loopback, off: usize) *u32 { + return @ptrCast(@alignCast(p.win.ptr + off)); + } + + /// Where this rank's part for exchange x goes (what the relay copies out, once `posted` passes x). + pub fn sendOffset(p: *const Loopback, x: u32) usize { + return p.l.send(x); + } + + /// Where the peer's part for exchange x goes (what the relay copies in, before setting `landed` to x). + pub fn recvOffset(p: *const Loopback, x: u32) usize { + return p.l.recv(x); + } + + /// Where the producer writes this rank's bf16 columns of the next exchange; the same shape as `Pair.part`. + pub fn nextX(p: *const Loopback) u32 { + return @truncate(p.x + 1); + } + + pub fn part(p: *const Loopback) Seam { + const x: u32 = @truncate(p.x + 1); + return .{ .x = x, .ref = .{ .buf = p.wbuf, .off = p.l.send(x) } }; + } + + /// Fill both slots with the seed and mark all arrived: arithmetic check only, timing must not use this. + pub fn seed(p: *Loopback, rows: u32, width: u32) void { + const bytes = stageBytes(rows, width); + if (bytes > p.l.slot) @panic("loopback: the seeded part does not fit the window's slot"); + for (0..2) |x| { + const recv = p.win[p.l.recv(@intCast(x))..][0..bytes]; + const send = p.win[p.l.send(@intCast(x))..][0..bytes]; + for (0..rows) |r| for (0..width) |i| { + const at = r * width + i; + writeHalf(recv, at, seedValue(@intCast(r), @intCast(i), 0x15)); + writeHalf(send, at, seedValue(@intCast(r), @intCast(i), 0x0A)); + }; + } + p.satisfy(); + } + + /// Release every wait with the sentinel, as a failed link does: arithmetic check only, timing runs `startRelay`. + pub fn satisfy(p: *Loopback) void { + p.landed.set(seam_failed); + } + + /// The in-process relay: wait for posted, copy in, release -- Pair's two events in order, memcpy for the wire. + pub fn startRelay(p: *Loopback) !void { + std.debug.assert(p.relay == null); + p.stop.store(false, .release); + p.relay = try std.Thread.spawn(.{}, relayLoop, .{p}); + } + + pub fn stopRelay(p: *Loopback) void { + p.stop.store(true, .release); + if (p.relay) |t| t.join(); + p.relay = null; + } + + fn relayLoop(p: *Loopback) void { + var x: u64 = 1; + while (!p.stop.load(.acquire)) { + if (p.posted.value() < x) continue; // spin: the host's side of the seam is part of what is measured + const at: u32 = @truncate(x); + const n: usize = p.bytesFor(at); + @memcpy(p.win[p.l.recv(at)..][0..n], p.win[p.l.send(at)..][0..n]); + p.landed.set(x); + x += 1; + } + } + + /// The part's bytes for exchange `x` as the last `gather`/`argmax` queued it (the relay's copy length). + pub fn bytesFor(p: *const Loopback, x: u32) u32 { + return @atomicLoad(u32, &p.bytes[x % JOBS], .acquire); + } + + /// Land a host-written peer part for exchange x (the relay's post-copy half): copy it in, then release the wait. + pub fn arrive(p: *Loopback, x: u32, bytes: []const u8) void { + std.debug.assert(bytes.len <= p.l.slot); + @memcpy(p.win[p.l.recv(x)..][0..bytes.len], bytes); + p.landed.set(x); + } + + /// The seam `Pair.stage` encodes (before its `gather`): this rank's `width` columns of the frame row into the slot. + pub fn stage(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, src: Ref, rows: u32, width: u32, stride: u32) void { + if (rows == 0 or width == 0 or width > stride) @panic("loopback: stage() wants rows > 0 and a width inside the source row"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + e.setPipeline(p.stage_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = stride, .rank = p.rank, .x = x, .eps = 0 }, 1); + e.setBuffer(src.buf, src.off, 2); + e.setBuffer(p.wbuf, p.l.send(x), 3); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The seam `Pair.gather` encodes: the relay, then the frame row as this rank's half and the peer's, both bf16. + pub fn gather(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, mine: Ref, out: Ref, rows: u32, width: u32, out_stride: u32) void { + if (rows == 0 or width == 0 or width % elems_per_word != 0) @panic("loopback: gather() wants rows > 0 and a width that fills whole window words"); + if (out_stride < 2 * width) @panic("loopback: gather()'s frame row is narrower than both halves"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + if (mine.buf.id != p.wbuf.id or mine.off != p.l.send(x)) @panic("loopback: gather()'s part is not the region part() returned"); + if (stageBytes(rows, width) > p.l.slot) @panic("loopback: this exchange's part does not fit the window's slot"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(stageBytes(rows, width)), .release); + seam(e, p.fence, p.posted, p.landed, x); + e.setPipeline(p.gather_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue(PairRows{ .width = width, .stride = out_stride, .rank = p.rank, .x = x, .eps = 0 }, 1); + e.setBuffer(mine.buf, mine.off, 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(out.buf, out.off, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// The loopback's head seam: this rank's argmax over n bf16 logits, the relay copy, then the pick over both parts. + pub fn argmaxBf16(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 1); + } + + /// The same seam over fp32 logits (GLM-5.3's frame): the loopback covers both widths the kernel reads. + pub fn argmax(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + p.argmaxWidth(e, x, logits, stride, lo, n, rows, picks, 0); + } + + fn argmaxWidth(p: *Loopback, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref, dtype: u32) void { + if (rows == 0 or n == 0) @panic("loopback: argmax() wants rows > 0 and n > 0"); + if (x != @as(u32, @truncate(p.x + 1))) @panic("loopback: an exchange was encoded out of order (part() not called first)"); + p.x += 1; + @atomicStore(u32, &p.bytes[x % JOBS], @intCast(pickBytes(rows)), .release); // two u32 a row: (value, id) + e.setPipeline(p.argmax_pipe); + e.setBuffer(logits.buf, logits.off, 0); + e.setValue([4]u32{ stride, lo, n, dtype }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + seam(e, p.fence, p.posted, p.landed, x); // the encoder boundary orders the argmax before the part is relayed + e.setPipeline(p.pick_pipe); + e.setBuffer(p.wbuf, p.l.sync(), 0); + e.setValue([2]u32{ rows, x }, 1); + e.setBuffer(p.wbuf, p.l.send(x), 2); + e.setBuffer(p.wbuf, p.l.recv(x), 3); + e.setBuffer(picks.buf, picks.off, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 32 }); + } + + /// Seams the GPU refused (must be 0 for a run to be valid). + pub fn gaveUp(p: *const Loopback) u32 { + return @atomicLoad(u32, p.word32(p.l.sync() + FAILED * 4), .acquire); + } + + /// There is no link to fail here, so this is `gaveUp` under the seam's name, and it reads 0. + pub fn failed(p: *const Loopback) bool { + return p.gaveUp() != 0; + } +}; diff --git a/zig/src/core/simd_qmm.zig b/zig/src/core/simd_qmm.zig index 844c2099e5..1856304072 100644 --- a/zig/src/core/simd_qmm.zig +++ b/zig/src/core/simd_qmm.zig @@ -8,6 +8,11 @@ pub fn splits(n: usize) u32 { return if (n <= 64) 32 else if (n <= 6144) 16 else 8; } +/// Chunks for `n`; a split rank passes the whole width as `whole` so it keeps that width's tree (0 keeps `n`). +pub fn chunks(n: usize, whole: usize) u32 { + return splits(if (whole != 0) whole else n); +} + /// Output tiles a simdgroup within the 16 KB reduction buffer; tiles never change a row's bits. pub fn tiles(n: usize, rt: u32, s: u32) u32 { var nt: u32 = if (n % 32 == 0) 4 else if (n % 16 == 0) 2 else 1; @@ -72,10 +77,10 @@ pub const Pipelines = struct { pub const Call = struct { pipeline: mtl.Pipeline, args: Args, groups: mtl.Size, threads: usize }; /// y = x W^T, 4-bit groups of 64: the scalar twin at 1-2 rows if `scalar`, chunk pairs if `paired`, else MMA tiles. -pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool) !Call { +pub fn call(p: *const Pipelines, n: usize, k: usize, rows: usize, scalar: bool, paired: bool, whole: usize) !Call { if (n == 0 or n % 8 != 0 or n > std.math.maxInt(i32) or k == 0 or k % 64 != 0 or k > std.math.maxInt(i32)) return error.UnsupportedSimdShape; if (rows == 0 or rows > std.math.maxInt(i32)) return error.NoRows; - const s = splits(n); + const s = chunks(n, whole); const si = std.mem.indexOfScalar(u32, &s_values, s).?; const args = Args{ .k = @intCast(k), .n = @intCast(n), .r = @intCast(rows) }; if (scalar and scalarBlock(rows, s) != 0) { @@ -127,7 +132,7 @@ pub fn pairedExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair), m, x, 0, o, 0); + for (outs, [_]bool{ false, true }) |o, pair| encode(e, try call(p, m.n, m.k, rows, false, pair, 0), m, x, 0, o, 0); e.end(); cb.commit(); cb.wait(); @@ -157,9 +162,9 @@ pub fn twinExact(p: *const Pipelines, device: mtl.Device, m: Matrix) !bool { defer queue.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - encode(e, try call(p, m.n, m.k, rows, false, false), m, x, 0, tiled, 0); - for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false), m, x, r * m.k * 2, twin, r * m.n * 2); - for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); + encode(e, try call(p, m.n, m.k, rows, false, false, 0), m, x, 0, tiled, 0); + for (0..rows) |r| encode(e, try call(p, m.n, m.k, 1, true, false, 0), m, x, r * m.k * 2, twin, r * m.n * 2); + for ([_]usize{ 0, rows - 2 }, 0..) |r, i| encode(e, try call(p, m.n, m.k, 2, true, false, 0), m, x, r * m.k * 2, twos, i * 2 * m.n * 2); e.end(); cb.commit(); cb.wait(); @@ -177,6 +182,8 @@ test "chunks follow output width only and tiles fit the reduction buffer" { try std.testing.expectEqual(@as(u32, 32), splits(48)); try std.testing.expectEqual(@as(u32, 16), splits(5120)); try std.testing.expectEqual(@as(u32, 8), splits(34816)); + try std.testing.expectEqual(splits(6240), chunks(3168, 6240)); + try std.testing.expectEqual(splits(6240), chunks(6240, 0)); for (s_values) |s| for ([_]u32{ 1, 2 }) |rt| { const nt = tiles(34816, rt, s); try std.testing.expect(s * rt * nt * 64 * 4 <= 16384); diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index bc96bb5a64..65231bb1f8 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -9,7 +9,7 @@ pub const Plan = planning.Plan; pub const Cache = struct { keys: Ref, values: Ref, capacity: u32, key_stride: u32 = 0, value_stride: u32 = 0 }; const CachePtrs = extern struct { keys: u64, values: u64 }; const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; -const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32 }; +const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32, kv_heads: i32 }; pub const Limits = struct { rows: u32 = 128, streams: u32 = 8, @@ -353,7 +353,7 @@ pub const Attention = struct { // before M5 (no shared prefix) the partials take the tail's buffers, which hold every chunk's packed rows const parts = if (g.shared_prefix) [3]Field{ .poa, .pma, .pla } else [3]Field{ .pob, .pmb, .plb }; if (!g.shared_prefix and @as(usize, chunks) * plan.tiles * 16 > @as(usize, scratch.tail_chunks) * scratch.limits.rows * 16) return error.AttentionLimits; - const args = PromptArgs{ .start = start, .rows = @intCast(plan.rows), .chunks = @intCast(chunks), .rp = @intCast(plan.tiles * 16), .key_stride = @intCast(binding.strides[0][0]), .value_stride = @intCast(binding.strides[0][1]), .scale2 = 0.0625 * std.math.log2e }; + const args = PromptArgs{ .start = start, .rows = @intCast(plan.rows), .chunks = @intCast(chunks), .rp = @intCast(plan.tiles * 16), .key_stride = @intCast(binding.strides[0][0]), .value_stride = @intCast(binding.strides[0][1]), .scale2 = 0.0625 * std.math.log2e, .kv_heads = @intCast(g.kv_heads) }; const pairs = prompt_threads / 64; // simdgroup pairs, each halving the head dim of one 16-row tile e.setPipeline(a.prompt.?[0]); bind(e, .{ scratch.get(.qa), binding.get(.caches) }); diff --git a/zig/src/families/qwen27/attention_plan.zig b/zig/src/families/qwen27/attention_plan.zig index 0308f18d6d..36afdf3696 100644 --- a/zig/src/families/qwen27/attention_plan.zig +++ b/zig/src/families/qwen27/attention_plan.zig @@ -46,7 +46,9 @@ pub const Geometry = struct { } pub fn check(g: Geometry) !void { - if (g.heads != 24 or g.kv_heads != 4 or g.dim != 256 or g.rotary != 64) return error.TargetAttentionShape; + // Shaders take the GQA group (6) as a constant, not head count: 12/2 on a rank equals 24/4 on one Mac. + if (g.dim != 256 or g.rotary != 64 or g.heads == 0 or g.kv_heads == 0 or g.heads % g.kv_heads != 0) return error.TargetAttentionShape; + if (g.heads / g.kv_heads != 6) return error.TargetAttentionShape; } pub fn promptRows(g: Geometry, requested: u32) u32 { return @min(requested, @as(u32, if (g.shared_prefix and g.fuse_chain) chain_rows else max_rows)); @@ -435,7 +437,13 @@ test "target attention reads every config dimension without widening kernel qual const g = try Geometry.fromConfig(c, false); try std.testing.expect(g.sameShape(.{})); try std.testing.expect(!g.shared_prefix); + c.heads = 12; // a rank of a two-way split: the same GQA6 over half the heads + c.kv_heads = 2; + const half = try Geometry.fromConfig(c, true); + try std.testing.expectEqual(@as(u32, 12), half.heads); + try std.testing.expectEqual(@as(u32, 6), half.group()); c.heads = 16; // the supplied 9B has GQA4; the unchanged shaders still require GQA6 + c.kv_heads = 4; try std.testing.expectError(error.TargetAttentionShape, Geometry.fromConfig(c, true)); c.heads = 24; c.kv_heads = 2; diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index e4122d631e..5e85c8d309 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -27,6 +27,9 @@ pub const Config = struct { eos_count: usize = 0, max_position: usize = 262144, mrope_interleaved: bool = true, + rank: u32 = 0, + ranks: u32 = 1, + split_attention: bool = false, pub fn kind(c: Config, layer: usize) Kind { return c.kinds[layer]; @@ -51,8 +54,178 @@ pub const Config = struct { pub fn isEos(c: Config, token: u32) bool { return std.mem.indexOfScalar(u32, c.eos[0..c.eos_count], token) != null; } + + // Rank/ranks: 0/1 on one Mac. The fused zba row holds this rank's z, then every rank's b, then every rank's a. + + /// This rank's first value head: where its b, a, A_log and dt_bias rows start; 0 on one Mac or a replicated shard. + pub fn valueHead0(c: Config) usize { + return if (c.split_attention) c.rank * c.v_heads else 0; + } + + /// The `b` block's width: every rank's, because the fused projection keeps it replicated. + fn zbaHeads(c: Config) usize { + return if (c.split_attention) c.ranks * c.v_heads else c.v_heads; + } + + /// The rank's `b` row inside the fused `zba` row. + pub fn zbaB(c: Config) usize { + return c.gdnValueDim() + c.valueHead0(); + } + + /// The rank's `a` row inside the fused `zba` row, after the `b` block. + pub fn zbaA(c: Config) usize { + return c.gdnValueDim() + c.zbaHeads() + c.valueHead0(); + } + + /// The fused `zba` row's width. A multiple of 32 keeps `projection.prepare` on its tiled kernel. + pub fn zbaDim(c: Config) usize { + return c.gdnValueDim() + 2 * c.zbaHeads(); + } +}; + +/// A rank's slice of one tensor: elements `[from, from+len)` of the checkpoint's rows, or of an input row. +pub const Span = struct { + from: usize = 0, + len: usize = 0, + + pub fn end(s: Span) usize { + return s.from + s.len; + } +}; + +/// A rank's share of the 27B: the checkpoint holds `global`'s shapes, `local` is what this rank runs. +pub const Shard = struct { + rank: u32, + ranks: u32, + global: Config, + local: Config, + + /// This rank's first output row in a module whose `local` rows each rank takes a block of. + pub fn row0(s: Shard, local: usize) usize { + return @as(usize, s.rank) * local; + } + + /// This rank's first input column in a module whose inputs each rank takes a block of. + pub fn col0(s: Shard, local: usize) usize { + return @as(usize, s.rank) * local; + } + + /// This rank's rows of a global-row module: local rows on a split dimension, the whole tensor on a replicated one. + fn rows(s: Shard, global: usize, local: usize) Span { + return if (local == global) .{ .from = 0, .len = global } else .{ .from = s.row0(local), .len = local }; + } + + /// This rank's input columns of a `global`-column module (the mirror of `rows`). + fn cols(s: Shard, global: usize, local: usize) Span { + return if (local == global) .{ .from = 0, .len = global } else .{ .from = s.col0(local), .len = local }; + } + + /// Whether attention and linear attention split across ranks; an mlp shard leaves them whole, so one seam a layer. + pub fn splitsAttention(s: Shard) bool { + return s.local.split_attention; + } + + /// self_attn.q_proj rows: q and its gate interleave every head_dim, so slices land on a 2*head_dim boundary. + pub fn qRows(s: Shard) Span { + return s.rows(s.global.qDim(), s.local.qDim()); + } + + /// `self_attn.k_proj` and `v_proj` rows (stacked k|v in that order): whole KV heads, one GQA group each. + pub fn kvRows(s: Shard) Span { + return s.rows(s.global.kvDim(), s.local.kvDim()); + } + + /// `self_attn.o_proj` input columns: this rank's q heads' outputs, `heads * head_dim` wide. + pub fn attentionCols(s: Shard) Span { + return s.cols(s.global.heads * s.global.head_dim, s.local.heads * s.local.head_dim); + } + + /// linear_attn.in_proj_qkv rows and conv1d channels: one q | k | v tensor, each block split on its own head count. + pub fn qkvSpans(s: Shard) [3]Span { + const gqk = s.global.k_heads * s.global.dk; + const q = s.rows(gqk, s.local.k_heads * s.local.dk); + const k = s.rows(gqk, s.local.k_heads * s.local.dk); + const v = s.rows(s.global.gdnValueDim(), s.local.gdnValueDim()); + return .{ q, .{ .from = gqk + k.from, .len = k.len }, .{ .from = 2 * gqk + v.from, .len = v.len } }; + } + + /// linear_attn.conv1d channels in qkvSpans order: the same three blocks, so local conv state indexes like qkv rows. + pub fn convSpans(s: Shard) [3]Span { + return s.qkvSpans(); + } + + /// `linear_attn.in_proj_z` rows: this rank's `v_heads * dv` value channels, one contiguous block. + pub fn zRows(s: Shard) Span { + return s.rows(s.global.gdnValueDim(), s.local.gdnValueDim()); + } + + /// linear_attn.in_proj_a and in_proj_b: replicated whole into zba; a 24-row split would drop the tiled projection. + pub fn abRows(s: Shard) Span { + return .{ .from = 0, .len = s.global.v_heads }; + } + + /// `linear_attn.A_log` and `dt_bias`: one constant a value head, so a rank takes its own heads' slice. + pub fn valueRows(s: Shard) Span { + return s.rows(s.global.v_heads, s.local.v_heads); + } + + /// `linear_attn.out_proj` input columns: this rank's value channels. + pub fn gdnCols(s: Shard) Span { + return s.cols(s.global.gdnValueDim(), s.local.gdnValueDim()); + } + + /// `mlp.gate_proj` and `up_proj` rows (stacked gate|up): this rank's share of the intermediate width. + pub fn intermediateRows(s: Shard) Span { + return s.rows(s.global.intermediate, s.local.intermediate); + } + + /// `mlp.down_proj` input columns: this rank's share of the same intermediate width. + pub fn downCols(s: Shard) Span { + return s.cols(s.global.intermediate, s.local.intermediate); + } + + /// `lm_head` rows: this rank's share of the vocabulary, one contiguous block of tokens. + pub fn vocabRows(s: Shard) Span { + return s.rows(s.global.vocab, s.local.vocab); + } + + /// `o_proj`/`out_proj`/`down_proj` rows: this rank's block over the whole K, or the width when replicated. + pub fn outputRows(s: Shard) Span { + return if (s.splitsAttention()) s.rows(s.global.hidden, s.global.hidden / s.ranks) else .{ .from = 0, .len = s.global.hidden }; + } }; +/// Mode: full halves attention, linear attention and MLP; mlp splits only the MLP and head, so one seam a layer. +pub const Mode = enum { full, mlp }; + +/// Split the even dimensions: heads, GQA groups, GDN heads, MLP width and vocabulary; embeddings and norms stay whole. +pub fn split(c: *Config, rank: u32, ranks: u32, mode: Mode) !Shard { + if (ranks < 2 or ranks > 8 or rank >= ranks) return error.BadShard; + if (c.intermediate % ranks != 0 or c.vocab % ranks != 0) return error.BadShard; + if (mode == .full and (c.heads % ranks != 0 or c.kv_heads % ranks != 0 or c.k_heads % ranks != 0 or c.v_heads % ranks != 0 or c.hidden % ranks != 0)) return error.BadShard; + const global = c.*; + if (mode == .full) { + c.heads /= ranks; + c.kv_heads /= ranks; + c.k_heads /= ranks; + c.v_heads /= ranks; + } + c.intermediate /= ranks; + c.vocab /= ranks; + c.rank = rank; + c.ranks = ranks; + c.split_attention = mode == .full; + // the local shapes must keep the group structure the attention, GDN and projection kernels run + if (c.heads == 0 or c.kv_heads == 0 or c.k_heads == 0 or c.v_heads == 0) return error.BadShard; + if (c.heads % c.kv_heads != 0 or c.v_heads % c.k_heads != 0) return error.BadShard; + if (c.heads / c.kv_heads != global.heads / global.kv_heads) return error.BadShard; + if (c.qDim() % 4 != 0 or c.kvDim() % 4 != 0 or c.gdnQkvDim() % 4 != 0 or c.gdnValueDim() % 4 != 0) return error.BadShard; + if (c.vocab % 4 != 0 or c.intermediate % 4 != 0) return error.BadShard; + if (c.layers != global.layers) return error.BadShard; + try shapeProducts(c.*); + return .{ .rank = rank, .ranks = ranks, .global = global, .local = c.* }; +} + pub fn object(v: std.json.Value) !std.json.ObjectMap { return if (v == .object) v.object else error.BadConfig; } @@ -184,121 +357,12 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) return c; } -const fixture = - \\{"model_type":"qwen3_5","hidden_size":64,"intermediate_size":128,"num_hidden_layers":4, - \\"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,"vocab_size":97, - \\"linear_num_key_heads":2,"linear_num_value_heads":4,"linear_key_head_dim":16, - \\"linear_value_head_dim":16,"linear_conv_kernel_dim":4,"eos_token_id":[2,11]} -; - -fn edit(a: std.mem.Allocator, key: []const u8, value: std.json.Value) ![]u8 { - var p = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); - defer p.deinit(); - try p.value.object.put(p.arena.allocator(), key, value); - return std.json.Stringify.valueAlloc(a, p.value, .{}); -} - -test "flat config preserves dimensions and interval schedule" { - const c = try parse(std.testing.allocator, fixture, null); - try std.testing.expectEqual(@as(usize, 128), c.qDim()); - try std.testing.expectEqual(@as(usize, 32), c.kvDim()); - try std.testing.expectEqual(@as(usize, 128), c.gdnQkvDim()); - try std.testing.expectEqual(Kind.linear, c.kind(0)); - try std.testing.expectEqual(Kind.attention, c.kind(3)); -} - -test "nested config merges generation EOS in stable order" { - const a = std.testing.allocator; - const text = try std.fmt.allocPrint(a, "{{\"model_type\":\"qwen3_5\",\"eos_token_id\":[11,2],\"text_config\":{s}}}", .{fixture}); - defer a.free(text); - const c = try parse(a, text, "{\"eos_token_id\":[2,3,11]}"); - try std.testing.expectEqualSlices(u32, &.{ 11, 2, 3 }, c.eos[0..c.eos_count]); - try std.testing.expect(c.isEos(3) and !c.isEos(4)); -} - -test "explicit layer types preserve their order" { - const a = std.testing.allocator; - var p = try std.json.parseFromSlice(std.json.Value, a, "[\"full_attention\",\"linear_attention\",\"full_attention\",\"linear_attention\"]", .{}); - defer p.deinit(); - const json = try edit(a, "layer_types", p.value); - defer a.free(json); - const c = try parse(a, json, null); - try std.testing.expectEqualSlices(Kind, &.{ .attention, .linear, .attention, .linear }, c.kinds[0..4]); -} - -test "rope parameters preserve precision and section metadata" { - const a = std.testing.allocator; - const p = try std.json.parseFromSlice(std.json.Value, a, "{\"rope_theta\":1234567.25,\"partial_rotary_factor\":0.5,\"mrope_section\":[2,1,1]}", .{}); - defer p.deinit(); - const json = try edit(a, "rope_parameters", p.value); - defer a.free(json); - const c = try parse(a, json, null); - try std.testing.expectEqual(@as(f64, 1234567.25), c.rope_theta); - try std.testing.expectEqual(@as(usize, 8), c.rope_dims); - try std.testing.expectEqualSlices(usize, &.{ 2, 1, 1 }, &c.mrope_section); -} - -test "dense config refuses MoE and tied embedding heads" { - const a = std.testing.allocator; - const moe = try edit(a, "num_experts", .{ .integer = 8 }); - defer a.free(moe); - try std.testing.expectError(error.NotQwenDense, parse(a, moe, null)); - const tied = try edit(a, "tie_word_embeddings", .{ .bool = true }); - defer a.free(tied); - try std.testing.expectError(error.TiedEmbeddingHead, parse(a, tied, null)); -} - -test "invalid scalar shapes and types fail without a cast trap" { - const a = std.testing.allocator; - for ([_]std.json.Value{ .{ .integer = -1 }, .{ .integer = 0 }, .{ .bool = true }, .{ .integer = 4294967296 } }) |v| { - const json = try edit(a, "hidden_size", v); - defer a.free(json); - try std.testing.expectError(error.BadConfig, parse(a, json, null)); - } - const json = try edit(a, "num_hidden_layers", .{ .integer = 129 }); - defer a.free(json); - try std.testing.expectError(error.BadConfig, parse(a, json, null)); -} - -test "missing and out of vocabulary EOS fail" { - const a = std.testing.allocator; - const missing = try edit(a, "eos_token_id", .null); - defer a.free(missing); - try std.testing.expectError(error.MissingEos, parse(a, missing, null)); - const c = try parse(a, missing, "{\"eos_token_id\":7}"); - try std.testing.expect(c.isEos(7)); - const wrong = try edit(a, "eos_token_id", .{ .integer = 97 }); - defer a.free(wrong); - try std.testing.expectError(error.BadEos, parse(a, wrong, null)); -} - -test "malformed layer and rope metadata fail" { - const a = std.testing.allocator; - const layers = try edit(a, "layer_types", .{ .bool = true }); - defer a.free(layers); - try std.testing.expectError(error.BadConfig, parse(a, layers, null)); - const rope = try edit(a, "partial_rotary_factor", .{ .float = 2 }); - defer a.free(rope); - try std.testing.expectError(error.BadConfig, parse(a, rope, null)); - try std.testing.expectError(error.BadConfig, parse(a, "[]", null)); -} - -test "head grouping requires complete key and value groups" { - const a = std.testing.allocator; - const attention = try edit(a, "num_key_value_heads", .{ .integer = 3 }); - defer a.free(attention); - try std.testing.expectError(error.BadConfig, parse(a, attention, null)); - const gdn = try edit(a, "linear_num_value_heads", .{ .integer = 3 }); - defer a.free(gdn); - try std.testing.expectError(error.BadConfig, parse(a, gdn, null)); +/// The embedding table stays whole on every rank, so its gather bound is the table's rows, not a shard's width. +pub fn embedBound(c: *const Config, rows: u64) !u32 { + if (rows < c.vocab) return error.BadEmbedding; + return std.math.cast(u32, rows) orelse error.BadEmbedding; } -test "text execution refuses unsupported gates and activation contracts" { - const a = std.testing.allocator; - const gate = try edit(a, "attn_output_gate", .{ .bool = false }); - defer a.free(gate); - try std.testing.expectError(error.UnsupportedAttentionGate, parse(a, gate, null)); - const act = try edit(a, "output_gate_type", .{ .string = "sigmoid" }); - defer a.free(act); - try std.testing.expectError(error.UnsupportedActivation, parse(a, act, null)); +test { + _ = @import("config_test.zig"); } diff --git a/zig/src/families/qwen27/config_test.zig b/zig/src/families/qwen27/config_test.zig new file mode 100644 index 0000000000..6f83ed12ba --- /dev/null +++ b/zig/src/families/qwen27/config_test.zig @@ -0,0 +1,287 @@ +const std = @import("std"); +const cfg = @import("config.zig"); +const parse = cfg.parse; +const Config = cfg.Config; +const Kind = cfg.Kind; +const Span = cfg.Span; +const Mode = cfg.Mode; +const split = cfg.split; +const object = cfg.object; + +const fixture = + \\{"model_type":"qwen3_5","hidden_size":64,"intermediate_size":128,"num_hidden_layers":4, + \\"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,"vocab_size":97, + \\"linear_num_key_heads":2,"linear_num_value_heads":4,"linear_key_head_dim":16, + \\"linear_value_head_dim":16,"linear_conv_kernel_dim":4,"eos_token_id":[2,11]} +; + +fn edit(a: std.mem.Allocator, key: []const u8, value: std.json.Value) ![]u8 { + var p = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + defer p.deinit(); + try p.value.object.put(p.arena.allocator(), key, value); + return std.json.Stringify.valueAlloc(a, p.value, .{}); +} + +test "flat config preserves dimensions and interval schedule" { + const c = try parse(std.testing.allocator, fixture, null); + try std.testing.expectEqual(@as(usize, 128), c.qDim()); + try std.testing.expectEqual(@as(usize, 32), c.kvDim()); + try std.testing.expectEqual(@as(usize, 128), c.gdnQkvDim()); + try std.testing.expectEqual(Kind.linear, c.kind(0)); + try std.testing.expectEqual(Kind.attention, c.kind(3)); +} + +test "nested config merges generation EOS in stable order" { + const a = std.testing.allocator; + const text = try std.fmt.allocPrint(a, "{{\"model_type\":\"qwen3_5\",\"eos_token_id\":[11,2],\"text_config\":{s}}}", .{fixture}); + defer a.free(text); + const c = try parse(a, text, "{\"eos_token_id\":[2,3,11]}"); + try std.testing.expectEqualSlices(u32, &.{ 11, 2, 3 }, c.eos[0..c.eos_count]); + try std.testing.expect(c.isEos(3) and !c.isEos(4)); +} + +test "explicit layer types preserve their order" { + const a = std.testing.allocator; + var p = try std.json.parseFromSlice(std.json.Value, a, "[\"full_attention\",\"linear_attention\",\"full_attention\",\"linear_attention\"]", .{}); + defer p.deinit(); + const json = try edit(a, "layer_types", p.value); + defer a.free(json); + const c = try parse(a, json, null); + try std.testing.expectEqualSlices(Kind, &.{ .attention, .linear, .attention, .linear }, c.kinds[0..4]); +} + +test "rope parameters preserve precision and section metadata" { + const a = std.testing.allocator; + const p = try std.json.parseFromSlice(std.json.Value, a, "{\"rope_theta\":1234567.25,\"partial_rotary_factor\":0.5,\"mrope_section\":[2,1,1]}", .{}); + defer p.deinit(); + const json = try edit(a, "rope_parameters", p.value); + defer a.free(json); + const c = try parse(a, json, null); + try std.testing.expectEqual(@as(f64, 1234567.25), c.rope_theta); + try std.testing.expectEqual(@as(usize, 8), c.rope_dims); + try std.testing.expectEqualSlices(usize, &.{ 2, 1, 1 }, &c.mrope_section); +} + +test "dense config refuses MoE and tied embedding heads" { + const a = std.testing.allocator; + const moe = try edit(a, "num_experts", .{ .integer = 8 }); + defer a.free(moe); + try std.testing.expectError(error.NotQwenDense, parse(a, moe, null)); + const tied = try edit(a, "tie_word_embeddings", .{ .bool = true }); + defer a.free(tied); + try std.testing.expectError(error.TiedEmbeddingHead, parse(a, tied, null)); +} + +test "invalid scalar shapes and types fail without a cast trap" { + const a = std.testing.allocator; + for ([_]std.json.Value{ .{ .integer = -1 }, .{ .integer = 0 }, .{ .bool = true }, .{ .integer = 4294967296 } }) |v| { + const json = try edit(a, "hidden_size", v); + defer a.free(json); + try std.testing.expectError(error.BadConfig, parse(a, json, null)); + } + const json = try edit(a, "num_hidden_layers", .{ .integer = 129 }); + defer a.free(json); + try std.testing.expectError(error.BadConfig, parse(a, json, null)); +} + +test "missing and out of vocabulary EOS fail" { + const a = std.testing.allocator; + const missing = try edit(a, "eos_token_id", .null); + defer a.free(missing); + try std.testing.expectError(error.MissingEos, parse(a, missing, null)); + const c = try parse(a, missing, "{\"eos_token_id\":7}"); + try std.testing.expect(c.isEos(7)); + const wrong = try edit(a, "eos_token_id", .{ .integer = 97 }); + defer a.free(wrong); + try std.testing.expectError(error.BadEos, parse(a, wrong, null)); +} + +test "malformed layer and rope metadata fail" { + const a = std.testing.allocator; + const layers = try edit(a, "layer_types", .{ .bool = true }); + defer a.free(layers); + try std.testing.expectError(error.BadConfig, parse(a, layers, null)); + const rope = try edit(a, "partial_rotary_factor", .{ .float = 2 }); + defer a.free(rope); + try std.testing.expectError(error.BadConfig, parse(a, rope, null)); + try std.testing.expectError(error.BadConfig, parse(a, "[]", null)); +} + +test "head grouping requires complete key and value groups" { + const a = std.testing.allocator; + const attention = try edit(a, "num_key_value_heads", .{ .integer = 3 }); + defer a.free(attention); + try std.testing.expectError(error.BadConfig, parse(a, attention, null)); + const gdn = try edit(a, "linear_num_value_heads", .{ .integer = 3 }); + defer a.free(gdn); + try std.testing.expectError(error.BadConfig, parse(a, gdn, null)); +} + +test "text execution refuses unsupported gates and activation contracts" { + const a = std.testing.allocator; + const gate = try edit(a, "attn_output_gate", .{ .bool = false }); + defer a.free(gate); + try std.testing.expectError(error.UnsupportedAttentionGate, parse(a, gate, null)); + const act = try edit(a, "output_gate_type", .{ .string = "sigmoid" }); + defer a.free(act); + try std.testing.expectError(error.UnsupportedActivation, parse(a, act, null)); +} + +/// The real 27B's shapes, so the shard arithmetic is pinned against the checkpoint rather than a toy config. +fn model27() Config { + return .{ + .hidden = 5120, + .intermediate = 17408, + .layers = 64, + .vocab = 248320, + .heads = 24, + .kv_heads = 4, + .head_dim = 256, + .k_heads = 16, + .v_heads = 48, + .dk = 128, + .dv = 128, + .conv_kernel = 4, + .eps = 1e-6, + .rope_dims = 64, + .rope_theta = 1e7, + }; +} + +fn expectSpan(span: Span, from: usize, len: usize) !void { + try std.testing.expectEqual(from, span.from); + try std.testing.expectEqual(len, span.len); +} + +test "a full shard takes whole heads and boundary-aligned row blocks" { + var d = model27(); + const s = try split(&d, 1, 2, .full); + try std.testing.expectEqual(@as(usize, 64), s.local.layers); + try std.testing.expectEqual(@as(usize, 12), s.local.heads); + try std.testing.expectEqual(@as(usize, 2), s.local.kv_heads); + try std.testing.expectEqual(@as(usize, 8), s.local.k_heads); + try std.testing.expectEqual(@as(usize, 24), s.local.v_heads); + try std.testing.expectEqual(@as(usize, 8704), s.local.intermediate); + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try std.testing.expect(s.splitsAttention()); + // q and its gate interleave every head: rank 1 starts at head 12, a 2*head_dim (512) boundary + try expectSpan(s.qRows(), 6144, 6144); + try std.testing.expectEqual(@as(usize, 0), s.qRows().from % (2 * s.global.head_dim)); + try expectSpan(s.kvRows(), 512, 512); + try expectSpan(s.attentionCols(), 3072, 3072); + // GDN q | k | v is one tensor of three blocks: q and k split 8 of 16 heads, v splits 24 of 48 + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 1024, 1024); + try expectSpan(blocks[1], 3072, 1024); + try expectSpan(blocks[2], 7168, 3072); + try std.testing.expectEqual(@as(usize, 5120), blocks[0].len + blocks[1].len + blocks[2].len); + const conv = s.convSpans(); + for (blocks, conv) |want, got| { + try expectSpan(got, want.from, want.len); + } + try expectSpan(s.zRows(), 3072, 3072); + try expectSpan(s.valueRows(), 24, 24); + try expectSpan(s.abRows(), 0, 48); + try expectSpan(s.gdnCols(), 3072, 3072); + try expectSpan(s.intermediateRows(), 8704, 8704); + try expectSpan(s.downCols(), 8704, 8704); + try expectSpan(s.vocabRows(), 124160, 124160); + // o_proj, out_proj and down_proj own rows of the hidden width, each over the whole input K + try expectSpan(s.outputRows(), 2560, 2560); +} + +test "the embedding bound is the whole table on a pair rank, not the shard's width" { + var d = model27(); + const s = try split(&d, 1, 2, .full); + // the table keeps 248320 rows on both ranks while the head's shard is 124160 wide + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try std.testing.expectEqual(@as(u32, 248320), try cfg.embedBound(&s.local, 248320)); + try std.testing.expectEqual(@as(u32, 248320), try cfg.embedBound(&s.global, 248320)); + // a bound below the shard's own width is a table that lost rows, so it stays an error + try std.testing.expectError(error.BadEmbedding, cfg.embedBound(&s.local, 124159)); +} + +test "rank 0's blocks start at the tensor's front" { + var d = model27(); + const s = try split(&d, 0, 2, .full); + try expectSpan(s.qRows(), 0, 6144); + try expectSpan(s.kvRows(), 0, 512); + try expectSpan(s.attentionCols(), 0, 3072); + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 0, 1024); + try expectSpan(blocks[1], 2048, 1024); + try expectSpan(blocks[2], 4096, 3072); + try expectSpan(s.zRows(), 0, 3072); + try expectSpan(s.valueRows(), 0, 24); + try expectSpan(s.outputRows(), 0, 2560); +} + +test "the GDN qkv blocks tile the global tensor exactly across ranks" { + var seen: [10240]bool = @splat(false); + for (0..2) |r| { + var d = model27(); + const s = try split(&d, @intCast(r), 2, .full); + for (s.qkvSpans()) |span| { + try std.testing.expect(span.end() <= seen.len); + for (span.from..span.end()) |i| { + try std.testing.expect(!seen[i]); + seen[i] = true; + } + } + } + for (seen) |hit| try std.testing.expect(hit); +} + +test "an mlp shard replicates attention and splits only the mlp and the head" { + var d = model27(); + const s = try split(&d, 1, 2, .mlp); + try std.testing.expect(!s.splitsAttention()); + try std.testing.expectEqual(@as(usize, 24), s.local.heads); + try std.testing.expectEqual(@as(usize, 4), s.local.kv_heads); + try std.testing.expectEqual(@as(usize, 16), s.local.k_heads); + try std.testing.expectEqual(@as(usize, 48), s.local.v_heads); + try std.testing.expectEqual(@as(usize, 8704), s.local.intermediate); + try std.testing.expectEqual(@as(usize, 124160), s.local.vocab); + try expectSpan(s.qRows(), 0, 12288); + try expectSpan(s.kvRows(), 0, 1024); + try expectSpan(s.attentionCols(), 0, 6144); + const blocks = s.qkvSpans(); + try expectSpan(blocks[0], 0, 2048); + try expectSpan(blocks[1], 2048, 2048); + try expectSpan(blocks[2], 4096, 6144); + try expectSpan(s.zRows(), 0, 6144); + try expectSpan(s.abRows(), 0, 48); + // attention is replicated, so every value head is this rank's: A_log and dt_bias stay whole too + try expectSpan(s.valueRows(), 0, 48); + try expectSpan(s.gdnCols(), 0, 6144); + try expectSpan(s.intermediateRows(), 8704, 8704); + try expectSpan(s.downCols(), 8704, 8704); + // an mlp shard keeps the mixers and down_proj replicated, so their output rows are the whole width + try expectSpan(s.outputRows(), 0, 5120); +} + +test "the fused zba row keeps the tiled projection and each rank's a and b" { + for ([_]Mode{ .full, .mlp }) |mode| { + var d = model27(); + const s = try split(&d, 1, 2, mode); + const c = s.local; + try std.testing.expectEqual(@as(usize, 0), c.zbaDim() % 32); + const whole_attention = mode == .mlp; + const ranks: usize = c.ranks; + const b_width = c.v_heads * (if (whole_attention) 1 else ranks); + const head0 = if (whole_attention) 0 else c.v_heads; + try std.testing.expectEqual(c.gdnValueDim() + 2 * b_width, c.zbaDim()); + try std.testing.expectEqual(c.gdnValueDim() + head0, c.zbaB()); + try std.testing.expectEqual(c.gdnValueDim() + b_width + head0, c.zbaA()); + try std.testing.expectEqual(c.zbaDim(), c.zbaA() + c.v_heads); + var e = model27(); + const z = try split(&e, 0, 2, mode); + try std.testing.expectEqual(z.local.gdnValueDim(), z.local.zbaB()); + try std.testing.expectEqual(z.local.gdnValueDim() + 48, z.local.zbaA()); + } + // one Mac is today's layout: z | b | a, each value head once + const u = model27(); + try std.testing.expectEqual(@as(usize, 6144), u.zbaB()); + try std.testing.expectEqual(@as(usize, 6192), u.zbaA()); + try std.testing.expectEqual(@as(usize, 6240), u.zbaDim()); +} diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index f074c0adc0..2b855427b3 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -4,6 +4,7 @@ const mtl = @import("metal"); const Model = @import("model.zig").Model; const Frame = @import("gpu_frame.zig").Frame; const forward = @import("forward.zig"); +const Ref = @import("projection.zig").Ref; /// Layers a verify encodes before committing them, so the GPU starts while the CPU encodes the rest. const commit_every = 4; const wts = @import("gpu_weights.zig"); @@ -57,6 +58,10 @@ pub const Runner = struct { images: ?image_rows.Images = null, /// A kept state's blit still reading a slot, by slot: the slot's next write waits for it (settleSlot). reading: [attn.planning.max_streams]?mtl.CommandBuffer = @splat(null), + /// A tensor-parallel rank's all-reduce, threaded straight into the layer program; null on one Mac. + reduce: ?forward.Reduce = null, + /// The pair's head-seam picks (one u32 a row), required with `reduce`: unexchanged halves make the ranks diverge. + picks: ?Ref = null, pub fn init(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32) !Runner { return initReady(a, m, slots, capacity, slots); @@ -243,11 +248,16 @@ pub const Runner = struct { if (r.capture_taps) r.taps.begin(); const pool = mtl.objc.Pool.push(); defer pool.pop(); - var stages = forward.Stages{ .model = r.model, .rows = @intCast(round.ids.len), .hooks = .{ .ptr = r, .gdn = gdnForward, .attention = attentionForward, .tap = if (r.capture_taps) tap else null }, .head = head, .images = placed.images }; + var stages = forward.Stages{ .model = r.model, .rows = @intCast(round.ids.len), .hooks = .{ .ptr = r, .gdn = gdnForward, .attention = attentionForward, .tap = if (r.capture_taps) tap else null }, .head = head, .images = placed.images, .reduce = r.reduce, .picks = r.picks }; r.last_gpu_seconds = @import("core").segments.runCommitting(r.model.device, &.{r.model.queue}, r.model.config.layers, .concurrent, &stages, commit_every) catch |err| { r.failed = true; return err; }; + // A seam cannot fail while encoding: the transport reports it out of band, and a broken pair fails the round. + if (r.reduce) |red| if (red.broken()) { + r.failed = true; + return error.PairFailed; + }; if (r.model.ane) |share| share.drain() catch |err| { r.failed = true; return err; diff --git a/zig/src/families/qwen27/dflash/runtime_callbacks.zig b/zig/src/families/qwen27/dflash/runtime_callbacks.zig index 46d252448c..bfea3d8552 100644 --- a/zig/src/families/qwen27/dflash/runtime_callbacks.zig +++ b/zig/src/families/qwen27/dflash/runtime_callbacks.zig @@ -29,7 +29,7 @@ fn project(ptr: *anyopaque, epoch: u64, x: op.View, w: ck.Tensor, mode: op.Mode) const input = try b.frame.ref(x); const output = try b.frame.ref(y); const s = @import("core").row_projection.simd; - s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false), t.matrix(), input.buf, input.off, output.buf, output.off); + s.encode(e, try s.call(&b.simd.?, t.n, t.k, x.rows, false, false, 0), t.matrix(), input.buf, input.off, output.buf, output.off); } else { const adapter = b.prepared.get(@intFromPtr(w.bytes.ptr)) orelse return error.DraftPreparationUnqualified; const input = try b.frame.ref(x); diff --git a/zig/src/families/qwen27/forward.zig b/zig/src/families/qwen27/forward.zig index 7cf8e564bc..c2e6d8964a 100644 --- a/zig/src/families/qwen27/forward.zig +++ b/zig/src/families/qwen27/forward.zig @@ -4,7 +4,8 @@ const mtl = @import("metal"); const Model = @import("model.zig").Model; const wts = @import("gpu_weights.zig"); const Config = @import("config.zig").Config; -const Ref = @import("projection.zig").Ref; +/// A tensor in a frame buffer: the `Reduce` callbacks pass refs, so the type the seam API exposes here. +pub const Ref = @import("projection.zig").Ref; const Frame = @import("gpu_frame.zig").Frame; const ane_mlp = @import("ane_mlp.zig"); @@ -15,6 +16,47 @@ pub const Mixer = struct { tap: ?*const fn (*anyopaque, mtl.ComputeEncoder, usize, Ref, u32) anyerror!void = null, }; +/// One exchange: `x` its index, `ref` this rank's `width` columns of the seam in the window's slot for x. +pub const Seam = struct { x: u32, ref: Ref }; + +/// A rank's per-layer all-gather behind one interface: transport is a model flag, not a fork in the layer program. +pub const Reduce = struct { + ptr: *anyopaque, + /// Pure, and taken BEFORE the producer kernel writes `ref`: this seam's index and where this rank's part goes. + part: *const fn (*anyopaque) Seam, + /// This rank's own columns of the seam's input rows into x's slot: a bf16 copy, so it carries the frame's bits. + stage: *const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, u32, u32, u32) void, + /// Relays x, waits for the peer's columns, and reopens the encoder: the caller uses the reopened copy. + gather: *const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, Ref, u32, u32, u32) void, + /// Out of band: a seam cannot fail while encoding, so `failed()` is checked after each seam and command buffer. + failed: *const fn (*anyopaque) bool, + /// Head seam: this rank's argmax over `lo..lo+n` (stride `stride`) meets the peer's in `picks`, one u32 a row. + argmax: ?*const fn (*anyopaque, *mtl.ComputeEncoder, u32, Ref, u32, u32, u32, u32, Ref) void = null, + pub fn next(r: Reduce) Seam { + return r.part(r.ptr); + } + /// The next exchange's index, for a seam whose part the kernel writes itself (`argmax` writes the picks). + pub fn nextX(r: Reduce) u32 { + return r.part(r.ptr).x; + } + /// A seam's input rows: stage this rank's columns of `src`, then take both halves of the frame into `out`. + pub fn rowsIn(r: Reduce, e: *mtl.ComputeEncoder, se: Seam, src: Ref, rows: u32, width: u32, stride: u32, out: Ref, out_stride: u32) void { + r.stage(r.ptr, e, se.x, src, rows, width, stride); + r.gather(r.ptr, e, se.x, se.ref, out, rows, width, out_stride); + } + /// A seam's output rows: the producer wrote this rank's rows to `se.ref`; both ranks' rows land in `out`. + pub fn rowsOut(r: Reduce, e: *mtl.ComputeEncoder, se: Seam, out: Ref, rows: u32, width: u32, out_stride: u32) void { + r.gather(r.ptr, e, se.x, se.ref, out, rows, width, out_stride); + } + pub fn pick(r: Reduce, e: *mtl.ComputeEncoder, x: u32, logits: Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: Ref) void { + r.argmax.?(r.ptr, e, x, logits, stride, lo, n, rows, picks); + } + /// Whether the peer link failed out of band: after a seam, before the next seam's kernel is encoded. + pub fn broken(r: Reduce) bool { + return r.failed(r.ptr); + } +}; + /// Serial decode and prompt segments own separate frames and order shared caches at the mixer. pub const Head = enum { all, last, none }; pub const Stages = struct { @@ -27,6 +69,10 @@ pub const Stages = struct { head: Head = .all, /// An image prompt's rows: the frame's `takes` rows take theirs over their placeholders' embeddings. images: ?Ref = null, + /// A tensor-parallel rank's all-reduce: null on one Mac, and then every seam is compiled out of the program. + reduce: ?Reduce = null, + /// Head seam picks (one u32 a row), required with `Reduce.argmax`: the ranks must pick the same token. + picks: ?Ref = null, pub fn begin(s: *Stages, lane: *@import("core").segments.Lane) !void { const m = s.model; if (s.rows == 0 or s.rows > m.frame.capacity or lane.k != 0) return error.BadFrameRows; @@ -67,21 +113,43 @@ pub const Stages = struct { pub fn post(s: *Stages, lane: *@import("core").segments.Lane, index: usize) !void { const m = s.model; const f = &m.frame; - const e = lane.enc; + var e = lane.enc; const layer = m.weights.layers[index]; const h: u32 = @intCast(m.config.hidden); const out = switch (layer.mixer) { .linear => |g| g.out, .attention => |a| a.out, }; - try m.kernels.quant(e, out, f.get(.mixer_out), null, f.get(.dims), f.get(.output), s.rows); - e.barrier(); + // Seam A: a rank owns output rows, not a K half: gather the mixer's columns before o_proj, its rows after. + const split_attention = s.reduce != null and m.shard.splitsAttention(); + var out_in = f.get(.mixer_out); + if (split_attention) { + std.debug.assert(m.shard.ranks == 2); // two Macs: gather's other half is the peer's half + const cols: u32 = @intCast(@divExact(out.k, 2)); // this rank's mixer columns, rows contiguous in `.mixer_out` + const first = s.reduce.?.next(); + s.reduce.?.rowsIn(&e, first, out_in, s.rows, cols, cols, f.get(.mixed), out.k); + lane.enc = e; // rowsIn reopened the encoder; the lane carries it to the next layer + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + out_in = f.get(.mixed); + } + const product = if (split_attention) s.reduce.?.next() else null; + if (product) |sm| try m.kernels.quant(lane.enc, out, out_in, null, f.get(.dims), sm.ref, s.rows) else try m.kernels.quant(lane.enc, out, out_in, null, f.get(.dims), f.get(.output), s.rows); + lane.enc.barrier(); + if (product) |sm| { + s.reduce.?.rowsOut(&e, sm, f.get(.output), s.rows, out.n, h); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + } try m.kernels.norm(e, s.hidden, f.get(.output), wts.ref(layer.post_norm), f.get(.eps), f.get(.dims), s.residual, f.get(.input), f.get(.sums), h, s.rows); e.barrier(); std.mem.swap(Ref, &s.hidden, &s.residual); var up = f.get(.mixed); up.offset += @as(usize, m.config.intermediate) * 2; const width: u32 = @intCast(m.config.intermediate); + // No ANE share on a pair (a rank's MLP columns are the GPU's); the assert keeps the branches from disagreeing. + std.debug.assert(s.reduce == null or m.ane == null); const shared = if (m.ane) |sh| m.kernels.prompt and s.rows <= sh.capacity() else false; if (if (shared) m.ane else null) |share| { // ANE slices run the shared columns; the GPU adds its remaining columns. @@ -111,14 +179,36 @@ pub const Stages = struct { try m.kernels.mlp(e, f.get(.mixed), up, f.get(.dims), f.get(.activated), f.get(.act_sums), width, s.rows); } lane.enc.barrier(); - try m.kernels.quant(lane.enc, layer.down, f.get(.activated), sums(f, .act_sums, s.rows), f.get(.dims), f.get(.output), s.rows); + // Seam B: gate|up gives each rank its columns, so down_proj's full K needs the peer's; `.mlp` keeps it whole. + var down_in = f.get(.activated); + const acts = if (s.reduce) |r| r.next() else null; + if (acts) |sm| { + s.reduce.?.rowsIn(&e, sm, down_in, s.rows, width, width, f.get(.mixed), layer.down.k); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + down_in = f.get(.mixed); + } + // The MLP is column-split on every pair (`.mlp` mode splits it too), so its seam is unconditional there. + const split_down = layer.down.n != h; + const seam_b = if (s.reduce != null and split_down) s.reduce.?.next() else null; + // A gathered frame's group sums are the peer's to write, so the kernel derives them where there is no peer. + const down_sums: ?@import("quant_gpu.zig").Sums = if (s.reduce == null) sums(f, .act_sums, s.rows) else null; + if (seam_b) |sm| try m.kernels.quant(lane.enc, layer.down, down_in, down_sums, f.get(.dims), sm.ref, s.rows) else try m.kernels.quant(lane.enc, layer.down, down_in, down_sums, f.get(.dims), f.get(.output), s.rows); lane.enc.barrier(); + if (seam_b) |sm| { + s.reduce.?.rowsOut(&e, sm, f.get(.output), s.rows, layer.down.n, h); + lane.enc = e; + if (s.reduce.?.broken()) return error.PairFailed; + e.barrier(); + } s.pending = f.get(.output); } pub fn finish(s: *Stages, lane: *@import("core").segments.Lane) !void { const m = s.model; const f = &m.frame; - const e = lane.enc; + // The head seam reopens the encoder, so it goes back in the lane; rank 0's half-vocab argmax is not a token. + var e = lane.enc; try m.kernels.norm(e, s.hidden, s.pending, wts.ref(m.weights.norm), f.get(.eps), f.get(.dims), s.residual, f.get(.input), f.get(.sums), @intCast(m.config.hidden), s.rows); e.barrier(); std.mem.swap(Ref, &s.hidden, &s.residual); @@ -134,12 +224,35 @@ pub const Stages = struct { try m.kernels.quant(e, m.weights.head, input, head_sums, f.get(.dims), f.get(.logits), if (s.head == .last) 1 else s.rows); } e.barrier(); + if (s.reduce) |r| { + if (s.head == .none) return; + const picks = s.picks orelse return error.HeadSeamMissing; + const seam = r.argmax orelse return error.HeadSeamMissing; + const span = m.shard.vocabRows(); + const local_rows: u32 = if (s.head == .last) 1 else s.rows; + seam(r.ptr, &e, r.nextX(), f.get(.logits), @intCast(m.config.vocab), @intCast(span.from), @intCast(span.len), local_rows, picks); + lane.enc = e; // the head seam reopened the encoder; the caller commits this one + if (r.broken()) return error.PairFailed; + } } }; -pub fn encode(m: *Model, e: mtl.ComputeEncoder, rows: u32, mixer: Mixer) !void { +/// `e` is the caller's encoder, but seams reopen it: the final one is written back through the pointer, so commit it. +pub fn encode(m: *Model, e: *mtl.ComputeEncoder, rows: u32, mixer: Mixer) !void { var stages = Stages{ .model = m, .rows = rows, .hooks = mixer }; - var lane = @import("core").segments.Lane{ .k = 0, .cb = undefined, .enc = e, .fence = undefined }; + try run(m, e, &stages); +} + +/// The tensor-parallel entry point: the one-Mac program with each layer's all-reduces inserted behind `reduce`. +pub fn encodeTp(m: *Model, e: *mtl.ComputeEncoder, rows: u32, mixer: Mixer, reduce: Reduce) !void { + var stages = Stages{ .model = m, .rows = rows, .hooks = mixer, .reduce = reduce }; + try run(m, e, &stages); +} + +fn run(m: *Model, e: *mtl.ComputeEncoder, stages: *Stages) !void { + // `cb`/`fence` stay undefined: a seam is the only closer, and it orders its writes through the transport's fence. + var lane = @import("core").segments.Lane{ .k = 0, .cb = undefined, .enc = e.*, .fence = undefined }; + defer e.* = lane.enc; try stages.begin(&lane); for (0..m.config.layers) |index| { try stages.pre(&lane, index); diff --git a/zig/src/families/qwen27/gdn_contract.zig b/zig/src/families/qwen27/gdn_contract.zig index 3dce6d0892..62ad130390 100644 --- a/zig/src/families/qwen27/gdn_contract.zig +++ b/zig/src/families/qwen27/gdn_contract.zig @@ -37,7 +37,7 @@ pub const Params = extern struct { slots: u32, weight_flags: u32, eps: f32, - /// 0: z, a and b in their own rows; else one fused zba row of this many values (z, then b, then a). + /// 0: z, a, b in own rows; else one fused zba row (z, b, a) whose rank blocks sit at zbaB/zbaA and is zbaDim wide. zba_stride: u32 = 0, advance: u32 = 0, // 1: a shared round's one-row windows step their state in the scan; the replay skips them }; @@ -84,7 +84,8 @@ pub const Bytes = struct { pub fn init(s: Shape, p: Params, k: WeightKinds) !Bytes { var expected = try s.params(p.rows, p.slots, k); - if (p.zba_stride != 0 and p.zba_stride < s.value + 2 * s.c.v_heads) return error.BadGdnDispatch; + // The row must span this rank's b and a reads past the other rank's padding, so the bound is the config stride. + if (p.zba_stride != 0 and p.zba_stride < s.c.zbaDim()) return error.BadGdnDispatch; expected.zba_stride = p.zba_stride; if (!std.mem.eql(u8, std.mem.asBytes(&p), std.mem.asBytes(&expected))) return error.BadGdnDispatch; return .{ diff --git a/zig/src/families/qwen27/gdn_hook.zig b/zig/src/families/qwen27/gdn_hook.zig index f2d6ce2e0f..944fd7aa30 100644 --- a/zig/src/families/qwen27/gdn_hook.zig +++ b/zig/src/families/qwen27/gdn_hook.zig @@ -33,12 +33,11 @@ pub const Hook = struct { } }; -/// z, b and a read in place from the fused zba rows. +/// z, b and a read in place from fused rows: a rank's b starts rank*24 after z, a 48 later (zbaDim = width + 96). fn refs(weight: Weights.Gdn, frame: *Frame) !pool_gpu.Refs { - const width = frame.config.gdnValueDim(); - const heads = frame.config.v_heads; + const c = frame.config; const mixed = frame.get(.mixed); - return .{ .z = mixed, .b = .{ .buffer = mixed.buffer, .offset = mixed.offset + width * 2 }, .a = .{ .buffer = mixed.buffer, .offset = mixed.offset + (width + heads) * 2 }, .output = frame.get(.mixer_out), .scalars = try pool_gpu.Scalars.from(weight), .zba_stride = @intCast(width + 2 * heads) }; + return .{ .z = mixed, .b = .{ .buffer = mixed.buffer, .offset = mixed.offset + c.zbaB() * 2 }, .a = .{ .buffer = mixed.buffer, .offset = mixed.offset + c.zbaA() * 2 }, .output = frame.get(.mixer_out), .scalars = try pool_gpu.Scalars.from(weight), .zba_stride = @intCast(c.zbaDim()) }; } fn devicePlan(plan: Plan, frame: *Frame) pool_gpu.DevicePlan { diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index d746877919..5fd13d67b1 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -1,7 +1,8 @@ //! Small operators use borrowed buffers and caller-owned encoders; copies never reinterpret activations. const mtl = @import("metal"); const source = @import("qwen_runtime_sources"); -const Config = @import("config.zig").Config; +const config = @import("config.zig"); +const Config = config.Config; const abi = @import("glue.zig"); const projection = @import("projection.zig"); const profile = @import("core").gpu_profile; @@ -51,12 +52,14 @@ pub const Kernels = struct { e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); } + /// The table is whole on every rank; only the head is sliced, so the bound is its row count, not `c.vocab`. pub fn embedding(k: Kernels, input_encoder: mtl.ComputeEncoder, c: Config, p: projection.Triple, ids: Ref, out: Ref, rows: u32) !void { + const vocab = try config.embedBound(&c, p.weight.shape[0]); var span = profile.begin(k.profiler, input_encoder, .embedding); errdefer span.cancel(); const e = span.encoder; - const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(c.vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; + const a = abi.Embed{ .width = @intCast(c.hidden), .rows = rows, .vocab = @intCast(vocab), .bits = p.bits, .group = p.group, .metadata = .bf16 }; try abi.checkEmbed(a); e.setPipeline(k.embed); e.setBuffer(ids.buffer, ids.offset, 0); diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index bb0fba68e3..abc577fa46 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -1,8 +1,8 @@ -//! The Metal graph borrows native scalar tensors and owns only byte-preserving affine layouts. +//! Borrows native scalar tensors and owns only byte-preserving affine layouts; rows and columns follow config.Shard. const std = @import("std"); const mtl = @import("metal"); const cp = @import("core").checkpoint_metal; -const Config = @import("config.zig").Config; +const cfg = @import("config.zig"); const projection = @import("projection.zig"); const lane = @import("core").lane_projection; const widths = @import("widths.zig"); @@ -89,31 +89,91 @@ const Builder = struct { return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; } - /// `parts` (stacked in order, `n` rows of `k` each, at their widths) staged from the files and repacked. + /// One part's staged tensors, shape-checked and viewed at the rank's rows; the stage must outlive these views. + fn viewed(p: Part, k: usize, tensors: [3]Tensor) !projection.Triple { + const weight = try checked(tensors[0], .u32, &.{ p.n, k * p.w.bits / 32 }); + const scales = try checked(tensors[1], .bf16, &.{ p.n, k / p.w.group }); + const biases = try checked(tensors[2], .bf16, &.{ p.n, k / p.w.group }); + return .{ .weight = try rowView(weight, p.span), .scales = try rowView(scales, p.span), .biases = try rowView(biases, p.span), .bits = p.w.bits, .group = p.w.group, .k = p.k, .k0 = p.k0, .gn = p.n }; + } + + /// Parts at their widths and K windows, all staged before any is viewed: the stage is one reused buffer. fn linear(b: Builder, parts: []const Part, k: usize) !Linear { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); - const names = try arena.allocator().alloc([]const u8, parts.len * 3); + if (parts.len > 3) return error.BadProjectionShape; + for (parts) |p| if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; + const a = arena.allocator(); + const names = try a.alloc([]const u8, parts.len * 3); for (parts, 0..) |p, i| for ([_][]const u8{ ".weight", ".scales", ".biases" }, 0..) |suffix, j| { - names[i * 3 + j] = try std.mem.concat(arena.allocator(), u8, &.{ b.prefix, p.module, suffix }); + names[i * 3 + j] = try std.mem.concat(a, u8, &.{ b.prefix, p.module, suffix }); }; try b.stage.reserve(b.checkpoint, names); defer b.stage.reset(); var triples: [3]projection.Triple = undefined; - if (parts.len > triples.len) return error.BadProjectionShape; for (parts, triples[0..parts.len], 0..) |p, *t, i| { - if (k * p.w.bits % 32 != 0 or k % p.w.group != 0) return error.BadNativeTensor; - t.* = .{ - .weight = try checked(try b.stage.read(b.checkpoint, names[i * 3]), .u32, &.{ p.n, k * p.w.bits / 32 }), - .scales = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 1]), .bf16, &.{ p.n, k / p.w.group }), - .biases = try checked(try b.stage.read(b.checkpoint, names[i * 3 + 2]), .bf16, &.{ p.n, k / p.w.group }), - .bits = p.w.bits, - .group = p.w.group, - }; + t.* = try viewed(p, k, .{ + try b.stage.read(b.checkpoint, names[i * 3]), + try b.stage.read(b.checkpoint, names[i * 3 + 1]), + try b.stage.read(b.checkpoint, names[i * 3 + 2]), + }); } return b.stack(triples[0..parts.len]); } + /// One tensor a rank takes in blocks (GDN q | k | v); `rows` are each block's checkpoint rows, which prepare sums. + fn blocks(b: Builder, module: []const u8, n: usize, k: usize, w: Width, spans: []const cfg.Span, rows: []const usize) !Linear { + var arena = std.heap.ArenaAllocator.init(b.allocator); + defer arena.deinit(); + if (k * w.bits % 32 != 0 or k % w.group != 0) return error.BadNativeTensor; + const suffixes = [_][]const u8{ ".weight", ".scales", ".biases" }; + const names = try arena.allocator().alloc([]const u8, suffixes.len); + for (suffixes, 0..) |suffix, j| names[j] = try std.mem.concat(arena.allocator(), u8, &.{ b.prefix, module, suffix }); + try b.stage.reserve(b.checkpoint, names); + defer b.stage.reset(); + const weight = try checked(try b.stage.read(b.checkpoint, names[0]), .u32, &.{ n, k * w.bits / 32 }); + const scales = try checked(try b.stage.read(b.checkpoint, names[1]), .bf16, &.{ n, k / w.group }); + const biases = try checked(try b.stage.read(b.checkpoint, names[2]), .bf16, &.{ n, k / w.group }); + var triples: [3]projection.Triple = undefined; + if (spans.len > triples.len or rows.len != spans.len) return error.BadProjectionShape; + for (spans, rows, triples[0..spans.len]) |span, gn, *t| { + t.* = .{ .weight = try rowView(weight, span), .scales = try rowView(scales, span), .biases = try rowView(biases, span), .bits = w.bits, .group = w.group, .gn = gn }; + } + return b.stack(triples[0..spans.len]); + } + + /// `conv1d.weight`'s row blocks fused in the same q | k | v order, as one bf16 copy (the file's is not repacked). + fn conv(b: Builder, module: []const u8, kernel: usize, n: usize, spans: []const cfg.Span) !Tensor { + const t = try b.tensor(module, .bf16, &.{ n, kernel, 1 }); + var at: usize = 0; + for (spans) |span| { + if (span.len == 0 or span.from != at) break; + at = span.end(); + } + if (at == n) return t; // the blocks are the whole tensor in order: no copy + var rows: usize = 0; + for (spans) |span| rows += span.len; + const out = try b.device.buffer(rows * kernel * @sizeOf(u16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + b.owned.append(b.allocator, out) catch |err| { + out.deinit(); + return err; + }; + const from = t.host(u16); + const to = out.slice(u16, rows * kernel); + var base: usize = 0; + for (spans) |span| { + if (span.end() > n) return error.BadNativeTensor; + @memcpy(to[base * kernel ..][0 .. span.len * kernel], from[span.from * kernel ..][0 .. span.len * kernel]); + base += span.len; + } + var view = t; + view.buffer = out; + view.offset = 0; + view.bytes = rows * kernel * @sizeOf(u16); + view.shape[0] = rows; + return view; + } + /// Fused members of one width take one matrix as before; a mixed stack becomes one matrix a member. fn stack(b: Builder, parts: []const projection.Triple) !Linear { var ws: [3]Width = undefined; @@ -133,12 +193,14 @@ const Builder = struct { } var fused = runs[0].linear; fused.n = first; + fused.gn = 0; + for (runs) |run| fused.gn += run.linear.gn; fused.raw = null; fused.runs = runs; return fused; } - fn layer(b: Builder, c: Config, i: usize) !Layer { + fn layer(b: Builder, s: *const cfg.Shard, i: usize) !Layer { var arena = std.heap.ArenaAllocator.init(b.allocator); defer arena.deinit(); const a = arena.allocator(); @@ -148,41 +210,86 @@ const Builder = struct { return std.mem.concat(allocator, u8, &.{ p, leaf }); } }; + const c = s.local; + const g = s.global; const w = b.widths.layers[i]; const M = widths.Member; + const qkv = s.qkvSpans(); + const qkvRows = [3]usize{ g.k_heads * g.dk, g.k_heads * g.dk, g.gdnValueDim() }; + const convRows = s.convSpans(); const mixer: @FieldType(Layer, "mixer") = switch (c.kind(i)) { .linear => .{ .linear = .{ - .qkv = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.in_proj_qkv"), .n = c.gdnQkvDim(), .w = w[@backingInt(M.qkv)] }}, c.hidden), - .zba = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = c.gdnValueDim(), .w = w[@backingInt(M.z)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = c.v_heads, .w = w[@backingInt(M.b)] }, .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = c.v_heads, .w = w[@backingInt(M.a)] } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = c.hidden, .w = w[@backingInt(M.out)] }}, c.gdnValueDim()), - .conv = try b.tensor(try Name.full(a, prefix, "linear_attn.conv1d.weight"), .bf16, &.{ c.gdnQkvDim(), c.conv_kernel, 1 }), - .a_log = try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{c.v_heads}), - .dt = try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{c.v_heads}), + .qkv = try b.blocks(try Name.full(a, prefix, "linear_attn.in_proj_qkv"), g.gdnQkvDim(), g.hidden, w[@backingInt(M.qkv)], &qkv, &qkvRows), + .zba = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_z"), .n = g.gdnValueDim(), .w = w[@backingInt(M.z)], .span = s.zRows() }, + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_b"), .n = g.v_heads, .w = w[@backingInt(M.b)], .span = s.abRows() }, + .{ .module = try Name.full(a, prefix, "linear_attn.in_proj_a"), .n = g.v_heads, .w = w[@backingInt(M.a)], .span = s.abRows() }, + }, g.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "linear_attn.out_proj"), .n = g.hidden, .w = w[@backingInt(M.out)], .span = s.outputRows() }}, g.gdnValueDim()), + .conv = try b.conv(try Name.full(a, prefix, "linear_attn.conv1d.weight"), g.conv_kernel, g.gdnQkvDim(), &convRows), + .a_log = try rowView(try b.tensor(try Name.full(a, prefix, "linear_attn.A_log"), .bf16, &.{g.v_heads}), s.valueRows()), + .dt = try rowView(try b.tensor(try Name.full(a, prefix, "linear_attn.dt_bias"), .bf16, &.{g.v_heads}), s.valueRows()), .norm = try b.tensor(try Name.full(a, prefix, "linear_attn.norm.weight"), .bf16, &.{c.dv}), } }, .attention => .{ .attention = .{ - .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = c.qDim(), .w = w[@backingInt(M.q)] }}, c.hidden), - .kv = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = c.kvDim(), .w = w[@backingInt(M.k)] }, .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = c.kvDim(), .w = w[@backingInt(M.v)] } }, c.hidden), - .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = c.hidden, .w = w[@backingInt(M.o)] }}, c.heads * c.head_dim), + .q = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.q_proj"), .n = g.qDim(), .w = w[@backingInt(M.q)], .span = s.qRows() }}, g.hidden), + .kv = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "self_attn.k_proj"), .n = g.kvDim(), .w = w[@backingInt(M.k)], .span = s.kvRows() }, + .{ .module = try Name.full(a, prefix, "self_attn.v_proj"), .n = g.kvDim(), .w = w[@backingInt(M.v)], .span = s.kvRows() }, + }, g.hidden), + .out = try b.linear(&.{.{ .module = try Name.full(a, prefix, "self_attn.o_proj"), .n = g.hidden, .w = w[@backingInt(M.o)], .span = s.outputRows() }}, g.heads * g.head_dim), .q_norm = try b.tensor(try Name.full(a, prefix, "self_attn.q_norm.weight"), .bf16, &.{c.head_dim}), .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); + const gu = try b.linear(&.{ + .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = g.intermediate, .w = w[@backingInt(M.gate)], .span = s.intermediateRows() }, + .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = g.intermediate, .w = w[@backingInt(M.up)], .span = s.intermediateRows() }, + }, g.hidden); + const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = g.hidden, .w = w[@backingInt(M.down)], .span = s.outputRows() }}, g.intermediate); return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .mixer = mixer }; } }; -const Part = struct { module: []const u8, n: usize, w: Width }; +const Part = struct { + /// The module's checkpoint name, without the shard's namespace prefix. + module: []const u8, + /// The rows the checkpoint's tensor has, before this rank takes its span. + n: usize, + w: Width, + /// The rows this rank takes: the whole tensor unless the rank splits this matrix. + span: cfg.Span = .{}, + /// The input columns this rank takes of each row, and their count, when the matrix is split on its input side. + k0: usize = 0, + k: usize = 0, +}; + +/// Rows `[span.from, end)` of `t`: rows are contiguous, so the same buffer at a later offset and a shorter shape. +fn rowView(t: Tensor, span: cfg.Span) !Tensor { + if (span.len == 0) return t; + if (span.from == 0 and span.len == t.shape[0]) return t; + if (span.end() > t.shape[0]) return error.BadNativeTensor; + var row: usize = 1; + for (t.shape[1..t.rank]) |d| row *= d; + const elem = t.dtype.size(); + const at = span.from * row * elem; + if (at % elem != 0) return error.BadNativeTensor; + var view = t; + view.offset = t.offset + at; + view.bytes = row * span.len * elem; + view.shape[0] = span.len; + return view; +} fn checked(t: Tensor, dtype: cp.DType, shape: []const usize) !Tensor { if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) return error.BadNativeTensor; return t; } -/// The graph from `checkpoint`, whose `repacked` tensors stay in their files: each Linear is staged, repacked, dropped. -pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, c: Config, ws: *const widths.Widths) !Weights { +/// Graph from `checkpoint`; each Linear is staged, repacked, dropped. `s` is the rank's share, `local` from `global`. +pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const cp.Checkpoint, s: *const cfg.Shard, ws: *const widths.Widths) !Weights { + const c = s.local; + const g = s.global; const prefix: []const u8 = if (checkpoint.has("language_model.model.embed_tokens.weight")) "language_model." else if (checkpoint.has("model.embed_tokens.weight")) "" else return error.MissingTextNamespace; if (checkpoint.has("language_model.model.embed_tokens.weight") and checkpoint.has("model.embed_tokens.weight")) return error.AmbiguousTextNamespace; var owned: std.ArrayList(mtl.Buffer) = .empty; @@ -200,10 +307,10 @@ pub fn load(allocator: std.mem.Allocator, device: mtl.Device, checkpoint: *const var stage = cp.Stage{ .device = device }; defer stage.deinit(); const b = Builder{ .allocator = allocator, .device = device, .checkpoint = checkpoint, .prefix = prefix, .owned = &owned, .stage = &stage, .runs = &runs, .widths = ws }; - for (layers, 0..) |*l, i| l.* = try b.layer(c, i); - const embed = try b.quant("model.embed_tokens", c.vocab, c.hidden, ws.embed); + for (layers, 0..) |*l, i| l.* = try b.layer(s, i); + const embed = try b.quant("model.embed_tokens", g.vocab, g.hidden, ws.embed); const norm = try b.tensor("model.norm.weight", .bf16, &.{c.hidden}); - const head = try b.linear(&.{.{ .module = "lm_head", .n = c.vocab, .w = ws.head }}, c.hidden); + const head = try b.linear(&.{.{ .module = "lm_head", .n = g.vocab, .w = ws.head, .span = s.vocabRows() }}, g.hidden); if (device.tensorUnits()) try laneOrder(allocator, device, layers, head); return .{ .allocator = allocator, .layers = layers, .embed = embed, .norm = norm, .head = head, .owned = owned, .runs = runs }; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 4c5e7e7212..28f6f5c9e4 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -18,6 +18,9 @@ pub const Model = struct { allocator: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, + /// This rank's share of the family's weights; `local` shapes are what this graph builds, `global` the file's. + shard: cfg.Shard, + /// The local config, so downstream shapes are the rank's own. config: cfg.Config, checkpoint: core.checkpoint_metal.Checkpoint, weights: wts.Weights, @@ -41,6 +44,16 @@ pub const Model = struct { return loadWith(a, io, dir, rows, null, true); } + /// One rank of a `ranks`-way tensor-parallel run: the same checkpoint, sliced in the engine. + pub fn loadShard(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, rank: u32, ranks: u32, mode: cfg.Mode, share_fits: ?*const fn (u64) bool) !*Model { + return loadShared(a, io, dir, rows, rank, ranks, mode, share_fits, false); + } + + /// One rank of a pair, its share decided after everything else is accounted for (the serving order). + pub fn loadShardDeferred(a: std.mem.Allocator, io: std.Io, dir: []const u8, rows: u32, rank: u32, ranks: u32, mode: cfg.Mode) !*Model { + return loadShared(a, io, dir, rows, rank, ranks, mode, null, true); + } + fn loadWith( a: std.mem.Allocator, io: std.Io, @@ -48,27 +61,44 @@ pub const Model = struct { rows: u32, share_fits: ?*const fn (u64) bool, later: bool, + ) !*Model { + return loadShared(a, io, dir, rows, 0, 1, .full, share_fits, later); + } + + fn loadShared( + a: std.mem.Allocator, + io: std.Io, + dir: []const u8, + rows: u32, + rank: u32, + ranks: u32, + mode: cfg.Mode, + share_fits: ?*const fn (u64) bool, + later: bool, ) !*Model { const path = try std.fs.path.join(a, &.{ dir, "config.json" }); defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); defer a.free(text); const c = try cfg.parse(a, text, null); - try supported(c); + var local = c; + const shard = if (ranks == 1) cfg.Shard{ .rank = 0, .ranks = 1, .global = c, .local = local } else try cfg.split(&local, rank, ranks, mode); + try supported(shard.global); var formats = try affine.Formats.init(a, text); defer formats.deinit(); - var ws = try widths.Widths.resolve(a, &formats, c); + var ws = try widths.Widths.resolve(a, &formats, shard.local); defer ws.deinit(); const m = try a.create(Model); errdefer a.destroy(m); m.allocator = a; - m.config = c; + m.shard = shard; + m.config = shard.local; m.widths_tag = ws.tag(); m.device = try mtl.Device.init(); errdefer m.device.deinit(); if (!qualifiedGeneration(m.device.generation())) return error.QwenMetalQualificationRequired; var why: widths.Why = .{}; - widths.admit(&ws, c, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { + widths.admit(&ws, shard.local, .{ .tensor_units = m.device.tensorUnits() }, &why) catch |err| { std.log.err("qwen27: {s} (layer {?d}) is {d}-bit group {d}; this chip's projection kernels do not run it ({s})", .{ why.module, why.layer, why.width.bits, why.width.group, @errorName(why.reason) }); return err; }; @@ -80,13 +110,13 @@ pub const Model = struct { defer core.checkpoint_host.freeShardFiles(a, files); // projections stay in the files until load repacks them: the weights live once, in their lane layouts for (files) |file| try m.checkpoint.addFileSplit(m.device, file, "", "language_model.", wts.repacked); - m.weights = try wts.load(a, m.device, &m.checkpoint, c, &ws); + m.weights = try wts.load(a, m.device, &m.checkpoint, &m.shard, &ws); errdefer m.weights.deinit(); m.kernels = .{ .allocator = a, .device = m.device, .raw_tiles = !m.device.tensorUnits() }; errdefer m.kernels.deinit(); m.glue = try glue.Kernels.init(m.device); errdefer m.glue.deinit(); - m.frame = try Frame.init(m.device, c, @max(rows, promptRows())); + m.frame = try Frame.init(m.device, shard.local, @max(rows, promptRows())); errdefer m.frame.deinit(); m.ane = null; m.share_pending = true; @@ -105,7 +135,9 @@ pub const Model = struct { defer m.checkpoint.closeFiles(); const c = m.config; const columns = try ane_policy.columns(if (platform.c.getenv("TF_ANE_COLUMNS")) |v| std.mem.span(v) else null); - const share = ane_policy.requestedOnChip(columns, m.device.tensorUnits()) and mtl.ane.available() and m.weights.uniformMlp(); + // The share is a whole MLP's columns: a rank holds half of them, so tensor parallel keeps prompts on the GPU. + const share = m.shard.ranks == 1 and ane_policy.requestedOnChip(columns, m.device.tensorUnits()) and mtl.ane.available() and m.weights.uniformMlp(); + if (m.shard.ranks != 1) std.log.info("qwen27: tensor parallel rank {d}/{d} runs prompts on the GPU alone", .{ m.shard.rank, m.shard.ranks }); const bytes = ane_mlp.Share.bytes(c, columns); const slices = (m.frame.capacity + ane_mlp.width - 1) / ane_mlp.width; const required: ?u64 = if (share) ane_mlp.Share.totalBytes(c, columns, slices, ane_policy.build_runtime_overhead) catch |err| switch (err) { diff --git a/zig/src/families/qwen27/projection.zig b/zig/src/families/qwen27/projection.zig index a3603ee3f2..60d5420d01 100644 --- a/zig/src/families/qwen27/projection.zig +++ b/zig/src/families/qwen27/projection.zig @@ -6,11 +6,12 @@ const row = @import("core").row_projection; const lane = @import("core").lane_projection; pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; -/// A device matrix; `runs` non-empty means fused members of different widths, each its own product into these rows. -pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{} }; +/// A device matrix; `runs` non-empty means fused members of different widths. `gn`: the checkpoint's rows (0 = `n`). +pub const Linear = struct { words: Ref, pairs: Ref, n: u32, k: u32, tile: u32, slices: u32, raw: ?row.Weights = null, bits: u8 = 4, group: u32 = 64, runs: []const Run = &.{}, gn: usize = 0 }; /// One product inside a mixed fused matrix, writing output columns [first, first + linear.n). pub const Run = struct { linear: Linear, first: u32 }; -pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64 }; +/// Staged tensors: `k`/`k0` the column window, `gn` the part's checkpoint rows when `weight.shape[0]` is a slice. +pub const Triple = struct { weight: checkpoint.Tensor, scales: checkpoint.Tensor, biases: checkpoint.Tensor, bits: u8 = 4, group: u32 = 64, k: usize = 0, k0: usize = 0, gn: usize = 0 }; pub fn split(n: usize, k: usize) !u32 { if (n == 0 or k == 0 or n % 4 != 0 or k % 64 != 0) return error.BadProjectionShape; @@ -33,24 +34,42 @@ pub fn promptSplit(n: usize, k: usize) ?u32 { return null; } +/// A part's rows in the checkpoint, which are the staged rows on one Mac. +fn globalRows(p: Triple) usize { + return if (p.gn != 0) p.gn else p.weight.shape[0]; +} + +/// The fused matrix's rows in the checkpoint: what the M5 and prompt tables key on. +fn splitRows(parts: []const Triple) usize { + var rows: usize = 0; + for (parts) |p| rows += globalRows(p); + return rows; +} + pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, tile: u32) !Linear { if (tile != 32 or parts.len == 0) return error.BadProjectionShape; const bits = parts[0].bits; const group = parts[0].group; if (bits == 0 or bits > 8 or group == 0) return error.BadAffineMetadata; - const k = parts[0].weight.shape[1] * 32 / bits; - var slices = try split(parts[0].weight.shape[0], k); + // the checkpoint's packed width, and the columns this rank takes of it (k == kg and k0 == 0 when whole) + const kg = parts[0].weight.shape[1] * 32 / bits; + const k = if (parts[0].k != 0) parts[0].k else kg; + const k0 = parts[0].k0; + if (k == 0 or k % group != 0 or k0 % group != 0 or k0 + k > kg) return error.BadAffineMetadata; + // A split is an fp32 order: the leading member's checkpoint rows key it, and every part's rows key the M5 tables. + const gn = splitRows(parts); + var slices = try split(globalRows(parts[0]), k); var n: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - if (p.bits != bits or p.group != group or k % group != 0) return error.BadAffineMetadata; - if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != k * bits or rows % 4 != 0 or try split(rows, k) != slices) return error.BadProjectionShape; - if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != k / group or p.biases.shape[1] != k / group) return error.BadAffineMetadata; + if (p.bits != bits or p.group != group or p.k != parts[0].k or p.k0 != k0) return error.BadAffineMetadata; + if (p.weight.rank != 2 or p.weight.dtype != .u32 or p.weight.shape[1] * 32 != kg * bits or rows % 4 != 0) return error.BadProjectionShape; + if (p.scales.dtype != .bf16 or p.biases.dtype != .bf16 or p.scales.rank != 2 or p.biases.rank != 2 or p.scales.shape[0] != rows or p.biases.shape[0] != rows or p.scales.shape[1] != kg / group or p.biases.shape[1] != kg / group) return error.BadAffineMetadata; n = try std.math.add(usize, n, rows); } - if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, tile, slices); - if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, tile, slices); - slices = m5Split(n, k) orelse slices; + if (!device.tensorUnits() or n % tile != 0) return prepareRows(device, allocator, owned, parts, n, k, kg, k0, tile, slices, gn); + if (bits != 4 or group != 64) return prepareLanes(device, allocator, owned, parts, n, k, kg, k0, tile, slices, gn); + slices = m5Split(gn, k) orelse slices; const w = try device.buffer(n * k / 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); owned.append(allocator, w) catch |err| { w.deinit(); @@ -64,21 +83,24 @@ pub fn prepare(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.Arr const dst = w.slice(u32, n * k / 8); const pair = sb.slice(u16, n * (k / 64) * 2); var base: usize = 0; + const groups = k / 64; + const src_groups = kg / 64; + const k0g = k0 / 64; for (parts) |p| { const rows = p.weight.shape[0]; const words = p.weight.host(u32); const scales = p.scales.host(u16); const biases = p.biases.host(u16); - for (0..rows) |r| for (0..k / 64) |g| { - const src = r * (k / 8) + g * 8; - const to = (((base + r) / tile) * (k / 64) + g) * tile * 8 + ((base + r) % tile) * 8; + for (0..rows) |r| for (0..groups) |g| { + const src = r * (kg / 8) + (k0g + g) * 8; + const to = (((base + r) / tile) * groups + g) * tile * 8 + ((base + r) % tile) * 8; @memcpy(dst[to..][0..8], words[src..][0..8]); - pair[(g * n + base + r) * 2] = scales[r * (k / 64) + g]; - pair[(g * n + base + r) * 2 + 1] = biases[r * (k / 64) + g]; + pair[(g * n + base + r) * 2] = scales[r * src_groups + k0g + g]; + pair[(g * n + base + r) * 2 + 1] = biases[r * src_groups + k0g + g]; }; base += rows; } - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .gn = gn }; } test "stacked gate-up keeps the member K split rather than recomputing it from stacked width" { @@ -89,29 +111,78 @@ test "stacked gate-up keeps the member K split rather than recomputing it from s try std.testing.expectEqual(@as(u32, 8), try split(5120, 17408)); } -fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { +/// A staged matrix of `rows` rows whose checkpoint matrix has `gn` (0: the rows are the whole matrix). +fn tri(rows: usize, gn: usize, blank: checkpoint.Tensor) Triple { + var t = blank; + t.shape[0] = rows; + return .{ .weight = t, .scales = t, .biases = t, .gn = gn }; +} + +test "a rank's row slice keeps the checkpoint matrix's K split" { + const blank: checkpoint.Tensor = .{ .buffer = undefined, .offset = 0, .bytes = 0, .dtype = .u32 }; + // gate | up: a rank stages 8,704 rows a member of the checkpoint's 17,408, so the leading member keys 2, not 4. + const gate = [2]Triple{ tri(8704, 17408, blank), tri(8704, 17408, blank) }; + try std.testing.expectEqual(@as(usize, 17408), globalRows(gate[0])); + try std.testing.expectEqual(@as(u32, 2), try split(globalRows(gate[0]), 5120)); + try std.testing.expectEqual(@as(u32, 4), try split(gate[0].weight.shape[0], 5120)); + // The fused matrix's checkpoint rows double, as the single Mac's stack: 34,816, the tables' key. + try std.testing.expectEqual(@as(usize, 34816), splitRows(&gate)); + // q: 12,288 checkpoint rows, 6,144 a rank's, 4 and 8 slices. + const q = tri(6144, 12288, blank); + try std.testing.expectEqual(@as(u32, 4), try split(globalRows(q), 5120)); + try std.testing.expectEqual(@as(u32, 8), try split(q.weight.shape[0], 5120)); + // down_proj: the M5 and prompt tables hold the checkpoint's 5,120 rows over 17,408, which a rank's 2,560 miss. + const down = tri(2560, 5120, blank); + try std.testing.expectEqual(@as(?u32, 4), m5Split(splitRows(&.{down}), 17408)); + try std.testing.expectEqual(@as(?u32, null), m5Split(down.weight.shape[0], 17408)); + try std.testing.expectEqual(@as(?u32, 2), promptSplit(splitRows(&.{down}), 17408)); + try std.testing.expectEqual(@as(?u32, null), promptSplit(down.weight.shape[0], 17408)); + // q | k | v: three blocks of one 10,240-row tensor, so the blocks' checkpoint rows sum to its width. + const qkv = [3]Triple{ tri(1024, 2048, blank), tri(1024, 2048, blank), tri(3072, 6144, blank) }; + try std.testing.expectEqual(@as(usize, 10240), splitRows(&qkv)); + try std.testing.expectEqual(@as(?u32, 2), m5Split(splitRows(&qkv), 5120)); + try std.testing.expectEqual(@as(?u32, 1), promptSplit(splitRows(&qkv), 5120)); + try std.testing.expectEqual(@as(u32, 8), try split(globalRows(qkv[0]), 5120)); + // A rank's three local blocks are 5,120 rows and key neither table, so the fallback split would stand. + var local: usize = 0; + for (qkv) |part| local += part.weight.shape[0]; + try std.testing.expectEqual(@as(usize, 5120), local); + try std.testing.expectEqual(@as(?u32, null), m5Split(local, 5120)); + try std.testing.expectEqual(@as(?u32, null), promptSplit(local, 5120)); +} + +fn prepareRows(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, kg: usize, k0: usize, tile: u32, slices: u32, gn: usize) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const bits: usize = parts[0].bits; const ng = k / parts[0].group; + const sng = kg / parts[0].group; const row_bytes = k * bits / 8; + const srow_words = kg * bits / 32; + const woff = k0 * bits / 32; + const goff = k0 / parts[0].group; const w = try own(device, allocator, owned, n * row_bytes, options); const scale = try own(device, allocator, owned, n * ng * 2, options); const bias = try own(device, allocator, owned, n * ng * 2, options); var base: usize = 0; for (parts) |p| { const rows = p.weight.shape[0]; - @memcpy(w.contents()[base * row_bytes .. (base + rows) * row_bytes], std.mem.sliceAsBytes(p.weight.host(u32))); - @memcpy(scale.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.scales.host(u16))); - @memcpy(bias.contents()[base * ng * 2 .. (base + rows) * ng * 2], std.mem.sliceAsBytes(p.biases.host(u16))); + const words = p.weight.host(u32); + const scales = p.scales.host(u16); + const biases = p.biases.host(u16); + for (0..rows) |r| { + @memcpy(w.contents()[(base + r) * row_bytes ..][0..row_bytes], std.mem.sliceAsBytes(words[r * srow_words + woff ..][0 .. k * bits / 32])); + @memcpy(scale.contents()[(base + r) * ng * 2 ..][0 .. ng * 2], std.mem.sliceAsBytes(scales[r * sng + goff ..][0..ng])); + @memcpy(bias.contents()[(base + r) * ng * 2 ..][0 .. ng * 2], std.mem.sliceAsBytes(biases[r * sng + goff ..][0..ng])); + } base += rows; } const raw = row.Weights{ .w = w, .scales = scale, .biases = bias, .n = n, .k = k, .group = parts[0].group, .bits = parts[0].bits, .sum = .f32 }; try raw.validate(); - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = scale }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .raw = raw, .bits = parts[0].bits, .group = parts[0].group, .gn = gn }; } /// Widths other than 4/64 on tensor units: lane_projection.weight_layout byte copies for the non-reg lane kernel. -fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, tile: u32, slices: u32) !Linear { +fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.ArrayList(mtl.Buffer), parts: []const Triple, n: usize, k: usize, kg: usize, k0: usize, tile: u32, slices: u32, gn: usize) !Linear { const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const layout = lane.Layout{ .n = n, .k = k, .format = .{ .bits = parts[0].bits, .group = parts[0].group }, .sk = slices }; try layout.validate(); @@ -119,24 +190,26 @@ fn prepareLanes(device: mtl.Device, allocator: std.mem.Allocator, owned: *std.Ar const sb = try own(device, allocator, owned, layout.metadataElements() * 2, options); var base: usize = 0; for (parts) |p| { - packTiles(.{ .words = p.weight.host(u32), .scales = p.scales.host(u16), .biases = p.biases.host(u16), .rows = p.weight.shape[0] }, base, n, k, parts[0].bits, parts[0].group, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); + packTiles(.{ .words = p.weight.host(u32), .scales = p.scales.host(u16), .biases = p.biases.host(u16), .rows = p.weight.shape[0], .kg = kg, .k0 = k0 }, base, n, k, parts[0].bits, parts[0].group, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); base += p.weight.shape[0]; } - return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group }; + return .{ .words = .{ .buffer = w }, .pairs = .{ .buffer = sb }, .n = @intCast(n), .k = @intCast(k), .tile = tile, .slices = slices, .bits = parts[0].bits, .group = parts[0].group, .gn = gn }; } -/// One member as stored: words [rows][k*bits/32], scales and biases [rows][k/group]. -pub const Member = struct { words: []const u32, scales: []const u16, biases: []const u16, rows: usize }; +/// One member as stored: words [rows][kg*bits/32], scales and biases [rows][kg/group], columns [k0, k0 + k). +pub const Member = struct { words: []const u32, scales: []const u16, biases: []const u16, rows: usize, kg: usize = 0, k0: usize = 0 }; -/// Copies a member into lane_projection.weight_layout rows [base, base + rows) of an n-row stack, bytes unchanged. +/// Copies a member's columns [k0, k0 + k) into lane_projection.weight_layout rows [base, base + rows), bytes unchanged. pub fn packTiles(p: Member, base: usize, n: usize, k: usize, bits: usize, group: usize, out_w: []u32, out_sb: []u16) void { const ng = k / group; + const sng = if (p.kg != 0) p.kg / group else ng; + const goff = p.k0 / group; const wp = group * bits / 32; for (0..p.rows) |r| for (0..ng) |g| { const at = base + r; - @memcpy(out_w[((at / 32 * ng + g) * 32 + at % 32) * wp ..][0..wp], p.words[(r * ng + g) * wp ..][0..wp]); - out_sb[(g * n + at) * 2] = p.scales[r * ng + g]; - out_sb[(g * n + at) * 2 + 1] = p.biases[r * ng + g]; + @memcpy(out_w[((at / 32 * ng + g) * 32 + at % 32) * wp ..][0..wp], p.words[(r * sng + goff + g) * wp ..][0..wp]); + out_sb[(g * n + at) * 2] = p.scales[r * sng + goff + g]; + out_sb[(g * n + at) * 2 + 1] = p.biases[r * sng + goff + g]; }; } diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index e38b913e2a..054aec2362 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -129,23 +129,26 @@ pub const Kernels = struct { return k.strided(input_encoder, raw, x, y, rows, @intCast(raw.n)); }; var span = profile.begin(k.profiler, input_encoder, .dense); - const sk = if (k.prompt) projection.promptSplit(p.n, p.k) orelse p.slices else p.slices; + // Prompt rows take the checkpoint matrix's split, as decoded rows already do through `p.slices`. + const gn = if (p.gn != 0) p.gn else p.n; + const sk = if (k.prompt) projection.promptSplit(gn, p.k) orelse p.slices else p.slices; span.shape = .{ .n = p.n, .k = p.k, .sk = sk, .bits = p.bits }; errdefer span.cancel(); const e = span.encoder; _ = dims; if (rows == 0 or rows > @import("gpu_frame.zig").max_rows or p.tile != 32) return error.BadCoreProjectionRows; - if (p.raw) |raw| if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { - const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; - const scalar = try k.twin(m); // loads the tiles before they're read - const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, try k.pairTwin(m)); - row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); - e.barrier(); - try span.end(); - return; - }; if (p.raw) |raw| { + if (raw.bits == 4 and raw.group == 64 and raw.sum == .f32 and raw.n % 8 == 0) { + const m = row.simd.Matrix{ .w = raw.w, .w_off = raw.w_off, .scales = raw.scales, .s_off = raw.s_off, .biases = raw.biases, .b_off = raw.b_off, .n = raw.n, .k = raw.k }; + const scalar = try k.twin(m); // loads the tiles before they're read + const paired = try k.pairTwin(m); + const call = try row.simd.call(&k.simd.?, raw.n, raw.k, rows, scalar, paired, gn); + row.simd.encode(e, call, m, x.buffer, x.offset, y.buffer, y.offset); + e.barrier(); + try span.end(); + return; + } if (k.row_pipelines == null) k.row_pipelines = try row.Pipelines.load(k.device); const call = try row.call(&k.row_pipelines.?, raw, rows, false); e.setPipeline(call.pipeline); diff --git a/zig/src/families/qwen27/session.zig b/zig/src/families/qwen27/session.zig index b0296cdf95..220c9b205d 100644 --- a/zig/src/families/qwen27/session.zig +++ b/zig/src/families/qwen27/session.zig @@ -34,17 +34,29 @@ pub const Session = struct { pub fn step(s: Session, token: u32) !void { try s.feed(&.{token}, .last); } + /// The last round's token: a pair reads the head seam's `picks`, whose `maxInt(u32)` sentinel marks a broken link. pub fn greedy(s: Session) !u32 { if (s.runner.failed or s.runner.active != null or s.runner.offsets[s.slot] == 0) return error.RoundNotReady; + if (s.runner.reduce) |reduce| { + const picks = s.runner.picks orelse return error.HeadSeamMissing; + const words: [*]const u32 = @ptrCast(@alignCast(picks.buffer.contents() + picks.offset)); + const token = words[0]; + if (token >= s.runner.model.shard.global.vocab) return if (reduce.broken()) error.PairFailed else error.NonfiniteLogits; + return token; + } const values = s.runner.model.frame.get(.logits).buffer.slice(u16, s.runner.model.config.vocab); return argmax(values); } - /// Hash committed recurrent bytes, visible KV prefixes and the last row's logits (not unused capacity). + /// Per-rank digest (same rank, never the peer): committed recurrent bytes, visible KV prefixes, the last logits. pub fn fingerprint(s: Session) ![32]u8 { const r = s.runner; if (r.active != null or r.failed or r.offsets[0] == 0) return error.RoundNotReady; var hash = std.crypto.hash.sha2.Sha256.init(.{}); hash.update(std.mem.asBytes(&r.offsets[0])); + if (r.reduce != null) { + const picks = r.picks orelse return error.HeadSeamMissing; + hash.update(picks.buffer.contents()[picks.offset..][0..@sizeOf(u32)]); + } for (0..r.gdn.budget.layers) |layer| { hash.update(try r.gdn.committed(layer, 0, true)); hash.update(try r.gdn.committed(layer, 0, false)); @@ -64,7 +76,8 @@ pub const Session = struct { } fn feed(s: Session, ids: []const u32, head: @import("forward.zig").Head) !void { - var round = try Round.init(s.runner.allocator, &.{.{ .slot = s.slot, .start = s.runner.offsets[s.slot], .capacity = s.runner.capacity, .ids = ids }}, @intCast(s.runner.model.config.conv_kernel), s.runner.slots, @intCast(s.runner.model.config.vocab)); + // Ids are global on every rank (only head rows split): validate against the whole vocab, not `config.vocab`. + var round = try Round.init(s.runner.allocator, &.{.{ .slot = s.slot, .start = s.runner.offsets[s.slot], .capacity = s.runner.capacity, .ids = ids }}, @intCast(s.runner.model.config.conv_kernel), s.runner.slots, @intCast(s.runner.model.shard.global.vocab)); defer round.deinit(); if (!s.reference_tree) { try s.runner.advance(&round, head); @@ -125,3 +138,35 @@ test "top-two margins retain first-ID ties and reject nonfinite scores" { try std.testing.expectEqualSlices(f32, &.{ 2, 2 }, &top.logits); try std.testing.expectError(error.NonfiniteLogits, topTwo(&.{ 0x3f80, 0x7fc0 })); } + +/// One rank's block of the vocabulary: `values` are its local logits rows, `lo` its first global token id. +pub const Block = struct { values: []const u16, lo: u32 }; + +/// Top two over the ranks' blocks, shifted by `lo`: a cross-block tie keeps the lower id, as `argmax`/`pair_pick` do. +pub fn topTwoBlocks(blocks: []const Block) !TopTwo { + if (blocks.len == 0) return error.BadLogits; + var result = TopTwo{ .ids = .{ 0, 1 }, .logits = .{ -std.math.inf(f32), -std.math.inf(f32) } }; + for (blocks) |block| { + const one = try topTwo(block.values); + for (one.ids, one.logits) |id, value| { + const global = block.lo + id; + if (value > result.logits[0]) { + result.logits[1] = result.logits[0]; + result.ids[1] = result.ids[0]; + result.logits[0] = value; + result.ids[0] = global; + } else if (value > result.logits[1]) { + result.logits[1] = value; + result.ids[1] = global; + } + } + } + return result; +} +test "a pair's merged top two shifts each block to global ids and a cross-block tie keeps the lower id" { + const top = try topTwoBlocks(&.{ .{ .values = &.{ 0x3f80, 0x4000, 0x4040 }, .lo = 0 }, .{ .values = &.{ 0x4040, 0xbf80 }, .lo = 4 } }); + try std.testing.expectEqualSlices(u32, &.{ 2, 4 }, &top.ids); + try std.testing.expectEqualSlices(f32, &.{ 3, 3 }, &top.logits); + try std.testing.expectError(error.NonfiniteLogits, topTwoBlocks(&.{.{ .values = &.{ 0x3f80, 0x7fc0 }, .lo = 0 }})); + try std.testing.expectError(error.BadLogits, topTwoBlocks(&.{})); +} diff --git a/zig/src/metal/command.zig b/zig/src/metal/command.zig index 49118d0b3d..45c2b29358 100644 --- a/zig/src/metal/command.zig +++ b/zig/src/metal/command.zig @@ -195,4 +195,14 @@ pub const ComputeEncoder = struct { pub fn end(self: ComputeEncoder) void { objc.msg(void, self.id, "endEncoding", .{}); } + + /// The command buffer this encoder was opened from, for the event ops a seam encodes BETWEEN encoders. + pub fn commandBuffer(self: ComputeEncoder) CommandBuffer { + return .{ .id = objc.msg(Id, self.id, "commandBuffer", .{}) }; + } + + /// Whether this encoder dispatches serially or concurrently: a seam reopens it in the same mode. + pub fn dispatchType(self: ComputeEncoder) types.DispatchType { + return @fromBackingInt(@intCast(objc.msg(usize, self.id, "dispatchType", .{}))); + } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index a42272fe56..04dda5a54c 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -335,7 +335,7 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { if (o.prompt_int8) return prompt_int8_refusal; if (o.slide and !std.mem.eql(u8, t, "nemotron_h")) return try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on Macs yet", .{t}); const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; - if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); + if (o.speed_up != null and !(std.mem.eql(u8, t, "qwen3_5") and dense)) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Qwen3.8-27B, Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; @@ -361,7 +361,7 @@ fn qwen27Engine(a: Allocator, io: std.Io, dir: []const u8) bool { return true; } -/// Qwen3.8-27B on the lane host: `--parallel` streams in shared rounds, DFlash2 drafts from --drafter for one of them. +/// 27B: --parallel shared rounds, DFlash2 from --drafter, --speed-up opens the two-Mac seam before the weights load. fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const context = o.context orelse 32768; if (context <= 0 or context > 262144) { @@ -375,7 +375,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ var why: []const u8 = ""; const drafter = if (o.drafts) o.drafter else null; const learn: ?@import("qwen27_disk.zig").Learn = if (o.learn_conversations) if (o.learn) |root| .{ .root = root, .dir = o.dir, .drafter = drafter, .cap = gibBytes(o.learn_gib), .floor = gibBytes(o.learn_min_free_gib), .daily = gibBytes(o.learn_daily_gib), .spill = o.spill, .slots = spill.slots(o) } else null else null; - const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { + const h = qwen27.open(gpa, io, o.dir, .{ .pair = o.speed_up, .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen3.8-27B engine cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return null; }; @@ -552,13 +552,13 @@ test "--spill-gib with --checkpoint-slots opens the four families with a disk st try std.testing.expectStringStartsWith((try refusal(a, .{ .dir = "d", .model_type = "glm5_next", .keep = 2 }, false)) orelse "", "--checkpoint-slots"); } -test "--speed-up is refused by Nemotron, the 27B and the 2B before any load, and not by the two-Mac families" { +test "--speed-up is refused by Nemotron and the 2B before any load, and opens the pair seam on the 27B" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); const two: api.Open = .{ .dir = "d", .model_type = "qwen3_5", .speed_up = "pair.json" }; const tail = " only: serve {s} without it"; - try std.testing.expectStringEndsWith((try refusal(a, two, true)) orelse "", try std.fmt.allocPrint(a, tail, .{"Qwen3.8-27B"})); + try std.testing.expect((try refusal(a, two, true)) == null); // the 27B pairs: the pair seam is its two-Mac transport try std.testing.expectStringEndsWith((try refusal(a, two, false)) orelse "", try std.fmt.allocPrint(a, tail, .{"Qwen3.5-2B"})); try std.testing.expectStringEndsWith((try refusal(a, .{ .dir = "d", .model_type = "nemotron_h", .speed_up = "pair.json" }, false)) orelse "", try std.fmt.allocPrint(a, tail, .{"Nemotron"})); for ([_][]const u8{ "qwen4_exp", "glm5_next", "glm_moe_dsa" }) |t| try std.testing.expect((try refusal(a, .{ .dir = "d", .model_type = t, .speed_up = "pair.json" }, false)) == null); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index 38a99f11e9..eb96786969 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -16,11 +16,19 @@ const Ring = q.tap_ring.Ring; const Context = q.dflash.runtime_model.Context; const Landing = @import("qwen27_landing.zig").Landing; const contract = tf.tree_round; +const pair_mod = @import("qwen27_pair.zig"); + +/// The pair transport's types and seam ABI, in their own unit: the lane backend and the link's adapter share them. +pub const Pair = pair_mod.Pair; +pub const Kind = pair_mod.Kind; +pub const Link = pair_mod.Link; +pub const batch_rows = pair_mod.batch_rows; +pub const pairReduce = pair_mod.pairReduce; +pub const pairBytes = pair_mod.pairBytes; /// A window's rows at most: the pending row and a DFlash2 tree. pub const max_rows = 16; -/// A shared round's rows at most: the frame's. -pub const batch_rows = 128; + /// Streams a runner holds at most (the attention plan's). pub const max_streams = q.attention_plan.max_streams; const ring = 64; @@ -73,6 +81,7 @@ pub const Metal = struct { fills: std.AutoHashMapUnmanaged(*lanes.Stream, slots.Fill) = .empty, // prompt passes the lane host runs by chunks image_fill: ?*const lanes.Stream = null, // the pass whose images the vision rows hold, until it closes promote_due: bool = false, // slot 0 freed: a stream moves in at the next backend call (slots.promoteDue) + link: ?*Link = null, // the pair: rank 0 sends each round here before it runs, rank 1 replays it pub fn init(gpa: std.mem.Allocator, runner: *q.decode_round.Runner, draft: ?*Draft) !*Metal { if (runner.slots == 0 or runner.slots > max_streams or runner.model.frame.capacity < batch_rows) return error.TargetBinding; @@ -122,7 +131,7 @@ pub const Metal = struct { const ringChunk = slots.ringChunk; const unbind = slots.unbind; const ready = slots.ready; - const refill = slots.refill; + pub const refill = slots.refill; pub fn backend(b: *Metal) be.Backend { return .{ .ptr = b, .vtable = &.{ .prefill = prefillFn, .prefill_step = slots.prefillStep, .cancel_prefill = slots.cancelPrefill, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draft_fns.draftFn, .probabilities = draft_fns.probabilitiesFn, .tree = draft_fns.treeFn, .release = releaseFn, .reports_fill = true } }; @@ -161,66 +170,9 @@ pub const Metal = struct { return slot < b.draft_slots and live <= b.draft_slots and x.caches.len > 0; } - /// Each window width's ms (verify, draws, keep) on a scratch context in slot 0; resets it. Kept per build. + /// Each window width's ms on a scratch slot 0, kept per build: warm.measure, for the A/B tool. pub fn measure(b: *Metal, io: std.Io) !void { - const r = b.runner; - if (b.owners[0] != null or r.active != null) return error.StreamBusy; - var shape: [128]u8 = undefined; - const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; - const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; - if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { - b.costs = kept; - b.timed = max_rows; - return; - }; - if (r.capacity <= max_rows) return error.ContextFull; - defer { - if (b.draft) |d| d.reset() catch {} else r.reset(0) catch {}; - } - var ids: [192]u32 = undefined; - for (&ids, 0..) |*t, i| t.* = @intCast((i * 7919 + 13) % r.model.config.vocab); - const prompt = ids[0..@min(ids.len, r.capacity - max_rows)]; // a small --context still holds a window past it - try b.refill(prompt); - var s = try lanes.Stream.init(b.gpa, .{ .id = "timing", .prompt = prompt, .max_new = 1 }); - defer s.deinit(b.gpa); - const timer: WindowTimer = .{ .b = b, .s = &s, .prompt = prompt }; - for (0..5) |_| _ = try timer.time(max_rows - 1); // the widest window first, while the GPU's clocks ramp up - var ms: [max_rows]f64 = undefined; - for (&ms, 0..) |*m, i| m.* = try lanes.cost_rule.fastest(timer, i); - try lanes.cost_rule.smooth(timer, &ms); - const ref = lanes.cost_cache.referenceKey(&parts); - if (lanes.cost_cache.load([max_rows]f64, b.gpa, io, ref)) |reference| if (lanes.cost_rule.drifted(&ms, &reference)) try lanes.cost_rule.again(timer, &ms); - for (&b.costs, ms, 0..) |*cost, m, i| cost.* = .{ .width = @intCast(i + 1), .ms = m }; - b.timed = max_rows; - if (k) |key| lanes.cost_cache.save([max_rows]lanes.config.Cost, b.gpa, io, key, b.costs); - lanes.cost_cache.save([max_rows]f64, b.gpa, io, ref, ms); - } - - /// Window entry `i` is `i + 1` rows, timed once; each keeps a row, so a full slot is refilled first. - const WindowTimer = struct { - b: *Metal, - s: *lanes.Stream, - prompt: []const u32, - pub fn time(t: WindowTimer, i: usize) !f64 { - if (t.b.runner.offsets[0] + i + 1 > t.b.runner.capacity) try t.b.refill(t.prompt); - return t.b.timeWindow(t.s, i + 1); - } - }; - - fn timeWindow(b: *Metal, s: *lanes.Stream, rows: usize) !f64 { - const t0 = mtl.clock.seconds(); - const positions: [max_rows]u64 = @splat(0); - const drafts: [max_rows]u32 = @splat(1); - var sampled: [max_rows]u32 = undefined; - var echoed: [max_rows]u32 = undefined; - b.owners[0] = s; - defer b.owners[0] = null; - const w = be.Window{ .stream = s, .pending = 0, .held = 0, .tokens = drafts[0 .. rows - 1], .parents = null, .positions = positions[0..rows] }; - var out = [_]be.Verified{.{ .sampled = sampled[0..rows], .drafts = echoed[0 .. rows - 1] }}; - try verifyFn(b, &.{w}, &out); - var path = [_]u32{0}; - try b.keepPaths(&.{&path}); - return (mtl.clock.seconds() - t0) * 1e3; + return @import("qwen27_warm.zig").measure(b, io); } pub fn self(ptr: *anyopaque) *Metal { @@ -268,6 +220,8 @@ pub const Metal = struct { /// Logits rows [first, first + rows) drawn at their positions: argmax on the GPU, or the stream's keyed sampler. fn draw(b: *Metal, s: *const lanes.Stream, first: usize, rows: usize, positions: []const u64, out: []u32, masks: []const u32) !void { + // a pair's token comes from the round's own head seam: a second local argmax would decide differently on rank 1 + if (b.runner.reduce != null) return error.PairGreedyOnly; const logits = b.runner.model.frame.get(.logits); const vocab: u32 = @intCast(b.runner.model.config.vocab); const width = (@as(usize, vocab) + 31) / 32; @@ -306,6 +260,16 @@ pub const Metal = struct { /// The tokens argmaxRows picked for rows [first, first + rows). fn picked(b: *const Metal, first: usize, rows: usize, out: []u32) !void { + if (b.runner.reduce != null) { + const seam_picks = b.runner.picks orelse return error.HeadSeamMissing; + const vocab: u32 = @intCast(b.runner.model.shard.global.vocab); + const ids: [*]const u32 = @ptrCast(@alignCast(seam_picks.buffer.contents() + seam_picks.offset)); + for (out[0..rows], ids[first..][0..rows]) |*t, id| { + if (id >= vocab) return error.NonfiniteLogits; + t.* = id; + } + return; + } const picks: [*]const contract.Pick = @ptrCast(@alignCast(b.picks.contents())); for (out[0..rows], picks[first..][0..rows]) |*t, p| { if (p.token == contract.invalid_token or p.nonfinite != 0) return error.NonfiniteLogits; @@ -314,7 +278,7 @@ pub const Metal = struct { } /// Keep the last verify's rows on `paths`: slot 0 alone keeps on the GPU with its taps, else the runner and rings. - fn keepPaths(b: *Metal, paths: []const []const u32) !void { + pub fn keepPaths(b: *Metal, paths: []const []const u32) !void { defer b.dropRound(); if (paths.len != b.count) return error.NothingToKeep; if (b.advanced) { // the verify kept each window's one row at its position; its logits already sit in row 0 @@ -417,8 +381,14 @@ pub const Metal = struct { // -- the vtable --------------------------------------------------------------------------------------------- + /// Rank 1 replays rank 0's rounds: refuse at the head of `prefill`, its first backend call (GLM's `beginPrompt`). + pub fn refuseFollower(m: *Metal) !void { + if (m.link) |l| if (l.pair.rank != 0) return error.FollowsPeer; + } + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { const b = self(ptr); + try refuseFollower(b); // before settle, dropRound, promoteDue and bind: rank 1 touched nothing yet b.settle() catch {}; b.dropRound(); slots.promoteDue(b); @@ -433,6 +403,11 @@ pub const Metal = struct { const b = self(ptr); _ = try b.slotFor(s); var token: [1]u32 = undefined; + if (b.runner.reduce != null) { + // the prompt's last chunk (head = .last) picked on both ranks: no second round, no unannounced exchange + try b.picked(0, 1, &token); + return b.push(token[0]); + } try b.draw(s, 0, 1, &.{position}, &token, &.{}); return b.push(token[0]); } @@ -465,7 +440,7 @@ pub const Metal = struct { } /// Every stream's window (pending token, then its drafts as a chain or tree) in one forward, every row drawn. - fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + pub fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const b = self(ptr); if (windows.len == 0 or windows.len > b.runner.slots or windows.len != out.len) return error.BadWindows; try b.settle(); @@ -501,16 +476,22 @@ pub const Metal = struct { if (windows.len > 1) b.dropTree(); const pool = mtl.objc.Pool.push(); defer pool.pop(); - b.round = try Round.init(b.gpa, inputs[0..windows.len], @intCast(r.model.config.conv_kernel), r.slots, @intCast(r.model.config.vocab)); + // Ids are global on every rank (only the head's rows are split), so the round validates the whole vocabulary. + b.round = try Round.init(b.gpa, inputs[0..windows.len], @intCast(r.model.config.conv_kernel), r.slots, @intCast(r.model.shard.global.vocab)); b.count = windows.len; errdefer b.dropRound(); + if (b.runner.reduce != null) { + // landing 1 is one greedy stream of one-row windows: the words are the slot and the pending token. + if (windows.len != 1 or total != 1) return error.CommandOutOfStep; + if (b.link) |l| try l.step(inputs[0].slot, b.ids[0]); + } if (plain) try r.advance(&b.round.?, .all) else try r.verifyHead(&b.round.?, .all); b.advanced = plain; b.rows = total; // every greedy window's rows from one argmax over the round; a sampled window draws on the host var greedy = false; for (windows) |w| greedy = greedy or (w.stream.sampling == null and w.masks.len == 0); - if (greedy) try b.argmaxRows(0, total); + if (greedy and b.runner.reduce == null) try b.argmaxRows(0, total); for (windows, out, b.round.?.firsts) |w, o, first| { const rows = w.rows(); if (w.stream.sampling != null or w.masks.len > 0) try b.draw(w.stream, first, rows, w.positions, o.sampled, w.masks) else try b.picked(first, rows, o.sampled); @@ -579,8 +560,10 @@ pub const Metal = struct { } fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { - slots.endFill(self(ptr), s); - self(ptr).unbind(s); + const b = self(ptr); + slots.endFill(b, s); + if (b.link) |l| if (b.slotOf(s)) |slot| l.end(slot); + b.unbind(s); } }; diff --git a/zig/src/native/qwen27_link.zig b/zig/src/native/qwen27_link.zig new file mode 100644 index 0000000000..82afc00bd1 --- /dev/null +++ b/zig/src/native/qwen27_link.zig @@ -0,0 +1,76 @@ +//! The pair link's liveness: rank 1 replays rank 0's rounds, and both ranks stop when the link dies. +const std = @import("std"); +const platform = @import("sys"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lb = @import("qwen27_lanes.zig"); + +/// Rank 1's follower: replay rank 0's rounds in order; a rank that cannot replay stops to be restarted. +pub fn follow(b: *lb.Metal) void { + const l = b.link orelse return; + while (true) { + const seen = l.pair.next(l.gpa, &l.words) catch |err| { + if (err != error.PairStopped) failFollow("a round", @errorName(err)); + return; + }; + const k = seen orelse return; + const kind = std.enums.fromInt(lb.Kind, k) orelse failFollow("an unknown command", k); + const w = l.words.items; + switch (kind) { + .begin => { + if (w.len != 1) failFollow("begin", w.len); + b.dropRound(); + b.runner.reset(w[0]) catch |err| return failFollow("begin", @errorName(err)); + }, + .chunk => { + if (w.len < 4 or w.len != 4 + w[1]) failFollow("chunk", w.len); + const session = q.session.Session{ .runner = b.runner, .slot = w[0] }; + const chunk = w[4..]; + if (w[3] != 0) { + session.decodeChunk(chunk, w[2] != 0) catch |err| return failFollow("decode chunk", @errorName(err)); + } else { + session.promptChunk(chunk, w[2] != 0) catch |err| return failFollow("prompt chunk", @errorName(err)); + } + }, + .step => { + if (w.len != 2) failFollow("step", w.len); + (q.session.Session{ .runner = b.runner, .slot = w[0] }).step(w[1]) catch |err| return failFollow("step", @errorName(err)); + }, + .end => {}, + } + } +} + +/// This rank can no longer hold rank 0's state: it stops rather than serve on a stale half, to be restarted. +fn failFollow(what: []const u8, detail: anytype) noreturn { + std.log.err("qwen27 rank 1: cannot replay rank 0's {s} ({any}): stopping so both ranks are restarted", .{ what, detail }); + std.process.exit(1); +} + +/// Both ranks stop when the link dies: a dead peer must not leave HTTP up on a half, so a supervisor restarts both. +pub fn watch(p: *lb.Pair, quit: *const std.atomic.Value(bool)) void { + var seen = p.heartbeats(); + var since = mtl.clock.seconds(); + while (!quit.load(.acquire)) { + if (p.failed()) down(p, "the pair link is down"); + p.beat(); + const beats = p.heartbeats(); + const now = mtl.clock.seconds(); + if (beats != seen) { + seen = beats; + since = now; + } else if (now - since > link_silence_s) down(p, "the peer stopped signalling"); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = @intCast(250 * std.time.ns_per_ms) }; + _ = platform.c.nanosleep(&ts, null); + } +} + +/// Beats the watchdog may miss before the peer counts as gone: twelve 250 ms intervals, longer than any seam. +const link_silence_s: f64 = 3.0; + +/// The link is not a serving state: name the half that is gone, then end the process the supervisor watches. +fn down(p: *lb.Pair, what: []const u8) noreturn { + std.log.err("qwen27 rank {d}: {s}: exiting so both ranks are restarted", .{ p.rank, what }); + std.process.exit(1); +} diff --git a/zig/src/native/qwen27_pair.zig b/zig/src/native/qwen27_pair.zig new file mode 100644 index 0000000000..3f2021d11f --- /dev/null +++ b/zig/src/native/qwen27_pair.zig @@ -0,0 +1,92 @@ +//! The pair transport's types, the seam ABI the engine calls, and the window one exchange may carry. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; + +/// The pair's transport type: the row-split seam the engine exports, so this file names no other family's type. +pub const Pair = tf.pair_rows.Pair; + +/// The command set rank 1 replays: the four GPU operations a pair round is made of (glm_dsa's `follow.zig` shape). +pub const Kind = enum(u32) { begin = 1, chunk, step, end }; + +/// Rank 0 puts each round on the wire before running it; rank 1 replays them in order here. Owned by the host. +pub const Link = struct { + pair: *Pair, + gpa: std.mem.Allocator, + words: std.ArrayList(u32) = .empty, + + pub fn deinit(l: *Link) void { + l.words.deinit(l.gpa); + } + + fn put(l: *Link, kind: Kind) !void { + try l.pair.send(@backingInt(kind), l.words.items); + l.words.clearRetainingCapacity(); + } + + /// A prompt's chunk boundaries are rank 0's memory policy's, so the send carries them: rank 1 cannot infer them. + pub fn begin(l: *Link, slot: u32) !void { + l.words.clearRetainingCapacity(); + try l.words.append(l.gpa, slot); + try l.put(.begin); + } + + pub fn chunk(l: *Link, slot: u32, ids: []const u32, last: bool, decoded: bool) !void { + l.words.clearRetainingCapacity(); + try l.words.appendSlice(l.gpa, &.{ slot, @intCast(ids.len), @intFromBool(last), @intFromBool(decoded) }); + try l.words.appendSlice(l.gpa, ids); + try l.put(.chunk); + } + + pub fn step(l: *Link, slot: u32, token: u32) !void { + l.words.clearRetainingCapacity(); + try l.words.appendSlice(l.gpa, &.{ slot, token }); + try l.put(.step); + } + + pub fn end(l: *Link, slot: u32) void { + l.words.clearRetainingCapacity(); + l.words.append(l.gpa, slot) catch return; + l.put(.end) catch {}; + } +}; + +/// The pair's seam ends in the engine's shape: the transport's `Ref` is the frame's, so the casts are field copies. +pub fn pairReduce(p: *Pair) q.forward.Reduce { + return .{ .ptr = p, .part = pairPart, .stage = pairStage, .gather = pairGather, .failed = pairFailed, .argmax = pairArgmax }; +} + +fn pairPart(ptr: *anyopaque) q.forward.Seam { + const p: *Pair = @ptrCast(@alignCast(ptr)); + const s = p.part(); + return .{ .x = s.x, .ref = .{ .buffer = s.ref.buf, .offset = s.ref.off } }; +} + +fn pairStage(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, src: q.forward.Ref, rows: u32, width: u32, stride: u32) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.stage(e, x, .{ .buf = src.buffer, .off = src.offset }, rows, width, stride); +} + +fn pairGather(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, mine: q.forward.Ref, out: q.forward.Ref, rows: u32, width: u32, out_stride: u32) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.gather(e, x, .{ .buf = mine.buffer, .off = mine.offset }, .{ .buf = out.buffer, .off = out.offset }, rows, width, out_stride); +} + +fn pairFailed(ptr: *anyopaque) bool { + const p: *Pair = @ptrCast(@alignCast(ptr)); + return p.failed(); +} + +fn pairArgmax(ptr: *anyopaque, e: *mtl.ComputeEncoder, x: u32, logits: q.forward.Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: q.forward.Ref) void { + const p: *Pair = @ptrCast(@alignCast(ptr)); + p.argmaxBf16(e, x, .{ .buf = logits.buffer, .off = logits.offset }, stride, lo, n, rows, .{ .buf = picks.buffer, .off = picks.offset }); +} + +/// The window an exchange may carry: a round's rows of the widest seam payload -- this rank's `intermediate / 2` bf16. +pub fn pairBytes() usize { + return @as(usize, batch_rows) * @divExact(17408, 2) * @sizeOf(u16); +} + +/// A shared round's rows at most: the frame's. +pub const batch_rows = 128; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 26946db9be..e48b8bbe92 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -12,6 +12,8 @@ const cache_fit = @import("cache_fit.zig"); const share_fit = @import("qwen27_share_fit.zig"); const Resident = @import("resident.zig").Resident; const lb = @import("qwen27_lanes.zig"); +const link = @import("qwen27_link.zig"); +const warm = @import("qwen27_warm.zig"); pub const lanes_backend = lb; const disk = @import("qwen27_disk.zig"); const images = @import("qwen27_images.zig"); @@ -28,6 +30,7 @@ pub const Options = struct { learn: ?disk.Learn = null, // --learn-conversations: kept states written behind to disk and read back vision: ?api.VisionOpen = null, // --vision: open loads the image tower, and requests may hold images images: ?*images.Vision = null, // the loaded tower attach serves (and the host then owns) + pair: ?[]const u8 = null, // --speed-up: this rank's settings file; `Pair.init` reads the rank out of it }; pub const Host = struct { @@ -47,6 +50,13 @@ pub const Host = struct { disk_stage: ?mtl.Buffer = null, // copies out of private kept states for the disk (qwen27_disk.zig) vision: ?*images.Vision = null, // --vision's tower and image rows spares: q.snapshot.Spares = .{}, // dropped kept states' buffers, for the next saves + pair: ?*lb.Pair = null, // --speed-up: the seam this rank owns; rank 1 follows through it + link: ?*lb.Link = null, // rank 0's round sender and rank 1's follower target; `Metal.link` points here + hw: ?mtl.Device = null, // the pair's device, opened before the model's so the seam exists for the load + picks: ?mtl.Buffer = null, // a round's head-seam picks (one u32 a row), read by the lane backend + follower: ?std.Thread = null, // rank 1's command loop + quit: std.atomic.Value(bool) = .init(false), // `close` sets it so the link watchdog never reads a freed pair + watcher: ?std.Thread = null, // the link watchdog both ranks run while the pair is up pub fn engine(h: *Host) api.Engine { return h.host.engine(); @@ -109,6 +119,8 @@ fn readMemory(_: ?*anyopaque, reset_peak: bool) ?api.Memory { fn words(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PromptTooLong => "the prompt and max_tokens do not fit this server's --context", + // rank 1 replays rank 0's rounds and owns no request (qwen27_lanes.prefillFn refuses first) + error.FollowsPeer => "speed-up mode: this Mac runs rank 0's requests; send requests to rank 0", else => null, }; } @@ -294,19 +306,37 @@ fn laterBytes(h: *const Host) u64 { pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, o: Options, a: Allocator, why: *[]const u8) !*Host { const pool = mtl.objc.Pool.push(); defer pool.pop(); - const model = try q.model.Model.loadDeferred(gpa, io, dir, 128); + // glm_dsa's order: the seam exists before the weights load and ready() comes once every GPU allocation is in. + var hw: ?mtl.Device = null; + var pair: ?*lb.Pair = null; + if (o.pair) |path| { + hw = try mtl.Device.init(); + pair = lb.Pair.init(gpa, io, hw.?, path, lb.pairBytes()) catch |err| { + why.* = try std.fmt.allocPrint(a, "the speed-up seam {s} would not open ({s})", .{ path, @errorName(err) }); + return err; + }; + } + errdefer if (pair) |p| p.deinit(); + errdefer if (hw) |d| d.deinit(); + const model = if (pair) |p| + try q.model.Model.loadShardDeferred(gpa, io, dir, 128, p.rank, 2, .full) + else + try q.model.Model.loadDeferred(gpa, io, dir, 128); errdefer model.deinit(); var with = o; // the tower before the slots: their room and the cache's budget count it as taken if (o.vision) |v| with.images = try images.open(gpa, io, model, dir, v, a, why); errdefer if (with.images) |v| v.deinit(); - const h = try attach(gpa, io, model, null, with, a, why); + const h = try attach(gpa, io, model, null, with, a, why, pair); h.owns_model = true; + h.pair = pair; + h.hw = hw; return h; } /// A loaded model served as `open` serves one, `drafter` bound in place of `o.drafter`; both stay the caller's. -pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.dflash.runtime_model.Model, o: Options, a: Allocator, why: *[]const u8) !*Host { +pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.dflash.runtime_model.Model, o: Options, a: Allocator, why: *[]const u8, pair: ?*lb.Pair) !*Host { + if (pair != null) try pairFlags(o, drafter != null, why); const pool = mtl.objc.Pool.push(); defer pool.pop(); const capacity = if (o.context == 0) @min(model.config.max_position, 32768) else o.context; @@ -326,9 +356,32 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.cache = null; h.resident = null; h.disk_stage = null; + // create() leaves unnamed Host fields undefined: a debug-filled optional is non-null, so name h.link first. + h.link = null; + h.hw = null; + h.picks = null; + h.follower = null; + h.quit = .init(false); + h.watcher = null; h.vision = o.images; // closed with the host; a failed attach leaves it to the caller h.runner = try q.decode_round.Runner.initReady(gpa, model, slots, capacity + lb.max_rows, 1); errdefer h.runner.deinit(); + h.pair = pair; + if (pair) |p| { + if (slots != 1) { + why.* = try std.fmt.allocPrint(a, "--speed-up serves one greedy stream: pass --parallel 1", .{}); + return error.PairOneStream; + } + // The picks live where the head seam writes them, so the lane backend reads them rather than a local argmax. + const picks = try model.device.buffer(lb.batch_rows * @sizeOf(u32), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer picks.deinit(); + h.picks = picks; + h.link = try gpa.create(lb.Link); + errdefer gpa.destroy(h.link.?); + h.link.?.* = .{ .pair = p, .gpa = gpa }; + h.runner.reduce = lb.pairReduce(p); + h.runner.picks = .{ .buffer = picks, .offset = 0 }; + } if (drafter) |d| h.draft = try lb.Draft.attach(gpa, &h.runner, d, false) else if (o.drafter) |d| h.draft = lb.Draft.open(gpa, io, model, &h.runner, d, if (o.drafter_bits == 4) .prepared_q4_reference else .bf16_reference) catch |err| { why.* = try std.fmt.allocPrint(a, "the native Qwen3.8-27B drafter cannot load {s} ({s})", .{ d, @errorName(err) }); return error.DrafterLoad; @@ -336,10 +389,11 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df errdefer if (h.draft) |d| d.deinit(); h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); + h.back.link = h.link; h.back.chunk = h.runner.promptRows(q.model.promptRows()); h.back.grow = .{ .ctx = h, .room = Host.roomFor, .added = Host.added }; h.back.vision = o.images; - if (h.draft != null) try h.back.measure(io); + if (h.draft != null) try warm.measure(h.back, io); h.cfg = try lanes.Config.init(gpa, h.back.facts(), lb.max_rows, lb.max_rows - 1); errdefer h.cfg.deinit(gpa); h.clock = .{ .io = io }; @@ -367,16 +421,59 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df }; if (model.ane) |share| share.footprint_growth = platform.memory_policy.growth(before, cache_fit.footprint()); } - try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); + // a pair's prompt states are two Macs' and its slots are one: no reuse in landing 1, and no disk learning + if (pair == null) try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); errdefer if (h.cache) |*store| store.deinit(); - if (o.learn) |l| try disk.attach(h, io, a, l); + if (pair == null) if (o.learn) |l| try disk.attach(h, io, a, l); errdefer if (h.cache) |*store| if (store.disk) |r| r.deinit(); h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + // both Macs are loaded and allocated before either serves: rank 1 may now replay rank 0's rounds + if (pair) |p| p.ready() catch |err| { + why.* = if (err == error.PairPeerSilent) + "the peer never signalled ready in five minutes: start rank 1 first, with the same build on both Macs and each rank's own settings file" + else + try std.fmt.allocPrint(a, "the pair link was not ready ({s}): restart both ranks", .{@errorName(err)}); + return err; + }; + // a dead link is not a serving state: exit so a supervisor restarts both ranks, not serve seam errors + if (pair) |p| { + // fail closed: a rank whose watchdog is missing would keep serving HTTP on a half once the peer died + h.watcher = std.Thread.spawn(.{}, link.watch, .{ p, &h.quit }) catch |err| { + why.* = try std.fmt.allocPrint(a, "the pair link's watchdog would not start ({s}): a dead peer would leave this rank up, so this rank stops", .{@errorName(err)}); + return error.PairNoWatchdog; + }; + } try h.host.start(); + if (pair) |p| if (p.rank == 1) { + h.follower = std.Thread.spawn(.{}, link.follow, .{h.back}) catch |err| { + why.* = try std.fmt.allocPrint(a, "rank 1's follower would not start ({s})", .{@errorName(err)}); + return err; + }; + }; std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); return h; } +/// The flags a pair refuses, named: quietly serving without one of them would misreport what the reply ran on. +fn pairFlags(o: Options, drafter: bool, why: *[]const u8) !void { + if (drafter) { + why.* = "--drafter does not run on a pair: rank 1 has no drafter checkpoint, so both ranks decode plain"; + return error.PairNoDrafter; + } + if (o.cache_gib != null) { + why.* = "--prompt-cache-gib does not run on a pair: each Mac keeps its own halves of the slots' states"; + return error.PairNoCache; + } + if (o.learn != null) { + why.* = "--learn does not run on a pair: each Mac learns its own halves of the slots' states"; + return error.PairNoLearn; + } + if (o.vision != null or o.images != null) { + why.* = "--vision does not run on a pair: rank 1 replays rank 0's token ids, so an image prompt would run different rows on the two ranks"; + return error.PairNoVision; + } +} + /// Reserve later buffers and the loaded share's priced bytes that its construction left outside the footprint. fn laterFootprint(h: *const Host) !u64 { const tower: u64 = if (h.back.vision) |v| v.firstUse() else 0; @@ -408,7 +505,18 @@ const holdResident = q.residency.hold; pub fn close(ptr: *anyopaque) void { const h = Host.of(ptr); + h.quit.store(true, .release); // the watchdog exits before the pair it reads is torn down h.host.stop(); + // the pair goes after the host stops (no round is mid-flight) and before the model frees the device + if (h.pair) |p| p.stopFollowing(); + if (h.follower) |t| t.join(); + if (h.watcher) |t| t.join(); + if (h.link) |l| { + l.deinit(); + h.gpa.destroy(l); + } + if (h.pair) |p| p.deinit(); + if (h.hw) |d| d.deinit(); const rows = if (h.cache) |*store| store.disk else null; if (h.cache) |*store| store.deinit(); // its states leave the set as they go h.spares.deinit(); @@ -420,6 +528,7 @@ pub fn close(ptr: *anyopaque) void { h.back.deinit(); if (h.draft) |d| d.deinit(); h.runner.deinit(); + if (h.picks) |b| b.deinit(); if (h.vision) |v| v.deinit(); if (h.owns_model) h.model.deinit(); h.gpa.destroy(h); @@ -450,3 +559,26 @@ test "Qwen27 interior state uses a guaranteed cold request cut" { try std.testing.expectEqual(@as(u32, 129024), followup.from); try std.testing.expect(store.held <= store.budget); } + +test "rank 1 refuses a request by name before its backend moves round state" { + try std.testing.expectEqualStrings("speed-up mode: this Mac runs rank 0's requests; send requests to rank 0", words(null, error.FollowsPeer).?); + // `link` is the only `Metal` field the guard reads: the rest staying undefined proves no other step ran + var pair: lb.Pair = undefined; + pair.rank = 1; + var l: lb.Link = undefined; + l.pair = &pair; + var m: lb.Metal = undefined; + m.link = &l; + try std.testing.expectError(error.FollowsPeer, lb.Metal.refuseFollower(&m)); + pair.rank = 0; + try lb.Metal.refuseFollower(&m); // rank 0's own requests run on it + m.link = null; + try lb.Metal.refuseFollower(&m); // no pair: nothing to refuse for +} + +test "--vision is refused on a pair, by name" { + var why: []const u8 = ""; + try std.testing.expectError(error.PairNoVision, pairFlags(.{ .context = 4096, .streams = 1, .vision = .{} }, false, &why)); + try std.testing.expect(std.mem.startsWith(u8, why, "--vision does not run on a pair")); + try pairFlags(.{ .context = 4096, .streams = 1 }, false, &why); // a single Mac reads images +} diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 1889a6c755..8d30d46531 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -28,6 +28,8 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { if (ids.len == 0 or ids.len + s.max_new + max_rows > b.runner.capacity) return error.PromptTooLong; if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; + // Announced before the first chunk: rank 1 opens where rank 0 does (glm_dsa order). + if (b.link) |l| try l.begin(slot); try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); @@ -115,6 +117,8 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill) !bool { }; const part = ids[at..end]; const last = end == ids.len; + // Rank 0 sends the chunk before running it here (glm_dsa order): the send carries this pass's boundary. + if (b.link) |l| try l.chunk(slot, part, last, decoded); if (d) |x| { if (decoded) try x.decodeChunk(part, last) else try x.promptChunk(part, last); } else { diff --git a/zig/src/native/qwen27_warm.zig b/zig/src/native/qwen27_warm.zig new file mode 100644 index 0000000000..085edeb4a8 --- /dev/null +++ b/zig/src/native/qwen27_warm.zig @@ -0,0 +1,73 @@ +//! The lane backend's window timing: each width's ms on a scratch slot 0, kept in the per-build cost cache. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lanes = tf.lanes; +const be = lanes.backend; + +const nl = @import("qwen27_lanes.zig"); // the native backend, not the core lane engine `lanes` names +const Metal = nl.Metal; +const max_rows = nl.max_rows; + +/// Each window width's ms (verify, draws, keep) on a scratch context in slot 0; resets it. Kept per build. +pub fn measure(b: *Metal, io: std.Io) !void { + const r = b.runner; + if (b.owners[0] != null or r.active != null) return error.StreamBusy; + var shape: [128]u8 = undefined; + const parts = [_][]const u8{ "qwen27-lanes", std.mem.span(r.model.device.name()), try q.widths.costShape(r.model.config, r.model.widths_tag, &shape, max_rows, b.draft != null) }; + const k = lanes.cost_cache.key(b.gpa, io, &parts) catch null; + if (k) |key| if (lanes.cost_cache.load([max_rows]lanes.config.Cost, b.gpa, io, key)) |kept| { + b.costs = kept; + b.timed = max_rows; + return; + }; + if (r.capacity <= max_rows) return error.ContextFull; + defer { + if (b.draft) |d| d.reset() catch {} else r.reset(0) catch {}; + } + var ids: [192]u32 = undefined; + for (&ids, 0..) |*t, i| t.* = @intCast((i * 7919 + 13) % r.model.config.vocab); + const prompt = ids[0..@min(ids.len, r.capacity - max_rows)]; // a small --context still holds a window past it + try b.refill(prompt); + var s = try lanes.Stream.init(b.gpa, .{ .id = "timing", .prompt = prompt, .max_new = 1 }); + defer s.deinit(b.gpa); + const timer: WindowTimer = .{ .b = b, .s = &s, .prompt = prompt }; + for (0..5) |_| _ = try timer.time(max_rows - 1); // the widest window first, while the GPU's clocks ramp up + var ms: [max_rows]f64 = undefined; + for (&ms, 0..) |*m, i| m.* = try lanes.cost_rule.fastest(timer, i); + try lanes.cost_rule.smooth(timer, &ms); + const ref = lanes.cost_cache.referenceKey(&parts); + if (lanes.cost_cache.load([max_rows]f64, b.gpa, io, ref)) |reference| if (lanes.cost_rule.drifted(&ms, &reference)) try lanes.cost_rule.again(timer, &ms); + for (&b.costs, ms, 0..) |*cost, m, i| cost.* = .{ .width = @intCast(i + 1), .ms = m }; + b.timed = max_rows; + if (k) |key| lanes.cost_cache.save([max_rows]lanes.config.Cost, b.gpa, io, key, b.costs); + lanes.cost_cache.save([max_rows]f64, b.gpa, io, ref, ms); +} + +/// Window entry `i` is `i + 1` rows, timed once; each keeps a row, so a full slot is refilled first. +const WindowTimer = struct { + b: *Metal, + s: *lanes.Stream, + prompt: []const u32, + pub fn time(t: WindowTimer, i: usize) !f64 { + if (t.b.runner.offsets[0] + i + 1 > t.b.runner.capacity) try t.b.refill(t.prompt); + return timeWindow(t.b, t.s, i + 1); + } +}; + +pub fn timeWindow(b: *Metal, s: *lanes.Stream, rows: usize) !f64 { + const t0 = mtl.clock.seconds(); + const positions: [max_rows]u64 = @splat(0); + const drafts: [max_rows]u32 = @splat(1); + var sampled: [max_rows]u32 = undefined; + var echoed: [max_rows]u32 = undefined; + b.owners[0] = s; + defer b.owners[0] = null; + const w = be.Window{ .stream = s, .pending = 0, .held = 0, .tokens = drafts[0 .. rows - 1], .parents = null, .positions = positions[0..rows] }; + var out = [_]be.Verified{.{ .sampled = sampled[0..rows], .drafts = echoed[0 .. rows - 1] }}; + try nl.Metal.verifyFn(b, &.{w}, &out); + var path = [_]u32{0}; + try b.keepPaths(&.{&path}); + return (mtl.clock.seconds() - t0) * 1e3; +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 47b6a354c4..6cd23eb757 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -30,6 +30,8 @@ pub const exlr8 = @import("core/exlr8.zig"); pub const exlr8_projection = @import("core/exlr8_projection.zig"); pub const exlr8_dense_gpu = @import("core/exlr8_dense_gpu.zig"); pub const exlr8_prompt = @import("core/exlr8_prompt.zig"); +/// The two-Mac pair seam's row split: the transport, its layout and the one-Mac Loopback the stage-2 checks use. +pub const pair_rows = @import("core/pair_rows.zig"); pub const glm_dsa = @import("families/glm_dsa/glm_dsa.zig"); pub const nemotron = @import("families/nemotron/nemotron.zig"); pub const flashnext_vision_contract = @import("families/flashnext/vision_contract.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index b367fa120a..691ecf11ce 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -272,7 +272,7 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: /// simd_qmm over rows [start, start + m), the scalar twin when `scalar` and it takes m rows. fn simdWindow(e: mtl.ComputeEncoder, p: *const sq.Pipelines, s: *const Synth, y: Out, scalar: bool, start: usize, m: usize) !void { const mat = sq.Matrix{ .w = s.w, .scales = s.s, .biases = s.b, .n = s.n, .k = s.k }; - sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); + sq.encode(e, try sq.call(p, s.n, s.k, m, scalar, true, 0), mat, s.x, start * s.k * 2, y.buf, start * s.n * 2); } /// simd_qmm (4-bit, group 64): tiles at every window against float64 and the wide call; the twin at 1 and 2 rows. diff --git a/zig/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 46f921e9f8..3aa37d5de0 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -198,7 +198,7 @@ pub const Proj = struct { e.dispatchGroups(mtl.Size.of(cl.groups, cl.slots, 1), mtl.Size.of(cl.threads, 1, 1)); }, .simd => { - const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true); + const cl = try sq.call(&p.sq_p.?, shape.n, shape.k, shape.rows, false, true, 0); sq.encode(e, cl, .{ .w = p.w.buf, .w_off = p.w.off, .scales = p.s.buf, .s_off = p.s.off, .biases = p.b.buf, .b_off = p.b.off, .n = shape.n, .k = shape.k }, p.x.buf, p.x.off, p.out.buf, p.out.off); }, .simd_bits => { diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ab64f33ed6..1bf3d51859 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -89,6 +89,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { text = if (has(d.name, " copy ")) tf.qwen27.runtime_sources.copy else tf.qwen27.runtime_sources.glue; + } else if (std.mem.startsWith(u8, d.name, "metal pair ")) { + text = sources.core_pair_rows; } else if (has(d.name, "tile_qmm")) text = sources.core_tile_qmm else if (has(d.name, "softmax")) text = sources.ops_softmax else if (has(d.name, "rms")) text = sources.ops_embed_norm else if (has(d.name, "top-k")) text = topk_text else if (has(d.name, "row projection")) text = sources.core_row_projection else if (has(d.name, "simd_qmm_bits")) text = sources.core_simd_qmm_bits else if (has(d.name, "simd_qmm")) text = sources.core_simd_qmm else if (has(d.name, "affine_mm")) { const bits: u8 = if (has(d.name, "q8")) 8 else 4; const group: u16 = if (has(d.name, "g32")) 32 else 64; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index f2a067de87..8fa7ebcdf4 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -17,6 +17,7 @@ const drama_combine = @import("metal_drama_combine.zig"); const drama_cfg = @import("metal_drama_cfg.zig"); const glm_eh_k = @import("metal_glm_eh.zig"); const glm_rows_k = @import("metal_glm_rows.zig"); +const pair_k = @import("metal_pair.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -163,6 +164,21 @@ pub fn main(init: std.process.Init) !u8 { try rows_list.append(a, r); try kernels.append(a, r.kernel()); } + var pair_list: std.ArrayList(*pair_k.Pair) = .empty; + defer { + for (pair_list.items) |p| { + p.deinit(); + a.destroy(p); + } + pair_list.deinit(a); + } + for (pair_k.all) |w| { + if (!wanted(o, pair_k.spec(w).name)) continue; + const p = try a.create(pair_k.Pair); + p.* = try pair_k.Pair.init(a, device, queue, w); + try pair_list.append(a, p); + try kernels.append(a, p.kernel()); + } var flash_list: [8]flash.Probe = undefined; var flash_made: usize = 0; defer for (flash_list[0..flash_made]) |*probe| probe.deinit(); diff --git a/zig/tests/contract/metal_pair.zig b/zig/tests/contract/metal_pair.zig new file mode 100644 index 0000000000..c38bb6c9c2 --- /dev/null +++ b/zig/tests/contract/metal_pair.zig @@ -0,0 +1,598 @@ +//! The pair seam's five kernels (core/pair_rows.metal), and the float64 rules the gates check. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; +const rows_k = contract.rows; + +/// The seam's five kernels; the tail of each registry name is the MSL symbol launchOf hands the source. +pub const Which = enum { add_rms, stage, gather, argmax, pick }; + +pub const all = [_]Which{ .add_rms, .stage, .gather, .argmax, .pick }; + +/// What pair_rows.zig dispatches per kernel, so `init` can compare it against Metal's own reflection. +pub const Launched = struct { threads: u32, tg_memory: u32 }; + +pub fn launch(w: Which) Launched { + return switch (w) { + .add_rms => .{ .threads = 256, .tg_memory = 32 }, + .stage => .{ .threads = 256, .tg_memory = 0 }, + .gather => .{ .threads = 256, .tg_memory = 0 }, + .argmax => .{ .threads = 1024, .tg_memory = 256 }, + .pick => .{ .threads = 32, .tg_memory = 0 }, + }; +} + +/// The MSL symbol pair_rows.metal defines. +pub fn symbol(w: Which) []const u8 { + return switch (w) { + .add_rms => "pair_add_rms", + .stage => "pair_stage", + .gather => "pair_gather", + .argmax => "pair_argmax", + .pick => "pair_pick", + }; +} + +/// The spec's name: `metal `, so launchOf's last word is the symbol and its prefix picks core_pair_rows. +pub fn name(w: Which) []const u8 { + return switch (w) { + .add_rms => "metal pair add_rms pair_add_rms", + .stage => "metal pair stage pair_stage", + .gather => "metal pair gather pair_gather", + .argmax => "metal pair argmax pair_argmax", + .pick => "metal pair pick pair_pick", + }; +} + +/// Qwen3.8-27B's rank-local slice: the seam reduces ids [124160, 248320) on rank 1. +const logit_lo: u32 = 124160; + +/// The RMS epsilon both callers pass (qwen27's frame and glm_dsa's 1e-5). +const eps: f32 = 1e-5; + +/// A thread's chain over every 256th square, five shuffle levels, then the eight partials. +fn chainSquares(s: c.Shape) usize { + return (s.k + 255) / 256 + 5 + 8; +} + +/// A thread's chain over every 1024th logit, five shuffle levels, then the 32 simdgroup winners. +fn chainLogits(s: c.Shape) usize { + return (s.k + 1023) / 1024 + 5 + 5; +} + +/// Only moves and compares values: no sums to carry an accumulator's error. +fn noSums(_: c.Shape) usize { + return 0; +} + +const add_why = "y follows the sum of squares as ss^-1/2 and stays fp32: a bf16 sum's error is 2^-8 of the norm"; +const move_why = "a move writes each bf16 element as it read it: no sum, no rounding, so no accumulator's error can enter"; + +const argmax_rule = "a winner is (value bits, id): NaN loses every comparison, so a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF"; +const pick_rule = "the pick copies one rank's (value, id) word pair; an id equal to 0xFFFFFFFF is how a broken link reads"; + +/// A kernel's spec. The raw kernels check themselves; the two moves' bound is exact, so gate 3 keeps every edge. +pub fn spec(w: Which) c.Spec { + return switch (w) { + .add_rms => .{ + .name = name(w), + .accum = .f32, + .depth = chainSquares, + .ops = 8, + .mids = 1, + .out = .f32, + .in_place = true, + .tiles = .{ .{}, .{}, .{ .size = 256, .step = 256, .min = 256 } }, + .served = &.{ + .{ .rows = 8, .n = 5120, .k = 5120 }, + .{ .rows = 8, .n = 6144, .k = 6144 }, + .{ .rows = 2, .n = 256, .k = 256 }, + }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = add_why }, + }, + // One spec for both moves; a gather's `n` is its frame row stride, and the 17,408 shape is seam B's. + .stage, .gather => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 256, .step = 256, .min = 256 } }, + .served = if (w == .stage) &.{.{ .rows = 8, .n = 5120, .k = 5120 }} else &.{ + .{ .rows = 8, .n = 10240, .k = 5120 }, .{ .rows = 8, .n = 17408, .k = 8704 }, + .{ .rows = 8, .n = 10496, .k = 5120 }, + }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = move_why }, + }, + .argmax => .{ + .name = name(w), + .accum = .none, + .depth = chainLogits, + .out = .raw, + .nan = .{ .stated = argmax_rule }, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &.{.{ .rows = 8, .n = 1, .k = 124160 }}, + }, + .pick => .{ + .name = name(w), + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = pick_rule }, + .served = &.{ .{ .rows = 32, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 } }, + }, + }; +} + +/// `constant PairRows&` at buffer 1: `width` then the frame row `stride`; argmax and pick take their own words. +const PairRows = extern struct { + width: u32, + stride: u32, + rank: u32, + x: u32, + eps: f32, + pad0: f32, + pad1: f32, + pad2: f32, +}; + +fn rowsArgs(width: usize, stride: usize, rank: u32, x: u32, e: f32) PairRows { + return .{ .width = @intCast(width), .stride = @intCast(stride), .rank = rank, .x = x, .eps = e, .pad0 = 0, .pad1 = 0, .pad2 = 0 }; +} + +fn bit(v: f32) u32 { + return @bitCast(v); +} + +fn floatOf(v: u32) f32 { + return @bitCast(v); +} + +/// One prepared call: the buffers the kernel is bound to, the host values its reference reads, and its own check. +pub const Pair = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + pipe: mtl.Pipeline, + /// The hooks this probe exposes: the raw kernels check their own rule, the three computing ones use the reference. + v: g.VTable, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + rows: usize = 0, + width: usize = 0, + /// The rank this prepare writes first: the gather's `mine` half is its payload. + rank: u32 = 0, + /// The row stride in PairRows, in bf16 elements: the stage's source frame row, the gather's destination row. + stride: usize = 0, + + // Host inputs of the last prepare, as the reference and the check read them (freed by `drop`). + xh: ?[]f32 = null, // add_rms: this rank's part + ph: ?[]f32 = null, // add_rms: the peer's part + wh: ?[]u16 = null, // add_rms: the RMS weights, bf16 + hh: ?[]f32 = null, // add_rms: h before the launch (the kernel folds into it in place) + sh: ?[]u16 = null, // stage: the src payload, bf16 + mh: ?[]u16 = null, // gather: this rank's half (mine), bf16 + nh: ?[]u16 = null, // gather: the peer's half (recv), bf16 + kh: ?[]u16 = null, // argmax: the bf16 logits + pw: ?[]u32 = null, // pick: this rank's (value bits, id) words + qw: ?[]u32 = null, // pick: the peer's words + + msg: [192]u8 = undefined, + sync: rig_.Ref = undefined, + mine: rig_.Ref = undefined, + peer: rig_.Ref = undefined, + h: rig_.Ref = undefined, + w: rig_.Ref = undefined, + y: rig_.Ref = undefined, + logits: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which) !Pair { + const pipe = try rig_.pipeline(device, sources.core_pair_rows, symbol(which)); + errdefer pipe.deinit(); + const want = launch(which); + const tg = rig_.tgMemory(pipe); + if (tg != want.tg_memory) { + std.debug.print("metal pair: {s} takes {d} bytes of threadgroup memory, its record says {d}\n", .{ symbol(which), tg, want.tg_memory }); + return error.ThreadgroupMemoryMismatch; + } + var v = g.hooks(Pair); + if (which != .argmax and which != .pick) v.check = null; + if (which != .stage and which != .gather) v.bounds = null; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .pipe = pipe, .v = v }; + } + + pub fn deinit(self: *Pair) void { + self.drop(); + self.rig.deinit(); + self.pipe.deinit(); + } + + pub fn kernel(self: *Pair) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = &self.v }; + } + + fn drop(self: *Pair) void { + if (self.xh) |s| self.a.free(s); + if (self.ph) |s| self.a.free(s); + if (self.wh) |s| self.a.free(s); + if (self.hh) |s| self.a.free(s); + if (self.sh) |s| self.a.free(s); + if (self.mh) |s| self.a.free(s); + if (self.nh) |s| self.a.free(s); + if (self.kh) |s| self.a.free(s); + if (self.pw) |s| self.a.free(s); + if (self.qw) |s| self.a.free(s); + self.xh = null; + self.ph = null; + self.wh = null; + self.hh = null; + self.sh = null; + self.mh = null; + self.nh = null; + self.kh = null; + self.pw = null; + self.qw = null; + } + + /// Ordinary values: what the seam carries in production, symmetric about zero. + fn ordinary(rng: std.Random) f32 { + return rng.float(f32) * 2 - 1; + } + + /// f32 values as bf16, for the buffers whose datum is a bf16 frame. + fn framed(rng: std.Random) f32 { + return c.fromBf16(c.toBf16(ordinary(rng))); + } + + /// A bf16 payload: what the two moves carry, with gate 1's non-finite plants where the reference looks for them. + fn payload(self: *Pair, len: usize, fill: g.Fill, rng: std.Random) ![]u16 { + const s = try self.a.alloc(u16, len); + for (s) |*v| v.* = c.toBf16(ordinary(rng)); + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for ([_]u16{ nan, inf, ninf, nan, ninf }, 0..) |plant, p| { + const at = (p + 1) * (len / 8 + 1) + p * 3; + if (at < len) s[at] = plant; + } + } + return s; + } + + /// An input the seam would give: this rank's words, the peer's words, or the weights. + fn part(self: *Pair, len: usize, bf16: bool, fill: g.Fill, rng: std.Random) ![]f32 { + const s = try self.a.alloc(f32, len); + for (s) |*v| v.* = if (bf16) framed(rng) else ordinary(rng); + if (fill == .nonfinite) { + const nan = std.math.nan(f32); + const inf = std.math.inf(f32); + for ([_]f32{ nan, inf, -inf, nan, -inf }, 0..) |plant, p| { + const at = (p + 1) * (len / 8 + 1) + p * 3; + if (at < len) s[at] = plant; + } + } + return s; + } + + /// Inputs for `shape` as `fill` asks, outputs poisoned and guarded; false when the kernel does not serve it. + pub fn prepare(self: *Pair, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 512) return false; + self.drop(); + self.rig.reset(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + switch (self.which) { + .add_rms => { + if (shape.n != shape.k or shape.k == 0 or shape.k > 32768 or shape.k % 256 != 0) return false; + const n = shape.rows * shape.k; + self.xh = try self.part(n, false, fill, rng); + self.ph = try self.part(n, false, fill, rng); + self.mine = try self.rig.input(std.mem.sliceAsBytes(self.xh.?)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(self.ph.?)); + const w = try self.a.alloc(u16, shape.k); + self.wh = w; + for (w) |*v| v.* = c.toBf16(0.5 + rng.float(f32)); + const h = try self.a.alloc(f32, n); + self.hh = h; + for (h) |*v| v.* = ordinary(rng); + self.sync = try self.rig.scratch(128); + self.w = try self.rig.input(std.mem.sliceAsBytes(w)); + self.h = try self.rig.output(n * 4); + self.y = try self.rig.output(n * 4); + @memcpy(self.rig.bytes(self.h, n * 4), std.mem.sliceAsBytes(h)); + }, + .stage, .gather => { + if (shape.rows == 0 or shape.k == 0 or shape.k > 32768 or shape.k % 256 != 0) return false; + const n = shape.rows * shape.k; + // The gather writes this rank's half first; alternating the rank with the seed covers both orders. + self.rank = @intCast(seed & 1); + const mine = try self.payload(n, fill, rng); + self.sync = try self.rig.scratch(128); + if (self.which == .stage) { + self.stride = shape.k; + self.sh = mine; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.out = try self.rig.output(n * 2); + } else { + // A frame row holds both halves, so its stride is at least 2*width: `n` sets a wider one. + self.stride = if (shape.n > 2 * shape.k) shape.n else 2 * shape.k; + self.mh = mine; + const peer = try self.payload(n, fill, rng); + self.nh = peer; + self.mine = try self.rig.input(std.mem.sliceAsBytes(mine)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(peer)); + self.out = try self.rig.output(2 * ((shape.rows - 1) * self.stride + 2 * shape.k)); + } + }, + .argmax => { + if (shape.k == 0 or shape.k > 262144) return false; + const n = shape.rows * shape.k; + const k = try self.a.alloc(u16, n); + self.kh = k; + for (k) |*v| v.* = c.toBf16(ordinary(rng)); + if (fill == .nonfinite) { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + for (0..shape.rows) |r| { + const row = k[r * shape.k ..][0..shape.k]; + row[shape.k / 3] = switch (r % 4) { + 0 => nan, + 1 => inf, + 2 => ninf, + else => row[shape.k / 3], + }; + if (r + 1 == shape.rows) for (row) |*v| { + v.* = nan; + }; + } + } + self.logits = try self.rig.input(std.mem.sliceAsBytes(k)); + self.out = try self.rig.output(shape.rows * 8); + }, + .pick => { + if (shape.n != 1 or shape.k != 1) return false; + self.pw = try self.a.alloc(u32, 2 * shape.rows); + self.qw = try self.a.alloc(u32, 2 * shape.rows); + for (0..shape.rows) |r| { + self.pw.?[2 * r] = bit(ordinary(rng)); + self.qw.?[2 * r] = bit(ordinary(rng)); + self.pw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + self.qw.?[2 * r + 1] = rng.uintLessThan(u32, 1 << 20); + } + if (fill == .nonfinite) { + const nan = bit(std.math.nan(f32)); + const inf = bit(std.math.inf(f32)); + const ninf = bit(-std.math.inf(f32)); + const plant = [_]struct { m: u32, p: u32 }{ + .{ .m = nan, .p = bit(ordinary(rng)) }, + .{ .m = bit(ordinary(rng)), .p = inf }, + .{ .m = ninf, .p = ninf }, + .{ .m = nan, .p = nan }, + .{ .m = inf, .p = nan }, + .{ .m = bit(ordinary(rng)), .p = ninf }, + }; + for (plant, 0..) |x, r| if (r < shape.rows) { + self.pw.?[2 * r] = x.m; + self.qw.?[2 * r] = x.p; + }; + } + self.sync = try self.rig.scratch(128); + self.mine = try self.rig.input(std.mem.sliceAsBytes(self.pw.?)); + self.peer = try self.rig.input(std.mem.sliceAsBytes(self.qw.?)); + self.out = try self.rig.output(shape.rows * 4); + }, + } + self.shape = shape; + self.rows = shape.rows; + self.width = shape.k; + return true; + } + + /// One launch on the prepared buffers, waited for. + pub fn run(self: *Pair) !void { + const p = self.rig.begin(); + const e = p.e; + const groups = mtl.Size.of(self.rows, 1, 1); + const threads = mtl.Size.of(launch(self.which).threads, 1, 1); + switch (self.which) { + .add_rms => { + @memset(self.rig.bytes(self.sync, 128), 0); + @memcpy(self.rig.bytes(self.h, self.rows * self.width * 4), std.mem.sliceAsBytes(self.hh.?)); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.width, 0, 1, eps), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.h.buf, self.h.off, 4); + e.setBuffer(self.w.buf, self.w.off, 5); + e.setBuffer(self.y.buf, self.y.off, 6); + e.dispatchGroups(groups, threads); + }, + .stage => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1, 0), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.out.buf, self.out.off, 3); + e.dispatchGroups(groups, threads); + }, + .gather => { + @memset(self.rig.bytes(self.sync, 128), 0); + const frame = self.rig.bytes(self.out, 2 * self.count()); + @memset(frame, 0); // the frame's gaps: no launch writes them, so their planted value is zero + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue(rowsArgs(self.width, self.stride, self.rank, 1, 0), 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(groups, threads); + }, + .argmax => { + e.setPipeline(self.pipe); + e.setBuffer(self.logits.buf, self.logits.off, 0); + const dims = [4]u32{ @intCast(self.width), logit_lo, @intCast(self.width), 1 }; + e.setValue(dims, 1); + e.setBuffer(self.out.buf, self.out.off, 2); + e.dispatchGroups(groups, threads); + }, + .pick => { + @memset(self.rig.bytes(self.sync, 128), 0); + e.setPipeline(self.pipe); + e.setBuffer(self.sync.buf, self.sync.off, 0); + e.setValue([2]u32{ @intCast(self.rows), 1 }, 1); + e.setBuffer(self.mine.buf, self.mine.off, 2); + e.setBuffer(self.peer.buf, self.peer.off, 3); + e.setBuffer(self.out.buf, self.out.off, 4); + e.dispatchGroups(mtl.Size.of(1, 1, 1), threads); + }, + } + try self.rig.finish(p); + } + + /// Output elements of the prepared call: bf16 elements for the three frames, the two raw kernels' bytes. + pub fn count(self: *Pair) usize { + return switch (self.which) { + .add_rms => 2 * self.rows * self.width, + .stage => self.rows * self.width, + .gather => (self.rows - 1) * self.stride + 2 * self.width, + .argmax => self.rows * 8, + .pick => self.rows * 4, + }; + } + + /// Float64 outputs of the prepared inputs and each one's sums' term magnitude (a move's term is the value itself). + pub fn reference(self: *Pair, want: []f64, mag: []f64) !void { + const width = self.width; + const total = self.rows * width; + switch (self.which) { + .add_rms => { + const xh = self.xh orelse return error.NoReference; + const ph = self.ph orelse return error.NoReference; + const hh = self.hh orelse return error.NoReference; + const wh = self.wh orelse return error.NoReference; + for (0..self.rows) |r| { + var ss: f64 = 0; + for (0..width) |i| { + const v = @as(f64, hh[r * width + i]) + xh[r * width + i] + ph[r * width + i]; + ss += v * v; + } + const rs = 1 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + for (0..width) |i| { + const v = @as(f64, hh[r * width + i]) + xh[r * width + i] + ph[r * width + i]; + const w = c.fromBf16(wh[i]); + const at = r * width + i; + const span = @abs(hh[r * width + i]) + @abs(xh[r * width + i]) + @abs(ph[r * width + i]); + want[at] = v; // x adds three terms, so its bound's magnitude is their span, not the sum + mag[at] = span; + want[total + at] = w * v * rs; // y is that span carried through w and 1/rms + mag[total + at] = @abs(w) * span * rs; + } + } + }, + .stage => { + const sh = self.sh orelse return error.NoReference; + for (want, mag, 0..) |*w, *m, i| { + const v = c.fromBf16(sh[i]); + w.* = v; + m.* = @abs(v); + } + }, + .gather => { + const mh = self.mh orelse return error.NoReference; + const nh = self.nh orelse return error.NoReference; + @memset(want, 0); // the frame's gaps: no launch writes them, so their reference value is zero + for (0..self.rows) |row| { + const at = row * self.stride; + for (0..width) |i| { + want[at + i] = c.fromBf16(if (self.rank == 0) mh[row * width + i] else nh[row * width + i]); + want[at + width + i] = c.fromBf16(if (self.rank == 0) nh[row * width + i] else mh[row * width + i]); + } + } + for (want, mag) |w, *m| m.* = @abs(w); + }, + else => return error.NoReference, + } + } + + /// The moves are exact: every output element must equal the reference's bf16 bit for bit. + pub fn bounds(self: *Pair, want: []const f64, mag: []const f64, allow: []f64) !void { + _ = want; + _ = mag; + if (self.which != .stage and self.which != .gather) return error.NotMove; + for (allow) |*l| l.* = 0; + } + + /// The outputs if every sum ran in `acc` in an order, the rest of the arithmetic exact (gate 4). + pub fn simulate(self: *Pair, acc: c.Accum, order: g.Order, out: []f64) !void { + if (self.which != .add_rms) return error.NoSums; + const width = self.width; + const hh = self.hh orelse return error.NoSums; + const xh = self.xh orelse return error.NoSums; + const ph = self.ph orelse return error.NoSums; + const wh = self.wh orelse return error.NoSums; + const terms = try self.a.alloc(f64, width); + defer self.a.free(terms); + for (0..self.rows) |r| { + for (terms, 0..) |*t, i| { + const v = @as(f64, hh[r * width + i]) + xh[r * width + i] + ph[r * width + i]; + t.* = v * v; + } + const ss = rows_k.sumIn(acc, order, terms); + const rs = 1 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + for (0..width) |i| { + const v = @as(f64, hh[r * width + i]) + xh[r * width + i] + ph[r * width + i]; + const at = r * width + i; + out[at] = v; + out[self.rows * width + at] = c.fromBf16(wh[i]) * v * rs; + } + } + } + + /// The seam's own checks: the raw kernels' rule at gate 1 (non-finite inputs) and gate 3 (every served edge). + pub fn check(self: *Pair, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + switch (self.which) { + .argmax => { + const k = self.kh orelse return error.NoReference; + const n = self.width; + for (0..self.rows) |r| { + var best: f32 = -std.math.inf(f32); + var at: u32 = 0xFFFFFFFF; + for (0..n) |i| { + const v = c.fromBf16(k[r * n + i]); + if (v > best) { + best = v; + at = logit_lo + @as(u32, @intCast(i)); + } + } + const got = std.mem.readInt(u32, out[r * 8 ..][0..4], .little); + const id = std.mem.readInt(u32, out[r * 8 + 4 ..][0..4], .little); + if (got != bit(best) or id != at) return try std.fmt.bufPrint(&self.msg, "row {d}: ({e} {d}) where its own rule gives ({e} {d})", .{ r, floatOf(got), id, best, at }); + } + }, + .pick => { + const pw = self.pw orelse return error.NoReference; + const qw = self.qw orelse return error.NoReference; + for (0..self.rows) |r| { + const vm = floatOf(pw[2 * r]); + const am = pw[2 * r + 1]; + const vp = floatOf(qw[2 * r]); + const ap = qw[2 * r + 1]; + const want = if (vp > vm or (vp == vm and ap < am)) ap else am; + const got = std.mem.readInt(u32, out[r * 4 ..][0..4], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "row {d} picks {d} where its own rule gives {d}", .{ r, got, want }); + } + }, + else => return error.NotRaw, + } + return null; + } +}; diff --git a/zig/tests/qwen27/dflash_served.zig b/zig/tests/qwen27/dflash_served.zig index c813567db6..f472c3389e 100644 --- a/zig/tests/qwen27/dflash_served.zig +++ b/zig/tests/qwen27/dflash_served.zig @@ -128,7 +128,7 @@ pub fn main(init: std.process.Init) !void { } var why: []const u8 = ""; // a context the window timing outgrows: a new build's first start refills the timing's slot - const served = host.attach(a, init.io, model, draft, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why) catch |err| { + const served = host.attach(a, init.io, model, draft, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why, null) catch |err| { std.debug.print("served host: {s}\n", .{why}); return err; }; diff --git a/zig/tests/qwen27/served.zig b/zig/tests/qwen27/served.zig index 858dabf5a3..d3714aebe7 100644 --- a/zig/tests/qwen27/served.zig +++ b/zig/tests/qwen27/served.zig @@ -58,7 +58,7 @@ pub fn main(init: std.process.Init) !void { } var box = Box{ .a = a, .io = init.io }; var why: []const u8 = ""; - const served = host.attach(a, init.io, model, null, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why) catch |err| { + const served = host.attach(a, init.io, model, null, .{ .context = 256, .streams = 1, .cache_gib = 0 }, a, &why, null) catch |err| { std.debug.print("served host: {s}\n", .{why}); return err; }; diff --git a/zig/tests/qwen27/synthetic_graph.zig b/zig/tests/qwen27/synthetic_graph.zig index 3da574b7d1..f0e707634f 100644 --- a/zig/tests/qwen27/synthetic_graph.zig +++ b/zig/tests/qwen27/synthetic_graph.zig @@ -95,7 +95,8 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { } var widths = try q.widths.Widths.uniformAll(a, c.layers); defer widths.deinit(); - var weights = try q.gpu_weights.load(a, device, &cp, c, &widths); + const shard = q.config.Shard{ .rank = 0, .ranks = 1, .global = c, .local = c }; + var weights = try q.gpu_weights.load(a, device, &cp, &shard, &widths); errdefer weights.deinit(); const glue = try q.glue_gpu.Kernels.init(device); errdefer glue.deinit(); @@ -105,6 +106,6 @@ pub fn createLayers(a: std.mem.Allocator, layers: usize) !*q.model.Model { f.deinit(); } const m = try a.create(q.model.Model); - m.* = .{ .allocator = a, .device = device, .queue = queue, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; + m.* = .{ .allocator = a, .device = device, .queue = queue, .shard = shard, .config = c, .checkpoint = cp, .weights = weights, .kernels = .{ .allocator = a, .device = device }, .glue = glue, .frame = frame }; return m; } diff --git a/zig/tests/qwen27/tp_check.zig b/zig/tests/qwen27/tp_check.zig new file mode 100644 index 0000000000..56494a7ef5 --- /dev/null +++ b/zig/tests/qwen27/tp_check.zig @@ -0,0 +1,572 @@ +//! The 27B TP gate on one Mac, no wire: ranks = 1 parity bit for bit, then two shards joined through `pair.Loopback`. +const std = @import("std"); +const q = @import("qwen27"); +const mtl = @import("metal"); +const pair = @import("core").pair; + +const Runner = q.decode_round.Runner; +const Session = q.session.Session; +const TopTwo = q.session.TopTwo; +const Model = q.model.Model; + +const help = + \\tf-qwen27-tp-check --model DIR [--steps N] [--ids N] [--chunk N] [--stage 1|2|3] [--mode full|mlp] + \\ stage 1: ranks = 1 through loadShard against the single-Mac load: logits bytes, state digest, greedy ids. + \\ stage 2: two rank-local shards, every layer seam through pair.Loopback (host memcpy link), ids vs stage 1. + \\ stage 3 (the default) runs both. + \\--chunk 1 (default) keeps the prompt on the decode path; > 1 uses the prompt kernels. Greedy only, no drafter. + \\--prompt-file F takes the prompt ids from F; --forced-file F feeds F's ids instead of each side's own picks. + \\--taps A,B,C,D,E hashes five layers' post-prompt hidden rows on every side (stage 2 divergence localisation). + \\Synthetic ids 1000.., no tokenizer, no weights written; JSON receipt on stdout, errors on stderr. + \\Run with no --model DIR to print SKIP and exit 0: the suite does that on hosts without the checkpoint. +; + +var dump = false; +/// Five layers whose post-block residuals the tap planes keep, for stage 2 divergence localisation (`--taps`). +var tap_ids: ?[5]usize = null; +fn hashBuffer(h: anytype, b: mtl.Buffer) void { + h.update(b.contents()[0..b.length()]); +} + +/// Five tapped planes of one round: post-block hidden rows, plane by plane, so a divergence names its layer. +fn tapDigests(r: *const Runner) ![5][32]u8 { + var out: [5][32]u8 = undefined; + for (0..out.len) |i| { + const plane = try r.taps.plane(i); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(plane.buffer.contents()[plane.offset..][0 .. @as(usize, r.taps.rows) * r.taps.width * 2]); + h.final(&out[i]); + } + return out; +} + +/// `--taps 0,1,2,3,4`: five layer indices, ascending, whose post-block residuals stage 2 keeps for comparison. +fn readTaps(text: []const u8) ![5]usize { + var out: [5]usize = undefined; + var it = std.mem.tokenizeScalar(u8, text, ','); + var n: usize = 0; + while (it.next()) |token| { + if (n == out.len) return error.BadTapList; + out[n] = try std.fmt.parseInt(usize, token, 10); + n += 1; + } + if (n != out.len) return error.BadTapList; + return out; +} + +/// Everything the next round reads back: both KV caches a layer slot, the GDN states, the slots' offsets. +fn stateDigest(r: *const Runner) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (r.caches) |cache| { + hashBuffer(&h, cache.keys.buffer); + hashBuffer(&h, cache.values.buffer); + } + for (r.gdn.buffers) |b| hashBuffer(&h, b); + h.update(std.mem.sliceAsBytes(r.offsets)); + var out: [32]u8 = undefined; + h.final(&out); + return out; +} + +/// The KV pages and recurrent state a round's arithmetic reads back, without the prompt kernels' scratch. +fn committedDigest(r: *const Runner) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (r.caches) |cache| { + hashBuffer(&h, cache.keys.buffer); + hashBuffer(&h, cache.values.buffer); + } + for (0..r.gdn.budget.layers) |layer| for ([_]bool{ true, false }) |recurrent| { + if (r.gdn.committed(layer, 0, recurrent)) |bytes| h.update(bytes) else |_| h.update("err"); + }; + h.update(std.mem.sliceAsBytes(r.offsets)); + var out: [32]u8 = undefined; + h.final(&out); + return out; +} + +fn dumpRow(tag: []const u8, m: *const Model, r: *const Runner) void { + if (!dump) return; + const values = m.frame.get(.logits).buffer.slice(u16, m.config.vocab); + const span = m.shard.vocabRows(); + const two = q.session.topTwo(values) catch |err| { + std.debug.print("[dump] {s} rank={d}/{d} vocab={d} span={d}+{d} top2=error.{s}\n", .{ tag, m.shard.rank, m.shard.ranks, m.config.vocab, span.from, span.len, @errorName(err) }); + return; + }; + std.debug.print("[dump] {s} rank={d}/{d} vocab={d} span={d}+{d} capacity={d} top2 ids={d},{d} logits={d},{d} first8={any}\n", .{ + tag, m.shard.rank, m.shard.ranks, m.config.vocab, span.from, span.len, m.frame.capacity, + two.ids[0], two.ids[1], two.logits[0], two.logits[1], values[0..8], + }); + const digest = std.fmt.bytesToHex(stateDigest(r), .lower); + const committed = std.fmt.bytesToHex(committedDigest(r), .lower); + std.debug.print("[state] {s} kv_gdn_offsets={s} committed={s} ready={d} offsets={any} heads={d} kv_heads={d} k_heads={d} v_heads={d} intermediate={d}\n", .{ + tag, &digest, &committed, r.ready, r.offsets, m.config.heads, m.config.kv_heads, m.config.k_heads, m.config.v_heads, m.config.intermediate, + }); +} + +fn shaLogits(m: *const Model) [32]u8 { + var hash = std.crypto.hash.sha2.Sha256.init(.{}); + hash.update(std.mem.sliceAsBytes(m.frame.get(.logits).buffer.slice(u16, m.config.vocab))); + var out: [32]u8 = undefined; + hash.final(&out); + return out; +} + +/// One rank's local top two, shifted to the global vocabulary, for the mismatch report. +fn mergeTop2(lo: [2]u32, local: [2]TopTwo) TopTwo { + var best = TopTwo{ .ids = .{ 0, 0 }, .logits = .{ -std.math.inf(f32), -std.math.inf(f32) } }; + for (local, 0..) |two, block| for (0..2) |k| { + const value = two.logits[k]; + const id = two.ids[k] + lo[block]; + // the old fields are read into locals first: assigning to `best` builds the literal in place + const old_id = best.ids[0]; + const old_value = best.logits[0]; + if (value > old_value) { + best = .{ .ids = .{ id, old_id }, .logits = .{ value, old_value } }; + } else if (value > best.logits[1]) { + best = .{ .ids = .{ old_id, id }, .logits = .{ old_value, value } }; + } + }; + return best; +} + +const Local = struct { + tokens: []u32, + top2: []TopTwo, + logits: [][32]u8, + state: [32]u8, + /// Post-block residuals of the prompt round at `tap_ids`, hashed on this side's own hidden state. + taps: [5][32]u8 = undefined, + /// `MTLDevice.currentAllocatedSize` with this one model resident (the engine's own memory accounting). + resident: u64, +}; + +/// A whole stream on one Mac: `shard1` takes `loadShard(..., rank = 0, ranks = 1, .full)`, else `Model.load`. +fn localRun(a: std.mem.Allocator, io: std.Io, dir: []const u8, ids: []const u32, steps: usize, chunk: usize, force: ?[]const u32, shard1: bool) !Local { + const m = if (shard1) try Model.loadShard(a, io, dir, 128, 0, 1, .full, null) else try Model.load(a, io, dir, 128, null); + defer m.deinit(); + var runner = try Runner.init(a, m, 1, rowsFor(ids.len, steps)); + defer runner.deinit(); + const s = Session{ .runner = &runner }; + if (tap_ids) |layers| { + runner.capture_taps = true; + runner.taps.ids = layers; + } + try s.prefill(ids, chunk); + const taps = try tapDigests(&runner); + runner.capture_taps = false; + const state = try s.fingerprint(); + dumpRow("single", m, &runner); + const resident = m.device.allocated(); + const tokens = try a.alloc(u32, steps); + errdefer a.free(tokens); + const top2 = try a.alloc(TopTwo, steps); + errdefer a.free(top2); + const logits = try a.alloc([32]u8, steps); + errdefer a.free(logits); + for (0..steps) |i| { + tokens[i] = try s.greedy(); + top2[i] = try q.session.topTwo(m.frame.get(.logits).buffer.slice(u16, m.config.vocab)); + logits[i] = shaLogits(m); + // a forced replay feeds the recorded history, so both sides stand at the same prefix when they compare + if (i + 1 < steps) try s.step(if (force) |f| f[i] else tokens[i]); + } + return .{ .tokens = tokens, .top2 = top2, .logits = logits, .state = state, .taps = taps, .resident = resident }; +} + +/// One rank of the pair: its shard, its seam transport, and the stream it produced. +const Side = struct { + rank: u32, + ids: []const u32, + steps: usize, + chunk: usize, + /// A recorded history both ranks step through instead of their own picks (the single Mac's reply, say). + force: ?[]const u32 = null, + m: *Model = undefined, + lb: pair.Loopback = undefined, + picks: mtl.Buffer = undefined, + runner: Runner = undefined, + tokens: []u32 = &.{}, + top2: []TopTwo = &.{}, + logits: [][32]u8 = &.{}, + /// When dumping, each rank copies its span of every step's logits here, so the whole row can be hashed. + whole: ?[]u8 = null, + state: [32]u8 = undefined, + /// Post-block residuals of the prompt round at `tap_ids`, hashed on this rank's own hidden state. + taps: [5][32]u8 = undefined, + err: ?anyerror = null, + + fn open(s: *Side, a: std.mem.Allocator, io: std.Io, dir: []const u8, mode: q.config.Mode, bytes: usize) !void { + s.m = try Model.loadShard(a, io, dir, 128, s.rank, 2, mode, null); + errdefer s.m.deinit(); + s.runner = try Runner.init(a, s.m, 1, rowsFor(s.ids.len, s.steps)); + errdefer s.runner.deinit(); + // one `pickBytes` a row, and a prompt round writes `chunk` rows before the decodes write one each + s.picks = try s.m.device.buffer(@max(s.steps, s.chunk) * 2 * @sizeOf(u32), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer s.picks.deinit(); + s.lb = try pair.Loopback.init(s.m.device, bytes, s.rank); + errdefer s.lb.deinit(); + s.runner.reduce = .{ .ptr = &s.lb, .part = loopPart, .stage = loopStage, .gather = loopGather, .failed = loopFailed, .argmax = loopArgmax }; + s.runner.picks = .{ .buffer = s.picks, .offset = 0 }; + s.tokens = try a.alloc(u32, s.steps); + s.top2 = try a.alloc(TopTwo, s.steps); + s.logits = try a.alloc([32]u8, s.steps); + } + + fn close(s: *Side) void { + s.runner.deinit(); + s.picks.deinit(); + s.lb.deinit(); + s.m.deinit(); + } + + fn run(s: *Side) void { + s.stream() catch |err| { + s.err = err; + // a rank that stops must not leave its peer's seam waiting: the sentinel is what a failed link does + s.lb.satisfy(); + }; + } + + fn stream(s: *Side) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const ses = Session{ .runner = &s.runner }; + if (tap_ids) |layers| { + s.runner.capture_taps = true; + s.runner.taps.ids = layers; + } + try ses.prefill(s.ids, s.chunk); + s.taps = try tapDigests(&s.runner); + s.runner.capture_taps = false; + s.state = try ses.fingerprint(); + trace(s, "prefill"); + dumpRow(if (s.rank == 0) "pair-rank0" else "pair-rank1", s.m, &s.runner); + for (0..s.steps) |i| { + const values = s.m.frame.get(.logits).buffer.slice(u16, s.m.config.vocab); + s.tokens[i] = try ses.greedy(); + s.top2[i] = try q.session.topTwo(values); + s.logits[i] = shaLogits(s.m); + if (s.whole) |w| { + const span = s.m.shard.vocabRows(); + const row = i * @as(usize, s.m.shard.global.vocab) * 2; + @memcpy(w[row + span.from * 2 ..][0 .. span.len * 2], std.mem.sliceAsBytes(values)); + } + trace(s, "step"); + if (i + 1 < s.steps) try ses.step(if (s.force) |f| f[i] else s.tokens[i]); + } + } +}; + +/// The seam's exchange counter after every round: both ranks must have encoded the same number of exchanges. +fn trace(s: *const Side, where: []const u8) void { + if (!dump) return; + std.debug.print("[x] rank={d} after={s} x={d} token={d}\n", .{ s.rank, where, s.lb.x, s.tokens[0] }); +} + +fn loopPart(p: *anyopaque) q.forward.Seam { + const lb: *pair.Loopback = @ptrCast(@alignCast(p)); + const seam = lb.part(); + return .{ .x = seam.x, .ref = .{ .buffer = seam.ref.buf, .offset = seam.ref.off } }; +} + +fn loopStage(p: *anyopaque, e: *mtl.ComputeEncoder, x: u32, src: q.forward.Ref, rows: u32, width: u32, stride: u32) void { + const lb: *pair.Loopback = @ptrCast(@alignCast(p)); + lb.stage(e, x, .{ .buf = src.buffer, .off = src.offset }, rows, width, stride); +} + +fn loopGather(p: *anyopaque, e: *mtl.ComputeEncoder, x: u32, mine: q.forward.Ref, out: q.forward.Ref, rows: u32, width: u32, out_stride: u32) void { + const lb: *pair.Loopback = @ptrCast(@alignCast(p)); + lb.gather(e, x, .{ .buf = mine.buffer, .off = mine.offset }, .{ .buf = out.buffer, .off = out.offset }, rows, width, out_stride); +} + +fn loopFailed(p: *anyopaque) bool { + const lb: *pair.Loopback = @ptrCast(@alignCast(p)); + return lb.failed(); +} + +fn loopArgmax(p: *anyopaque, e: *mtl.ComputeEncoder, x: u32, logits: q.forward.Ref, stride: u32, lo: u32, n: u32, rows: u32, picks: q.forward.Ref) void { + const lb: *pair.Loopback = @ptrCast(@alignCast(p)); + lb.argmaxBf16(e, x, .{ .buf = logits.buffer, .off = logits.offset }, stride, lo, n, rows, .{ .buf = picks.buffer, .off = picks.offset }); +} + +/// The wire, factored out: one thread per direction relays the producer's posted part into the consumer's slot. +fn pump(from: *pair.Loopback, to: *pair.Loopback, stop: *std.atomic.Value(bool)) void { + var x: u32 = 1; + while (!stop.load(.acquire)) { + if (from.posted.value() < x) { + std.atomic.spinLoopHint(); + continue; + } + const n: usize = from.bytesFor(x); + to.arrive(x, from.win[from.sendOffset(x)..][0..n]); + x +%= 1; + } +} + +fn runSide(s: *Side) void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + s.run(); +} + +/// KV rows one stream needs: every prompt id plus every decode step, padded so a round never runs out mid-way. +fn rowsFor(ids: usize, steps: usize) u32 { + return @intCast(@min(ids + steps + 8, 32 * 1024)); +} + +/// Ids from a dump (`[248045, 846, …]`) or any run of decimals: brackets, commas and whitespace all separate. +fn readIds(a: std.mem.Allocator, io: std.Io, path: []const u8) ![]u32 { + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 20)); + var out: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, text, "[](), \t\r\n"); + while (it.next()) |token| try out.append(a, try std.fmt.parseInt(u32, token, 10)); + return out.toOwnedSlice(a); +} + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len == 1) { + try std.Io.File.stdout().writeStreamingAll(io, "SKIP qwen27-tp-check: no --model DIR, so this host has no checkpoint to check\n"); + return; + } + if (args.len == 2 and std.mem.eql(u8, args[1], "--help")) { + try std.Io.File.stdout().writeStreamingAll(io, help ++ "\n"); + return; + } + var dir: ?[]const u8 = null; + var steps: usize = 8; + var count: usize = 64; + var chunk: usize = 1; + var stage: u8 = 3; + var mode: q.config.Mode = .full; + var prompt_file: ?[]const u8 = null; + var forced_file: ?[]const u8 = null; + var i: usize = 1; + while (i < args.len) : (i += 1) { + const arg = args[i]; + if (std.mem.eql(u8, arg, "--dump")) { + dump = true; + continue; + } + if (i + 1 >= args.len) return error.MissingValue; + i += 1; + const value = args[i]; + if (std.mem.eql(u8, arg, "--model")) dir = value else if (std.mem.eql(u8, arg, "--steps")) { + steps = try std.fmt.parseInt(usize, value, 10); + } else if (std.mem.eql(u8, arg, "--ids")) { + count = try std.fmt.parseInt(usize, value, 10); + } else if (std.mem.eql(u8, arg, "--chunk")) { + chunk = try std.fmt.parseInt(usize, value, 10); + } else if (std.mem.eql(u8, arg, "--stage")) { + stage = try std.fmt.parseInt(u8, value, 10); + } else if (std.mem.eql(u8, arg, "--mode")) { + mode = if (std.mem.eql(u8, value, "mlp")) .mlp else if (std.mem.eql(u8, value, "full")) .full else return error.UnknownMode; + } else if (std.mem.eql(u8, arg, "--prompt-file")) { + prompt_file = value; + } else if (std.mem.eql(u8, arg, "--forced-file")) { + forced_file = value; + } else if (std.mem.eql(u8, arg, "--taps")) { + tap_ids = try readTaps(value); + } else return error.UnknownOption; + } + if (dir == null or steps == 0 or steps > 256 or count == 0 or count > 512 or chunk == 0 or chunk > 128 or stage == 0 or stage > 3) return error.BadOptions; + // synthetic ids: no tokenizer, deterministic, every id a valid embedding row + const ids = if (prompt_file) |f| try readIds(a, io, f) else blk: { + const out = try a.alloc(u32, count); + for (out, 0..) |*id, j| id.* = @intCast(1000 + j); + break :blk out; + }; + if (ids.len == 0 or ids.len > 512) return error.BadOptions; + count = ids.len; + const force: ?[]const u32 = if (forced_file) |f| try readIds(a, io, f) else null; + if (force) |f| if (f.len < steps) return error.ShortForcedFile; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const out = std.Io.File.stdout(); + + var reference: ?Local = null; + if (stage & 1 != 0) { + const single = try localRun(a, io, dir.?, ids, steps, chunk, force, false); + const shard1 = try localRun(a, io, dir.?, ids, steps, chunk, force, true); + const same_tokens = std.mem.eql(u32, single.tokens, shard1.tokens); + const same_state = std.mem.eql(u8, &single.state, &shard1.state); + var first_logits: ?usize = null; + for (single.logits, shard1.logits, 0..) |one, other, at| if (first_logits == null and !std.mem.eql(u8, &one, &other)) { + first_logits = at; + }; + const single_hex = std.fmt.bytesToHex(single.state, .lower); + const shard1_hex = std.fmt.bytesToHex(shard1.state, .lower); + const receipt = try std.json.Stringify.valueAlloc(a, .{ + .kind = "qwen27_tp_check", + .stage = 1, + .model = dir.?, + .ids = count, + .steps = steps, + .chunk = chunk, + .single_mac_tokens = single.tokens, + .shard1_tokens_equal = same_tokens, + .shard1_logits_byte_equal = first_logits == null, + .shard1_first_logits_difference = first_logits, + .shard1_prompt_state_equal = same_state, + .single_mac_resident_bytes = single.resident, + .shard1_resident_bytes = shard1.resident, + .single_mac_prompt_state = @as([]const u8, &single_hex), + .shard1_prompt_state = @as([]const u8, &shard1_hex), + }, .{}); + try out.writeStreamingAll(io, receipt); + try out.writeStreamingAll(io, "\n"); + if (!same_tokens or first_logits != null or !same_state) return error.ShardAtRanksOneDiffers; + reference = single; + } + + if (stage & 2 == 0) return; + const ref = reference orelse return error.StageTwoNeedsTheSingleMacStream; + + // The window holds the widest prompt round: the MLP seam stages a bf16 row of this rank's half of `intermediate`. + const bytes = pair.stageBytes(128, 17408 / 2); + const zero = try a.create(Side); + const one = try a.create(Side); + zero.* = .{ .rank = 0, .ids = ids, .steps = steps, .chunk = chunk, .force = force }; + one.* = .{ .rank = 1, .ids = ids, .steps = steps, .chunk = chunk, .force = force }; + try zero.open(a, io, dir.?, mode, bytes); + defer zero.close(); + const rank0_alone = zero.m.device.allocated(); + try one.open(a, io, dir.?, mode, bytes); + defer one.close(); + const both = zero.m.device.allocated(); + const working = zero.m.device.maxWorkingSet(); + const global_vocab = zero.m.shard.global.vocab; + + // Dumping also keeps every step's logits row whole, so the pair's row can be hashed against the single Mac's row. + const whole: ?[]u8 = if (dump) try a.alloc(u8, steps * @as(usize, global_vocab) * 2) else null; + zero.whole = whole; + one.whole = whole; + + var stop = std.atomic.Value(bool).init(false); + const a2b = try std.Thread.spawn(.{}, pump, .{ &zero.lb, &one.lb, &stop }); + const b2a = try std.Thread.spawn(.{}, pump, .{ &one.lb, &zero.lb, &stop }); + const t0 = try std.Thread.spawn(.{}, runSide, .{zero}); + const t1 = try std.Thread.spawn(.{}, runSide, .{one}); + t0.join(); + t1.join(); + stop.store(true, .release); + a2b.join(); + b2a.join(); + + if (zero.err) |err| { + try out.writeStreamingAll(io, try std.fmt.allocPrint(a, "{{\"kind\":\"qwen27_tp_check\",\"stage\":2,\"rank0_error\":\"{s}\"}}\n", .{@errorName(err)})); + return err; + } + if (one.err) |err| { + try out.writeStreamingAll(io, try std.fmt.allocPrint(a, "{{\"kind\":\"qwen27_tp_check\",\"stage\":2,\"rank1_error\":\"{s}\"}}\n", .{@errorName(err)})); + return err; + } + + if (dump) { + // both ranks recorded a length per exchange: a disagreement means the layer loops entered the seam differently + const x0 = zero.lb.x; + const x1 = one.lb.x; + var bad: usize = 0; + var first_bad: ?usize = null; + for (1..@min(x0, x1) + 1) |at| { + const want = zero.lb.bytesFor(@intCast(at)); + const got = one.lb.bytesFor(@intCast(at)); + if (want == got) continue; + bad += 1; + if (first_bad == null) first_bad = at; + if (bad <= 8) std.debug.print("[audit] x={d} rank0_bytes={d} rank1_bytes={d}\n", .{ at, want, got }); + } + std.debug.print("[audit] x_rank0={d} x_rank1={d} mismatched_exchanges={d} first_bad={?d}\n", .{ x0, x1, bad, first_bad }); + } + + // A tapped layer whose hidden rows differ names where a seam stops being bit-exact; the logits only show the drift. + var first_tap_diff: ?usize = null; + if (tap_ids) |layers| for (layers, 0..) |layer, at| { + const zero_hex = std.fmt.bytesToHex(zero.taps[at], .lower); + const one_hex = std.fmt.bytesToHex(one.taps[at], .lower); + const want_hex = std.fmt.bytesToHex(ref.taps[at], .lower); + std.debug.print("[taps] layer={d} single={s} rank0={s} rank1={s}\n", .{ layer, want_hex[0..], zero_hex[0..], one_hex[0..] }); + if (first_tap_diff == null and !std.mem.eql(u8, &ref.taps[at], &zero.taps[at]) and !std.mem.eql(u8, &ref.taps[at], &one.taps[at])) first_tap_diff = layer; + }; + + // The token comparison only sees a tie when one flips; hashing every step's whole row shows the first bit of drift. + var first_logit_step: ?usize = null; + if (whole) |w| { + const row_bytes = @as(usize, global_vocab) * 2; + for (0..steps) |at| { + var hash = std.crypto.hash.sha2.Sha256.init(.{}); + hash.update(w[at * row_bytes ..][0..row_bytes]); + var got: [32]u8 = undefined; + hash.final(&got); + if (std.mem.eql(u8, &got, &ref.logits[at])) continue; + first_logit_step = at; + const merged = mergeTop2(.{ 0, @intCast(zero.m.config.vocab) }, .{ zero.top2[at], one.top2[at] }); + const want = ref.top2[at]; + std.debug.print("[logits] first_differing_step={d} single={d}@{d},{d}@{d} pair={d}@{d},{d}@{d}\n", .{ + at, want.ids[0], want.logits[0], want.ids[1], want.logits[1], merged.ids[0], merged.logits[0], merged.ids[1], merged.logits[1], + }); + break; + } + std.debug.print("[logits] steps_compared={d} first_differing_step={?d}\n", .{ steps, first_logit_step }); + } + + var first_difference: ?usize = null; + for (ref.tokens, zero.tokens, 0..) |want, got, at| if (first_difference == null and want != got) { + first_difference = at; + }; + const ranks_agree = std.mem.eql(u32, zero.tokens, one.tokens); + var bad_pick: ?usize = null; + for (zero.tokens, one.tokens, 0..) |x, y, at| if (bad_pick == null and (x >= global_vocab or y >= global_vocab or x != y)) { + bad_pick = at; + }; + var merged: ?TopTwo = null; + var single: ?TopTwo = null; + if (first_difference) |at| { + merged = mergeTop2(.{ 0, @intCast(zero.m.config.vocab) }, .{ zero.top2[at], one.top2[at] }); + single = ref.top2[at]; + } + const zero_state = std.fmt.bytesToHex(zero.state, .lower); + const one_state = std.fmt.bytesToHex(one.state, .lower); + const receipt = try std.json.Stringify.valueAlloc(a, .{ + .kind = "qwen27_tp_check", + .stage = 2, + .model = dir.?, + .mode = @tagName(mode), + .ranks = 2, + .ids = count, + .steps = steps, + .chunk = chunk, + .rank0_tokens = zero.tokens, + .rank1_tokens = one.tokens, + .reference_tokens = ref.tokens, + .ranks_token_equal = ranks_agree, + .reference_token_equal = first_difference == null, + .first_difference = first_difference, + .single_top2 = single, + .rank0_top2 = if (first_difference) |at| zero.top2[at] else null, + .rank1_top2 = if (first_difference) |at| one.top2[at] else null, + .pair_merged_top2 = merged, + .pair_logits_compared = whole != null, + .pair_logits_byte_equal = if (whole == null) null else first_logit_step == null, + .pair_first_logits_difference = first_logit_step, + .bad_pick_at = bad_pick, + .seam_gave_up_rank0 = zero.lb.gaveUp(), + .seam_gave_up_rank1 = one.lb.gaveUp(), + .rank0_prompt_state = @as([]const u8, &zero_state), + .rank1_prompt_state = @as([]const u8, &one_state), + .global_vocab = global_vocab, + .local_vocab = zero.m.config.vocab, + .rank0_alone_resident_bytes = rank0_alone, + .both_ranks_resident_bytes = both, + .second_rank_increment_bytes = both - rank0_alone, + .device_max_working_set = working, + }, .{}); + try out.writeStreamingAll(io, receipt); + try out.writeStreamingAll(io, "\n"); + if (!ranks_agree or bad_pick != null) return error.RanksDisagree; + if (zero.lb.gaveUp() != 0 or one.lb.gaveUp() != 0) return error.SeamGaveUp; + if (first_difference != null) return error.PairTokensDifferFromSingleMac; + if (first_logit_step != null) return error.PairLogitsDifferFromSingleMac; +}