diff --git a/CHANGELOG.md b/CHANGELOG.md index 53e64cda2c..984be25440 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the GLM CUDA recipe says what `TF_GLM_CHUNK_SOLO` gains at TP=3 with `TF_GLM_ROCE_MULTI` (a little under half of its gain without it) (#1996). Thanks to @Bizuayeu. - DeepSeek-V4.1 on CUDA reads `--prompt-cache-gib` as the kept states' resident limit (0: none kept) unless `TF_DSV41_SESSION_RESIDENT_GIB` is set, which wins; with TP > 1 the variable is still needed on every rank (#1625). - DeepSeek-V4.1 on CUDA reads `--prompt-cache-gib` as the kept states' resident limit on every rank (0: none resident in RAM; with `TF_DSV41_SESSION_DISK` an evicted chat is parked on disk): with TP > 1 rank 0 sends its limit to the other ranks at startup. A set `TF_DSV41_SESSION_RESIDENT_GIB` wins on rank 0; set on another rank, it must equal rank 0's limit or every rank refuses to start (#1625). +- The native HIP engine serves `Qwen/Qwen3-Embedding-8B` as Qwen publishes it (BF16) through `/v1/embeddings` on a Radeon AI PRO R9700 (gfx1201): BF16 dense weights stay BF16 on a new dot2 GEMM tile, inputs of up to 32,768 tokens batch to the same vectors as alone, and `dimensions` takes any MRL width from 32 to 4,096. Every family's `/v1/embeddings` now refuses a whitespace-only text as SGLang does, and names an input longer than the context with its token count (`context_length_exceeded`) before the batch limit; see [the API notes](docs/api.md#embeddings) (#1978, part of #1977). Thanks to @laudney. + - Eight remaining host BF16 roundings in test oracles (GLM's Metal checks and CUDA FP8 KV reference, Nemotron's training scan, LFM2's DSpark taps, Flash Next's HIP kernel tests) and the XPU FP64 GDN reference use core's shared conversion, so a negative NaN gives a NaN instead of overflowing in Debug builds; every other non-NaN input gives the same bits. No kernel or served path changes (part of #1863). - Living Weights on Qwen4 dense, step 1: the QSA backward pass on the host, checked against a binary64 reference; nothing served changes (#2021, part of #1804). diff --git a/docs/api.md b/docs/api.md index 1586675e0b..e044f3cff0 100644 --- a/docs/api.md +++ b/docs/api.md @@ -498,6 +498,22 @@ On a discrete card (an RTX 3090, an A100) the server sizes its window from the c reports it, less a reserve of a 32nd of the card and at least 1 GiB. `TENSORFOLD_MEMORY_RESERVE_GIB` sets that one from 0.5 GiB up. +## Embeddings + +`POST /v1/embeddings` takes OpenAI's body: `input` is one text, a list of texts, one list of token IDs or a list of +them, with `encoding_format` (`float` or `base64`) and `dimensions`. Every family that serves embeddings shares these +refusals, each an HTTP 400 `invalid_request_error`: + +| Input | Refusal | +| --- | --- | +| a text that is empty or whitespace only (Python's whitespace set, which SGLang strips before it refuses) | "embedding text must not be empty or whitespace only"; in a list, "embedding input at index N must not be empty or whitespace only" | +| one input longer than the model's context | code `context_length_exceeded`: "the input (N tokens) is longer than the model's context length (C tokens)" | +| inputs that each fit but together reach the server's batch limit | "the embedding batch exceeds the admitted N tokens" | + +An over-long input is named before the batch limit is checked, so a client learns that the input itself is too long, +not that the batch is too full. `dimensions` crops each vector to that many leading values and renormalizes it; a +model with Matryoshka (MRL) widths takes any width in its declared range (32 to 4,096 for Qwen3-Embedding-8B). + ## Responses `choices[0].message.content` holds the answer. Reasoning uses `reasoning_content`, or diff --git a/docs/recipes/qwen3-embedding-hip.md b/docs/recipes/qwen3-embedding-hip.md new file mode 100644 index 0000000000..364864d4d8 --- /dev/null +++ b/docs/recipes/qwen3-embedding-hip.md @@ -0,0 +1,100 @@ +# Qwen3-Embedding-8B on native HIP + +The native HIP engine serves `Qwen/Qwen3-Embedding-8B` as published (BF16 +safetensors, revision `1d8ad4ca9b3dd8059ad90a75d4983776a23d44af`) through +`POST /v1/embeddings`. It is qualified on one Radeon AI PRO R9700 (gfx1201). +gfx11 builds run the same code; FP16 builds (RDNA2) refuse the checkpoint. + +| Responsibility | Files | +| --- | --- | +| Config, BF16 weights, batched prompt pass | `zig/src/families/qwen3/` (`config.zig`, `model.zig`, `forward.zig`) | +| Prompt-only engine, lane backend, embedding head | `zig/src/families/qwen3/` (`engine.zig`, `prompt.zig`, `native.zig`) | +| BF16 dense GEMM (dot2 tile, BF16 weight decoder) | `zig/kernels/hip/quant/bf16_decoder.hpp`, `dense_tiles.hpp`, `zig/src/hip/launches/dense.zig` | +| Head-major attention rows | `tf_heads_rows` in `zig/kernels/hip/ops/attention.hpp` | +| Stored-dtype dense weights | `zig/src/core/quant/dense.zig` (`readStored`) | +| Device contracts | `zig/tests/hip/contract_dense.zig`, `contract_heads_rows.zig` | +| Pinned model-card oracle | `tools/zig/oracle_qwen3_embedding.py` | + +The weights stay BF16 on the GPU (14.1 GiB). Activations are BF16 and every +product sums in fp32. The head comes from the sentence-transformers folder: +`1_Pooling/config.json` must pool the last token only, and `modules.json`'s +`Normalize` makes every full-width vector unit length. A folder without +`1_Pooling` is refused. + +## Build and serve + +```sh +zig build native-hip -Dhip=true -Dhip-arch=gfx1201 -Dhipcc=/opt/rocm/core/bin/hipcc +zig-out/native-hip/bin/tensorfold-native serve MODEL_DIRECTORY \ + --name Qwen3-Embedding-8B --host 127.0.0.1 --port 8090 +``` + +With ROCm 10.1, pass the hipcc under `/opt/rocm/core`: the build takes the +ROCm root from the hipcc path, and `/opt/rocm` lacks the device math headers. + +The server sizes one pass to `--context` (32,768 tokens by default, at most +the config's 40,960) and shrinks it by halves to fit the memory budget, down +to 2,048; the startup line names the window. A request holds at most 128 +inputs and the window's tokens in all. `dimensions` crops to any width from +32 to 4,096 (the card's MRL range) and renormalizes. The chat and completion +routes are refused. + +Queries take the card's instruction prefix; documents are sent as they are: + +```sh +curl -s http://127.0.0.1:8090/v1/embeddings -H 'Content-Type: application/json' -d '{ + "input": ["Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:Explain gravity", + "Gravity is a force that attracts two bodies towards each other."]}' +``` + +Text inputs end in `<|endoftext|>` through the tokenizer's post-processor, +which last-token pooling reads. Token-ID inputs are used as sent, so a client +must append 151643 itself. + +## Gates + +The oracle runs the model card's Transformers example in a venv with +Transformers 5.22.1. It refuses a snapshot or a Transformers install whose +files differ from its pins before any model code runs. Its fidelity set is +the card's 4 inputs and 5 long inputs of 2,044 to 32,759 tokens, each the +longest whole-line prefix of a pinned file within 2K, 4K, 8K, 16K or 32K +tokens (the card's `LICENSE` and `README.md`, three Transformers sources); +every input's text and token IDs are pinned by SHA256. The card's inputs run +as the card runs them, one padded batch; each long input runs alone. The +fp32 reference takes about 20 minutes on 32 CPU cores: + +```sh +python3 tools/zig/oracle_qwen3_embedding.py sources --model MODEL_DIRECTORY +python3 tools/zig/oracle_qwen3_embedding.py ref --model MODEL_DIRECTORY --dtype float32 --device cpu --out ref32.json +python3 tools/zig/oracle_qwen3_embedding.py ref --model MODEL_DIRECTORY --dtype bfloat16 --device cuda --out ref16.json +python3 tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8090 +``` + +`check` sends the inputs in requests that fit the server's window (the card's +and the 2K to 16K inputs together, the 32K one alone). It passes when every +vector lies within transformers' own BF16 distance of the fp32 vector, +rounded up to two digits and at most 0.001 (distance = 1 - cosine), has unit +norm, is the same fp32 bytes alone, in a batch and from token IDs, and the +card's 2 x 2 scores are within 0.01. At MRL widths 32, 1024 and 4096 (`dimensions`) every vector +must have unit norm and lie within the same bound, taken at that width, of +the fp32 vector cropped and renormalized. The full-width gate runs on the +GPU without a server too: + +```sh +TF_QWEN3_EMBED_DIR=MODEL_DIRECTORY TF_QWEN3_EMBED_REF=ref32.json TF_QWEN3_EMBED_BAR=ref16.json \ + zig build qwen3-hip-gpu-test -Dhip=true -Dhip-arch=gfx1201 -Dhipcc=/opt/rocm/core/bin/hipcc +``` + +Results on the R9700 (ROCm 10.1): + +| Check | Result | +| --- | --- | +| Card example, distance to fp32 (4 inputs) | 1.32e-4, 9.30e-5, 7.76e-5, 9.27e-5 | +| Long inputs, distance to fp32 (2,044, 4,043, 8,177, 16,383, 32,759 tokens) | 1.16e-4, 5.24e-5, 7.00e-5, 9.90e-5, 9.22e-5 | +| Full-width limit (transformers BF16's worst, 1.572e-4, rounded up) | 1.6e-4 | +| MRL width 32, 1024, 4096: worst distance (limit; transformers BF16's worst) | 7.92e-5 (1.9e-4; 1.884e-4), 1.18e-4 (1.7e-4; 1.624e-4), 1.32e-4 (1.6e-4; 1.572e-4) | +| Card 2 x 2 scores | [[0.7479, 0.0759], [0.0886, 0.6324]] (card: [[0.7493, 0.0751], [0.0880, 0.6318]]) | +| Batch, alone and token-ID vectors (all 9 inputs) | identical fp32 bytes | +| `hip-kernel-test`, `dense_gemm_block` | fp32 within the dot2 chain's bound against float64; BF16 out is the fp32 sum rounded once | +| `hip-kernel-contract` | dense GEMM (fp32 and BF16 out) and `tf_heads_rows` pass gates 1 to 4; their gate 5 records are current (gfx1030 does not serve the dense GEMM) | +| Load | 4 s from the page cache | diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py new file mode 100644 index 0000000000..25795a79a6 --- /dev/null +++ b/tools/zig/oracle_qwen3_embedding.py @@ -0,0 +1,419 @@ +"""Pinned offline Qwen3-Embedding-8B oracle: the model card's Transformers example, and a server checked against it.""" + +import argparse +import base64 +import hashlib +import importlib.metadata +import importlib.util +import json +import math +import sys +import urllib.request +from pathlib import Path + +USAGE = """ +Pins checked before any model code runs (the snapshot's revision, every file's SHA256, Transformers 5.22.1's tree): + python tools/zig/oracle_qwen3_embedding.py sources --model DIR +Reference (writes JSON with token ids, vectors and the card's 2 x 2 scores; the card's 4 inputs, then 5 long inputs +of about 2K to 32K tokens cut from pinned files): + python tools/zig/oracle_qwen3_embedding.py ref --model DIR --dtype float32 --device cpu --out ref32.json +Compare a running server's /v1/embeddings (same texts in batches that fit its window, and one by one) with the fp32 +target and the transformers BF16 bar, at full width and at MRL widths 32, 1024 and 4096: + python tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8080 +Gates, at each width: cosine distance (1 - cos) to fp32 (cropped and renormalized for MRL) at most 0.001 and at most +the BF16 bar's worst distance at that width, rounded up to two digits; unit norm. At full width also: the card's scores +within 0.01; batch, alone and token-id vectors equal as fp32 bytes (base64). +""" + +CHECKPOINT = ("Qwen/Qwen3-Embedding-8B", "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af") +# Every non-weight file of the pinned snapshot: card, configs, sentence-transformers modules, tokenizer, licence. +SOURCES = { + ".gitattributes": "34448b82c17d60fec9b65b1f093c115ddbaadc04beb1b0140b6bfed2e012a930", + "1_Pooling/config.json": "2e1da26b3fd65cf7e370d2fabf28a8c59efa7edb525d1b8b50be8e5ca1048ea6", + "LICENSE": "832dd9e00a68dd83b3c3fb9f5588dad7dcf337a0db50f7d9483f310cd292e92e", + "README.md": "1fd6f3583d642a3319efb3eb8c22b6aa6320e6ab92263447bd84fd92615206d4", + "config.json": "dcccf7c7890c8debc4af8dace8d6acd8dd50bd56d50d4d8b60529381ba3de3d7", + "config_sentence_transformers.json": "10667c72ddb772627bf1780cb7f86af8e2ae0032b8c243c731172064105c6961", + "generation_config.json": "28396d421a2108acce96383f6a7de78008f7f1b17f807958f3c14c51dbfb65fb", + "merges.txt": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5", + "model.safetensors.index.json": "ceaf1429bfa18eecc75d079737e8d1fe8dc0627f5ec7e6a985514626d1ac6a77", + "modules.json": "84e40c8e006c9b1d6c122e02cba9b02458120b5fb0c87b746c41e0207cf642cf", + "tokenizer.json": "83cdf8c3a34f68862319cb1810ee7b1e2c0a44e0864ae930194ddb76bb7feb8d", + "tokenizer_config.json": "2f58f4bbd7bbce15d683f525954ef3a92cd82f5e06415a9c513859bf8ab72436", + "vocab.json": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910", +} +# The shards' SHA256 (the Hub's LFS object ids). +WEIGHTS = { + "model-00001-of-00004.safetensors": "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451", + "model-00002-of-00004.safetensors": "dff635b0f6dbbaad2a2d633ef037ec0a39bc165cc1806c712fbd6fcbcb4526c0", + "model-00003-of-00004.safetensors": "30b1d4c53d84eb018f642cad7b373f0aabf79699872d8702c1f38577c0a59a2f", + "model-00004-of-00004.safetensors": "36cbc9c60375693629f25743c1e77ebb1724af58e671b2376463193c7fd21ef6", +} +VERSION = "5.22.1" +# tree_sha256 of the whole installed Transformers package, so the shared modeling and tokenizer code is pinned too. +TRANSFORMERS_TREE = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b" +# The Qwen3 code the card's example runs, and the files the long inputs are cut from (package-relative). +LIBRARY = { + "models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b", + "models/qwen3/configuration_qwen3.py": "49d9ccd0f29ccd977b93dba2004f84b867891b3d0509b71fd8ba3ed01054e64d", + "models/qwen2/tokenization_qwen2.py": "fac4e6576bfe2369731be147a4e530f262bdf32f2ac50436f96f0d8bdd2fc628", + "tokenization_utils_base.py": "a70a4eda47e15756cb4ddd10ecd3a4c2c92d00b556fc00c47e7ce8a6bdb15dbd", + "generation/utils.py": "966f70ba83d94c9ce07b21150be49ff84b24eadbf9c47de052a1fb0f8ba38679", + "modeling_utils.py": "251aaae7733affbb7cdb5aa05a84157c46cac02518d868898cd6d6282cc14571", +} +# Long inputs (most tokens with the closing <|endoftext|>, pin record, file): see long_texts. +LONG = [ + (2048, "checkpoint", "LICENSE"), + (4096, "checkpoint", "README.md"), + (8192, "library", "tokenization_utils_base.py"), + (16384, "library", "generation/utils.py"), + (32768, "library", "modeling_utils.py"), +] +# Every input in order (the card's 4, then LONG) as fingerprint() gives it; check_inputs refuses any other. +INPUTS = [ + (27, "b11f8d3ae64c3c7029275c5d8d82172ef31a9a354223c12845b61047a0c972e8", "3c1537f6fc0a496f2b68c1aaa38484d8a82030338ae52bbd955b45b50c9efc24"), + (23, "bd6e4e8bb1b5c2e15db6003175c53f084690aeb6c09b34f1808161ef3353d524", "a08f60368f1da49a05dd929f406c503bcd0864ee464f25807866753bffe3062c"), + (8, "b6d0a6ed3a94d5e532ae83127502407102e51342d1db011caf4c5b9808a1e131", "c1e33d8973efc767e95f15a7b31af3c68eccc5ca89e0c2648f5337fec3c8f2c6"), + (31, "ba5bd6b5de9d9f827336e76d982b5dbae8aa3cc075c115f57405ba370bba8dbc", "4bd9685251a5e3587a645c0bf96dbc76ca70998073523b5537a1513c15f4c8ea"), + (2044, "f60c57fd5911f2c10704c555aef3a6e38543afa78604385ad5d8dde87050e0ff", "e10c109c63916e2cacd0143d8295fa5fa449c9eb09b7d3126405bfec9a6c9d92"), + (4043, "93c6978d85521f7b15d47dfedd386b72f063fd25ea5ada5ccf38fefdc88e7766", "a87db7cfe361a40e15c9b209796f6260bfb74017bc1e897e6e5be3d1a557cdbf"), + (8177, "3f5c76527fb4d032b31c827e5ec4d171a9feb94e7ab2ab96341028b1dd412505", "17c5b8a1cb0f288b40b9810549881b901eea58c83edb6e6d1b89ac1beab603eb"), + (16383, "c0a043846083a03ec205a34af84d7b1053a8d1a4ea4f3d3c028bb4c048ef1366", "528fda3905d2fc0340f8d5ca34f4bd91703613d76b40e5ac0835fab641acf173"), + (32759, "009256a9df975a6f67e5984760078c1dbcb0bd0b9aefbce0bdb8373ebd386dd7", "1dfcec470654b993960be543f8e84c461bc28d743aa5a50ee63bd81a9b8e15a8"), +] +# MRL widths checked besides the full one (the card's range is 32 to 4096). +WIDTHS = (32, 1024, 4096) + +TASK = "Given a web search query, retrieve relevant passages that answer the query" +CARD_SCORES = [[0.7493016123771667, 0.0750647559762001], [0.08795969933271408, 0.6318399906158447]] + + +def detailed_instruct(task: str, query: str) -> str: + return f"Instruct: {task}\nQuery:{query}" + + +def card_texts() -> list: + queries = [detailed_instruct(TASK, "What is the capital of China?"), detailed_instruct(TASK, "Explain gravity")] + documents = [ + "The capital of China is Beijing.", + "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and " + "is responsible for the movement of planets around the sun.", + ] + return queries + documents + + +def long_texts(model: Path, package: Path, count) -> list: + """Each LONG file's longest whole-line prefix whose `count` (with the closing <|endoftext|>) is within its bound.""" + roots = {"checkpoint": model, "library": package} + out = [] + for most, where, name in LONG: + lines = (roots[where] / name).read_text(encoding="utf-8").splitlines(keepends=True) + lo, hi = 0, len(lines) # the longest k with count(lines[:k]) <= most, by bisection on whole lines + while lo < hi: + k = (lo + hi + 1) // 2 + lo, hi = (k, hi) if count("".join(lines[:k])) <= most else (lo, k - 1) + text = "".join(lines[:lo]) + if count(text) < most // 2: + raise SystemExit(f"FAIL: {name} gives only {count(text)} of {most} tokens") + out.append(text) + return out + + +def fingerprint(text: str, ids: list) -> tuple: + """An input's tokens, text SHA256 and token-ID SHA256 (little-endian uint32), as INPUTS pins them.""" + import struct + raw = struct.pack(f"<{len(ids)}I", *ids) + return len(ids), hashlib.sha256(text.encode()).hexdigest(), hashlib.sha256(raw).hexdigest() + + +def check_inputs(name: str, texts, ids) -> None: + """Refuse inputs that are not the pinned fidelity set, text and token IDs both.""" + if not isinstance(texts, list) or not isinstance(ids, list) or len(texts) != len(INPUTS) or len(ids) != len(INPUTS): + raise SystemExit(f"FAIL: {name} does not hold the {len(INPUTS)} pinned inputs") + for i, (text, one, pin) in enumerate(zip(texts, ids, INPUTS)): + if not isinstance(text, str) or not isinstance(one, list) or not all(type(x) is int and 0 <= x <= 0xFFFFFFFF for x in one): + raise SystemExit(f"FAIL: {name} input {i} is not a text and its token IDs") + if fingerprint(text, one) != pin: + raise SystemExit(f"FAIL: {name} input {i} differs from its pin") + + +def groups(ids: list, window: int, most: int = 128) -> list: + """Consecutive input indices in requests of at most `most` inputs and `window` tokens in all.""" + out, part, used = [], [], 0 + for i, one in enumerate(ids): + if len(one) > window: + raise SystemExit(f"FAIL: input {i} has {len(one)} tokens, more than the {window}-token window") + if part and (used + len(one) > window or len(part) == most): + out.append(part) + part, used = [], 0 + part.append(i) + used += len(one) + return out + [part] if part else out + + +def round_up(x: float) -> float: + """`x` in [1e-20, 10) rounded up to two significant digits: the BF16 bar's distance as a gate, never below x.""" + scale = float(10 ** max(0, 1 - math.floor(math.log10(x)))) + up = math.ceil(x * scale) + while (up - 1) / scale >= x: + up -= 1 + while up / scale < x: + up += 1 + return up / scale + + +def crop(v, width: int) -> list: + """An MRL vector: the first `width` values, unit length again.""" + head = list(v[:width]) + norm = math.sqrt(sum(x * x for x in head)) + return [x / norm for x in head] + + +def digest(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def hashed(root: Path, pins: dict, what: str) -> dict: + """The pinned files' hashes under root, or a refusal naming every file that differs from its pin.""" + found = {name: digest(root / name) if (root / name).is_file() else None for name in pins} + if wrong := sorted(name for name, value in found.items() if value != pins[name]): + raise SystemExit(f"FAIL: {what} differs from its pin: {', '.join(wrong)}") + return found + + +def tree_sha256(package: Path) -> str: + """Every file of an installed package (bytecode caches excluded) as sorted "path NUL sha256" lines, hashed.""" + if links := sorted(p.relative_to(package).as_posix() for p in package.rglob("*") if p.is_symlink() and p.is_dir()): + raise SystemExit(f"FAIL: {package.name} has linked directories: {', '.join(links)}") + files = sorted(p for p in package.rglob("*") if p.is_file() and "__pycache__" not in p.relative_to(package).parts) + return hashlib.sha256("".join(f"{p.relative_to(package).as_posix()}\0{digest(p)}\n" for p in files).encode()).hexdigest() + + +def pinned(model: Path) -> dict: + """Refuse, before any model code runs, a snapshot or Transformers install whose sources differ from the pins.""" + repo, revision = CHECKPOINT + model = model.resolve() + if model.parent.name == "snapshots" and model.name != revision: + raise SystemExit(f"FAIL: snapshot directory is not the pinned revision {revision}") + if links := sorted(p.relative_to(model).as_posix() for p in model.rglob("*") if p.is_symlink() and p.is_dir()): + raise SystemExit(f"FAIL: snapshot has linked directories: {', '.join(links)}") + files = {p.relative_to(model).as_posix() for p in model.rglob("*") if p.is_file()} + if extra := sorted(files - set(SOURCES) - set(WEIGHTS)): + raise SystemExit(f"FAIL: snapshot has unpinned files: {', '.join(extra)}") + if (version := importlib.metadata.version("transformers")) != VERSION: + raise SystemExit(f"FAIL: expected Transformers {VERSION}, found {version}") + package = Path(importlib.util.find_spec("transformers").origin).parent + if (tree := tree_sha256(package)) != TRANSFORMERS_TREE: + raise SystemExit(f"FAIL: Transformers {VERSION} package tree differs from its pin: {tree}") + return {"repo": repo, "revision": revision, "checkpoint": hashed(model, SOURCES, f"{repo}@{revision}"), + "weights": hashed(model, WEIGHTS, f"{repo}@{revision} weights"), + "transformers": VERSION, "transformers_tree": tree, + "library": hashed(package, LIBRARY, f"Transformers {VERSION}"), + "torch": importlib.metadata.version("torch")} + + +def sources(args) -> None: + print(json.dumps(pinned(args.model), indent=1)) + + +def expected_pins() -> dict: + """The pin record every reference must carry (torch's version is recorded, not pinned).""" + return {"repo": CHECKPOINT[0], "revision": CHECKPOINT[1], "checkpoint": SOURCES, "weights": WEIGHTS, + "transformers": VERSION, "transformers_tree": TRANSFORMERS_TREE, "library": LIBRARY} + + +def check_pins(name: str, made: dict) -> None: + """Refuse a reference whose pin record is missing a field or differs from the pins in any of them.""" + pins = made.get("pins") + if not isinstance(pins, dict): + raise SystemExit(f"FAIL: {name} carries no pin record") + if wrong := sorted(k for k, v in expected_pins().items() if pins.get(k) != v): + raise SystemExit(f"FAIL: {name} was not made under the pins: {', '.join(wrong)}") + + +def reference(args) -> None: + # the pins first: no Transformers or torch code runs before its tree and the snapshot are checked + pins = pinned(args.model) + import torch + import torch.nn.functional as F + from transformers import AutoModel, AutoTokenizer + + def last_token_pool(last_hidden_states, attention_mask): + # the card's pooling, verbatim in effect + left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] + if left_padding: + return last_hidden_states[:, -1] + lengths = attention_mask.sum(dim=1) - 1 + return last_hidden_states[torch.arange(last_hidden_states.shape[0], device=last_hidden_states.device), lengths] + + dtype = getattr(torch, args.dtype) + tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left", local_files_only=True) + package = Path(importlib.util.find_spec("transformers").origin).parent + longs = long_texts(args.model, package, lambda text: len(tokenizer(text)["input_ids"])) + texts = card_texts() + longs + ids = [tokenizer(t)["input_ids"] for t in texts] + check_inputs("the tokenized inputs", texts, ids) # before the model loads + model = AutoModel.from_pretrained(args.model, dtype=dtype, local_files_only=True).to(args.device).eval() + # the card's 4 inputs as the card runs them: one left-padded batch, truncated at 8,192 tokens + batch = tokenizer(card_texts(), padding=True, truncation=True, max_length=8192, return_tensors="pt").to(args.device) + with torch.no_grad(): + out = model(**batch) + emb = [F.normalize(last_token_pool(out.last_hidden_state, batch["attention_mask"]).float(), p=2, dim=1)] + # the long inputs one by one, whole: no padding, no truncation + for text in longs: + one = tokenizer(text, return_tensors="pt").to(args.device) + with torch.no_grad(): + out = model(**one) + emb.append(F.normalize(last_token_pool(out.last_hidden_state, one["attention_mask"]).float(), p=2, dim=1)) + print(f"long input: {one['input_ids'].shape[1]} tokens", flush=True) + emb = torch.cat(emb) + scores = (emb[:2] @ emb[2:4].T).tolist() + print("scores", scores) + print("card ", CARD_SCORES) + json.dump({"dtype": args.dtype, "device": args.device, "pins": pins, "texts": texts, "ids": ids, + "vectors": emb.cpu().tolist(), "scores": scores}, open(args.out, "w")) + + +def post(url: str, body: dict) -> dict: + req = urllib.request.Request(url + "/v1/embeddings", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=600) as r: + return json.load(r) + + +def vectors(resp: dict, count: int, width: int) -> list: + """The response's fp32 vectors, in index order, after count, index, width and finiteness checks.""" + data = resp.get("data", []) + if len(data) != count or [d.get("index") for d in data] != list(range(count)): + raise SystemExit(f"FAIL: expected {count} vectors with indices 0..{count - 1}, got {len(data)}") + out = [] + for d in data: + raw = base64.b64decode(d["embedding"]) + if len(raw) != 4 * width: + raise SystemExit(f"FAIL: vector {d['index']} has {len(raw) // 4} values, expected {width}") + out.append(raw) + for i, raw in enumerate(out): + if not all(math.isfinite(x) for x in floats(raw)): + raise SystemExit(f"FAIL: vector {i} is not finite") + return out + + +def floats(raw: bytes) -> list: + import struct + return list(struct.unpack(f"<{len(raw) // 4}f", raw)) + + +def cosine(a, b) -> float: + dot = sum(x * y for x, y in zip(a, b)) + return dot / math.sqrt(sum(x * x for x in a) * sum(y * y for y in b)) + + +def check_vectors(name: str, made: dict, count: int) -> int: + """The reference's width, after checking it holds `count` finite, nonzero vectors of one width.""" + vectors = made.get("vectors") + if not isinstance(vectors, list) or len(vectors) != count or not vectors: + raise SystemExit(f"FAIL: {name} does not hold {count} vectors") + width = len(vectors[0]) + for i, v in enumerate(vectors): + if len(v) != width or not all(isinstance(x, (int, float)) and math.isfinite(x) for x in v): + raise SystemExit(f"FAIL: {name} vector {i} is not {width} finite values") + if sum(x * x for x in v) == 0: + raise SystemExit(f"FAIL: {name} vector {i} is zero") + return width + + +def check(args) -> None: + ref = json.load(open(args.ref)) + bar = json.load(open(args.bar)) + if ref["dtype"] != "float32" or bar["dtype"] != "bfloat16" or ref["texts"] != bar["texts"]: + raise SystemExit("FAIL: --ref must be the float32 reference and --bar the bfloat16 one, of the same texts") + if ref["texts"][:4] != card_texts(): + raise SystemExit("FAIL: the references do not start with the card's 4 inputs") + check_inputs("--ref", ref.get("texts"), ref.get("ids")) + check_inputs("--bar", bar.get("texts"), bar.get("ids")) + check_pins("--ref", ref) + check_pins("--bar", bar) + texts = ref["texts"] + n = len(texts) + width = check_vectors("--ref", ref, n) + if check_vectors("--bar", bar, n) != width or ref.get("ids") != bar.get("ids"): + raise SystemExit("FAIL: --ref and --bar differ in width or token ids") + parts = groups(ref["ids"], args.window) + enc = {"encoding_format": "base64"} + + def served(inputs: list, dims) -> list: + """Every input's vector, from requests of the inputs in `parts`, cropped by the server to `dims` if given.""" + out = [b""] * n + for part in parts: + body = {"input": [inputs[i] for i in part], **enc} | ({"dimensions": dims} if dims else {}) + for i, raw in zip(part, vectors(post(args.url, body), len(part), dims or width)): + out[i] = raw + return out + + def limit_at(dims: int) -> float: + worst = max(1 - cosine(crop(b, dims), crop(r, dims)) for b, r in zip(bar["vectors"], ref["vectors"])) + if not (math.isfinite(worst) and worst > 0): + raise SystemExit(f"FAIL: the bar's distance at width {dims} is not finite and positive") + limit = min(args.max_distance, round_up(worst)) + print(f"width {dims}: transformers bf16 worst distance {worst:.2e}; limit {limit:.2e}") + return limit + + print(f"requests: {[[i for i in part] for part in parts]} (window {args.window})") + limit = limit_at(width) + batch = served(texts, None) + alone = [vectors(post(args.url, {"input": t, **enc}), 1, width)[0] for t in texts] + ids = served(ref["ids"], None) + ok = True + for i in range(n): + v = floats(batch[i]) + dist = 1 - cosine(v, ref["vectors"][i]) + norm = math.sqrt(sum(x * x for x in v)) + same, by_ids = batch[i] == alone[i], batch[i] == ids[i] + print(f"input {i} ({len(ref['ids'][i])} tokens): distance {dist:.2e}, norm {norm:.6f}, " + f"batch == alone {same}, ids == text {by_ids}") + ok &= math.isfinite(dist) and dist <= limit and same and by_ids and abs(norm - 1) < 1e-4 + for dims in args.widths: + limit = limit_at(dims) + for i, raw in enumerate(served(texts, dims)): + v = floats(raw) + dist = 1 - cosine(v, crop(ref["vectors"][i], dims)) + norm = math.sqrt(sum(x * x for x in v)) + print(f"width {dims} input {i}: distance {dist:.2e}, norm {norm:.6f}") + ok &= math.isfinite(dist) and dist <= limit and abs(norm - 1) < 1e-4 + vs = [floats(b) for b in batch[:4]] + scores = [[sum(x * y for x, y in zip(vs[q], vs[2 + d])) for d in range(2)] for q in range(2)] + worst = max(abs(scores[q][d] - CARD_SCORES[q][d]) for q in range(2) for d in range(2)) + print("scores", scores, "card", CARD_SCORES, f"max diff {worst:.4f}") + ok &= worst <= args.max_score_diff + print("PASS" if ok else "FAIL") + sys.exit(0 if ok else 1) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + so = sub.add_parser("sources") + so.add_argument("--model", required=True, type=Path) + r = sub.add_parser("ref") + r.add_argument("--model", required=True, type=Path) + r.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"]) + r.add_argument("--device", default="cpu") + r.add_argument("--out", required=True) + c = sub.add_parser("check") + c.add_argument("--ref", required=True, help="float32 reference JSON (the target)") + c.add_argument("--bar", required=True, help="bfloat16 transformers reference JSON (the bar)") + c.add_argument("--url", default="http://127.0.0.1:8080") + c.add_argument("--max-distance", type=float, default=0.001) + c.add_argument("--max-score-diff", type=float, default=0.01) + c.add_argument("--window", type=int, default=32768, help="the server's window: the most tokens in one request") + c.add_argument("--widths", type=int, nargs="*", default=list(WIDTHS), help="MRL widths to check") + args = p.parse_args() + {"sources": sources, "ref": reference, "check": check}[args.cmd](args) + + +if __name__ == "__main__": + main() diff --git a/tools/zig/test_oracle_qwen3_embedding.py b/tools/zig/test_oracle_qwen3_embedding.py new file mode 100644 index 0000000000..dea021f051 --- /dev/null +++ b/tools/zig/test_oracle_qwen3_embedding.py @@ -0,0 +1,116 @@ +"""The Qwen3-Embedding oracle refuses references not made under its pins, and checks pins before model code.""" +import copy +import inspect +import math +import tempfile +import unittest +from pathlib import Path + +import oracle_qwen3_embedding as oracle + + +def made() -> dict: + return {"pins": copy.deepcopy(oracle.expected_pins()) | {"torch": "any"}} + + +class Gate(unittest.TestCase): + def test_the_bar_distance_rounds_up_to_two_digits(self): + for x, want in [(1.572e-4, 1.6e-4), (1.624e-4, 1.7e-4), (1.884e-4, 1.9e-4), (1.6e-4, 1.6e-4), + (0.1099, 0.11), (9.96e-5, 1.0e-4)]: + self.assertAlmostEqual(want, oracle.round_up(x), delta=want * 1e-12) + self.assertGreaterEqual(oracle.round_up(x), x) + + def test_two_digit_values_stay_and_the_next_float_rounds_past_them(self): + for k in range(2, 10): + for d in range(10, 100): + x = d / 10**k + self.assertEqual(x, oracle.round_up(x)) + self.assertEqual((d + 1) / 10**k, oracle.round_up(math.nextafter(x, 1))) + + +class PinRecord(unittest.TestCase): + def test_a_reference_made_under_the_pins_passes(self): + oracle.check_pins("--ref", made()) + + def test_a_changed_or_missing_record_is_refused(self): + changes = [ + lambda p: p.update(revision="0" * 40), + lambda p: p["checkpoint"].update({"tokenizer.json": "0" * 64}), + lambda p: p.update(weights={}), + lambda p: p.update(library={}), + lambda p: p.update(transformers="5.22.0"), + lambda p: p.update(transformers_tree="0" * 64), + lambda p: p.pop("checkpoint"), + ] + for change in changes: + ref = made() + change(ref["pins"]) + with self.assertRaises(SystemExit): + oracle.check_pins("--ref", ref) + with self.assertRaises(SystemExit): + oracle.check_pins("--ref", {}) + + def test_a_reference_with_a_nonfinite_zero_or_ragged_vector_is_refused(self): + good = {"vectors": [[0.6, 0.8], [1.0, 0.0]]} + self.assertEqual(2, oracle.check_vectors("--bar", good, 2)) + for vectors in ([[float("nan"), 1.0], [1.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]], [[0.6, 0.8], [1.0]], + [[0.6, 0.8]], [[float("inf"), 0.0], [1.0, 0.0]]): + with self.assertRaises(SystemExit): + oracle.check_vectors("--bar", {"vectors": vectors}, 2) + + def test_pins_are_checked_before_torch_or_transformers_is_imported(self): + body = inspect.getsource(oracle.reference) + self.assertLess(body.index("pinned(args.model)"), body.index("import torch")) + self.assertLess(body.index("pinned(args.model)"), body.index("from transformers")) + self.assertNotIn("transformers", {m.split(".")[0] for m in oracle.__dict__ if not m.startswith("_")}) + + +class FidelitySet(unittest.TestCase): + def test_every_long_input_is_cut_from_a_pinned_file_and_fits_the_window(self): + pins = {"checkpoint": oracle.SOURCES, "library": oracle.LIBRARY} + for most, where, name in oracle.LONG: + self.assertIn(name, pins[where]) + self.assertLessEqual(most, 32768) + self.assertEqual([2048, 4096, 8192, 16384, 32768], [most for most, _, _ in oracle.LONG]) + self.assertEqual((32, 1024, 4096), oracle.WIDTHS) + + def test_a_long_input_is_the_longest_prefix_of_whole_lines_within_its_bound(self): + with tempfile.TemporaryDirectory() as root: + lines = [f"line {i} " + "word " * i + "\n" for i in range(400)] + for _, where, name in oracle.LONG: + path = Path(root, where, name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(lines)) + texts = oracle.long_texts(Path(root, "checkpoint"), Path(root, "library"), len) + for (most, _, _), text in zip(oracle.LONG, texts): + self.assertLessEqual(len(text), most) + k = text.count("\n") + self.assertEqual(text, "".join(lines[:k])) + self.assertGreater(len(text) + len(lines[k]), most) + + def test_requests_hold_at_most_the_window_and_128_inputs_in_order(self): + self.assertEqual([[0, 1], [2], [3]], oracle.groups([[1] * 3, [1] * 5, [1] * 9, [1]], 9)) + self.assertEqual([[0], [1, 2]], oracle.groups([[1] * 4, [1] * 2, [1] * 2], 4)) + self.assertEqual([list(range(128)), [128]], oracle.groups([[1]] * 129, 1000)) + with self.assertRaises(SystemExit): + oracle.groups([[1] * 10], 9) + + def test_inputs_other_than_the_pinned_set_are_refused(self): + texts, ids = ["a"] * 9, [[1, 2]] * 9 + for bad in ([[1, 2]] * 9, [[1, 2]] * 8 + [[1 << 32]], [[1, 2]] * 8 + [[-1]], [[1, 2]] * 8 + [[1.0]]): + with self.assertRaises(SystemExit): + oracle.check_inputs("--ref", texts, bad) + self.assertEqual((2, "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", + oracle.hashlib.sha256(bytes([1, 0, 0, 0, 2, 0, 0, 0])).hexdigest()), + oracle.fingerprint("a", [1, 2])) + self.assertEqual(4 + len(oracle.LONG), len(oracle.INPUTS)) + for (most, _, _), (tokens, _, _) in zip(oracle.LONG, oracle.INPUTS[4:]): + self.assertTrue(most // 2 <= tokens <= most) + self.assertLessEqual(sum(tokens for tokens, _, _ in oracle.INPUTS[:-1]), 32768) + + def test_an_mrl_vector_is_the_cropped_head_at_unit_length(self): + self.assertEqual([0.6, 0.8], oracle.crop([3.0, 4.0, 12.0], 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 9c98b050d6..1fb6ae2eaf 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -141,6 +141,9 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.flashnext })).step); b.step("test-flashnext-proposal-hip-host", "Original Flash Next HIP IDs and head-row spans without device runtime").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.flashnext, .filters = &.{"Flash Next proposal:"} })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.nemotron })).step); + const qwen3_host = b.addRunArtifact(b.addTest(.{ .root_module = host.qwen3 })); + test_step.dependOn(&qwen3_host.step); + b.step("test-qwen3-hip-host", "Qwen3 dense config, checkpoint prefix and lm_head rows on the CPU").dependOn(&qwen3_host.step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.serve })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.native })).step); // the GPU programs compile on any host against stub kernels, so `test` catches what only they reach @@ -150,6 +153,17 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const check = engineModules(b, check_target, .Debug, stub_kernels, null); b.step("check-native-hip-host", "Compile the real native HIP server with inert kernels, without GPU execution").dependOn(&server(b, check_target, check, build_options).step); const gpu = engineModules(b, target, optimize, objects, hip_include); + const qwen3_gpu = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/families/qwen3/forward_test.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "hip", .module = gpu.hip }, .{ .name = "qwen3", .module = gpu.qwen3 } }, + }) }); + const qwen3_run = b.addRunArtifact(qwen3_gpu); + qwen3_run.has_side_effects = true; // the checkpoints and oracle files named by TF_QWEN3_* are not cached inputs + b.step("qwen3-hip-gpu-build", "Compile the Qwen3 dense GPU checks without running them").dependOn(&qwen3_gpu.step); + b.step("qwen3-hip-gpu-test", "Qwen3 lastRows on the GPU against transformers (TF_QWEN3_EMBED_*, TF_QWEN3_RERANK_*)").dependOn(&qwen3_run.step); b.step("tf-qwen35-test", "Qwen3.5 / 3.6 model checks on the GPU (-Dhip-arch)").dependOn(&b.addInstallArtifact(checksExe(b, target, optimize, gpu), .{}).step); const fx_test = b.addExecutable(.{ .name = "tf-flashnext-hip-test", .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/flashnext_hip/main.zig"), @@ -212,6 +226,7 @@ pub const Engine = struct { family: *std.Build.Module, flashnext: *std.Build.Module, nemotron: *std.Build.Module, + qwen3: *std.Build.Module, native: *std.Build.Module, }; @@ -280,6 +295,18 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: }); const ocr_kernels = b.createModule(.{ .root_source_file = b.addWriteFiles().add("glm_ocr_kernels.zig", "const hip = @import(\"hip\"); pub const arch = hip.kernels.arch; pub const bytes = hip.kernels.images[@intFromEnum(hip.kernels.Group.glm_ocr)];\n"), .imports = &.{.{ .name = "hip", .module = hip }} }); const glm_ocr = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm_ocr/hip.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "hip", .module = hip }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes }, .{ .name = "engine_api", .module = api }, .{ .name = "hip_serve", .module = serve }, .{ .name = "glm_ocr_kernels", .module = ocr_kernels } } }); + const qwen3 = b.createModule(.{ + .root_source_file = b.path("zig/src/families/qwen3/qwen3.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "hip", .module = hip }, + .{ .name = "core", .module = core }, + .{ .name = "lanes", .module = lanes }, + .{ .name = "engine_api", .module = api }, + }, + }); const native = b.createModule(.{ .root_source_file = b.path("zig/src/native/hip.zig"), .target = target, @@ -294,6 +321,7 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: .{ .name = "nemotron", .module = nemotron }, .{ .name = "glm_ocr", .module = glm_ocr }, .{ .name = "hip_serve", .module = serve }, + .{ .name = "qwen3", .module = qwen3 }, }, }); return .{ @@ -304,6 +332,7 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: .family = family, .flashnext = flashnext, .nemotron = nemotron, + .qwen3 = qwen3, .native = native, }; } @@ -335,15 +364,15 @@ const groups = [_][]const u8{ /// What the group sources include, so an edit to one rebuilds them. const headers = [_][]const u8{ - "attention/attention_fa.hip", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", - "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", - "ops/linear.hpp", "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", - "common/arch.hpp", "common/dot2.hpp", "common/vec.hpp", "common/wmma.hpp", - "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", "quant/mlx_pieces.hpp", - "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", - "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", - "flashnext/streams.hpp", "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", - "ops/nemotron_moe.hpp", + "attention/attention_fa.hip", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", + "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", + "ops/linear.hpp", "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", + "common/arch.hpp", "common/dot2.hpp", "common/vec.hpp", "common/wmma.hpp", + "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", "quant/mlx_pieces.hpp", + "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", + "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", + "flashnext/streams.hpp", "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", + "ops/nemotron_moe.hpp", "quant/bf16_decoder.hpp", "quant/dense_tiles.hpp", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. diff --git a/zig/kernels/hip/ops/attention.hpp b/zig/kernels/hip/ops/attention.hpp index c97bcf67f5..6e305c42d2 100644 --- a/zig/kernels/hip/ops/attention.hpp +++ b/zig/kernels/hip/ops/attention.hpp @@ -19,6 +19,16 @@ extern "C" __global__ void tf_attn_gate(const float* att, const void* qg, void* store(out, kind, i, a * sigmoid(g)); } +// out[r, h * d + j] = att[h, r, j] rounded once to kind: the causal kernel's head-major rows, row-major for o_proj. +extern "C" __global__ void tf_heads_rows(const float* att, void* out, int kind, int len, int heads, int d) { + long long i = gid(); + if (i >= static_cast(len) * heads * d) return; + int j = static_cast(i % d); + int h = static_cast((i / d) % heads); + int r = static_cast(i / (static_cast(d) * heads)); + store(out, kind, i, att[(static_cast(h) * len + r) * d + j]); +} + // keys or values (len, kv_heads, d) of kind into the cache (kv_heads, total, d) at slot pos0. extern "C" __global__ void tf_kv_write(const void* src, void* cache, int kind, int len, int kv_heads, int d, int total, int pos0) { diff --git a/zig/kernels/hip/ops/elementwise.hpp b/zig/kernels/hip/ops/elementwise.hpp index 0b1d5bb7ae..13037942b0 100644 --- a/zig/kernels/hip/ops/elementwise.hpp +++ b/zig/kernels/hip/ops/elementwise.hpp @@ -32,11 +32,13 @@ extern "C" __global__ void tf_embed_rows(const uint32_t* words, const void* scal } // The float gather: rows of an fp32 table (an embedding kept unquantized), written in the activation type. -extern "C" __global__ void tf_embed_dense(const float* table, const int* ids, int n, int k, void* out, int out_kind) { +// Rows `ids` of a float table stored as `table_kind` (fp32, or BF16 / FP16 as a checkpoint stores them). +extern "C" __global__ void tf_embed_dense(const void* table, int table_kind, const int* ids, int n, int k, void* out, + int out_kind) { long long i = gid(); if (i >= static_cast(n) * k) return; int r = static_cast(i / k), c = static_cast(i % k); - store(out, out_kind, i, table[static_cast(ids[r]) * k + c]); + store(out, out_kind, i, load(table, table_kind, static_cast(ids[r]) * k + c)); } extern "C" __global__ void tf_cast(const void* src, int skind, void* dst, int dkind, long long n) { diff --git a/zig/kernels/hip/quant/bf16_decoder.hpp b/zig/kernels/hip/quant/bf16_decoder.hpp new file mode 100644 index 0000000000..77fe35e686 --- /dev/null +++ b/zig/kernels/hip/quant/bf16_decoder.hpp @@ -0,0 +1,65 @@ +#pragma once + +// The BF16 dense WeightDecoder: a column's K weights stored as BF16, two a 32-bit word; no scales, no biases. + +#include + +#include "common/vec.hpp" +#include "quant/mlx.hpp" + +namespace tf { +namespace rocm { + +// The GEMM tile's decoder over raw BF16 weights (Affine args with bits 16; its tables are unused). Every group of +// `a.group` K folds into the accumulator with scale 1, so the sums are the fp32 dot2 chain, group after group. +struct Bf16Decoder { + using Args = Affine; + static constexpr int kBits = 16; + // Words of a 32-weight chunk of a column. + static constexpr int kWords = 16; + // No bias term: the tile neither sums x nor folds a bias. + static constexpr bool kBias = false; + + // A plain product only: a routed plan's items are not this decoder's. + __device__ static bool take_item(Args& a, int) { return a.route.items == nullptr; } + + __device__ static const uint32_t* words(const Args& a) { return a.words; } + + template + __device__ static Idx row_words(const Args& a) { + return static_cast(a.k) / 2; + } + + // A chunk is 64 bytes: four 16-byte loads (the host checks the alignment). + template + __device__ static void load(const uint32_t* src, uint32_t (&w)[kWords]) { +#pragma unroll + for (int i = 0; i < kWords / 4; ++i) { + const u32x4 v = reinterpret_cast(src)[i]; + w[4 * i] = v.x; + w[4 * i + 1] = v.y; + w[4 * i + 2] = v.z; + w[4 * i + 3] = v.w; + } + } + + // Weights t and t + 1 of a chunk (t even): the word that holds them, as stored. + template + __device__ static typename T::pair pair(const uint32_t (&w)[kWords], int t) { + static_assert(std::is_same_v, "BF16 weights take BF16 activations"); + return __builtin_bit_cast(typename T::pair, w[t >> 1]); + } + + struct Term {}; + + __device__ static Term term(const Args&, long long) { return {}; } + __device__ static float scale_value(const Args&, const Term&) { return 1.f; } + __device__ static float bias_value(const Args&, const Term&) { return 0.f; } + + // A group's dot added as it is: the scale is 1. + __device__ static float fold_scale(float acc, float dot, float) { return acc + dot; } + __device__ static float fold_bias(float acc, float, float) { return acc; } +}; + +} // namespace rocm +} // namespace tf diff --git a/zig/kernels/hip/quant/dense_tiles.hpp b/zig/kernels/hip/quant/dense_tiles.hpp new file mode 100644 index 0000000000..dccb4266fe --- /dev/null +++ b/zig/kernels/hip/quant/dense_tiles.hpp @@ -0,0 +1,25 @@ +#pragma once + +// BF16 dense kernels: the shared prefill GEMM tile over the BF16 decoder, named as the launchers look up. + +#include "quant/act.hpp" +#include "quant/bf16_decoder.hpp" +#include "tiles/epilogue.hpp" +#include "tiles/gemm.hpp" + +namespace tf { +namespace rocm { + +// y = x . W^T for BF16 x and BF16 W at any row count; a build without the BF16 dot2 keeps the symbol and traps. +template +__global__ void __launch_bounds__(256) dense_gemm_block(Affine a) { +#if TF_DEVICE_BF16_DOT2 + gemm_tile, T, GemmOut, RT>(a); +#else + (void)a; + __builtin_trap(); +#endif +} + +} // namespace rocm +} // namespace tf diff --git a/zig/kernels/hip/quant/mlx_decoder.hpp b/zig/kernels/hip/quant/mlx_decoder.hpp index 6f94a58974..f35b257654 100644 --- a/zig/kernels/hip/quant/mlx_decoder.hpp +++ b/zig/kernels/hip/quant/mlx_decoder.hpp @@ -21,6 +21,8 @@ struct MlxDecoder { static constexpr int kBits = BITS; // Words of a 32-code chunk of a column. static constexpr int kWords = BITS; + // A bias a group: the tiles sum x and fold the bias with it. + static constexpr bool kBias = true; // ---- the plan: a routed item's expert and rows ---- diff --git a/zig/kernels/hip/tiles/dot2_tiles.hip b/zig/kernels/hip/tiles/dot2_tiles.hip index 65ef2d4e85..13fd252479 100644 --- a/zig/kernels/hip/tiles/dot2_tiles.hip +++ b/zig/kernels/hip/tiles/dot2_tiles.hip @@ -6,6 +6,7 @@ #include "tiles/stream.hpp" #include "tiles/matrix_gemm.hpp" #include "quant/mlx_tiles.hpp" +#include "quant/dense_tiles.hpp" namespace tf { namespace rocm { @@ -376,5 +377,8 @@ TF_TILES(5) TF_TILES(6) TF_TILES(8) +// The BF16 dense GEMM on every build (its body traps where there is no BF16 dot2), so every image has the symbol. +template __global__ void dense_gemm_block::rt>(Affine); + } // namespace rocm } // namespace tf diff --git a/zig/kernels/hip/tiles/gemm.hpp b/zig/kernels/hip/tiles/gemm.hpp index 6a95340979..b67bedd55e 100644 --- a/zig/kernels/hip/tiles/gemm.hpp +++ b/zig/kernels/hip/tiles/gemm.hpp @@ -41,6 +41,7 @@ struct GemmShape { }; // One 128 x 128 block of y = x . W^T: Dec reads weights and group terms, Act the rows of x, T multiplies, Epi stores. +// A decoder without a bias term (Dec::kBias false) leaves out the x sums and the bias folds; the MLX one keeps them. template __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { if (!Dec::take_item(a, blockIdx.z)) return; @@ -122,7 +123,7 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { } if (side == 0 && st % per == per - 1) { sc[(st / per) % 3][line] = Dec::scale_value(a, term[J]); - bi[(st / per) % 3][line] = Dec::bias_value(a, term[J]); + if constexpr (Dec::kBias) bi[(st / per) % 3][line] = Dec::bias_value(a, term[J]); } }; // Side 1 sums a row's x of the stage in LDS, lo then hi of each pair in order, on through the group. @@ -180,7 +181,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { // Two stages a trip so a stage's buffer and parity are known when the code is built. auto step = [&](std::integral_constant, int s) { constexpr int cur = P; - if (live && s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1); + if constexpr (Dec::kBias) { + if (live && s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1); + } if constexpr (P == 1) fetch(s + 1); if (live) { constexpr int kXRegs = RT / 4; // x words (4 pairs each) a lane loads a chunk @@ -235,7 +238,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { } } } - if (side == 1 && sxlive) sumx(s, cur); + if constexpr (Dec::kBias) { + if (side == 1 && sxlive) sumx(s, cur); + } stash(std::integral_constant{}, s + 1, cur ^ 1); __syncthreads(); }; @@ -243,7 +248,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { step(std::integral_constant{}, s); if (s + 1 < stages) step(std::integral_constant{}, s + 1); } - if (live && (stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1); + if constexpr (Dec::kBias) { + if (live && (stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1); + } #pragma unroll for (int r = 0; r < RT; ++r) { const int row = m0 + rg * 16 + hrow + r; diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index c0fc468210..3991c3c210 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -40,6 +40,7 @@ test { _ = cubin; _ = hsaco; _ = hip_records; + _ = @import("contract/hip_records_qwen3.zig"); _ = @import("contract/fakes.zig"); } diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 899a403a9d..986be387fe 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -496,10 +496,8 @@ const own_records = [_]Record{ } ++ @import("hip_records_pair.zig").records ++ @import("hip_records_prompt.zig").records ++ @import("hip_records_score.zig").records; -pub const records = own_records ++ @import("hip_records_stream.zig").records ++ - @import("hip_records_router.zig").records ++ - @import("hip_records_glm_ocr.zig").records ++ - @import("hip_records_item.zig").records; +pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ + @import("hip_records_glm_ocr.zig").records ++ @import("hip_records_item.zig").records ++ @import("hip_records_qwen3.zig").records; /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64. pub const controls = [_]Record{ @@ -540,7 +538,9 @@ pub fn check(a: std.mem.Allocator, list: []const Record, quiet: bool) !usize { defer out.deinit(a); var n: usize = 0; for (list) |r| for (limits.rdna, r.usage) |t, u| { - if (u.vgprs == 0) { + if (!u.served) { + if (!quiet) std.debug.print("N/A {s} on {s}: not served there\n", .{ r.launch.name, t.name }); + } else if (u.vgprs == 0) { try out.print(a, "{s} on {s}: no compiled usage recorded\n", .{ r.launch.name, t.name }); n += 1; } else n += try limits.hip(r.launch, u, t, &out, a); @@ -590,6 +590,10 @@ test "a control cannot fit or lack compiled usage on any one chip" { var missing = r; missing.usage[i].vgprs = 0; try std.testing.expectEqual(@as(usize, 1), try controlProblems(a, missing, true)); + // a control is never excused as unserved + var unserved = r; + unserved.usage[i] = .{ .vgprs = 0, .served = false }; + try std.testing.expectEqual(@as(usize, 1), try controlProblems(a, unserved, true)); } } } diff --git a/zig/src/core/contract/hip_records_qwen3.zig b/zig/src/core/contract/hip_records_qwen3.zig new file mode 100644 index 0000000000..565fd22ece --- /dev/null +++ b/zig/src/core/contract/hip_records_qwen3.zig @@ -0,0 +1,38 @@ +//! Gate 5 for qwen3's tf_heads_rows and BF16 dense GEMM; the GEMM does not serve gfx1030 (only a trap body there). +const std = @import("std"); +const limits = @import("limits.zig"); +const hip_records = @import("hip_records.zig"); + +const Record = hip_records.Record; +const check = hip_records.check; + +pub const records = [_]Record{ + .{ + .launch = .{ .name = "attention tf_heads_rows", .symbol = "tf_heads_rows", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 30, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + }, + }, + .{ + .launch = .{ .name = "dense gemm dense_gemm_block", .symbol = "_ZN2tf4rocm16dense_gemm_blockINS0_7DotBF16ELi8EEEvNS0_6AffineE", .image = "affine_tiles", .threads = 256 }, + .usage = .{ + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + }, + }, +}; + +test "a chip a record does not serve is skipped; a served chip without usage is still a problem" { + const a = std.testing.allocator; + const launch: limits.HipLaunch = .{ .name = "t", .symbol = "t", .image = "ops", .threads = 256 }; + const one: limits.HipUsage = .{ .vgprs = 16, .bound = 256 }; + const skip: Record = .{ .launch = launch, .usage = .{ .{ .vgprs = 0, .served = false }, one, one, one } }; + try std.testing.expectEqual(@as(usize, 0), try check(a, &.{skip}, true)); + const missing: Record = .{ .launch = launch, .usage = .{ .{ .vgprs = 0 }, one, one, one } }; + try std.testing.expectEqual(@as(usize, 1), try check(a, &.{missing}, true)); +} diff --git a/zig/src/core/contract/limits.zig b/zig/src/core/contract/limits.zig index 3089d01d31..b624135929 100644 --- a/zig/src/core/contract/limits.zig +++ b/zig/src/core/contract/limits.zig @@ -200,6 +200,8 @@ pub const HipLaunch = struct { /// What hipcc made of one kernel for one chip, from the code object's metadata (vgprs 0: none recorded). pub const HipUsage = struct { vgprs: u32, + /// False: the kernel is not launched on this chip (its image keeps a trap body there); no usage is checked. + served: bool = true, sgprs: u32 = 0, lds: u32 = 0, private: u32 = 0, diff --git a/zig/src/core/lanes/embedding.zig b/zig/src/core/lanes/embedding.zig index 0490730668..ef4573e41a 100644 --- a/zig/src/core/lanes/embedding.zig +++ b/zig/src/core/lanes/embedding.zig @@ -7,12 +7,15 @@ pub const Info = struct { max_batch_tokens: u32, max_inputs: u32 = 128, mrl_dimensions: []const u32 = &.{}, + /// Every width from this one to `dimensions` crops (a model trained for any MRL width); 0: only the list's. + mrl_min: u32 = 0, pub fn dimension(i: Info, requested: ?u32) !u32 { const n = requested orelse return i.dimensions; if (n == 0 or n > i.dimensions) return error.UnsupportedEmbeddingDimension; - if (i.mrl_dimensions.len == 0) return error.MrlNotDeclared; + if (i.mrl_dimensions.len == 0 and i.mrl_min == 0) return error.MrlNotDeclared; if (n == i.dimensions or std.mem.indexOfScalar(u32, i.mrl_dimensions, n) != null) return n; + if (i.mrl_min > 0 and n >= i.mrl_min) return n; return error.UnsupportedEmbeddingDimension; } }; @@ -47,3 +50,16 @@ fn overlap(a: []const u8, b: []const u8) bool { const bp = @intFromPtr(b.ptr); return if (ap <= bp) bp - ap < a.len else ap - bp < b.len; } + +test "an MRL range takes every width from its least to the full one; a list takes its own widths" { + const range: Info = .{ .dimensions = 4096, .vocabulary = 10, .max_batch_tokens = 8, .mrl_min = 32 }; + for ([_]u32{ 32, 384, 2049, 4096 }) |n| try std.testing.expectEqual(n, try range.dimension(n)); + for ([_]u32{ 0, 31, 4097 }) |n| try std.testing.expectError(error.UnsupportedEmbeddingDimension, range.dimension(n)); + try std.testing.expectEqual(@as(u32, 4096), try range.dimension(null)); + const list: Info = .{ .dimensions = 8, .vocabulary = 10, .max_batch_tokens = 8, .mrl_dimensions = &.{ 2, 4 } }; + try std.testing.expectEqual(@as(u32, 4), try list.dimension(4)); + try std.testing.expectError(error.UnsupportedEmbeddingDimension, list.dimension(3)); + const none: Info = .{ .dimensions = 8, .vocabulary = 10, .max_batch_tokens = 8 }; + try std.testing.expectError(error.MrlNotDeclared, none.dimension(4)); + try std.testing.expectEqual(@as(u32, 8), try none.dimension(null)); +} diff --git a/zig/src/core/quant/dense.zig b/zig/src/core/quant/dense.zig index a36f114e6f..c989a48d83 100644 --- a/zig/src/core/quant/dense.zig +++ b/zig/src/core/quant/dense.zig @@ -1,4 +1,4 @@ -//! Weights kept in float: fp32 [N, K], the unquantized projections a conversion leaves (an MTP head's fc). +//! Float weights [N, K]: fp32 (an MTP head's fc a conversion leaves), or BF16 / FP16 as stored (`readStored`). const std = @import("std"); const types = @import("types.zig"); @@ -36,6 +36,13 @@ pub fn read(a: Allocator, t: anytype, key: []const u8) !Host { return .{ .weight = try convert.float32(a, try t.get(join(&buf, &.{ key, ".weight" }))) }; } +/// The float weight in its stored dtype when that is BF16 or FP16 (shared, not copied), else widened as `read` does. +pub fn readStored(a: Allocator, t: anytype, key: []const u8) !Host { + var buf: [256]u8 = undefined; + const w = try t.get(join(&buf, &.{ key, ".weight" })); + return .{ .weight = if (w.dtype == .bf16 or w.dtype == .f16) w else try convert.float32(a, w) }; +} + pub fn rows(h: Host) usize { return h.weight.shape[0]; } @@ -72,24 +79,26 @@ pub fn upload(u: types.Uploader, h: Host) !Device { return .{ .weight = try u.tensor(h.weight) }; } -/// What the kernels read: the fp32 weight's address. -pub const Matrix = struct { weight: u64 }; +/// What the kernels read: the weight's address and its element type (fp32, or BF16 / FP16 from `readStored`). +pub const Matrix = struct { weight: u64, dtype: types.DType = .f32 }; pub const View = struct { n: u32, k: u32, matrix: Matrix }; pub fn view(d: Device, stacked: bool) error{UnsupportedTables}!View { const lead: usize = @intFromBool(stacked); - return .{ .n = @intCast(d.weight.dim(lead)), .k = @intCast(d.weight.dim(lead + 1)), .matrix = .{ .weight = d.weight.ptr } }; + return .{ .n = @intCast(d.weight.dim(lead)), .k = @intCast(d.weight.dim(lead + 1)), .matrix = .{ .weight = d.weight.ptr, .dtype = d.weight.dtype } }; } pub fn elements(h: Host) usize { return h.weight.numel(); } -/// The weight as fp64, row by row. +/// The weight as fp64, row by row, read in its own dtype. pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void { - if (out.len != h.weight.numel()) return error.UnexpectedTensor; - for (out, 0..) |*o, i| o.* = convert.load(.f32, h.weight.bytes, i); + const dtype = h.weight.dtype; + if (dtype != .f32 and dtype != .bf16 and dtype != .f16) return error.UnexpectedTensor; + if (out.len != h.weight.numel() or h.weight.bytes.len != out.len * dtype.size()) return error.UnexpectedTensor; + for (out, 0..) |*o, i| o.* = convert.load(dtype, h.weight.bytes, i); } test "a float weight halves along K and cuts its rows" { @@ -105,3 +114,83 @@ test "a float weight halves along K and cuts its rows" { try reference(low, &out); try std.testing.expectEqualSlices(f64, &.{ 5, 6, 7, 8 }, &out); } + +/// A one-tensor table for the tests: `get` and `has` over a single `.weight`. +const OneTable = struct { + weight: Tensor, + pub fn get(t: *const OneTable, name: []const u8) !Tensor { + if (!std.mem.eql(u8, name, "p.weight")) return error.MissingTensor; + return t.weight; + } + pub fn has(_: *const OneTable, name: []const u8) bool { + return std.mem.eql(u8, name, "p.weight"); + } + /// What the affine reader asks of a table; this one declares no quantization. + pub fn width(_: *const OneTable, _: []const u8) error{UnsupportedQuantization}!@import("mlx.zig").Width { + return error.UnsupportedQuantization; + } +}; + +test "readStored keeps a BF16 or FP16 weight as stored; read still widens it to fp32" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // 1.5, -2, 0.25, 3 in BF16 and in FP16 + const bf = [_]u16{ 0x3FC0, 0xC000, 0x3E80, 0x4040 }; + const hf = [_]u16{ 0x3E00, 0xC000, 0x3400, 0x4200 }; + for ([_]struct { dtype: types.DType, bits: []const u16 }{ .{ .dtype = .bf16, .bits = &bf }, .{ .dtype = .f16, .bits = &hf } }) |c| { + const t: OneTable = .{ .weight = .{ .dtype = c.dtype, .rank = 2, .shape = .{ 2, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(c.bits) } }; + const stored = try readStored(a, &t, "p"); + try std.testing.expectEqual(c.dtype, stored.weight.dtype); + try std.testing.expectEqual(@as(usize, 8), stored.weight.bytes.len); + var got: [4]f64 = undefined; + try reference(stored, &got); + try std.testing.expectEqualSlices(f64, &.{ 1.5, -2, 0.25, 3 }, &got); + const wide = try read(a, &t, "p"); + try std.testing.expectEqual(types.DType.f32, wide.weight.dtype); + try reference(wide, &got); + try std.testing.expectEqualSlices(f64, &.{ 1.5, -2, 0.25, 3 }, &got); + // a BF16 host halves along K and cuts rows by its own element size + const pair = try halves(a, stored); + try std.testing.expectEqual(@as(usize, 4), pair[1].weight.bytes.len); + var right: [2]f64 = undefined; + try reference(pair[1], &right); + try std.testing.expectEqualSlices(f64, &.{ -2, 3 }, &right); + } +} + +test "the view of an uploaded 16-bit weight reports its dtype" { + var b: types.Buf = undefined; + b.ptr = 0x1000; + b.dtype = .bf16; + b.rank = 2; + b.shape = .{ 3, 128, 1, 1, 1 }; + b.len = 3 * 128 * 2; + const v = try view(.{ .weight = b }, false); + try std.testing.expectEqual(types.DType.bf16, v.matrix.dtype); + try std.testing.expectEqual(@as(u32, 3), v.n); + try std.testing.expectEqual(@as(u32, 128), v.k); +} + +test "a reference refuses bytes that do not match the dtype's size" { + const bits = [_]u16{ 0x3FC0, 0xC000 }; + const h: Host = .{ .weight = .{ .dtype = .f32, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bits)[0..2] } }; + var out: [1]f64 = undefined; + try std.testing.expectError(error.UnexpectedTensor, reference(h, &out)); +} + +test "quant.readStored dispatches a float tensor to the dense format in its stored dtype; quant.read widens it" { + const quant = @import("quant.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bf = [_]u16{ 0x3FC0, 0xC000, 0x3E80, 0x4040 }; + const t: OneTable = .{ .weight = .{ .dtype = .bf16, .rank = 2, .shape = .{ 2, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bf) } }; + const stored = try quant.readStored(a, &t, "p"); + try std.testing.expectEqual(types.DType.bf16, stored.dense.weight.dtype); + const wide = try quant.read(a, &t, "p"); + try std.testing.expectEqual(types.DType.f32, wide.dense.weight.dtype); + const f32s = [_]f32{ 1, 2 }; + const t32: OneTable = .{ .weight = .{ .dtype = .f32, .rank = 2, .shape = .{ 1, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&f32s) } }; + try std.testing.expectEqual(types.DType.f32, (try quant.readStored(a, &t32, "p")).dense.weight.dtype); +} diff --git a/zig/src/core/quant/quant.zig b/zig/src/core/quant/quant.zig index bf46e04dd5..26581a9980 100644 --- a/zig/src/core/quant/quant.zig +++ b/zig/src/core/quant/quant.zig @@ -110,6 +110,14 @@ pub fn read(a: Allocator, t: anytype, key: []const u8) !Host { return types.refuse(key, "no weight format reads this tensor"); } +/// `read`, keeping a dense tensor stored in BF16 or FP16 in that dtype (every other format reads as `read` does). +pub fn readStored(a: Allocator, t: anytype, key: []const u8) !Host { + inline for (read_order) |f| { + if (Module(f).matches(t, key)) return wrap(f, try if (f == .dense) dense.readStored(a, t, key) else Module(f).read(a, t, key)); + } + return types.refuse(key, "no weight format reads this tensor"); +} + /// A fused [embedding | hidden] projection as its two halves along K. pub fn halves(a: Allocator, h: Host) ![2]Host { switch (h) { diff --git a/zig/src/core/quant_geometry_host.zig b/zig/src/core/quant_geometry_host.zig index 41aa2ab82d..f96f791889 100644 --- a/zig/src/core/quant_geometry_host.zig +++ b/zig/src/core/quant_geometry_host.zig @@ -1,6 +1,7 @@ -//! Standalone host entry for the actual affine reader, expert stack and bit-preserving slicing regressions. +//! Standalone host entry for the actual affine reader, dense stored dtypes, expert stack and bit-exact slicing. test { _ = @import("quant/mlx.zig"); _ = @import("quant/experts.zig"); _ = @import("quant/exact_test.zig"); + _ = @import("quant/dense.zig"); } diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 9e537a0e20..5fd3441d3e 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -12,6 +12,7 @@ pub const eaq_container = @import("eaq_container.zig"); pub const digest_cache = @import("digest_cache.zig"); pub const checkpoint_identity = @import("checkpoint_identity.zig"); pub const checkpoint = @import("checkpoint.zig"); +pub const checkpoint_fixture = @import("checkpoint_fixture.zig"); pub const exl3_load = @import("exl3_load.zig"); pub const direct_io = @import("direct_io.zig"); pub const direct_gather = @import("direct_gather.zig"); diff --git a/zig/src/families/qwen3/config.zig b/zig/src/families/qwen3/config.zig new file mode 100644 index 0000000000..35e6e15136 --- /dev/null +++ b/zig/src/families/qwen3/config.zig @@ -0,0 +1,153 @@ +//! A Qwen3 dense checkpoint's config.json: the geometry the HIP forward runs, and refusals of what it does not. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +pub const Error = error{ InvalidConfig, UnsupportedConfig } || Allocator.Error; + +/// The forward's geometry and constants, as config.json names them. +pub const Spec = struct { + layers: usize, + hidden: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + inter: usize, + vocab: usize, + eps: f64, + theta: f64, + max_positions: usize, + /// `tie_word_embeddings`: the output head is the embedding table. + tied: bool, +}; + +/// The spec of config.json's `bytes`; a missing or malformed key is InvalidConfig, an unrun feature UnsupportedConfig. +pub fn parse(gpa: Allocator, bytes: []const u8) Error!Spec { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const doc = std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), bytes, .{}) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + else => { + std.log.warn("qwen3 config.json is not JSON ({s})", .{@errorName(err)}); + return error.InvalidConfig; + }, + }; + if (doc != .object) return invalid("the document", "is not an object"); + const o = doc.object; + const model_type = o.get("model_type") orelse return invalid("model_type", "is missing"); + if (model_type != .string or !std.mem.eql(u8, model_type.string, "qwen3")) return unsupported("model_type", "is not qwen3"); + const s: Spec = .{ + .layers = try count(o, "num_hidden_layers"), + .hidden = try count(o, "hidden_size"), + .heads = try count(o, "num_attention_heads"), + .kv_heads = try count(o, "num_key_value_heads"), + // Qwen3Config's own default when an older config.json leaves it out + .head_dim = if (o.get("head_dim")) |v| (if (v == .null) 128 else try count(o, "head_dim")) else 128, + .inter = try count(o, "intermediate_size"), + .vocab = try count(o, "vocab_size"), + .eps = try number(o, "rms_norm_eps"), + .theta = try number(o, "rope_theta"), + .max_positions = try count(o, "max_position_embeddings"), + .tied = try flag(o, "tie_word_embeddings", false), + }; + if (try flag(o, "attention_bias", false)) return unsupported("attention_bias", "is true: the forward has no projection biases"); + if (try flag(o, "use_sliding_window", false)) return unsupported("use_sliding_window", "is true: the forward attends to the whole input"); + if (o.get("rope_scaling")) |r| if (r != .null) return unsupported("rope_scaling", "is set: the forward runs plain RoPE"); + if (s.head_dim > 256) return unsupported("head_dim", "is above 256, the attention kernel's limit"); + if (s.heads % s.kv_heads != 0) return unsupported("num_attention_heads", "is not a multiple of num_key_value_heads"); + const act = o.get("hidden_act") orelse return invalid("hidden_act", "is missing"); + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return unsupported("hidden_act", "is not silu"); + if (!(s.eps > 0) or !(s.theta > 0)) return invalid(if (s.eps > 0) "rope_theta" else "rms_norm_eps", "is not positive"); + return s; +} + +fn invalid(key: []const u8, why: []const u8) error{InvalidConfig} { + std.log.warn("qwen3 config.json: {s} {s}", .{ key, why }); + return error.InvalidConfig; +} + +fn unsupported(key: []const u8, why: []const u8) error{UnsupportedConfig} { + std.log.warn("qwen3 config.json: {s} {s}", .{ key, why }); + return error.UnsupportedConfig; +} + +/// A positive integer key. +fn count(o: std.json.ObjectMap, key: []const u8) Error!usize { + const v = o.get(key) orelse return invalid(key, "is missing"); + if (v != .integer or v.integer < 1) return invalid(key, "is not a positive integer"); + return @intCast(v.integer); +} + +/// A number key, integer or float. +fn number(o: std.json.ObjectMap, key: []const u8) Error!f64 { + const v = o.get(key) orelse return invalid(key, "is missing"); + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => invalid(key, "is not a number"), + }; +} + +/// A boolean key, `default` when absent or null. +fn flag(o: std.json.ObjectMap, key: []const u8, default: bool) Error!bool { + const v = o.get(key) orelse return default; + return switch (v) { + .null => default, + .bool => |b| b, + else => invalid(key, "is not a boolean"), + }; +} + +/// Qwen3-Embedding-8B's config.json (the reranker's differs only in vocab_size 151669). +const embedding = + \\{"architectures": ["Qwen3ForCausalLM"], "attention_bias": false, "attention_dropout": 0.0, + \\ "bos_token_id": 151643, "eos_token_id": 151645, "head_dim": 128, "hidden_act": "silu", + \\ "hidden_size": 4096, "initializer_range": 0.02, "intermediate_size": 12288, + \\ "max_position_embeddings": 40960, "max_window_layers": 36, "model_type": "qwen3", + \\ "num_attention_heads": 32, "num_hidden_layers": 36, "num_key_value_heads": 8, "rms_norm_eps": 1e-06, + \\ "rope_scaling": null, "rope_theta": 1000000, "sliding_window": null, "tie_word_embeddings": false, + \\ "torch_dtype": "bfloat16", "transformers_version": "4.51.2", "use_cache": true, + \\ "use_sliding_window": false, "vocab_size": 151665} +; + +test "the 8B embedding and reranker configs give their geometry" { + const a = std.testing.allocator; + const s = try parse(a, embedding); + try std.testing.expectEqual(Spec{ .layers = 36, .hidden = 4096, .heads = 32, .kv_heads = 8, .head_dim = 128, .inter = 12288, .vocab = 151665, .eps = 1e-6, .theta = 1e6, .max_positions = 40960, .tied = false }, s); + const reranker = try std.mem.replaceOwned(u8, a, embedding, "\"vocab_size\": 151665", "\"vocab_size\": 151669"); + defer a.free(reranker); + try std.testing.expectEqual(@as(usize, 151669), (try parse(a, reranker)).vocab); + const older = try std.mem.replaceOwned(u8, a, embedding, "\"head_dim\": 128, ", ""); + defer a.free(older); + try std.testing.expect(!std.mem.eql(u8, older, embedding)); + try std.testing.expectEqual(@as(usize, 128), (try parse(a, older)).head_dim); +} + +test "each feature the forward does not run is refused by name, and each malformed key is invalid" { + const a = std.testing.allocator; + const Case = struct { []const u8, []const u8, Error }; + const cases = [_]Case{ + .{ "\"model_type\": \"qwen3\"", "\"model_type\": \"qwen3_moe\"", error.UnsupportedConfig }, + .{ "\"attention_bias\": false", "\"attention_bias\": true", error.UnsupportedConfig }, + .{ "\"use_sliding_window\": false", "\"use_sliding_window\": true", error.UnsupportedConfig }, + .{ "\"rope_scaling\": null", "\"rope_scaling\": {\"type\": \"yarn\", \"factor\": 4}", error.UnsupportedConfig }, + .{ "\"head_dim\": 128", "\"head_dim\": 512", error.UnsupportedConfig }, + .{ "\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.UnsupportedConfig }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 7", error.UnsupportedConfig }, + .{ "\"model_type\": \"qwen3\",", "", error.InvalidConfig }, + .{ "\"hidden_size\": 4096", "\"hidden_size\": 0", error.InvalidConfig }, + .{ "\"num_hidden_layers\": 36", "\"num_hidden_layers\": \"36\"", error.InvalidConfig }, + .{ "\"intermediate_size\": 12288,", "", error.InvalidConfig }, + .{ "\"rms_norm_eps\": 1e-06", "\"rms_norm_eps\": 0", error.InvalidConfig }, + .{ "\"rope_theta\": 1000000", "\"rope_theta\": \"1e6\"", error.InvalidConfig }, + .{ "\"tie_word_embeddings\": false", "\"tie_word_embeddings\": 0", error.InvalidConfig }, + .{ "\"hidden_act\": \"silu\",", "", error.InvalidConfig }, + }; + for (cases) |case| { + const bytes = try std.mem.replaceOwned(u8, a, embedding, case[0], case[1]); + defer a.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, embedding)); + try std.testing.expectError(case[2], parse(a, bytes)); + } + try std.testing.expectError(error.InvalidConfig, parse(a, "[1]")); + try std.testing.expectError(error.InvalidConfig, parse(a, "{not json")); +} diff --git a/zig/src/families/qwen3/engine.zig b/zig/src/families/qwen3/engine.zig new file mode 100644 index 0000000000..197968f27a --- /dev/null +++ b/zig/src/families/qwen3/engine.zig @@ -0,0 +1,127 @@ +//! One Qwen3 dense model on one HIP device for prompt-only heads: weights, a stream and scratch sized to a window. +const std = @import("std"); +const hip = @import("hip"); +const model = @import("model.zig"); +const forward = @import("forward.zig"); + +const Allocator = std.mem.Allocator; + +/// Inputs one pass pools at most (the embedding lane's default batch). +pub const max_inputs: usize = 128; +/// The least window the budget may shrink a load to. +pub const least_window: usize = 2048; + +pub const Engine = struct { + gpa: Allocator, + driver: hip.Runtime, + ctx: hip.Context, + lib: hip.Launcher, + stream: hip.Stream, + m: model.Model, + /// Device bytes the weights took. + weights: usize = 0, + /// Tokens a pass takes at most, over all its inputs; set by `size`. + window: usize = 0, + arena: hip.Arena = undefined, + ids_host: hip.HostBuffer = undefined, + ids_dev: hip.DeviceBuffer = undefined, + rows_dev: hip.DeviceBuffer = undefined, + rows_host: hip.HostBuffer = undefined, + sized: bool = false, + + /// The model on device `index`, its scratch not yet sized. + pub fn load(gpa: Allocator, io: std.Io, dir: []const u8, index: c_int, choices: hip.Launcher.Choices) !*Engine { + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = .{ .gpa = gpa, .driver = undefined, .ctx = undefined, .lib = undefined, .stream = undefined, .m = undefined }; + e.driver = try hip.Runtime.open(); + errdefer e.driver.close(); + e.ctx = try hip.Context.init(&e.driver, index); + errdefer e.ctx.deinit(); + e.lib = try hip.Launcher.load(&e.driver, choices, hip.kernels.images); + errdefer e.lib.unload(); + e.stream = try hip.Stream.init(&e.driver); + errdefer e.stream.deinit(); + const before = hip.usage(false).device; + e.m = try model.Model.load(gpa, io, &e.driver, dir); + e.weights = hip.usage(false).device - before; + return e; + } + + pub fn ops(e: *Engine, arena: *hip.Arena) hip.ops.Ops { + return .{ .l = &e.lib, .bf16 = true, .stream = e.stream.handle, .arena = arena }; + } + + /// The pass's scratch for `window` tokens (one input of them: the longest cache) plus the ids and pooled rows. + pub fn scratchBytes(e: *Engine, window: usize) !usize { + var arena = hip.Arena.counting(); + var o = e.ops(&arena); + o.stream = hip.counting; + const out: hip.ops.Tensor = .{ .ptr = try arena.take(max_inputs * e.m.spec.hidden * 2), .kind = .bf16 }; + try forward.lastRows(o, &e.m, 0, &[_]u32{@intCast(window)}, out); + return arena.peak + Engine.fixedBytes(e, window); + } + + /// The buffers beside the arena: ids for `window` tokens and `max_inputs` pooled BF16 rows. + fn fixedBytes(e: *const Engine, window: usize) usize { + return window * 4 + max_inputs * e.m.spec.hidden * 2; + } + + /// Allocates the scratch for passes of `window` tokens; refuses when the bytes held differ from the count. + pub fn size(e: *Engine, window: usize) !void { + const want = try e.scratchBytes(window); + const arena_bytes = want - e.fixedBytes(window); + const before = hip.usage(false).device; + e.arena = try hip.Arena.init(&e.driver, arena_bytes); + errdefer e.arena.deinit(); + e.ids_dev = try hip.DeviceBuffer.alloc(&e.driver, window * 4); + errdefer e.ids_dev.free(); + e.rows_dev = try hip.DeviceBuffer.alloc(&e.driver, max_inputs * e.m.spec.hidden * 2); + errdefer e.rows_dev.free(); + e.ids_host = try hip.HostBuffer.alloc(&e.driver, window * 4); + errdefer e.ids_host.free(); + e.rows_host = try hip.HostBuffer.alloc(&e.driver, max_inputs * e.m.spec.hidden * 2); + errdefer e.rows_host.free(); + const held = hip.usage(false).device - before; + if (held != want) { + std.log.err("qwen3: sizing held {d} device bytes, the count says {d}", .{ held, want }); + return error.SizeMismatch; + } + e.window = window; + e.sized = true; + } + + /// Each input's final-normed last row as fp32 into `out` (ends.len rows of hidden); inputs end at `ends`. + pub fn run(e: *Engine, tokens: []const u32, ends: []const u32, out: []f32) !void { + const hidden = e.m.spec.hidden; + if (!e.sized or tokens.len == 0 or tokens.len > e.window or ends.len == 0 or ends.len > max_inputs or + ends[ends.len - 1] != tokens.len or out.len != ends.len * hidden) return error.InvalidBatch; + for (tokens) |t| if (t >= e.m.spec.vocab) return error.InvalidBatch; + @memcpy(e.ids_host.slice(u32)[0..tokens.len], tokens); + try e.ids_dev.uploadAsync(0, e.ids_host.view(0, tokens.len * 4), e.stream); + const mark = e.arena.mark(); + defer e.arena.release(mark); + const rows: hip.ops.Tensor = .{ .ptr = e.rows_dev.base(), .kind = .bf16 }; + try forward.lastRows(e.ops(&e.arena), &e.m, e.ids_dev.base(), ends, rows); + try e.rows_dev.downloadAsync(0, e.rows_host.view(0, ends.len * hidden * 2), e.stream); + try e.stream.synchronize(); + for (out, e.rows_host.slice(u16)[0..out.len]) |*o, b| o.* = @bitCast(@as(u32, b) << 16); + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + if (e.sized) { + e.rows_host.free(); + e.ids_host.free(); + e.rows_dev.free(); + e.ids_dev.free(); + e.arena.deinit(); + } + e.m.deinit(); + e.stream.deinit(); + e.lib.unload(); + e.ctx.deinit(); + e.driver.close(); + e.gpa.destroy(e); + } +}; diff --git a/zig/src/families/qwen3/forward.zig b/zig/src/families/qwen3/forward.zig new file mode 100644 index 0000000000..8e62817640 --- /dev/null +++ b/zig/src/families/qwen3/forward.zig @@ -0,0 +1,134 @@ +//! A batch of independent prompts through every layer; each input's final-normed last row comes out. +const std = @import("std"); +const hip = @import("hip"); +const model = @import("model.zig"); + +const Ops = hip.ops.Ops; +const Tensor = hip.ops.Tensor; +const Kind = hip.ops.Kind; + +/// Rows of one block of row-wise work, and of one attention query block. +pub const SPAN = 2048; + +pub const Error = hip.ops.Error; + +fn take(o: Ops, kind: Kind, n: usize) Error!Tensor { + return .{ .ptr = try o.arena.take(n * kind.size()), .kind = kind }; +} + +fn at(t: Tensor, n: usize) Tensor { + return .{ .ptr = t.ptr + n * t.kind.size(), .kind = t.kind }; +} + +/// head_dim ** -0.5 in double, as Python has it, passed as fp32. +pub fn scaleOf(head_dim: usize) f32 { + return @floatCast(std.math.pow(f64, @floatFromInt(head_dim), -0.5)); +} + +/// Each input's final-normed last row into `out` (ends.len, hidden); inputs end at `ends`, start at position 0. +pub fn lastRows(ops: Ops, m: *const model.Model, ids: u64, ends: []const u32, out: Tensor) Error!void { + var o = ops; + o.prefill = true; // prefill's kernels at any row count: an input's rows have the same bits alone or batched + const s = m.spec; + const total: usize = if (ends.len == 0) return else ends[ends.len - 1]; + var longest: usize = 0; + var begin: usize = 0; + for (ends) |e| { + if (e <= begin) return error.BadShape; + longest = @max(longest, e - begin); + begin = e; + } + const x = try take(o, m.act, total * s.hidden); + var start: usize = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + try o.embedRows(m.embed, ids + start * 4, rows, at(x, start * s.hidden)); + } + for (m.layers) |*layer| { + const mark = o.arena.mark(); + defer o.arena.release(mark); + try layerRows(o, m, layer, x, ends, total, longest); + } + for (ends, 0..) |e, i| try o.rms(at(x, (e - 1) * s.hidden), m.final_norm, at(out, i * s.hidden), 1, s.hidden, @floatCast(s.eps)); +} + +/// One decoder layer over every row: norms, projections and MLP in SPAN blocks, attention input by input. +fn layerRows(o: Ops, m: *const model.Model, l: *const model.Layer, x: Tensor, ends: []const u32, total: usize, longest: usize) Error!void { + const s = m.spec; + const hd = s.head_dim; + const qw = s.heads * hd; + const kvw = s.kv_heads * hd; + const q = try take(o, m.act, total * qw); + const k = try take(o, m.act, total * kvw); + const v = try take(o, m.act, total * kvw); + var start: usize = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const normed = try take(o, m.act, rows * s.hidden); + try o.rms(at(x, start * s.hidden), l.input_norm, normed, rows, s.hidden, @floatCast(s.eps)); + try copyRows(o, try o.project(normed, l.q, rows, false), at(q, start * qw), rows * qw); + try copyRows(o, try o.project(normed, l.k, rows, false), at(k, start * kvw), rows * kvw); + try copyRows(o, try o.project(normed, l.v, rows, false), at(v, start * kvw), rows * kvw); + } + const attn = try take(o, m.act, total * qw); + // one input's flat cache, reused by each input in turn: (kv_heads, longest, d) + const cache: Ops.Cache = .{ .k = try o.arena.take(kvw * longest * m.act.size()), .v = try o.arena.take(kvw * longest * m.act.size()), .kind = m.act, .kv_heads = s.kv_heads, .total = longest, .d = hd }; + var a: usize = 0; + for (ends) |b| { + defer a = b; + try attention(o, m, l, q, k, v, attn, cache, a, b); + } + start = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const xs = at(x, start * s.hidden); + const y = try o.project(at(attn, start * qw), l.o, rows, false); + try o.add(xs, y, xs, rows * s.hidden); + const normed = try take(o, m.act, rows * s.hidden); + try o.rms(xs, l.post_norm, normed, rows, s.hidden, @floatCast(s.eps)); + const gate = try o.project(normed, l.gate, rows, false); + const up = try o.project(normed, l.up, rows, false); + const act = try take(o, m.act, rows * s.inter); + try o.siluMul(gate, up, act, rows * s.inter); + const z = try o.project(act, l.down, rows, false); + try o.add(xs, z, xs, rows * s.hidden); + } +} + +/// Input rows [a, b): head norms, RoPE from position 0, its cache, causal attention, head-major rows back. +fn attention(o: Ops, m: *const model.Model, l: *const model.Layer, q: Tensor, k: Tensor, v: Tensor, attn: Tensor, c: Ops.Cache, a: usize, b: usize) Error!void { + const s = m.spec; + const hd = s.head_dim; + const qw = s.heads * hd; + const kvw = s.kv_heads * hd; + const theta: f32 = @floatCast(s.theta); + const len = b - a; + var p: usize = 0; + while (p < len) : (p += SPAN) { + const rows = @min(SPAN, len - p); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const qn = try take(o, m.act, rows * qw); + try o.rms(at(q, (a + p) * qw), l.q_norm, qn, rows * s.heads, hd, @floatCast(s.eps)); + const kn = try take(o, m.act, rows * kvw); + try o.rms(at(k, (a + p) * kvw), l.k_norm, kn, rows * s.kv_heads, hd, @floatCast(s.eps)); + // q rotated, rounded to the activation dtype, then widened: (heads, rows, d) fp32 for the kernel + const q32 = try o.arena.of(f32, rows * qw); + try o.ropePrefill(qn, .{ .ptr = q32, .kind = .f32 }, rows * hd, hd, rows, s.heads, hd, hd, p, theta); + // k rotated into the cache (kv_heads, total, d) from slot p (the kernel writes row r at r * s_row); v as it is + try o.ropePrefill(kn, .{ .ptr = c.k + p * hd * c.kind.size(), .kind = c.kind }, c.total * hd, hd, rows, s.kv_heads, hd, hd, p, theta); + try o.kvWrite(at(v, (a + p) * kvw), c.v, rows, s.kv_heads, hd, c.total, p); + const att = try o.arena.of(f32, rows * qw); + try o.causalPrefill(q32, c, att, rows, p + rows, s.heads, scaleOf(hd), p); + try o.headsRows(att, at(attn, (a + p) * qw), rows, s.heads, hd); + } +} + +/// `n` values of `src` into `dst` (same kind), as a one-row column copy. +fn copyRows(o: Ops, src: Tensor, dst: Tensor, n: usize) Error!void { + try o.copyCols(src, n, 0, dst.ptr, 1, n); +} diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig new file mode 100644 index 0000000000..8d29e1c5e8 --- /dev/null +++ b/zig/src/families/qwen3/forward_test.zig @@ -0,0 +1,494 @@ +//! lastRows on the GPU against transformers: the embedding's pooled rows and the reranker's yes/no logits. + +// Skipped without TF_QWEN3_EMBED_DIR / _REF / _BAR (fp32, BF16 oracles) or TF_QWEN3_RERANK_DIR / _REF. +const std = @import("std"); +const hip = @import("hip"); +const qwen3 = @import("qwen3"); + +const t = std.testing; +const gpa = t.allocator; + +/// The device, its kernels and an arena: what the engine holds around `lastRows`. +const Gpu = struct { + r: hip.Runtime, + ctx: hip.Context, + stream: hip.Stream, + l: hip.Launcher, + arena: hip.Arena, + + fn open(g: *Gpu, arena_bytes: usize) !void { + g.r = try hip.Runtime.open(); + errdefer g.r.close(); + g.ctx = try hip.Context.init(&g.r, 0); + errdefer g.ctx.deinit(); + g.stream = try hip.Stream.init(&g.r); + errdefer g.stream.deinit(); + g.l = try hip.Launcher.load(&g.r, .{}, hip.kernels.images); + errdefer g.l.unload(); + g.arena = try hip.Arena.init(&g.r, arena_bytes); + } + + fn close(g: *Gpu) void { + g.arena.deinit(); + g.l.unload(); + g.stream.deinit(); + g.ctx.deinit(); + g.r.close(); + } + + fn ops(g: *Gpu) hip.ops.Ops { + return .{ .l = &g.l, .bf16 = true, .stream = g.stream.handle, .arena = &g.arena }; + } + + /// Each input's final-normed last row, as BF16 bits (`ends.len` rows of `hidden`). + fn lastRows(g: *Gpu, m: *const qwen3.Model, ids: []const i32, ends: []const u32) ![]u16 { + const mark = g.arena.mark(); + defer g.arena.release(mark); + var dev = try hip.DeviceBuffer.fromHost(&g.r, std.mem.sliceAsBytes(ids)); + defer dev.free(); + const hidden = m.spec.hidden; + var out = try hip.DeviceBuffer.alloc(&g.r, ends.len * hidden * 2); + defer out.free(); + try qwen3.lastRows(g.ops(), m, @intFromPtr(dev.ptr), ends, .{ .ptr = @intFromPtr(out.ptr), .kind = .bf16 }); + try g.stream.synchronize(); + const rows = try gpa.alloc(u16, ends.len * hidden); + errdefer gpa.free(rows); + try out.download(0, std.mem.sliceAsBytes(rows)); + return rows; + } +}; + +fn env(a: std.mem.Allocator, key: []const u8) !?[]u8 { + return t.environ.getAlloc(a, key) catch |err| switch (err) { + error.EnvironmentVariableMissing => null, + else => err, + }; +} + +fn widen(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +fn readJson(a: std.mem.Allocator, path: []const u8) !std.json.Value { + const bytes = try std.Io.Dir.cwd().readFileAlloc(t.io, path, a, .limited(256 << 20)); + return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}); +} + +fn idsOf(a: std.mem.Allocator, v: std.json.Value) ![]i32 { + const out = try a.alloc(i32, v.array.items.len); + for (v.array.items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn number(v: std.json.Value) f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => std.math.nan(f64), + }; +} + +fn cosine(x: []const f64, y: []const f64) f64 { + var xy: f64 = 0; + var xx: f64 = 0; + var yy: f64 = 0; + for (x, y) |p, q| { + xy += p * q; + xx += p * p; + yy += q * q; + } + return xy / @sqrt(xx * yy); +} + +/// The pins tools/zig/oracle_qwen3_embedding.py writes into each reference: the checkpoint and the Transformers tree. +const embed_revision = "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af"; +const transformers_tree = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b"; + +/// One fidelity input as the oracle's INPUTS pins it: tokens, SHA256 of the UTF-8 text, SHA256 of the IDs as LE u32. +const InputPin = struct { tokens: usize, text: *const [64]u8, ids: *const [64]u8 }; + +/// The card's 4 inputs, then 5 long inputs of 2K to 32K tokens cut from pinned files. +const embed_inputs = [_]InputPin{ + .{ .tokens = 27, .text = "b11f8d3ae64c3c7029275c5d8d82172ef31a9a354223c12845b61047a0c972e8", .ids = "3c1537f6fc0a496f2b68c1aaa38484d8a82030338ae52bbd955b45b50c9efc24" }, + .{ .tokens = 23, .text = "bd6e4e8bb1b5c2e15db6003175c53f084690aeb6c09b34f1808161ef3353d524", .ids = "a08f60368f1da49a05dd929f406c503bcd0864ee464f25807866753bffe3062c" }, + .{ .tokens = 8, .text = "b6d0a6ed3a94d5e532ae83127502407102e51342d1db011caf4c5b9808a1e131", .ids = "c1e33d8973efc767e95f15a7b31af3c68eccc5ca89e0c2648f5337fec3c8f2c6" }, + .{ .tokens = 31, .text = "ba5bd6b5de9d9f827336e76d982b5dbae8aa3cc075c115f57405ba370bba8dbc", .ids = "4bd9685251a5e3587a645c0bf96dbc76ca70998073523b5537a1513c15f4c8ea" }, + .{ .tokens = 2044, .text = "f60c57fd5911f2c10704c555aef3a6e38543afa78604385ad5d8dde87050e0ff", .ids = "e10c109c63916e2cacd0143d8295fa5fa449c9eb09b7d3126405bfec9a6c9d92" }, + .{ .tokens = 4043, .text = "93c6978d85521f7b15d47dfedd386b72f063fd25ea5ada5ccf38fefdc88e7766", .ids = "a87db7cfe361a40e15c9b209796f6260bfb74017bc1e897e6e5be3d1a557cdbf" }, + .{ .tokens = 8177, .text = "3f5c76527fb4d032b31c827e5ec4d171a9feb94e7ab2ab96341028b1dd412505", .ids = "17c5b8a1cb0f288b40b9810549881b901eea58c83edb6e6d1b89ac1beab603eb" }, + .{ .tokens = 16383, .text = "c0a043846083a03ec205a34af84d7b1053a8d1a4ea4f3d3c028bb4c048ef1366", .ids = "528fda3905d2fc0340f8d5ca34f4bd91703613d76b40e5ac0835fab641acf173" }, + .{ .tokens = 32759, .text = "009256a9df975a6f67e5984760078c1dbcb0bd0b9aefbce0bdb8373ebd386dd7", .ids = "1dfcec470654b993960be543f8e84c461bc28d743aa5a50ee63bd81a9b8e15a8" }, +}; + +/// Refuses references whose texts and token IDs are not exactly the pinned inputs, in order. +fn checkInputs(ref: std.json.Value, pins: []const InputPin) !void { + const texts = ref.object.get("texts") orelse return error.IncompleteReference; + const ids = ref.object.get("ids") orelse return error.IncompleteReference; + try holds(texts == .array and ids == .array and texts.array.items.len == pins.len and ids.array.items.len == pins.len); + for (texts.array.items, ids.array.items, pins) |x, one, pin| { + try holds(x == .string and one == .array and one.array.items.len == pin.tokens); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (one.array.items) |id| { + try holds(id == .integer and id.integer >= 0 and id.integer <= std.math.maxInt(u32)); + var le: [4]u8 = undefined; + std.mem.writeInt(u32, &le, @intCast(id.integer), .little); + h.update(&le); + } + try holds(std.mem.eql(u8, &std.fmt.bytesToHex(h.finalResult(), .lower), pin.ids)); + var th: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(x.string, &th, .{}); + try holds(std.mem.eql(u8, &std.fmt.bytesToHex(th, .lower), pin.text)); + } +} + +test "Qwen3-Embedding references: inputs other than the pinned ones are refused" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + // "a" and the IDs [1, 2] (bytes 01 00 00 00 02 00 00 00) + const pin = [_]InputPin{.{ .tokens = 2, .text = "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", .ids = "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f" }}; + const parse = struct { + fn of(al: std.mem.Allocator, body: []const u8) !std.json.Value { + return std.json.parseFromSliceLeaky(std.json.Value, al, body, .{}); + } + }.of; + try checkInputs(try parse(a, "{\"texts\": [\"a\"], \"ids\": [[1, 2]]}"), &pin); + for ([_][]const u8{ + "{\"texts\": [\"b\"], \"ids\": [[1, 2]]}", + "{\"texts\": [\"a\"], \"ids\": [[1, 3]]}", + "{\"texts\": [\"a\"], \"ids\": [[1]]}", + "{\"texts\": [\"a\", \"a\"], \"ids\": [[1, 2], [1, 2]]}", + "{\"texts\": [\"a\"], \"ids\": [[1, -2]]}", + "{\"texts\": [\"a\"]}", + }) |body| try t.expect(std.meta.isError(checkInputs(try parse(a, body), &pin))); + try t.expectEqual(@as(usize, 9), embed_inputs.len); +} + +fn text(v: std.json.Value, key: []const u8) ?[]const u8 { + const x = (if (v == .object) v.object.get(key) else null) orelse return null; + return if (x == .string) x.string else null; +} + +fn same(expected: []const u8, got: ?[]const u8) !void { + if (!std.mem.eql(u8, expected, got orelse "")) return error.ReferenceMismatch; +} + +/// A silent check: a refused reference is an error, without the expect helpers' diff on stderr. +fn holds(ok: bool) !void { + if (!ok) return error.ReferenceMismatch; +} + +/// Refuses, before any GPU work, references that are not the fp32 target and BF16 bar of the same pinned inputs. +fn checkReferences(ref: std.json.Value, bar: std.json.Value) !void { + try same("float32", text(ref, "dtype")); + try same("bfloat16", text(bar, "dtype")); + // the oracle checks every hash of the pin record; here both references carry one complete, identical record + for ([_]std.json.Value{ ref, bar }) |r| { + const pins = r.object.get("pins") orelse return error.UnpinnedReference; + try same("Qwen/Qwen3-Embedding-8B", text(pins, "repo")); + try same(embed_revision, text(pins, "revision")); + try same("5.22.1", text(pins, "transformers")); + try same(transformers_tree, text(pins, "transformers_tree")); + for ([_][]const u8{ "checkpoint", "weights", "library" }) |key| { + const files = pins.object.get(key) orelse return error.UnpinnedReference; + try holds(files == .object and files.object.count() > 0); + for (files.object.values()) |hash| try holds(hash == .string and hash.string.len == 64); + } + } + try holds(jsonEqual(ref.object.get("pins").?, bar.object.get("pins").?)); + for ([_][]const u8{ "texts", "ids" }) |key| { + const x = ref.object.get(key) orelse return error.IncompleteReference; + const y = bar.object.get(key) orelse return error.IncompleteReference; + try holds(std.meta.activeTag(x) == .array and std.meta.activeTag(y) == .array and x.array.items.len == y.array.items.len); + for (x.array.items, y.array.items) |p, q| try holds(jsonEqual(p, q)); + } + for ([_]std.json.Value{ ref, bar }) |r| for (r.object.get("vectors").?.array.items) |v| { + var square: f64 = 0; + for (v.array.items) |x| { + const f = number(x); + try holds(std.math.isFinite(f)); + square += f * f; + } + try holds(square > 0); + }; +} + +/// The SHA256 of a file, as lowercase hex. +fn fileSha256(io: std.Io, path: []const u8) ![64]u8 { + const f = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer f.close(io); + const size = (try f.stat(io)).size; + var h = std.crypto.hash.sha2.Sha256.init(.{}); + var buf: [1 << 22]u8 = undefined; + var at: u64 = 0; + while (at < size) { + const n = try f.readPositionalAll(io, buf[0..@intCast(@min(buf.len, size - at))], at); + if (n == 0) return error.ShortRead; + h.update(buf[0..n]); + at += n; + } + return std.fmt.bytesToHex(h.finalResult(), .lower); +} + +/// The references' pinned files are the checkpoint the test loads: each one's SHA256, and every indexed shard pinned. +fn checkFiles(a: std.mem.Allocator, io: std.Io, dir: []const u8, pins: std.json.Value) !void { + for ([_][]const u8{ "checkpoint", "weights" }) |key| { + const files = pins.object.get(key).?.object; + for (files.keys(), files.values()) |name, hash| { + const got = try fileSha256(io, try std.fs.path.join(a, &.{ dir, name })); + try holds(std.mem.eql(u8, &got, hash.string)); + } + } + const index = try readJson(a, try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" })); + const shards = index.object.get("weight_map").?.object; + const weights = pins.object.get("weights").?.object; + for (shards.values()) |file| try holds(file == .string and weights.get(file.string) != null); + try holds(pins.object.get("checkpoint").?.object.get("model.safetensors.index.json") != null); +} + +fn jsonEqual(p: std.json.Value, q: std.json.Value) bool { + if (std.meta.activeTag(p) != std.meta.activeTag(q)) return false; + return switch (p) { + .string => |x| std.mem.eql(u8, x, q.string), + .integer => |x| x == q.integer, + .array => |x| x.items.len == q.array.items.len and for (x.items, q.array.items) |a, b| { + if (!jsonEqual(a, b)) break false; + } else true, + .object => |x| x.count() == q.object.count() and for (x.keys(), x.values()) |k, v| { + const w = q.object.get(k) orelse break false; + if (!jsonEqual(v, w)) break false; + } else true, + else => false, + }; +} + +test "Qwen3-Embedding references: a bar of other inputs, an unpinned one or a non-finite vector is refused" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const good = + \\{"dtype": "float32", "pins": {"repo": "Qwen/Qwen3-Embedding-8B", + \\ "revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", "transformers": "5.22.1", + \\ "transformers_tree": "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b", + \\ "checkpoint": {"config.json": "dcccf7c7890c8debc4af8dace8d6acd8dd50bd56d50d4d8b60529381ba3de3d7"}, + \\ "weights": {"model-00001-of-00004.safetensors": "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451"}, + \\ "library": {"models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b"}}, + \\ "texts": ["a"], "ids": [[1, 2]], "vectors": [[0.6, 0.8]]} + ; + const ref = try std.json.parseFromSliceLeaky(std.json.Value, a, good, .{}); + const bar_text = try std.mem.replaceOwned(u8, a, good, "float32", "bfloat16"); + try checkReferences(ref, try std.json.parseFromSliceLeaky(std.json.Value, a, bar_text, .{})); + const bad = [_][2][]const u8{ + .{ "[[1, 2]]", "[[1, 3]]" }, + .{ "\"texts\": [\"a\"]", "\"texts\": [\"b\"]" }, + .{ "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", "0000000000000000000000000000000000000000" }, + .{ "[[0.6, 0.8]]", "[[0.0, 0.0]]" }, + .{ "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451", "00b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451" }, + .{ "\"weights\": {\"model-00001-of-00004.safetensors\": \"99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451\"},", "" }, + .{ "\"transformers\": \"5.22.1\"", "\"transformers\": \"5.22.0\"" }, + }; + for (bad) |swap| { + const other = try std.mem.replaceOwned(u8, a, bar_text, swap[0], swap[1]); + try t.expect(std.meta.isError(checkReferences(ref, try std.json.parseFromSliceLeaky(std.json.Value, a, other, .{})))); + } +} + +test "Qwen3-Embedding references: pins that do not hash to the checkpoint's files are refused, whoever made both" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(t.io, .{ .sub_path = "s.safetensors", .data = "weights" }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\": {\"x\": \"s.safetensors\"}}" }); + const dir = try tmp.dir.realPathFileAlloc(t.io, ".", a); + const shard = try fileSha256(t.io, try std.fs.path.join(a, &.{ dir, "s.safetensors" })); + const index = try fileSha256(t.io, try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" })); + const good = try std.fmt.allocPrint(a, "{{\"checkpoint\": {{\"model.safetensors.index.json\": \"{s}\"}}, \"weights\": {{\"s.safetensors\": \"{s}\"}}}}", .{ &index, &shard }); + try checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, good, .{})); + const wrong = try std.mem.replaceOwned(u8, a, good, &shard, "0000000000000000000000000000000000000000000000000000000000000000"); + try t.expect(std.meta.isError(checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, wrong, .{})))); + const unpinned = try std.fmt.allocPrint(a, "{{\"checkpoint\": {{\"model.safetensors.index.json\": \"{s}\"}}, \"weights\": {{}}}}", .{&index}); + try t.expect(std.meta.isError(checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, unpinned, .{})))); +} + +/// The oracle's requests: consecutive inputs, at most 128 and `window` tokens in all (`ids` lengths in order). +fn groupEnds(a: std.mem.Allocator, lengths: []const usize, window: usize) ![]usize { + var out: std.ArrayList(usize) = .empty; + var used: usize = 0; + var count: usize = 0; + for (lengths, 0..) |n, i| { + if (n > window) return error.InputPastWindow; + if (count > 0 and (used + n > window or count == 128)) { + try out.append(a, i); + used = 0; + count = 0; + } + used += n; + count += 1; + } + try out.append(a, lengths.len); + return out.items; +} + +test "Qwen3-Embedding: the GPU test batches inputs as the oracle's check does" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + try t.expectEqualSlices(usize, &.{ 2, 3, 4 }, try groupEnds(a, &.{ 3, 5, 9, 1 }, 9)); + try t.expectEqualSlices(usize, &.{ 1, 3 }, try groupEnds(a, &.{ 4, 2, 2 }, 4)); + try t.expectError(error.InputPastWindow, groupEnds(a, &.{10}, 9)); +} + +/// `x` in [1e-20, 10) rounded up to two significant digits, as the oracle's round_up: never below x. +fn roundUp(x: f64) f64 { + var scale: f64 = 1; + var k = 1 - @floor(std.math.log10(x)); + while (k > 0) : (k -= 1) scale *= 10; + var up = @ceil(x * scale); + while ((up - 1) / scale >= x) up -= 1; + while (up / scale < x) up += 1; + return up / scale; +} + +test "Qwen3-Embedding gate: the BF16 bar's distance rounds up to two digits" { + const cases = [_][2]f64{ + .{ 1.572e-4, 1.6e-4 }, .{ 1.624e-4, 1.7e-4 }, .{ 1.884e-4, 1.9e-4 }, + .{ 1.6e-4, 1.6e-4 }, .{ 0.1099, 0.11 }, .{ 9.96e-5, 1.0e-4 }, + }; + for (cases) |c| { + try t.expectApproxEqRel(c[1], roundUp(c[0]), 1e-12); + try t.expect(roundUp(c[0]) >= c[0]); + } + var tens: f64 = 100; + while (tens <= 1e9) : (tens *= 10) for (10..100) |d| { + const x = @as(f64, @floatFromInt(d)) / tens; + try t.expectEqual(x, roundUp(x)); + try t.expectEqual(@as(f64, @floatFromInt(d + 1)) / tens, roundUp(std.math.nextAfter(f64, x, 1))); + }; +} + +test "Qwen3-Embedding: the card's and the long inputs batched equal them alone, within the BF16 bar of fp32" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try env(a, "TF_QWEN3_EMBED_DIR") orelse return error.SkipZigTest; + const ref = try readJson(a, try env(a, "TF_QWEN3_EMBED_REF") orelse return error.SkipZigTest); + const bar = try readJson(a, try env(a, "TF_QWEN3_EMBED_BAR") orelse return error.SkipZigTest); + try checkReferences(ref, bar); + try checkInputs(ref, &embed_inputs); + try checkFiles(a, t.io, dir, ref.object.get("pins").?); + var g: Gpu = undefined; + try g.open(4 << 30); + defer g.close(); + // the card's 2 queries and 2 documents, then 5 long inputs of 2K to 32K tokens, each vector full width + const inputs = ref.object.get("ids").?.array.items; + const refs = ref.object.get("vectors").?.array.items; + const bars = bar.object.get("vectors").?.array.items; + try t.expectEqual(@as(usize, 9), inputs.len); + try t.expectEqual(inputs.len, refs.len); + try t.expectEqual(inputs.len, bars.len); + const lengths = try a.alloc(usize, inputs.len); + for (inputs, lengths) |one, *n| { + try t.expect(one.array.items.len > 0); + n.* = one.array.items.len; + } + try t.expect(lengths[inputs.len - 1] > 16384); + const window = 32768; + var m = try qwen3.Model.load(gpa, t.io, &g.r, dir); + defer m.deinit(); + for (refs, bars) |r, b| { + try t.expectEqual(m.spec.hidden, r.array.items.len); + try t.expectEqual(m.spec.hidden, b.array.items.len); + } + const hidden = m.spec.hidden; + var bar_worst: f64 = 0; + for (bars, refs) |b, r| { + const bv = try a.alloc(f64, hidden); + const rv = try a.alloc(f64, hidden); + for (b.array.items, r.array.items, bv, rv) |x, y, *p, *q| { + p.* = number(x); + q.* = number(y); + } + bar_worst = @max(bar_worst, 1 - cosine(bv, rv)); + } + try t.expect(std.math.isFinite(bar_worst) and bar_worst > 0); + const limit = @min(0.001, roundUp(bar_worst)); + var first: usize = 0; + for (try groupEnds(a, lengths, window)) |last| { + defer first = last; + var ids: std.ArrayList(i32) = .empty; + const ends = try a.alloc(u32, last - first); + for (inputs[first..last], ends) |one, *e| { + try ids.appendSlice(a, try idsOf(a, one)); + e.* = @intCast(ids.items.len); + } + const batch = try g.lastRows(&m, ids.items, ends); + defer gpa.free(batch); + var begin: u32 = 0; + for (ends, 0..) |e, j| { + defer begin = e; + const i = first + j; + const alone = try g.lastRows(&m, ids.items[begin..e], &.{e - begin}); + defer gpa.free(alone); + try t.expectEqualSlices(u16, batch[j * hidden ..][0..hidden], alone); + const got = try a.alloc(f64, hidden); + const want = try a.alloc(f64, hidden); + for (batch[j * hidden ..][0..hidden], refs[i].array.items, got, want) |bits, r, *p, *q| { + p.* = widen(bits); + q.* = number(r); + } + const distance = 1 - cosine(got, want); + std.debug.print("embedding input {d} ({d} tokens): cosine distance {e:.3} (limit {e:.3})\n", .{ i, e - begin, distance, limit }); + try t.expect(std.math.isFinite(distance) and distance <= limit); + } + } +} + +test "Qwen3-Reranker: each case's yes - no logit from lastRows and lm_head rows matches transformers" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try env(a, "TF_QWEN3_RERANK_DIR") orelse return error.SkipZigTest; + const ref = try readJson(a, try env(a, "TF_QWEN3_RERANK_REF") orelse return error.SkipZigTest); + var g: Gpu = undefined; + try g.open(8 << 30); + defer g.close(); + var m = try qwen3.Model.load(gpa, t.io, &g.r, dir); + defer m.deinit(); + const yes: u32 = @intCast(ref.object.get("yes").?.integer); + const no: u32 = @intCast(ref.object.get("no").?.integer); + const head = try qwen3.readRows(gpa, t.io, dir, m.head, &.{ yes, no }); + defer gpa.free(head); + const hidden = m.spec.hidden; + var worst_diff: f64 = 0; + var worst_score: f64 = 0; + var cases: usize = 0; + for (ref.object.get("cases").?.array.items) |c| { + const want_diff = c.object.get("diff") orelse continue; // ids-only cases carry no score + const ids = try idsOf(a, c.object.get("ids").?); + const row = try g.lastRows(&m, ids, &.{@intCast(ids.len)}); + defer gpa.free(row); + var ly: f64 = 0; + var ln: f64 = 0; + for (row, head[0..hidden], head[hidden..][0..hidden]) |bits, wy, wn| { + ly += widen(bits) * wy; + ln += widen(bits) * wn; + } + const diff = ly - ln; + const score = 1 / (1 + @exp(-diff)); + const ref_score = number(c.object.get("score").?); + // a NaN would slip through @max: every value must be finite before it counts + for ([_]f64{ diff, score, number(want_diff), ref_score }) |v| try t.expect(std.math.isFinite(v)); + const d = @abs(diff - number(want_diff)); + const ds = @abs(score - ref_score); + std.debug.print("rerank {s} {d} tokens: diff {d:.4} ref {d:.4} score {d:.6} ref {d:.6}\n", .{ c.object.get("kind").?.string, ids.len, diff, number(want_diff), score, number(c.object.get("score").?) }); + worst_diff = @max(worst_diff, d); + worst_score = @max(worst_score, ds); + cases += 1; + } + std.debug.print("rerank: {d} cases, worst |diff| {d:.4} (gate 0.25), worst |score| {d:.5} (gate 0.02)\n", .{ cases, worst_diff, worst_score }); + try t.expect(cases >= 32); + try t.expect(worst_diff <= 0.25); + try t.expect(worst_score <= 0.02); +} diff --git a/zig/src/families/qwen3/model.zig b/zig/src/families/qwen3/model.zig new file mode 100644 index 0000000000..d299b66621 --- /dev/null +++ b/zig/src/families/qwen3/model.zig @@ -0,0 +1,221 @@ +//! A Qwen3 dense checkpoint on the GPU: BF16 projections as stored, fp32 norms, and host rows of lm_head on request. +const std = @import("std"); +const core = @import("core"); +const hip = @import("hip"); +const config = @import("config.zig"); + +const quant = core.quant; +const Checkpoint = core.checkpoint.Checkpoint; +const Allocator = std.mem.Allocator; + +pub const Projection = quant.Projection; + +/// One decoder layer: norms are fp32 device addresses, projections BF16 (N, K) as the checkpoint stores them. +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + q_norm: u64, + k_norm: u64, + q: Projection, + k: Projection, + v: Projection, + o: Projection, + gate: Projection, + up: Projection, + down: Projection, +}; + +pub const Model = struct { + gpa: Allocator, + spec: config.Spec, + /// The activation dtype: BF16, as the checkpoint's weights. + act: hip.ops.Kind = .bf16, + embed: Projection, + /// The stored name of the output head's rows for `readRows`: lm_head, or the embedding when the model ties them. + head: []const u8, + final_norm: u64, + layers: []Layer, + buffers: std.ArrayList(hip.memory.DeviceBuffer) = .empty, + + /// The checkpoint in `dir` on the device of `rt`; its mapped pages are dropped layer by layer. + pub fn load(gpa: Allocator, io: std.Io, rt: *const hip.Runtime, dir: []const u8) !Model { + const spec = try readSpec(gpa, io, dir); + var ck = try Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + ck.scope = scopeOf(&ck); + var m: Model = .{ .gpa = gpa, .spec = spec, .embed = undefined, .head = headKey(spec.tied, ck.scope), .final_norm = 0, .layers = &.{} }; + errdefer m.deinit(); + const up: hip.Upload = .{ .r = rt, .buffers = &m.buffers, .gpa = gpa }; + var scratch: std.heap.ArenaAllocator = .init(gpa); + defer scratch.deinit(); + var t: Table = .{ .ck = &ck }; + m.embed = try projection(scratch.allocator(), &t, &up, "embed_tokens", spec.vocab, spec.hidden); + m.final_norm = try norm(scratch.allocator(), &t, &up, "norm.weight", spec.hidden); + m.layers = try gpa.alloc(Layer, spec.layers); + for (m.layers, 0..) |*layer, i| { + _ = scratch.reset(.retain_capacity); + // no ck.release() between layers: its fadvise would drop the shards from the page cache mid-load + layer.* = try readLayer(scratch.allocator(), &t, &up, spec, i); + } + return m; + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.layers); + m.* = undefined; + } +}; + +/// The reranker's save keeps the base model under `model.` (beside `lm_head`); the embedding's has no prefix. +pub fn scopeOf(ck: *const Checkpoint) []const u8 { + return if (ck.has("model.embed_tokens.weight")) "model." else ""; +} + +/// The head rows' stored name: `lm_head.weight`, or the embedding's (prefixed as the checkpoint stores it) when tied. +pub fn headKey(tied: bool, scope: []const u8) []const u8 { + if (!tied) return "lm_head.weight"; + return if (scope.len > 0) "model.embed_tokens.weight" else "embed_tokens.weight"; +} + +/// `quant.read`'s table over a checkpoint: dense tensors only, so no format asks for an affine width. +const Table = struct { + ck: *Checkpoint, + + pub fn get(t: *const Table, name: []const u8) !quant.Tensor { + return t.ck.get(name); + } + + pub fn has(t: *const Table, name: []const u8) bool { + return t.ck.has(name); + } + + pub fn width(_: *const Table, _: []const u8) error{UnsupportedQuantization}!quant.mlx.Width { + return error.UnsupportedQuantization; + } +}; + +fn readSpec(gpa: Allocator, io: std.Io, dir: []const u8) !config.Spec { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return config.parse(gpa, bytes); +} + +/// A BF16 projection `key` of shape (n, k), refused in any other dtype or shape. +fn projection(a: Allocator, t: *const Table, up: *const hip.Upload, key: []const u8, n: usize, k: usize) !Projection { + const host = try quant.readStored(a, t, key); + const w = switch (host) { + .dense => |d| d.weight, + else => return refuse(key, "is not a dense tensor"), + }; + if (w.dtype != .bf16 or w.rank != 2 or w.shape[0] != n or w.shape[1] != k) { + std.log.err("{s}: {t} {any}, expected bf16 [{d}, {d}]", .{ key, w.dtype, w.shape[0..w.rank], n, k }); + return error.UnexpectedTensor; + } + return quant.view(try quant.upload(up.uploader(), host), false); +} + +/// An RMS norm weight `key` of `n` values, widened to fp32 on the device. +fn norm(a: Allocator, t: *const Table, up: *const hip.Upload, key: []const u8, n: usize) !u64 { + const src = try t.get(key); + if (src.rank != 1 or src.shape[0] != n) { + std.log.err("{s}: shape {any}, expected [{d}]", .{ key, src.shape[0..src.rank], n }); + return error.UnexpectedTensor; + } + const wide = try quant.convert.float32(a, src); + return (try up.uploader().tensor(wide)).ptr; +} + +fn readLayer(a: Allocator, t: *const Table, up: *const hip.Upload, s: config.Spec, i: usize) !Layer { + var buf: [128]u8 = undefined; + const q_width = s.heads * s.head_dim; + const kv_width = s.kv_heads * s.head_dim; + return .{ + .input_norm = try norm(a, t, up, try layerKey(&buf, i, "input_layernorm.weight"), s.hidden), + .post_norm = try norm(a, t, up, try layerKey(&buf, i, "post_attention_layernorm.weight"), s.hidden), + .q_norm = try norm(a, t, up, try layerKey(&buf, i, "self_attn.q_norm.weight"), s.head_dim), + .k_norm = try norm(a, t, up, try layerKey(&buf, i, "self_attn.k_norm.weight"), s.head_dim), + .q = try projection(a, t, up, try layerKey(&buf, i, "self_attn.q_proj"), q_width, s.hidden), + .k = try projection(a, t, up, try layerKey(&buf, i, "self_attn.k_proj"), kv_width, s.hidden), + .v = try projection(a, t, up, try layerKey(&buf, i, "self_attn.v_proj"), kv_width, s.hidden), + .o = try projection(a, t, up, try layerKey(&buf, i, "self_attn.o_proj"), s.hidden, q_width), + .gate = try projection(a, t, up, try layerKey(&buf, i, "mlp.gate_proj"), s.inter, s.hidden), + .up = try projection(a, t, up, try layerKey(&buf, i, "mlp.up_proj"), s.inter, s.hidden), + .down = try projection(a, t, up, try layerKey(&buf, i, "mlp.down_proj"), s.hidden, s.inter), + }; +} + +fn layerKey(buf: *[128]u8, i: usize, rest: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "layers.{d}.{s}", .{ i, rest }); +} + +fn refuse(key: []const u8, why: []const u8) error{UnexpectedTensor} { + std.log.err("{s} {s}", .{ key, why }); + return error.UnexpectedTensor; +} + +/// Rows `rows` of the 2-D tensor `key` (named as stored, no prefix) as fp32 host values, row after row. +pub fn readRows(gpa: Allocator, io: std.Io, dir: []const u8, key: []const u8, rows: []const u32) ![]f32 { + var ck = try Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const t = try ck.get(key); + if (t.rank != 2) return refuse(key, "is not a matrix"); + const width = t.shape[1]; + const out = try gpa.alloc(f32, rows.len * width); + errdefer gpa.free(out); + const size = t.dtype.size(); + for (rows, 0..) |r, i| { + if (r >= t.shape[0]) { + std.log.warn("{s} has no row {d} of {d}", .{ key, r, t.shape[0] }); + return error.UnexpectedTensor; + } + var shape = t.shape; + shape[0] = width; + const one: quant.Tensor = .{ .bytes = t.bytes[r * width * size ..][0 .. width * size], .dtype = t.dtype, .rank = 1, .shape = shape }; + var scratch: std.heap.ArenaAllocator = .init(gpa); + defer scratch.deinit(); + const wide = try quant.convert.float32(scratch.allocator(), one); + // little-endian fp32 words, read one by one: a mapped fp32 tensor need not be 4-byte aligned + for (out[i * width ..][0..width], 0..) |*x, j| x.* = @bitCast(std.mem.readInt(u32, wide.bytes[4 * j ..][0..4], .little)); + } + return out; +} + +fn writeFixture(dir: std.Io.Dir, prefixed: bool) !void { + const fixture = core.checkpoint_fixture; + const a = std.testing.allocator; + // lm_head: 3 rows of 2 bf16 values: row r is (r + 1, -(r + 1)) as 1.0, 2.0, 3.0 and their negatives + const head = [_]u16{ 0x3f80, 0xbf80, 0x4000, 0xc000, 0x4040, 0xc040 }; + const embed = [_]u16{ 0, 0, 0, 0, 0, 0 }; + try fixture.write(a, std.testing.io, dir, "model.safetensors", &.{ + .{ .name = if (prefixed) "model.embed_tokens.weight" else "embed_tokens.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .bytes = std.mem.sliceAsBytes(&embed) }, + .{ .name = "lm_head.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .bytes = std.mem.sliceAsBytes(&head) }, + }); +} + +test "the reranker's model. prefix is found, the embedding's bare names are not, and lm_head rows read as fp32" { + const a = std.testing.allocator; + const io = std.testing.io; + for ([_]bool{ true, false }) |prefixed| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeFixture(tmp.dir, prefixed); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try Checkpoint.openModel(a, io, dir); + defer ck.close(); + try std.testing.expectEqualStrings(if (prefixed) "model." else "", scopeOf(&ck)); + const rows = try readRows(a, io, dir, "lm_head.weight", &.{ 2, 0 }); + defer a.free(rows); + try std.testing.expectEqualSlices(f32, &.{ 3, -3, 1, -1 }, rows); + try std.testing.expectError(error.UnexpectedTensor, readRows(a, io, dir, "lm_head.weight", &.{3})); + // a tied model reads its head rows from the embedding, as stored + const tied = try readRows(a, io, dir, headKey(true, scopeOf(&ck)), &.{1}); + defer a.free(tied); + try std.testing.expectEqualSlices(f32, &.{ 0, 0 }, tied); + try std.testing.expectEqualStrings("lm_head.weight", headKey(false, scopeOf(&ck))); + } +} diff --git a/zig/src/families/qwen3/native.zig b/zig/src/families/qwen3/native.zig new file mode 100644 index 0000000000..7efa2b9540 --- /dev/null +++ b/zig/src/families/qwen3/native.zig @@ -0,0 +1,232 @@ +//! Qwen3 dense checkpoints for the native HIP server: the head the folder declares, sized to the memory budget. +const std = @import("std"); +const hip = @import("hip"); +const lanes = @import("lanes"); +const api = @import("engine_api"); +const engine = @import("engine.zig"); +const prompt = @import("prompt.zig"); + +const Engine = engine.Engine; +const Allocator = std.mem.Allocator; + +/// The registry entry: `native/hip.zig` opens `qwen3` checkpoints through it. +pub const model_type = "qwen3"; +pub const formats: []const []const u8 = &.{"bf16"}; +/// The window when neither --context nor the checkpoint's config names one (the Qwen3 8B serving default). +pub const default_context: i64 = 32768; +/// The engine takes a stream's whole prompt in one step. +pub const prefill_step: u32 = 0; +pub const slide_refusal = "--slide: Qwen3 dense checkpoints do not learn on HIP"; +pub const confidence_refusal = "--mtp-confidence: Qwen3 dense checkpoints have no draft head"; +/// An embedding engine writes no reply text, so no reasoning block (Info.reasoning). +pub const reasoning = false; +/// Checked by native/hip.zig before any rank joins a group. +pub const tp_refusal = "--tp: Qwen3 dense checkpoints run on one HIP GPU; serve with --tp 1"; +pub const open = openEngine; +pub const follow = followRank; + +/// What the server drives, as the other HIP families' `Loaded`, plus what this engine serves. +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + streams: usize, + startup: []const u8, + /// No chat or completions: the server refuses the generation routes. + text_generation: bool = false, + embeddings: ?api.EmbeddingInfo = null, + /// The window the engine was sized to (the asked one, or less where the budget shrank it). + window: ?usize = null, +}; + +pub const Refused = error{Refused}; + +fn refuse(a: Allocator, problem: *[]const u8, comptime fmt: []const u8, args: anytype) Refused { + problem.* = std.fmt.allocPrint(a, fmt, args) catch fmt; + return error.Refused; +} + +/// The head a checkpoint folder declares. +pub const Head = union(enum) { + /// sentence-transformers last-token pooling, unit-normed when modules.json lists Normalize. + embedding: struct { normalize: bool }, +}; + +const pooling_modes = [_][]const u8{ + "pooling_mode_cls_token", "pooling_mode_mean_tokens", "pooling_mode_max_tokens", + "pooling_mode_mean_sqrt_len_tokens", "pooling_mode_weightedmean_tokens", "pooling_mode_lasttoken", +}; + +/// The head from `1_Pooling/config.json` and `modules.json`, or a refusal naming what the folder lacks. +pub fn headOf(a: Allocator, io: std.Io, dir: []const u8, problem: *[]const u8) (Refused || Allocator.Error)!Head { + const pooling = try readJson(a, io, dir, "1_Pooling/config.json", problem) orelse return refuse( + a, + problem, + "qwen3 on HIP serves sentence-transformers embedding checkpoints: {s} has no 1_Pooling/config.json", + .{dir}, + ); + if (pooling != .object) return refuse(a, problem, "1_Pooling/config.json is not a JSON object", .{}); + for (pooling_modes) |mode| { + const on = if (pooling.object.get(mode)) |v| v == .bool and v.bool else false; + const want = std.mem.eql(u8, mode, "pooling_mode_lasttoken"); + if (on != want) return refuse(a, problem, "1_Pooling/config.json: {s} is {s}; qwen3 on HIP pools the last token only", .{ mode, if (on) "on" else "off" }); + } + // no modules.json: the pooled rows as they are; one that is there must be a list of modules with string types + var normalize = false; + if (try readJson(a, io, dir, "modules.json", problem)) |mods| { + if (mods != .array) return refuse(a, problem, "modules.json is not a JSON list of modules", .{}); + for (mods.array.items) |m| { + const kind = if (m == .object) m.object.get("type") else null; + if (kind == null or kind.? != .string) return refuse(a, problem, "modules.json: a module without a string \"type\"", .{}); + if (std.mem.eql(u8, kind.?.string, "sentence_transformers.models.Normalize")) normalize = true; + } + } + return .{ .embedding = .{ .normalize = normalize } }; +} + +/// The folder's JSON file `name`, null when it is absent; unreadable or malformed JSON is refused. +fn readJson(a: Allocator, io: std.Io, dir: []const u8, name: []const u8, problem: *[]const u8) (Refused || Allocator.Error)!?std.json.Value { + const path = try std.fs.path.join(a, &.{ dir, name }); + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 20)) catch |err| switch (err) { + error.FileNotFound => return null, + error.OutOfMemory => return error.OutOfMemory, + else => return refuse(a, problem, "{s}: cannot read it ({s})", .{ name, @errorName(err) }), + }; + return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch |err| switch (err) { + error.OutOfMemory => error.OutOfMemory, + else => refuse(a, problem, "{s} is not valid JSON ({s})", .{ name, @errorName(err) }), + }; +} + +fn gibs(bytes: usize) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + +/// The window the budget holds: `asked`, halved while its scratch passes `room`, down to `engine.least_window`. +pub fn fitWindow(asked: usize, room: usize, ctx: anytype, bytesFor: fn (@TypeOf(ctx), usize) anyerror!usize) !?usize { + var window = asked; + while (true) { + if (try bytesFor(ctx, window) <= room) return window; + if (window <= engine.least_window) return null; + window = @max(window / 2, engine.least_window); + } +} + +fn scratchOf(e: *Engine, window: usize) anyerror!usize { + return e.scratchBytes(window); +} + +const Owned = struct { + gpa: Allocator, + e: *Engine, + p: prompt.Prompt, + startup: []const u8, +}; + +fn release(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.p.deinit(); + own.e.deinit(); + own.gpa.free(own.startup); + own.gpa.destroy(own); +} + +/// The engine for `dir` on `dev`, its head from the folder and its window from the memory budget. +fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, dev: hip.Device, dir: []const u8, o: anytype, problem: *[]const u8) anyerror!Loaded { + if (o.world > 1) return refuse(a, problem, "--tp {d}: Qwen3 dense checkpoints run on one HIP GPU", .{o.world}); + // the BF16 weights run on the BF16 dot2 GEMM tile: gfx11 and gfx12, not an FP16 build (RDNA2) + if (dev.caps.act != .bf16) return refuse(a, problem, "Qwen3 dense checkpoints on HIP need a GPU with BF16 dot2 (gfx11 or gfx12); {s} runs FP16 activations", .{hip.kernels.arch}); + const head = try headOf(a, io, dir, problem); + const e = Engine.load(gpa, io, dir, dev.index, dev.policy.choices()) catch |err| + return refuse(a, problem, "the native HIP engine cannot load {s} ({s})", .{ dir, @errorName(err) }); + errdefer e.deinit(); + const pool = hip.admission.budget(gpa, io, e.ctx) catch |err| + return refuse(a, problem, "the native HIP engine cannot read the GPU's memory ({s})", .{@errorName(err)}); + const room = pool.room(e.weights); + const window = (fitWindow(o.window, room, e, scratchOf) catch |err| + return refuse(a, problem, "the native HIP engine cannot count its scratch ({s})", .{@errorName(err)})) orelse { + const least = e.scratchBytes(@min(o.window, engine.least_window)) catch 0; + return refuse(a, problem, "the HIP memory budget fits no {d}-token pass: it takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; free device memory or lower --context", .{ @min(o.window, engine.least_window), gibs(least), gibs(room), gibs(e.weights), gibs(pool.reserve) }); + }; + e.size(window) catch |err| + return refuse(a, problem, "the native HIP engine cannot allocate its scratch ({s})", .{@errorName(err)}); + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .e = e, .p = prompt.Prompt.init(gpa, e, head.embedding.normalize), .startup = undefined }; + const shrunk: []const u8 = if (window < o.window) " (shrunk from the asked window to fit the budget)" else ""; + own.startup = try std.fmt.allocPrint(gpa, "HIP {s} device {d}: qwen3 embedding, model {d:.2} GiB; {d}-token passes{s}, {d:.2} GiB scratch, of {d:.1} GiB left after a {d:.1} GiB reserve; vectors {s}", .{ + hip.kernels.arch, dev.index, gibs(e.weights), window, shrunk, gibs(e.scratchBytes(window) catch 0), gibs(room), gibs(pool.reserve), if (head.embedding.normalize) "unit-normed" else "as pooled", + }); + const spec = e.m.spec; + return .{ + .backend = own.p.backend(), + .facts = .{}, + .rows = 1, + .ctx = own, + .deinit = release, + .streams = 1, + .startup = own.startup, + .text_generation = false, + .embeddings = .{ .dimensions = @intCast(spec.hidden), .vocabulary = @intCast(spec.vocab), .max_batch_tokens = @intCast(window), .max_inputs = @intCast(engine.max_inputs), .mrl_min = @min(32, @as(u32, @intCast(spec.hidden))) }, + .window = window, + }; +} + +/// Ranks above 0 have nothing to hold: the engine runs on one GPU. +fn followRank(a: Allocator, _: Allocator, _: std.Io, _: hip.Device, _: []const u8, o: anytype, problem: *[]const u8) anyerror!void { + return refuse(a, problem, "--tp {d}: Qwen3 dense checkpoints run on one HIP GPU", .{o.world}); +} + +test "the window halves to fit the room, down to the least window" { + const Count = struct { + per: usize, + fn bytes(c: @This(), window: usize) anyerror!usize { + return window * c.per; + } + }; + const c: Count = .{ .per = 10 }; + try std.testing.expectEqual(@as(?usize, 32768), try fitWindow(32768, 327680, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 8192), try fitWindow(32768, 100000, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 2048), try fitWindow(32768, 20480, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, null), try fitWindow(32768, 20479, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 1000), try fitWindow(1000, 10000, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, null), try fitWindow(1000, 9999, c, Count.bytes)); +} + +test "a pooling folder declares the embedding head: last token only, normalized when modules.json says so" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + var problem: []const u8 = ""; + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "1_Pooling") != null); + try tmp.dir.createDirPath(io, "1_Pooling"); + const last = + \\{"word_embedding_dimension": 4096, "pooling_mode_cls_token": false, "pooling_mode_mean_tokens": false, + \\ "pooling_mode_max_tokens": false, "pooling_mode_mean_sqrt_len_tokens": false, + \\ "pooling_mode_weightedmean_tokens": false, "pooling_mode_lasttoken": true, "include_prompt": true} + ; + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = last }); + try std.testing.expect(!(try headOf(a, io, dir, &problem)).embedding.normalize); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"type\": \"sentence_transformers.models.Pooling\"}, {\"type\": \"sentence_transformers.models.Normalize\"}]" }); + try std.testing.expect((try headOf(a, io, dir, &problem)).embedding.normalize); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"type\": \"sentence_transformers.models.Normalize\"" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "not valid JSON") != null); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "{\"type\": \"sentence_transformers.models.Normalize\"}" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"path\": \"1_Pooling\"}]" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[]" }); + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = "{\"pooling_mode_lasttoken\": tru" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = "{\"pooling_mode_mean_tokens\": true, \"pooling_mode_lasttoken\": true}" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "pooling_mode_mean_tokens") != null); +} diff --git a/zig/src/families/qwen3/prompt.zig b/zig/src/families/qwen3/prompt.zig new file mode 100644 index 0000000000..2ff8e34887 --- /dev/null +++ b/zig/src/families/qwen3/prompt.zig @@ -0,0 +1,106 @@ +//! The lane backend of a prompt-only Qwen3 engine: a stream's whole prompt in one step, pooled rows until release. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api"); +const Engine = @import("engine.zig").Engine; + +const be = lanes.backend; +const Stream = lanes.Stream; +const Allocator = std.mem.Allocator; + +/// What a prompt-only engine refuses: drawing, verifying or drafting tokens. +pub const Refusal = error{PromptOnly}; + +pub const Prompt = struct { + gpa: Allocator, + engine: *Engine, + /// The checkpoint's Normalize module: unit L2 rows. + normalize: bool, + /// Each stream's pooled fp32 rows, from its prompt pass to its release. + slots: std.AutoHashMapUnmanaged(*Stream, []f32) = .empty, + table: be.Backend.VTable = .{ + .embedding = embedding, + .prefill = prefill, + .prefill_step = prefillStep, + .first = first, + .queue = queue, + .read = read, + .verify = verify, + .keep = keep, + .draft = draft, + .release = release, + }, + + pub fn init(gpa: Allocator, engine: *Engine, normalize: bool) Prompt { + return .{ .gpa = gpa, .engine = engine, .normalize = normalize }; + } + + pub fn deinit(p: *Prompt) void { + var it = p.slots.valueIterator(); + while (it.next()) |rows| p.gpa.free(rows.*); + p.slots.deinit(p.gpa); + } + + /// The backend over this engine; heads add their hooks to `table` before it serves. + pub fn backend(p: *Prompt) be.Backend { + return .{ .ptr = p, .vtable = &p.table, .chunk_invariant = true }; + } + + fn of(ptr: *anyopaque) *Prompt { + return @ptrCast(@alignCast(ptr)); + } + + /// The whole prompt in one step: every input's last row, kept for `embedding`. + fn prefillStep(ptr: *anyopaque, s: *Stream, _: u32) anyerror!bool { + const p = of(ptr); + const batch = s.embedding orelse return error.PromptOnly; + const hidden = p.engine.m.spec.hidden; + const rows = try p.gpa.alloc(f32, batch.ends.len * hidden); + errdefer p.gpa.free(rows); + try p.engine.run(s.prompt(), batch.ends, rows); + const slot = try p.slots.getOrPut(p.gpa, s); + if (slot.found_existing) p.gpa.free(slot.value_ptr.*); + slot.value_ptr.* = rows; + return true; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + _ = try prefillStep(ptr, s, 0); + } + + /// The stream's pooled rows into `out`, each unit-normed when the checkpoint normalizes. + fn embedding(ptr: *anyopaque, s: *Stream, out: []f32) anyerror!void { + const p = of(ptr); + const rows = p.slots.get(s) orelse return error.EmbeddingNotFilled; + if (rows.len != out.len) return error.EmbeddingShape; + @memcpy(out, rows); + if (!p.normalize) return; + const hidden = p.engine.m.spec.hidden; + var at: usize = 0; + while (at < out.len) : (at += hidden) try api.embedding_vector.normalize(out[at..][0..hidden]); + } + + fn release(ptr: *anyopaque, s: *Stream) void { + const p = of(ptr); + if (p.slots.fetchRemove(s)) |kv| p.gpa.free(kv.value); + } + + fn first(_: *anyopaque, _: *Stream, _: u64) anyerror!u64 { + return error.PromptOnly; + } + fn queue(_: *anyopaque, _: *Stream, _: be.Feed, _: u64) anyerror!u64 { + return error.PromptOnly; + } + fn read(_: *anyopaque, _: u64) anyerror!u32 { + return error.PromptOnly; + } + fn verify(_: *anyopaque, _: []const be.Window, _: []be.Verified) anyerror!void { + return error.PromptOnly; + } + fn keep(_: *anyopaque, _: []const be.Window, _: []const []const u32) anyerror!void { + return error.PromptOnly; + } + fn draft(_: *anyopaque, _: []const be.DraftRequest) anyerror!void { + return error.PromptOnly; + } +}; diff --git a/zig/src/families/qwen3/qwen3.zig b/zig/src/families/qwen3/qwen3.zig new file mode 100644 index 0000000000..775e417391 --- /dev/null +++ b/zig/src/families/qwen3/qwen3.zig @@ -0,0 +1,21 @@ +//! Qwen3 dense checkpoints on the native HIP engine: config, BF16 weights on the device and the batched prompt pass. + +pub const config = @import("config.zig"); +pub const model = @import("model.zig"); +pub const forward = @import("forward.zig"); +pub const engine = @import("engine.zig"); +pub const prompt = @import("prompt.zig"); +pub const native = @import("native.zig"); +pub const Spec = config.Spec; +pub const Model = model.Model; +pub const readRows = model.readRows; +pub const lastRows = forward.lastRows; + +test { + _ = config; + _ = model; + _ = forward; + _ = engine; + _ = prompt; + _ = native; +} diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index ff98d00f33..440b1062df 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -26,6 +26,7 @@ test { _ = @import("kernel_tests/attention_pin.zig"); _ = @import("kernel_tests/stats.zig"); _ = @import("kernel_tests/ops.zig"); + _ = @import("kernel_tests/dense.zig"); _ = @import("kernel_tests/argmax.zig"); _ = @import("kernel_tests/tp.zig"); _ = @import("kernel_tests/affine/affine.zig"); diff --git a/zig/src/hip/kernel_tests/dense.zig b/zig/src/hip/kernel_tests/dense.zig new file mode 100644 index 0000000000..6d93a2d2aa --- /dev/null +++ b/zig/src/hip/kernel_tests/dense.zig @@ -0,0 +1,106 @@ +//! The BF16 dense GEMM against float64: within the dot2 chain's bound, BF16 out rounded once, rows batch-independent. +const std = @import("std"); +const rig = @import("rig.zig"); +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; +const ops = @import("../ops/ops.zig"); +const Tensor = ops.Tensor; + +const Case = struct { m: usize, n: usize, k: usize }; + +/// Edges of both tile dimensions, one and several groups of K, and Qwen3's widest K. +const cases = [_]Case{ + .{ .m = 1, .n = 2, .k = 128 }, + .{ .m = 7, .n = 130, .k = 4096 }, + .{ .m = 128, .n = 256, .k = 128 }, + .{ .m = 129, .n = 300, .k = 4096 }, + .{ .m = 300, .n = 130, .k = 12288 }, +}; + +/// The fp32 value as BF16, rounded to nearest even (the tile's epilogue; the rig's bf16Bits truncates). +fn nearest(v: f32) u16 { + const u: u32 = @bitCast(v); + if (std.math.isNan(v)) return 0x7fc0; + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf16Of(rng: *rig.Rng) u16 { + return rig.bf16Bits(@floatCast(rng.normal())); +} + +test "dense GEMM: BF16 weights against float64, fp32 within the dot2 chain's bound, BF16 out rounded once" { + const t = try Rig.open(64 << 20); + defer t.close(); + if (!t.bf16) return error.SkipZigTest; + const o = t.ops(&t.on); + var rng: rig.Rng = .{ .state = 0xDE45E }; + for (cases) |c| { + const hx = try gpa.alloc(u16, c.m * c.k); + defer gpa.free(hx); + const hw = try gpa.alloc(u16, c.n * c.k); + defer gpa.free(hw); + for (hx) |*v| v.* = bf16Of(&rng); + for (hw) |*v| v.* = bf16Of(&rng); + var x = try t.upload(hx); + defer x.free(); + var w = try t.upload(hw); + defer w.free(); + var y32 = try t.alloc(c.m * c.n * 4); + defer y32.free(); + var y16 = try t.alloc(c.m * c.n * 2); + defer y16.free(); + const xt: Tensor = .{ .ptr = at(x), .kind = .bf16 }; + try o.denseGemm(xt, at(w), .{ .ptr = at(y32), .kind = .f32 }, c.m, c.n, c.k); + try o.denseGemm(xt, at(w), .{ .ptr = at(y16), .kind = .bf16 }, c.m, c.n, c.k); + try t.stream.synchronize(); + const g32 = try rig.download(f32, y32); + defer gpa.free(g32); + const g16 = try rig.download(u16, y16); + defer gpa.free(g16); + const depth: f64 = @floatFromInt(@import("../launches/dense.zig").Kernels.chain(c.k)); + var worst: f64 = 0; + for (0..c.m) |r| for (0..c.n) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (0..c.k) |i| { + const p = rig.bf16Value(hx[r * c.k + i]) * rig.bf16Value(hw[j * c.k + i]); + want += p; + mag += @abs(p); + } + const got = g32[r * c.n + j]; + const err = @abs(@as(f64, got) - want); + worst = @max(worst, err / @max(mag, 1e-30)); + try rig.expect(err <= depth * std.math.pow(f64, 2, -24) * mag, "dense m {d} n {d} k {d} at ({d}, {d}): {d} vs {d} (mag {d})", .{ c.m, c.n, c.k, r, j, got, want, mag }); + try rig.expect(g16[r * c.n + j] == nearest(got), "dense bf16 out at ({d}, {d}) is {x}, not {x}: the fp32 sum {d} rounded once", .{ r, j, g16[r * c.n + j], nearest(got), got }); + }; + std.log.info("dense m {d} n {d} k {d}: worst error {e:.2} of the magnitude", .{ c.m, c.n, c.k, worst }); + // the last row alone has the bits it has in the batch + var lone = try t.alloc(c.n * 4); + defer lone.free(); + const last: Tensor = .{ .ptr = at(x) + (c.m - 1) * c.k * 2, .kind = .bf16 }; + try o.denseGemm(last, at(w), .{ .ptr = at(lone), .kind = .f32 }, 1, c.n, c.k); + try t.stream.synchronize(); + const g1 = try rig.download(f32, lone); + defer gpa.free(g1); + try rig.expect(std.mem.eql(u8, std.mem.sliceAsBytes(g1), std.mem.sliceAsBytes(g32[(c.m - 1) * c.n ..][0..c.n])), "dense m {d}: the last row alone differs from its batch bits", .{c.m}); + } +} + +test "dense GEMM: misaligned x, K off a group and FP16 operands are refused before launch" { + const t = try Rig.open(1 << 20); + defer t.close(); + if (!t.bf16) return error.SkipZigTest; + const o = t.ops(&t.on); + var x = try t.alloc(4 * 256 * 2 + 16); + defer x.free(); + var w = try t.alloc(4 * 256 * 2); + defer w.free(); + var y = try t.alloc(4 * 4 * 4); + defer y.free(); + const out: Tensor = .{ .ptr = at(y), .kind = .f32 }; + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x) + 2, .kind = .bf16 }, at(w), out, 1, 4, 256)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .bf16 }, at(w), out, 1, 4, 200)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .f16 }, at(w), out, 1, 4, 256)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .bf16 }, at(w), .{ .ptr = at(y), .kind = .f16 }, 1, 4, 256)); +} diff --git a/zig/src/hip/kernel_tests/ops.zig b/zig/src/hip/kernel_tests/ops.zig index 1c427945f3..c9358329a9 100644 --- a/zig/src/hip/kernel_tests/ops.zig +++ b/zig/src/hip/kernel_tests/ops.zig @@ -104,6 +104,38 @@ test "casts, sums, column copies and cache writes give torch's bits" { } } +test "head-major attention rows become row-major o_proj input, rounded once to the activation or kept in fp32" { + const t = try Rig.open(1 << 20); + defer t.close(); + const o = t.ops(&t.on); + var rng: rig.Rng = .{ .state = 0x4EAD }; + const len = 37; + const heads = 5; + const d = 24; + const n = len * heads * d; + const ha = try gpa.alloc(f32, n); + defer gpa.free(ha); + for (ha) |*x| x.* = @floatCast(rng.normal() * 9.1); + var att = try t.upload(ha); + defer att.free(); + var half = try t.alloc(n * 2); + defer half.free(); + var wide = try t.alloc(n * 4); + defer wide.free(); + try o.headsRows(at(att), tensor(half, act(t)), len, heads, d); + try o.headsRows(at(att), tensor(wide, .f32), len, heads, d); + try t.stream.synchronize(); + const gh = try rig.download(u16, half); + defer gpa.free(gh); + const gw = try rig.download(f32, wide); + defer gpa.free(gw); + for (0..len) |r| for (0..heads) |h| for (0..d) |j| { + const x = ha[(h * len + r) * d + j]; + try std.testing.expectEqual(round(t, x), gh[(r * heads + h) * d + j]); + try std.testing.expectEqual(x, gw[(r * heads + h) * d + j]); + }; +} + test "argmax and top-k follow torch's order: the larger value, then the lower index" { const t = try Rig.open(1 << 20); defer t.close(); diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index a2aae7237b..06608d5731 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -9,6 +9,7 @@ const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; const caps = @import("caps.zig"); const affine = @import("launches/affine.zig"); +const dense = @import("launches/dense.zig"); const util = @import("launches/util.zig"); @@ -86,6 +87,8 @@ pub const Launcher = struct { wide: bool, /// The MLX affine products and the registry that picks their tiles. affine: affine.Kernels, + /// The BF16 dense GEMM (raw BF16 weights through the dot2 GEMM tile). + dense: dense.Kernels, /// Flash Next's kernels: its residual streams, n-gram rows, indexer and sparse attention. fnx: flashnext.Fns, /// Nemotron's Mamba-2 mixer kernels. @@ -129,6 +132,7 @@ pub const Launcher = struct { silu_mul: Function, add: Function, attn_gate: Function, + heads_rows: Function, gnorm_silu: Function, copy_cols: Function, rope_prefill: Function, @@ -203,6 +207,7 @@ pub const Launcher = struct { .silu_mul = try ops.function("tf_silu_mul"), .add = try ops.function("tf_add"), .attn_gate = try ops.function("tf_attn_gate"), + .heads_rows = try ops.function("tf_heads_rows"), .gnorm_silu = try ops.function("tf_gnorm_silu"), .copy_cols = try ops.function("tf_copy_cols"), .rope_prefill = try ops.function("tf_rope_prefill"), @@ -291,7 +296,9 @@ pub const Launcher = struct { l.mamba = try mamba.Fns.load(l.mods[@backingInt(kernels.Group.mamba)]); const tiles = l.mods[@backingInt(kernels.Group.affine_tiles)]; const dot2 = l.mods[@backingInt(kernels.Group.affine_dot2)]; - l.affine = try affine.Kernels.load(tiles, dot2, caps.Caps.of(kernels.arch) orelse return error.Invalid, choices.products); + const chip = caps.Caps.of(kernels.arch) orelse return error.Invalid; + l.affine = try affine.Kernels.load(tiles, dot2, chip, choices.products); + l.dense = try dense.Kernels.load(tiles, chip.act == .bf16); try l.affine.fillByteLut(r, dot2); l.fnx = try flashnext.Fns.load(l.mods[@backingInt(kernels.Group.flashnext)]); return l; @@ -325,6 +332,7 @@ pub const Launcher = struct { pub const tf_moe_route = moe.tf_moe_route; pub const tf_causal = attention.tf_causal; pub const tf_attn_gate = attention.tf_attn_gate; + pub const tf_heads_rows = attention.tf_heads_rows; pub const tf_kv_write = attention.tf_kv_write; pub const tf_kv_write_at = attention.tf_kv_write_at; pub const tf_argmax_rows = draw.tf_argmax_rows; diff --git a/zig/src/hip/launches/attention.zig b/zig/src/hip/launches/attention.zig index ea31be1c9f..00e6ea71ba 100644 --- a/zig/src/hip/launches/attention.zig +++ b/zig/src/hip/launches/attention.zig @@ -134,6 +134,17 @@ pub fn tf_attn_gate(l: *const Launcher, att: CF, qg: C, out: P, kind: c_int, len try l.flat(l.op.attn_gate, @as(i64, len) * heads * d, s, &a); } +pub fn tf_heads_rows(l: *const Launcher, att: CF, out: P, kind: c_int, len: c_int, heads: c_int, d: c_int, s: S) Error!void { + var a: Args = .{}; + try a.add(ad(att)); + try a.add(ad(out)); + try a.add(kind); + try a.add(len); + try a.add(heads); + try a.add(d); + try l.flat(l.op.heads_rows, @as(i64, len) * heads * d, s, &a); +} + pub fn tf_kv_write(l: *const Launcher, src: C, cache: P, kind: c_int, len: c_int, kv_heads: c_int, d: c_int, total: c_int, pos0: c_int, s: S) Error!void { var a: Args = .{}; try a.add(ad(src)); diff --git a/zig/src/hip/launches/dense.zig b/zig/src/hip/launches/dense.zig new file mode 100644 index 0000000000..ecbb867cc2 --- /dev/null +++ b/zig/src/hip/launches/dense.zig @@ -0,0 +1,89 @@ +//! The BF16 dense GEMM of quant/dense_tiles.hpp (the dot2 GEMM tile over raw BF16 weights) launched from Zig. + +const std = @import("std"); +const abi = @import("../abi.zig"); +const runtime = @import("../runtime.zig"); +const hl = @import("../launch.zig"); +const Module = @import("../module.zig").Module; +const Function = @import("../module.zig").Function; +const affine = @import("affine.zig"); +const util = @import("util.zig"); + +const Error = runtime.Error; + +/// K a group folds into the accumulator (the tile's group; the weights carry no scales). +pub const group: c_int = 128; +/// The tile's rows and columns a block. +const side: c_int = 128; + +const symbol = "_ZN2tf4rocm16dense_gemm_blockINS0_7DotBF16ELi8EEEvNS0_6AffineE"; + +pub const Kernels = struct { + tile: Function, + /// The build's activations are BF16 (the tile's dot2); on an FP16 build the symbol only traps. + bf16: bool, + + /// `tiles` is dot2_tiles.hip's module, where every build keeps the symbol. + pub fn load(tiles: Module, bf16: bool) Error!Kernels { + return .{ .tile = try tiles.function(symbol), .bf16 = bf16 }; + } + + /// The product's refusal, or null when the tile takes it: BF16 build, K in whole groups, 16-byte rows and bases. + pub fn refusal(k: *const Kernels, x: u64, w: u64, m: c_int, n: c_int, kk: c_int) ?[]const u8 { + if (!k.bf16) return "the BF16 dense GEMM needs a BF16 build"; + if (m <= 0 or n <= 0 or kk <= 0) return "empty product"; + if (@rem(kk, group) != 0) return "K is not a multiple of 128"; + if (x % 16 != 0 or w % 16 != 0) return "x or the weights are not 16-byte aligned"; + return null; + } + + /// y = x . W^T for `m` BF16 rows of x and the [n, k] BF16 weights: fp32 into `out`, or BF16 into `out16`. + pub fn gemm(k: *const Kernels, r: *const runtime.Runtime, x: u64, w: u64, out: u64, out16: u64, m: c_int, n: c_int, kk: c_int, s: abi.Stream) Error!void { + if (k.refusal(x, w, m, n, kk)) |why| { + std.log.err("dense gemm: {s} (m {d}, n {d}, k {d})", .{ why, m, n, kk }); + return error.Invalid; + } + if ((out == 0) == (out16 == 0)) { + std.log.err("dense gemm: exactly one of out and out16", .{}); + return error.Invalid; + } + const a: affine.Arg = .{ + .x = x, + .words = w, + .scale = .{ .p = 0, .kind = 0 }, + .bias = .{ .p = 0, .kind = 0 }, + .out = out, + .m = m, + .n = n, + .k = kk, + .bits = 16, + .group = group, + .fp16 = 0, + .out16 = out16, + }; + if (s == util.counting) return; + var args: hl.Args = .{}; + try args.add(a); + try hl.launch(k.tile, .{ .grid = .{ .x = cdiv(n), .y = cdiv(m) }, .block = .{ .x = 256 } }, .{ .r = r, .handle = s }, &args); + } + + /// The longest rounding chain of an output: a group's 64 dot2s (two roundings each at most), one add a group. + pub fn chain(k: usize) usize { + return 2 * (@as(usize, @intCast(group)) / 2) + k / @as(usize, @intCast(group)) + 2; + } + + fn cdiv(v: c_int) c_uint { + return @intCast(@divTrunc(v + side - 1, side)); + } +}; + +test "the dense GEMM refuses what its tile does not take" { + const k: Kernels = .{ .tile = undefined, .bf16 = true }; + try std.testing.expect(k.refusal(0, 0, 3, 5, 256) == null); + try std.testing.expect(k.refusal(0, 0, 3, 5, 200) != null); + try std.testing.expect(k.refusal(8, 0, 3, 5, 256) != null); + try std.testing.expect(k.refusal(0, 8, 3, 5, 256) != null); + try std.testing.expect(k.refusal(0, 0, 0, 5, 256) != null); + const half: Kernels = .{ .tile = undefined, .bf16 = false }; + try std.testing.expect(half.refusal(0, 0, 3, 5, 256) != null); +} diff --git a/zig/src/hip/launches/elementwise.zig b/zig/src/hip/launches/elementwise.zig index 4f767f37fd..4dcdb77a22 100644 --- a/zig/src/hip/launches/elementwise.zig +++ b/zig/src/hip/launches/elementwise.zig @@ -28,9 +28,10 @@ pub fn tf_embed_rows(l: *const Launcher, words: C, scale: C, bias: C, scale_kind try l.flat(l.op.embed_rows, @as(i64, n) * k, s, &a); } -pub fn tf_embed_dense(l: *const Launcher, table: CF, ids: CI, n: c_int, k: c_int, out: P, out_kind: c_int, s: S) Error!void { +pub fn tf_embed_dense(l: *const Launcher, table: C, table_kind: c_int, ids: CI, n: c_int, k: c_int, out: P, out_kind: c_int, s: S) Error!void { var a: Args = .{}; try a.add(ad(table)); + try a.add(table_kind); try a.add(ad(ids)); try a.add(n); try a.add(k); diff --git a/zig/src/hip/ops/attention.zig b/zig/src/hip/ops/attention.zig index 7e356c4e4c..132e2d29e3 100644 --- a/zig/src/hip/ops/attention.zig +++ b/zig/src/hip/ops/attention.zig @@ -14,6 +14,11 @@ pub fn attnGate(o: Ops, att: u64, qg: Tensor, out: Tensor, len: usize, heads: us try o.l.tf_attn_gate(f(att), p(qg.ptr), p(out.ptr), @backingInt(out.kind), int(len), int(heads), int(d), @intFromBool(rows_major), o.stream); } +/// The causal kernel's fp32 (heads, len, d) rows as (len, heads * d) of `out`'s kind, rounded once: o_proj's input. +pub fn headsRows(o: Ops, att: u64, out: Tensor, len: usize, heads: usize, d: usize) Error!void { + try o.l.tf_heads_rows(f(att), p(out.ptr), @backingInt(out.kind), int(len), int(heads), int(d), o.stream); +} + /// Prefill RoPE of x (len, heads, d), rounded to x's kind, stored as `out` kind at out[h * s_head + r * s_row + j]. pub fn ropePrefill(o: Ops, x: Tensor, out: Tensor, s_head: usize, s_row: usize, len: usize, heads: usize, d: usize, rotary: usize, pos0: usize, theta: f32) Error!void { try o.l.tf_rope_prefill(p(x.ptr), @backingInt(x.kind), p(out.ptr), @backingInt(out.kind), @intCast(s_head), @intCast(s_row), int(len), int(heads), int(d), int(rotary), int(pos0), theta, o.stream); diff --git a/zig/src/hip/ops/ops.zig b/zig/src/hip/ops/ops.zig index 8c842f4e2f..98907749f1 100644 --- a/zig/src/hip/ops/ops.zig +++ b/zig/src/hip/ops/ops.zig @@ -48,6 +48,7 @@ pub const Ops = struct { pub const affineRouted = products.affineRouted; pub const embedAffine = products.embedAffine; pub const denseRows = products.denseRows; + pub const denseGemm = products.denseGemm; /// Whether decode.hip's merged launches are on. pub fn fused(o: Ops) bool { @@ -55,6 +56,7 @@ pub const Ops = struct { } pub const attnGate = attention.attnGate; + pub const headsRows = attention.headsRows; pub const ropePrefill = attention.ropePrefill; pub const qkRope = attention.qkRope; pub const ropeDecode = attention.ropeDecode; diff --git a/zig/src/hip/ops/project.zig b/zig/src/hip/ops/project.zig index a4d7e38914..7638a75337 100644 --- a/zig/src/hip/ops/project.zig +++ b/zig/src/hip/ops/project.zig @@ -30,11 +30,21 @@ fn affineOf(proj: Projection, h: mlx.Matrix) Affine { pub fn project(o: Ops, x: Tensor, proj: Projection, m: usize, f32_out: bool) Error!Tensor { switch (proj.handle) { .mlx => |h| return affine(o, x, affineOf(proj, h), m, f32_out), - .dense => |h| { - if (f32_out or m == 0 or x.kind == .f32) return error.BadShape; - const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * x.kind.size()), .kind = x.kind }; - try denseRows(o, x, h.weight, out, m, proj.n, proj.k); - return out; + .dense => |h| switch (h.dtype) { + .f32 => { + if (f32_out or m == 0 or x.kind == .f32) return error.BadShape; + const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * x.kind.size()), .kind = x.kind }; + try denseRows(o, x, h.weight, out, m, proj.n, proj.k); + return out; + }, + .bf16 => { + if (m == 0 or x.kind != .bf16) return error.BadShape; + const kind: t.Kind = if (f32_out) .f32 else .bf16; + const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * kind.size()), .kind = kind }; + try denseGemm(o, x, h.weight, out, m, proj.n, proj.k); + return out; + }, + else => return error.BadShape, }, } } @@ -200,7 +210,7 @@ pub fn projectRouted(o: Ops, x: Tensor, proj: Projection, items: u64, count: usi pub fn embedRows(o: Ops, table: Projection, ids: u64, n: usize, out: Tensor) Error!void { switch (table.handle) { .mlx => |h| return embedAffine(o, affineOf(table, h), ids, n, out), - .dense => |h| try o.l.tf_embed_dense(f(h.weight), @ptrFromInt(ids), int(n), int(table.k), p(out.ptr), @backingInt(out.kind), o.stream), + .dense => |h| try o.l.tf_embed_dense(p(h.weight), @backingInt(try kindOf(h.dtype)), @ptrFromInt(ids), int(n), int(table.k), p(out.ptr), @backingInt(out.kind), o.stream), } } @@ -351,3 +361,21 @@ pub fn embedAffine(o: Ops, table: Affine, ids: u64, n: usize, out: Tensor) Error pub fn denseRows(o: Ops, x: Tensor, w: u64, out: Tensor, rows: usize, n: usize, k: usize) Error!void { try o.l.tf_dense_rows(p(x.ptr), @backingInt(x.kind), f(w), p(out.ptr), int(rows), int(n), int(k), o.stream); } + +/// x (rows, k) BF16 times BF16 weights [n, k] on the dot2 GEMM tile: `out` fp32, or BF16 rounded once. +pub fn denseGemm(o: Ops, x: Tensor, w: u64, out: Tensor, rows: usize, n: usize, k: usize) Error!void { + if (x.kind != .bf16 or out.kind == .f16) return error.BadShape; + const wide = out.kind == .f32; + if (o.l.dense.refusal(x.ptr, w, int(rows), int(n), int(k)) != null) return error.BadShape; + try o.l.dense.gemm(o.l.r, x.ptr, w, if (wide) out.ptr else 0, if (wide) 0 else out.ptr, int(rows), int(n), int(k), o.stream); +} + +/// The kernels' element kind of a dense weight's dtype. +fn kindOf(dtype: quant.DType) Error!t.Kind { + return switch (dtype) { + .f32 => .f32, + .bf16 => .bf16, + .f16 => .f16, + else => error.BadShape, + }; +} diff --git a/zig/src/native/hip.zig b/zig/src/native/hip.zig index 8a922cdcc6..3a74134408 100644 --- a/zig/src/native/hip.zig +++ b/zig/src/native/hip.zig @@ -9,12 +9,13 @@ const flashnext = @import("flashnext"); const nemotron = @import("nemotron"); const serve = @import("hip_serve"); const glm_ocr = @import("glm_ocr"); +const qwen3 = @import("qwen3"); const Allocator = std.mem.Allocator; const admission = hip.admission; const modelContext = admission.modelContext; /// The HIP families: namespaces with `model_type`, `formats`, `default_context`, `prefill_step`, `open` and `follow`. -const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native }; +const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native, qwen3.native }; pub const backends: []const []const u8 = &.{"hip"}; /// No per-thread alternate signal stack: its thread-local leaves RCCL's and HIP's threads too little stack to start. @@ -310,7 +311,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O h.clock = .{ .io = io }; h.core = lanes.Engine.init(gpa, &h.cfg, loaded.backend, h.clock.clock()); errdefer h.core.deinit(); - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .benchmark_ranks = if (r.dev.group) |g| @intCast(g.world) else 1, .lanes = @intCast(loaded.streams), .context_window = @intCast(r.options.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = true, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); + const what = Serves.of(loaded, r.options.window); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .benchmark_ranks = if (r.dev.group) |g| @intCast(g.world) else 1, .lanes = @intCast(loaded.streams), .context_window = @intCast(what.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = what.text_generation, .text_generation = what.text_generation, .embeddings = what.embeddings, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); h.host.memory = .{ .read = readMemory }; if (@hasField(@TypeOf(loaded), "cache")) { h.host.cache = loaded.cache; @@ -322,6 +324,24 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// What a family's `Loaded` says it serves; a family that names none generates text in the window asked. +const Serves = struct { + text_generation: bool = true, + embeddings: ?api.EmbeddingInfo = null, + window: usize, + + fn of(loaded: anytype, asked: usize) Serves { + const L = @TypeOf(loaded); + var s: Serves = .{ .window = asked }; + if (@hasField(L, "text_generation")) s.text_generation = loaded.text_generation; + if (@hasField(L, "embeddings")) s.embeddings = loaded.embeddings; + if (@hasField(L, "window")) if (loaded.window) |w| { + s.window = w; + }; + return s; + } +}; + /// A rank above 0: holds its share and runs rank 0's steps until rank 0 stops; false with `problem` if it cannot start. pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { const report = @import("sys").optional_libraries.atBoundary(); @@ -366,6 +386,7 @@ test "every registered family is listed for capabilities" { try std.testing.expectEqualStrings("qwen4_exp", families[2].model_type); try std.testing.expectEqualStrings("nemotron_h", families[3].model_type); try std.testing.expectEqualStrings("glm_ocr", families[4].model_type); + try std.testing.expectEqualStrings("qwen3", families[5].model_type); } test "GLM-OCR native flags: MTP and tensor ranks refuse before readiness; copy drafts remain" { @@ -386,6 +407,17 @@ test "Nemotron refuses --slide and --mtp-confidence on HIP before it loads, and try std.testing.expect(flagRefused(nemotron.native, .{ .dir = "d", .model_type = "nemotron_h" }) == null); } +test "a one-GPU family refuses --tp before the group forms; the others are not refused for it" { + const OneGpu = struct { + pub const tp_refusal = "--tp: one GPU only"; + }; + try std.testing.expectEqualStrings("--tp: one GPU only", flagRefused(OneGpu, .{ .dir = "d", .model_type = "m", .tp = 2 }).?); + try std.testing.expect(flagRefused(OneGpu, .{ .dir = "d", .model_type = "m" }) == null); + try std.testing.expect(flagRefused(nemotron.native, .{ .dir = "d", .model_type = "nemotron_h", .tp = 2 }) == null); + try std.testing.expectEqualStrings(qwen3.native.tp_refusal, flagRefused(qwen3.native, .{ .dir = "d", .model_type = "qwen3", .tp = 2, .rank = 1 }).?); + try std.testing.expect(!reasoning(qwen3.native)); +} + test "a family that names a --slide or --mtp-confidence refusal is refused before it loads; the others are not" { const Refuses = struct { pub const slide_refusal = "--slide: this family does not learn on HIP"; @@ -482,3 +514,14 @@ test "slide admission: HIP followers refuse before joining ranks" { try std.testing.expectEqualStrings(try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on HIP yet", .{f.model_type}), why); } } + +test "a family's Loaded names what it serves; one without the fields generates text in the asked window" { + const Plain = struct { streams: usize }; + const p = Serves.of(Plain{ .streams = 1 }, 4096); + try std.testing.expect(p.text_generation and p.embeddings == null and p.window == 4096); + const Embed = struct { text_generation: bool, embeddings: ?api.EmbeddingInfo, window: ?usize }; + const e = Serves.of(Embed{ .text_generation = false, .embeddings = .{ .dimensions = 8, .vocabulary = 4, .max_batch_tokens = 2048 }, .window = 2048 }, 32768); + try std.testing.expect(!e.text_generation and e.embeddings.?.dimensions == 8 and e.window == 2048); + const Unsized = struct { window: ?usize }; + try std.testing.expectEqual(@as(usize, 32768), Serves.of(Unsized{ .window = null }, 32768).window); +} diff --git a/zig/src/server/embedding_protocol_test.zig b/zig/src/server/embedding_protocol_test.zig index c2f175dd9b..af3dba7369 100644 --- a/zig/src/server/embedding_protocol_test.zig +++ b/zig/src/server/embedding_protocol_test.zig @@ -117,3 +117,34 @@ test "embeddings protocol: float and little-endian base64 preserve MRL rows and var nan = [_]f32{ 3, 4, std.math.nan(f32) }; try t.expectError(error.Refused, response.encode(&cx, "m", parsed, &nan)); } + +test "embeddings protocol: whitespace-only text is refused with its index; an over-long input names its length" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + var cx: errors.Cx = .{ .a = arena.allocator() }; + try t.expectError(error.Refused, parse(&cx, "{\"input\":\" \\t\\n\"}", cap)); + try t.expectEqualStrings("embedding text must not be empty or whitespace only", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"a\",\" \"]}", cap)); + try t.expectEqualStrings("embedding input at index 1 must not be empty or whitespace only", cx.message); + // the context is 4 tokens here: a 5-token input names itself, not the batch + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"abcde\"}", cap)); + try t.expectEqual(errors.Kind.context_length, cx.kind); + try t.expectEqualStrings("the input (5 tokens) is longer than the model's context length (4 tokens)", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[1,2,3,4,5]}", cap)); + try t.expectEqual(errors.Kind.context_length, cx.kind); + // past both the context (4) and the batch (8): the input's own length is what is refused, counted whole + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"abcdefghi\"}", cap)); + try t.expectEqualStrings("the input (9 tokens) is longer than the model's context length (4 tokens)", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[1,2,3,4,5,6,7,8,9]}", cap)); + try t.expectEqualStrings("the input (9 tokens) is longer than the model's context length (4 tokens)", cx.message); + // Unicode spaces are spaces, as Python's strip() takes them; a letter among them is text + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"\\u00a0\\u2003\\u3000\"}", cap)); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"a\",\"\\u2003\"]}", cap)); + try t.expectEqualStrings("embedding input at index 1 must not be empty or whitespace only", cx.message); + _ = try parse(&cx, "{\"input\":\"\\u00a0a\"}", cap); + // inputs within the context that together pass the batch's 8 tokens + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"abcd\",\"efgh\",\"i\"]}", cap)); + try t.expectEqualStrings("the embedding batch exceeds the admitted 8 tokens", cx.message); + const full = try parse(&cx, "{\"input\":[\"abcd\",\"efgh\"]}", cap); + try t.expectEqual(@as(usize, 8), full.tokens.len); +} diff --git a/zig/src/server/embedding_request.zig b/zig/src/server/embedding_request.zig index 20dd00fef2..82120ba7f4 100644 --- a/zig/src/server/embedding_request.zig +++ b/zig/src/server/embedding_request.zig @@ -41,8 +41,9 @@ pub fn parse(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, body: Va try sequence(cx, text, cap, context, input, &tokens, &ends); } else if (first == .string or first == .array) { if (input.array.len > cap.max_inputs) return cx.refuse("too many embedding inputs"); - for (input.array) |item| { + for (input.array, 0..) |item, i| { if (@as(std.meta.Tag(Value), item) != @as(std.meta.Tag(Value), first)) return cx.refuse("embedding input forms cannot be mixed"); + if (item == .string and blank(item.string)) return cx.fail(.request, "embedding input at index {d} must not be empty or whitespace only", .{i}); try sequence(cx, text, cap, context, item, &tokens, &ends); } } else return cx.refuse("input must contain strings or token IDs"); @@ -56,14 +57,16 @@ fn sequence(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, input: Va if (ends.items.len >= cap.max_inputs) return cx.refuse("too many embedding inputs"); const at = tokens.items.len; if (input == .string) { - if (input.string.len == 0) return cx.refuse("embedding text must not be empty"); + if (blank(input.string)) return cx.refuse("embedding text must not be empty or whitespace only"); const ids = text.encode(cx.a, input.string, true) catch |err| return switch (err) { error.OutOfMemory => error.OutOfMemory, error.Template => cx.refuse("the tokenizer cannot encode this embedding input"), }; defer cx.a.free(ids); + try fits(cx, context, ids.len); for (ids) |id| try token(cx, cap, id, tokens); } else if (input == .array) { + try fits(cx, context, input.array.len); for (input.array) |item| { if (item != .int) return cx.refuse("token IDs must be integers"); const id = std.fmt.parseInt(u32, item.int, 10) catch return cx.refuse("token ID is outside the vocabulary"); @@ -72,12 +75,32 @@ fn sequence(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, input: Va } else return cx.refuse("each embedding input must be text or token IDs"); const count = tokens.items.len - at; if (count == 0) return cx.refuse("embedding input must contain tokens"); - if (context > 0 and count > context) return cx.fail(.context_length, "embedding input exceeds the {d}-token context", .{context}); try ends.append(cx.a, @intCast(tokens.items.len)); } +/// One input longer than the context is refused as such, before the batch's own limit. +fn fits(cx: *Cx, context: u32, count: usize) errors.Refused!void { + if (context > 0 and count > context) return cx.fail(.context_length, "the input ({d} tokens) is longer than the model's context length ({d} tokens)", .{ count, context }); +} + +/// Empty or whitespace only, by Python's str.isspace (as SGLang strips): nothing to embed; invalid UTF-8 is not blank. +fn blank(s: []const u8) bool { + const view = std.unicode.Utf8View.init(s) catch return false; + var it = view.iterator(); + while (it.nextCodepoint()) |c| if (!space(c)) return false; + return true; +} + +fn space(c: u21) bool { + return switch (c) { + 0x09...0x0d, 0x1c...0x20, 0x85, 0xa0, 0x1680, 0x2000...0x200a, 0x2028, 0x2029, 0x202f, 0x205f, 0x3000 => true, + else => false, + }; +} + fn token(cx: *Cx, cap: api.EmbeddingInfo, id: u32, tokens: *std.ArrayList(u32)) errors.Refused!void { if (id >= cap.vocabulary) return cx.refuse("token ID is outside the vocabulary"); - if (tokens.items.len >= cap.max_batch_tokens) return cx.refuse("embedding batch exceeds the admitted token limit"); + // the input's own length was checked before: here only the whole batch can pass the admitted tokens + if (tokens.items.len >= cap.max_batch_tokens) return cx.fail(.request, "the embedding batch exceeds the admitted {d} tokens", .{cap.max_batch_tokens}); try tokens.append(cx.a, id); } diff --git a/zig/tests/hip/contract_dense.zig b/zig/tests/hip/contract_dense.zig new file mode 100644 index 0000000000..28a4f7e8a6 --- /dev/null +++ b/zig/tests/hip/contract_dense.zig @@ -0,0 +1,114 @@ +//! dense_gemm_block under the contract on HIP: BF16 x times raw BF16 weights on the dot2 GEMM tile, fp32 and BF16 out. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; + +/// The kernel's own chain (launches/dense.zig): a group's dot2s, then one add a group into the accumulator. +fn depth(s: c.Shape) usize { + return @FieldType(hip.Launcher, "dense").chain(s.k); +} + +/// Qwen3-8B's projections as float64 runs them: 2,048-row passes, 12,288-wide gate/up, K to 12,288, both tile edges. +const served = [_]c.Shape{ + .{ .rows = 1, .n = 4096, .k = 4096 }, + .{ .rows = 7, .n = 1024, .k = 12288 }, + .{ .rows = 16, .n = 12288, .k = 4096 }, + .{ .rows = 129, .n = 130, .k = 4096 }, + .{ .rows = 300, .n = 256, .k = 4096 }, + .{ .rows = 2048, .n = 128, .k = 4096 }, +}; + +pub fn spec(wide: bool) c.Spec { + return .{ + .name = if (wide) "hip dense_gemm_block (bf16 weights, fp32 out)" else "hip dense_gemm_block (bf16 weights, bf16 out)", + .accum = .f32, + .depth = depth, + .out = if (wide) .f32 else .bf16, + .served = &served, + .tiles = .{ .{ .size = 128 }, .{ .size = 128 }, .{ .size = 12288, .step = 128, .min = 128 } }, + }; +} + +pub const DenseGemm = struct { + rig: rig_.Rig, + l: *const hip.Launcher, + /// fp32 out, or BF16 rounded once in the epilogue. + wide: bool, + d: ?Dense = null, + x: u64 = 0, + w: u64 = 0, + out: u64 = 0, + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, wide: bool) DenseGemm { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .wide = wide }; + } + + pub fn deinit(self: *DenseGemm) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *DenseGemm) g.Kernel { + return .{ .spec = spec(self.wide), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(DenseGemm); + + /// x from the dense family; W as BF16 in [-1, 1), or the stress fill's mirrored integers; FP16 builds serve none. + pub fn prepare(self: *DenseGemm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 128 != 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.a; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 1); + const bits = try a.alloc(u16, d.w.len); + defer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed ^ 0xD3); + const rng = prng.random(); + const k = shape.k; + for (bits) |*b| b.* = c.toBf16(rng.float(f32) * 2 - 1); + if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + bits[j * k + i] = c.toBf16(v); + bits[j * k + k / 2 + i] = c.toBf16(-v); + }; + if (fill == .nonfinite and shape.n > 5) { + bits[1 * k + k / 4] = c.toBf16(std.math.nan(f32)); + bits[5 * k] = c.toBf16(std.math.inf(f32)); + } + for (d.w, bits) |*w, b| w.* = c.fromBf16(b); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.w = try self.rig.input(std.mem.sliceAsBytes(bits)); + self.out = try self.rig.output(d.outputs() * @as(usize, if (self.wide) 4 else 2)); + return true; + } + + pub fn run(self: *DenseGemm) !void { + const s = self.d.?.shape; + const out32: u64 = if (self.wide) self.out else 0; + const out16: u64 = if (self.wide) 0 else self.out; + try self.l.dense.gemm(self.l.r, self.x, self.w, out32, out16, @intCast(s.rows), @intCast(s.n), @intCast(s.k), self.rig.s.handle); + try self.rig.s.synchronize(); + } + + pub fn count(self: *DenseGemm) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *DenseGemm, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + } + + pub fn simulate(self: *DenseGemm, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/hip/contract_heads_rows.zig b/zig/tests/hip/contract_heads_rows.zig new file mode 100644 index 0000000000..0f319f43cc --- /dev/null +++ b/zig/tests/hip/contract_heads_rows.zig @@ -0,0 +1,90 @@ +//! tf_heads_rows under the contract on HIP: the causal kernel's fp32 (heads, len, d) rows to row-major o_proj input. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const kit = @import("contract_kit.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +fn none(_: c.Shape) usize { + return 0; +} + +pub fn spec(k: kit.Kind) c.Spec { + return .{ + .name = if (k == .bf16) "hip tf_heads_rows (bf16)" else "hip tf_heads_rows (fp16)", + .accum = .none, + .depth = none, + .out = kit.elem(k), + // rows: a prompt block (1, a mid one, a whole 2048); n: query heads; k: head width (Qwen3-8B: 32 of 128) + .served = &.{ .{ .rows = 1, .n = 32, .k = 128 }, .{ .rows = 300, .n = 32, .k = 128 }, .{ .rows = 2048, .n = 32, .k = 128 } }, + }; +} + +/// shape.n heads of shape.rows rows of shape.k fp32 values, each to (row, head, j) of the activation kind. +pub const HeadsRows = struct { + rig: rig_.Rig, + l: *const hip.Launcher, + k: kit.Kind, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + att: []f64 = &.{}, + in: u64 = 0, + out: u64 = 0, + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher) HeadsRows { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .k = kit.act(), .arena = .init(a) }; + } + + pub fn deinit(self: *HeadsRows) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *HeadsRows) g.Kernel { + return .{ .spec = spec(self.k), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(HeadsRows); + + pub fn prepare(self: *HeadsRows, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const n = shape.rows * shape.n * shape.k; + // fp32 attention outputs: spread past the activation's step so the rounding is exercised everywhere + self.att = try kit.floats(a, prng.random(), n, -40, 40); + if (fill == .nonfinite) for ([_]f64{ std.math.nan(f64), std.math.inf(f64), -std.math.inf(f64) }, 0..) |v, i| { + self.att[(i * 97 + 3) % n] = v; + }; + self.in = try self.rig.input(try kit.f32Bytes(a, self.att)); + self.out = try self.rig.output(n * kit.size(self.k)); + return true; + } + + pub fn run(self: *HeadsRows) !void { + const s = self.shape; + try self.l.tf_heads_rows(kit.ptr(?[*]const f32, self.in), kit.ptr(?*anyopaque, self.out), @backingInt(self.k), @intCast(s.rows), @intCast(s.n), @intCast(s.k), self.rig.s.handle); + try self.rig.s.synchronize(); + } + + pub fn count(self: *HeadsRows) usize { + return self.shape.rows * self.shape.n * self.shape.k; + } + + /// Output (r, h, j) is input (h, r, j), exactly; the bound adds the output's one rounding. + pub fn reference(self: *HeadsRows, want: []f64, mag: []f64) !void { + const s = self.shape; + for (0..s.rows) |r| for (0..s.n) |h| for (0..s.k) |j| { + const v = self.att[(h * s.rows + r) * s.k + j]; + want[(r * s.n + h) * s.k + j] = v; + mag[(r * s.n + h) * s.k + j] = @abs(v); + }; + } +}; diff --git a/zig/tests/hip/contract_limits.zig b/zig/tests/hip/contract_limits.zig index 23e11b25ee..56ae3c123b 100644 --- a/zig/tests/hip/contract_limits.zig +++ b/zig/tests/hip/contract_limits.zig @@ -72,6 +72,11 @@ pub fn live(a: std.mem.Allocator, r: *const hip.Runtime, ctx: hip.Context, arch: std.debug.print(say, .{ word, arch, d[0], got, d[2] }); } for (all) |rec| { + // a planted control is checked on every chip: only a registered launch may be unserved somewhere + if (!rec.usage[at].served and !isControl(rec)) { + std.debug.print("N/A limits live {s}: not served on {s}\n", .{ rec.launch.name, arch }); + continue; + } const bytes = try image(rec.launch.image); const listed = try hsaco.kernels(a, bytes); defer a.free(listed); @@ -108,6 +113,11 @@ pub fn objects(a: std.mem.Allocator, io: std.Io, dir: []const u8) !usize { var line: std.ArrayList(u8) = .empty; defer line.deinit(a); for (limits.rdna, 0..) |t, i| { + if (!rec.usage[i].served and !control) { + std.debug.print("N/A limits {s} {s}: not served there\n", .{ t.name, rec.launch.name }); + try line.appendSlice(a, ".{ .vgprs = 0, .served = false }, "); + continue; + } const path = try file(a, dir, t.name, rec.launch.image); defer a.free(path); const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 28)) catch |err| { diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 32657d7a93..e2b159baf2 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -18,6 +18,8 @@ const relu2_k = @import("contract_relu2.zig"); const kv_k = @import("contract_kv_write.zig"); const argmax_k = @import("contract_argmax.zig"); const pair_k = @import("contract_relu2_pair.zig"); +const heads_rows_k = @import("contract_heads_rows.zig"); +const dense_k = @import("contract_dense.zig"); const gate_k = @import("contract_moe_gate.zig"); const router_k = @import("contract_router.zig"); const kit = @import("contract_kit.zig"); @@ -152,6 +154,16 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { var pair = pair_k.Pair.init(a, &r, s, &l); defer pair.deinit(); try kernels.append(a, pair.kernel()); + var heads_rows = heads_rows_k.HeadsRows.init(a, &r, s, &l); + defer heads_rows.deinit(); + try kernels.append(a, heads_rows.kernel()); + var dense: [2]dense_k.DenseGemm = undefined; + for (&dense, [_]bool{ true, false }) |*d, fp32_out| d.* = .init(a, &r, s, &l, fp32_out); + defer for (&dense) |*d| d.deinit(); + // the BF16 dense GEMM runs on BF16 builds only; an FP16 build keeps its trap body and serves no shape of it + if (l.dense.bf16) { + for (&dense) |*d| try kernels.append(a, d.kernel()); + } else std.debug.print("N/A hip dense_gemm_block: not served on this FP16 build\n", .{}); var gates: [3]gate_k.Gate = undefined; for (&gates, std.enums.values(gate_k.Which)) |*m, w| m.* = .init(a, &r, s, &l, w); defer for (&gates) |*m| m.deinit();