diff --git a/tools/dsv41-zig/eagle/README.md b/tools/dsv41-zig/eagle/README.md new file mode 100644 index 0000000000..e669375a13 --- /dev/null +++ b/tools/dsv41-zig/eagle/README.md @@ -0,0 +1,116 @@ +# The TFE draft head (DeepSeek-V4.1, Zig engine) + +The TFE head is an EAGLE-style draft head. It is trained on the target's hidden states (the EAGLE study's +`eagle_train.py` Head; the ideas come from the public EAGLE papers, not their code). The engine code is in +`zig/src/families/dsv41/eagle.zig` (the pass), `eagle_file.zig` (the head file and its load) and `eagle_check.zig` +(the self-test). The kernels are in `zig/kernels/cuda/dsv41/eagle.cu`. + +## Switches + +| Variable | Effect | +|---|---| +| `TF_DS_DRAFTER=eagle` | The TFE head drafts and no rank loads DSpark. Unset or `dspark`: DSpark. All ranks must use the same value. | +| `TF_DS_EAGLE_HEAD=FILE` | The head file (`eagle_convert.py` output). | +| `TF_DS_EAGLE_FP8` | Unset: FP8 linears when the file has the FP8 forms. `1`: FP8 (the forms must be there). `0`: bf16. | +| `TF_DS_EAGLE_GEMV` | Linears on the own GEMV: `wo` (default: wo and w2), `all`, or `cublas`. | +| `TF_DS_EAGLE_GRAPH=1` | The steps of a pass as one CUDA graph for each (streams, steps). | +| `TF_DS_EAGLE_TOPK=K` | two_stage heads: K candidates (at most 512) in place of the file's `cand_k`. | +| `TF_DS_EAGLE_CONF=a,b` | The confidence `conf = a * log(p / (1 - p)) + b`. | +| `TF_DS_EAGLE_CHECK=FILE` | Before serving: the head against `eagle_ref.py vectors` output. | +| `TF_DS_EAGLE_TAPDUMP=FILE` | Every absorb writes its tap rows (see "Tap dump"). | + +## The head + +The head runs on rank 0 only. Rank 0 sends every draft to rank 1 in the verify frame, so rank 1 runs no head and a +pass has no collective. + +For a row at position p, with the target's taps at p (the tap layers side by side, the final state last) and the next +token x_{p+1}: + +1. g = RMSNorm(fc(taps)). +2. x = w_in [RMSNorm(embed(x_{p+1})) ; RMSNorm(g)]. +3. One decoder layer: heads x 128 queries, one key and value head. With qk_norm, an RMSNorm on each query and key + head. RoPE on 64 of the 128 dims. Attention over the stream's newest window + 1 positions. Then a SwiGLU MLP. +4. logits = up(down(RMSNorm(o))). + +A two_stage head takes the top K of those logits (K = `cand_k`, or `TF_DS_EAGLE_TOPK`). It drafts the argmax of their +exact logits from the target's own head (a dense bf16 copy in the head file, on rank 0). + +## The chain + +A stream's chain starts from its pending token at q0: + +- The real rows below q0 have their keys and values in the stream's ring. +- Draft 1 comes from the row q0 - 1. +- Draft j > 1 comes from a row at position q0 + j - 2, with g = the step before's o and the token draft j - 1. Its key + and value also go into the ring. The next pass's real rows write over them. + +`eagle_ref.py` is this order of work in torch (the self-test's reference). It is checked against the training code's +chain (`eagle_ref.py cross`). + +The g rows come from each round (absorbRound: the drafting streams' window rows) and from each prompt chunk that +reaches the head (its last ring rows). A pass first runs the rows from the stream's last extended position to q0 - 1 +through the layer's input part (their keys and values). Then it runs `steps` single-row steps for all streams at once. + +A prompt chunk drops its rows below the replay floor (the prompt length - the model window) at the middle layer. A tap +of a layer below layers / 2 keeps the chunk's rows from its first row, so the absorb reads it at the offset +`ch.start - start0`. + +## Linears and FP8 + +The linears are bf16 cuBLAS calls or, for up to 16 rows, the own GEMV of `eagle.cu`. Both sum in fp32. + +With FP8, each linear (fc, w_in, wq, wkv, wo, w13, w2, down, up) is e4m3 with an fp32 scale for each output row (the +file's `NAME.fp8` and `NAME.scale`, made by `eagle_convert.py`). FP8 linears always use the own GEMV. The norms, the +target head's copy and the logits do not change. + +The FP8 rule: s = amax(|W| of the row) / 448 in fp32 (1 if the row is 0); q = e4m3 round-to-nearest-even of +fp32(W) / s; the weight is fp32(q) * s. + +## Kernels (eagle.cu) + +- RMSNorm: a row a block (512 threads). +- RoPE: in place on the first rd features of each head of row r at position pos[r]. The halves are rotated with the + frequency theta^(-i / (rd / 2)). A row a block. +- Attention: one row and one query head a block (grid rows x heads, 128 threads; head_dim 128). It reads the keys and + values of the row's stream ring (K, V bf16 [ring rows, kv_heads * 128], the stream's rows from base[r]) at the + positions max(lo[r], pos[r] - window) .. pos[r] (ring row base + position % ring). The scores are fp32 times + `scale`, then softmax, then the fp32 sum of the values, bf16 out [rows, heads * 128]. A thread computes the score of + one key. For the values, each warp sums a quarter of the keys (a lane 4 features), then the 4 warp sums are added. +- top1: a row a block (1024 threads). The largest logit of [n] fp32 (a tie goes to the lower token id), its softmax + probability p over the n, and conf = a * log(p / (1 - p)) + b (the lane code reads sigmoid(conf) as the draft's + survival). map: null (logit i is token i), or the row's token ids int64 [n] (a two_stage head's candidates). +- top1_part and top1_join: the top-1 of long rows (the low-rank head's 129280 logits). The part kernel (grid rows x + parts, 256 threads) finds each part's largest logit (a tie: the lower id) and its sum of exp(logit - that max). The + join kernel (grid rows, 32 threads) joins the parts and writes the same outputs as top1. Part buffers: pm, pz fp32 + and pi int32 [rows, parts]. +- gemv_bf16 and gemv_fp8: y[r, o] = scale[o] * sum_i x[r, i] * W[o, i] for up to 16 rows (fp32 sums). A group of G + lanes computes one output (G = 32, or the number of 16-byte chunks in a row when that is less). A lane reads every + G-th chunk. Mode 0: out bf16 = y. Mode 1: out fp32 = y. Mode 2: out bf16 = out + y with one rounding (as cuBLAS beta + 1). x is bf16 [rows, k] with 16-byte aligned rows. 256 threads. +- topk: a two_stage head's candidates, a row a block (1024 threads). The k largest of [n] fp32 logits; a tie at the + k-th value goes to the lower token ids. The output is int64 ids in cand [rows, k] (the order in a row is not fixed). + A radix select on the keys (4 passes of 8 bits), then the keys above the k-th and the first of the equal ones in id + order. +- rescore: a two_stage head's exact logits, out[r, j] = W[cand[r, j]] . x[r] (W = the target head bf16 [vocab, d], + fp32 sums), fp32 out [rows, k]. Grid rows x ceil(k / 8), 256 threads: a warp a candidate. + +## Self-test (TF_DS_EAGLE_CHECK) + +The self-test uses slot 0 before serving and resets the slot after. For each case it absorbs the rows, runs one pass, +and compares each step's o, draft and top-1 probability with the vector file. It writes one JSON line for each case +and a verdict on stderr. Pass: every step-1 o is within 3% (max abs difference over max abs), and every draft is equal +up to the first step whose reference top-2 logit gap is less than 0.05 (there, a rounding can select either token). + +## Tap dump (TF_DS_EAGLE_TAPDUMP) + +Each absorb writes one record: a u32 header length, the JSON header `{kind, slot, start, n, cut, taps, dim}`, then the +rows' taps as bf16 [n, taps * dim] in the engine's tap order, before tap_norm. `kind` is `prompt` or `round`. `cut` is +the number of rows that a prompt chunk dropped at its replay floor. `tfe_tapcheck.py` reads these files. + +## Limit: head windows past the model window + +A prompt gives the head only its last `sliding_window` rows (128 for V4.1): the rows below the replay floor do not +run the layers that the taps come from. A head with a longer window (for example 511) starts each prompt with 128 keys +and gets its full window only after more tokens. The engine logs a warning at load for such a head. Heads trained with +`--serve-floor 128` see the same history in training. diff --git a/tools/dsv41-zig/eagle/eagle_convert.py b/tools/dsv41-zig/eagle/eagle_convert.py new file mode 100644 index 0000000000..13d0750464 --- /dev/null +++ b/tools/dsv41-zig/eagle/eagle_convert.py @@ -0,0 +1,218 @@ +"""A TFE draft head checkpoint (eagle_train.py's HEAD.pt: {"cfg", "state"}) as the file the Zig engine loads +(TF_DS_DRAFTER=eagle TF_DS_EAGLE_HEAD=FILE): every tensor in bf16, row-major, after one JSON header, and each linear +also in e4m3 with an fp32 scale a row (the engine's TF_DS_EAGLE_FP8=1 form). + + python3 eagle_convert.py HEAD.pt OUT.tfe [--target-head PREFIX] [--no-fp8] + +--target-head: a two_stage head's rescoring table, the target's own vocabulary head from eagle_collect.py --head-dump +(PREFIX-r0.pt, PREFIX-r1.pt: {"rank", "wt" [dim, vocab / 2], "norm"}), as tensor "target_head" bf16 [vocab, dim] +(the ranks' halves in rank order). The header's "target_head" names the files, their sha256 and shapes. + +File: b"TFEHEAD1", u64 little-endian header length, the JSON header (padded with spaces so the data starts at a +multiple of 4096), then the tensors (each at a multiple of 256 bytes from the data start). Header fields: the head's +shape (dim, heads, kv_heads, head_dim, inter, rank, vocab, window, rope_dim, rope_theta, eps), tap_norm, qk_norm, +head_type ("lowrank": up(down(o)) picks the draft; "two_stage": its top cand_k ids rescored by the target's own head), +cand_k, `taps` (the +tap layers in the engine's order: layers ascending, the final hidden state 1000 last) and `tensors` {name: {offset, +shape, dtype}}. + +The FP8 form of a linear NAME.weight [n, k] (bf16 W): s = amax(|W| a row) / 448 in fp32 (1 where the row is zero), +q = e4m3 (round to nearest even) of fp32(W) / s, tensors NAME.fp8 (dtype f8e4m3, [n, k]) and NAME.scale (f32 [n]); the +weight it stands for is fp32(q) * s. The header's "fp8" records the rule, the linears and the sha256 of each q and s. + +The engine writes its taps in that order. The training order (cfg["tap_layers"], the order of --layers) can differ, so +fc's input column blocks (and tap_norms' rows) are put in the engine's order here; fc's output is the same. +Only the base head (eagle_train.py Head) is supported; a "wide" head is refused. +""" + +import argparse +import glob +import hashlib +import json +import os +import struct + +import torch + +FINAL = 1000 +MAGIC = b"TFEHEAD1" +LINEARS = ("fc", "w_in", "wq", "wkv", "wo", "w13", "w2", "down", "up") +DTYPES = {"bf16": (torch.bfloat16, torch.int16), "f8e4m3": (torch.float8_e4m3fn, torch.uint8), "f32": (torch.float32, torch.int32)} +NAMES = ("fc_norm", "e_norm", "g_norm", "attn_norm", "mlp_norm", "out_norm", "fc.weight", "w_in.weight", "wq.weight", + "wkv.weight", "wo.weight", "w13.weight", "w2.weight", "down.weight", "up.weight") + + +def engine_order(taps: list[int]) -> list[int]: + """The engine's tap order: layers ascending, the final state last.""" + + return sorted(t for t in taps if t != FINAL) + ([FINAL] if FINAL in taps else []) + + +def sha256(path: str) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for b in iter(lambda: f.read(1 << 24), b""): + h.update(b) + return h.hexdigest() + + +def target_head(prefix: str, vocab: int, dim: int) -> tuple[torch.Tensor, dict]: + """The head dump files of PREFIX as one bf16 [vocab, dim] table and their record.""" + + paths = sorted(glob.glob(f"{prefix}-r*.pt")) + dumps = [(p, torch.load(p, map_location="cpu", weights_only=True)) for p in paths] + dumps.sort(key=lambda x: x[1]["rank"]) + if [d["rank"] for _, d in dumps] != list(range(len(dumps))) or not dumps: + raise SystemExit(f"head dump ranks of {prefix}: {[d['rank'] for _, d in dumps]}") + for p, d in dumps: + if d["wt"].shape[0] != dim: + raise SystemExit(f"{p}: wt {tuple(d['wt'].shape)}, want [{dim}, *]") + lm = torch.cat([d["wt"].T for _, d in dumps]).contiguous().to(torch.bfloat16) + if tuple(lm.shape) != (vocab, dim): + raise SystemExit(f"target head {tuple(lm.shape)}, want ({vocab}, {dim})") + rec = {"source": "eagle_collect.py --head-dump", "shape": [vocab, dim], + "files": [{"name": os.path.basename(p), "rank": d["rank"], "wt": list(d["wt"].shape), "sha256": sha256(p)} + for p, d in dumps], + "sha256_bf16": hashlib.sha256(lm.view(torch.int16).numpy().tobytes()).hexdigest()} + return lm, rec + + +def fp8_rows(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The FP8 form of a bf16 linear [n, k]: (q e4m3 [n, k], s fp32 [n]) with fp32(q) * s standing for w.""" + + wf = w.to(torch.bfloat16).float() + s = wf.abs().amax(1) / 448.0 + s = torch.where(s == 0, torch.ones_like(s), s) + return (wf / s[:, None]).to(torch.float8_e4m3fn), s + + +def raw(t: torch.Tensor) -> bytes: + for name, (dt, iv) in DTYPES.items(): + if t.dtype == dt: + return t.contiguous().view(iv).numpy().tobytes() + raise ValueError(f"dtype {t.dtype}") + + +def dtype_name(t: torch.Tensor) -> str: + return next(n for n, (dt, _) in DTYPES.items() if t.dtype == dt) + + +def convert(ck: dict, target: str | None = None, fp8: bool = True) -> tuple[dict, dict]: + """-> (header, {name: tensor}) of a checkpoint; `target`: the head dump prefix (a two_stage head); fp8: add each + linear's FP8 form.""" + + cfg = ck["cfg"] + if cfg.get("kind", "base") != "base": + raise SystemExit(f"head kind {cfg.get('kind')!r} is not supported (base only)") + sd = ck["state"] + taps = list(cfg["tap_layers"]) + if len(taps) != cfg["layers"] or len(set(taps)) != len(taps): + raise SystemExit(f"bad tap_layers {taps} for {cfg['layers']} layers") + order = engine_order(taps) + perm = [taps.index(t) for t in order] + D = cfg["dim"] + out = {} + for n in NAMES: + out[n] = sd[n].float() + fc = out["fc.weight"].view(D, len(taps), D) + out["fc.weight"] = fc[:, perm].reshape(D, len(taps) * D) + tap_norm = bool(cfg.get("tap_norm", False)) + if tap_norm: + out["tap_norms"] = sd["tap_norms"].float()[perm] + out["fb_norm"] = sd["fb_norm"].float() + qk_norm = bool(cfg.get("qk_norm", False)) + if qk_norm: + out["q_norm"] = sd["q_norm"].float() + out["k_norm"] = sd["k_norm"].float() + head_type = cfg.get("head_type", "lowrank") + cand_k = int(cfg.get("cand_k", 0)) + if head_type not in ("lowrank", "two_stage") or (head_type == "two_stage") != (cand_k > 0): + raise SystemExit(f"head_type {head_type!r} with cand_k {cand_k} is not supported") + if (head_type == "two_stage") != (target is not None): + raise SystemExit("a two_stage head needs --target-head (and only a two_stage head takes it)") + rec = None + if target is not None: + out["target_head"], rec = target_head(target, cfg["vocab"], D) + hd = cfg["head_dim"] + if out["wkv.weight"].shape != (2 * hd, D) or out["wq.weight"].shape != (cfg["heads"] * hd, D): + raise SystemExit("unexpected attention shapes") + head = {"format": 1, "kind": "base", "dim": D, "heads": cfg["heads"], "kv_heads": 1, "head_dim": hd, + "inter": cfg["inter"], "rank": cfg["rank"], "vocab": cfg["vocab"], "window": cfg["window"], + "rope_dim": cfg["rope_dim"], "rope_theta": 10000.0, "eps": 1e-6, "tap_norm": tap_norm, "qk_norm": qk_norm, + "head_type": head_type, "cand_k": cand_k, "taps": order, + "train_taps": taps} + if rec is not None: + head["target_head"] = rec + tensors = {k: v.to(torch.bfloat16).contiguous() for k, v in out.items()} + if fp8: + shas = {} + for n in LINEARS: + q, sc = fp8_rows(tensors[n + ".weight"]) + tensors[n + ".fp8"], tensors[n + ".scale"] = q, sc + shas[n] = {"fp8": hashlib.sha256(raw(q)).hexdigest(), "scale": hashlib.sha256(raw(sc)).hexdigest()} + head["fp8"] = {"rule": "s = amax(|bf16 W| a row) / 448 (1 if 0); q = e4m3 RNE of fp32(W) / s; W ~ fp32(q) * s", + "linears": list(LINEARS), "sha256": shas} + return head, tensors + + +def write(path: str, head: dict, tensors: dict) -> None: + table, off = {}, 0 + for k, t in tensors.items(): + off = (off + 255) // 256 * 256 + table[k] = {"offset": off, "shape": list(t.shape), "dtype": dtype_name(t)} + off += t.numel() * t.element_size() + head = {**head, "tensors": table, "data_bytes": off} + text = json.dumps(head, separators=(",", ":")).encode() + start = (len(MAGIC) + 8 + len(text) + 4095) // 4096 * 4096 + text += b" " * (start - len(MAGIC) - 8 - len(text)) + with open(path, "wb") as f: + f.write(MAGIC + struct.pack(" tuple[dict, dict]: + """The header and the tensors of a converted file.""" + + with open(path, "rb") as f: + b = f.read() + if b[:8] != MAGIC: + raise ValueError("not a TFE head file") + n = struct.unpack(" None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("head") + p.add_argument("out") + p.add_argument("--target-head", default=None, help="a two_stage head: the head dump prefix (PREFIX-r*.pt)") + p.add_argument("--no-fp8", action="store_true", help="leave out the linears' FP8 forms") + a = p.parse_args() + ck = torch.load(a.head, map_location="cpu", weights_only=True) + head, tensors = convert(ck, a.target_head, not a.no_fp8) + write(a.out, head, tensors) + print(json.dumps({"out": a.out, "taps": head["taps"], "train_taps": head["train_taps"], + "head_type": head["head_type"], "cand_k": head["cand_k"], + "target_head": head.get("target_head"), + "fp8": head.get("fp8", {}).get("sha256"), + "params": sum(t.numel() for k, t in tensors.items() if k.endswith(".weight") or k.endswith("_norm") or k == "tap_norms")})) + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/eagle_kernels_test.py b/tools/dsv41-zig/eagle/eagle_kernels_test.py new file mode 100644 index 0000000000..4f7c7afd5c --- /dev/null +++ b/tools/dsv41-zig/eagle/eagle_kernels_test.py @@ -0,0 +1,250 @@ +"""The TFE head's row kernels (zig/kernels/cuda/dsv41/eagle.cu) against torch on the CPU, on any CUDA GPU (the +fatbin built for it: nvcc -fatbin -arch=sm_XX eagle.cu). Driver API through ctypes; one line a kernel. + + python3 eagle_kernels_test.py EAGLE.fatbin +""" + +import ctypes +import math +import sys + +import torch +import torch.nn.functional as F + +cu = ctypes.CDLL("nvcuda.dll" if sys.platform == "win32" else "libcuda.so.1") + + +def ok(r, what): + if r != 0: + raise RuntimeError(f"{what}: CUresult {r}") + + +ok(cu.cuInit(0), "cuInit") +dev = ctypes.c_int() +ok(cu.cuDeviceGet(ctypes.byref(dev), 0), "cuDeviceGet") +ctx = ctypes.c_void_p() +ok(cu.cuDevicePrimaryCtxRetain(ctypes.byref(ctx), dev), "retain") +ok(cu.cuCtxSetCurrent(ctx), "current") +mod = ctypes.c_void_p() +img = open(sys.argv[1], "rb").read() +ok(cu.cuModuleLoadData(ctypes.byref(mod), ctypes.c_char_p(img)), "load") + + +def fn(name): + f = ctypes.c_void_p() + ok(cu.cuModuleGetFunction(ctypes.byref(f), mod, name.encode()), name) + return f + + +def dev_t(t: torch.Tensor): + t = t.contiguous() + p = ctypes.c_uint64() + n = t.numel() * t.element_size() + ok(cu.cuMemAlloc_v2(ctypes.byref(p), ctypes.c_size_t(max(n, 4))), "alloc") + ok(cu.cuMemcpyHtoD_v2(p, ctypes.c_void_p(t.data_ptr()), ctypes.c_size_t(n)), "h2d") + return p + + +def back(p, like: torch.Tensor): + out = torch.empty_like(like).contiguous() + ok(cu.cuMemcpyDtoH_v2(ctypes.c_void_p(out.data_ptr()), p, ctypes.c_size_t(out.numel() * out.element_size())), "d2h") + return out + + +def launch(f, grid, block, args): + vals = [] + for kind, v in args: + vals.append({"p": ctypes.c_uint64, "ll": ctypes.c_longlong, "i": ctypes.c_int, "f": ctypes.c_float}[kind](v)) + arr = (ctypes.c_void_p * len(vals))(*[ctypes.cast(ctypes.byref(v), ctypes.c_void_p) for v in vals]) + g = grid + (1,) * (3 - len(grid)) + ok(cu.cuLaunchKernel(f, g[0], g[1], g[2], block, 1, 1, 0, None, arr, None), "launch") + ok(cu.cuCtxSynchronize(), "sync") + + +r = lambda x: x.to(torch.bfloat16).float() +gen = torch.Generator().manual_seed(0) +res = {} + +# rmsnorm: rows with a stride, into the second half of a wider row +R, D = 7, 5120 +x = (torch.randn(R, 2 * D, generator=gen) * 3).to(torch.bfloat16) +w = (torch.rand(D, generator=gen) + 0.5).to(torch.bfloat16) +o = torch.zeros(R, 2 * D, dtype=torch.bfloat16) +px, pw, po = dev_t(x), dev_t(w), dev_t(o) +launch(fn("tf_eagle_rmsnorm_kernel"), (R,), 512, [("p", px.value), ("ll", 2 * D), ("p", pw.value), ("p", po.value + D * 2), ("ll", 2 * D), ("i", D), ("f", 1e-6)]) +got = back(po, o)[:, D:].float() +xf = x[:, :D].float() +want = r(xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) * w.float()) +res["rmsnorm"] = float((got - want).abs().max() / want.abs().max()) + +# rope: 32 heads x 128, the first 64 rotated, positions up to 300K +H, hd, rd = 32, 128, 64 +q = torch.randn(R, H * hd, generator=gen).to(torch.bfloat16) +pos = torch.tensor([0, 1, 5, 127, 128, 4096, 300001], dtype=torch.long) +pq, pp = dev_t(q), dev_t(pos) +launch(fn("tf_eagle_rope_kernel"), (R,), 256, [("p", pq.value), ("ll", H * hd), ("i", H), ("i", hd), ("i", rd), ("p", pp.value), ("f", 10000.0)]) +got = back(pq, q).float().view(R, H, hd) +half = rd // 2 +f = 10000.0 ** (-torch.arange(half, dtype=torch.float32) / half) +a = pos.float()[:, None] * f[None] +c, s = a.cos()[:, None], a.sin()[:, None] +qf = q.float().view(R, H, hd) +want = torch.cat([qf[..., :half] * c - qf[..., half:rd] * s, qf[..., :half] * s + qf[..., half:rd] * c, qf[..., rd:]], -1) +res["rope"] = float((got - r(want)).abs().max()) + +# attention: 3 streams' rings of 144 rows (+1 sink), windows clipped by lo and by 128 +ring, W, slots = 144, 128, 3 +K = torch.randn(slots * ring + 1, hd, generator=gen).to(torch.bfloat16) +V = torch.randn(slots * ring + 1, hd, generator=gen).to(torch.bfloat16) +rows = [(0, 10, 0), (1, 200, 50), (2, 140, 0), (1, 201, 50), (0, 0, 0)] # (slot, position, lo) +N = len(rows) +qa = torch.randn(N, H * hd, generator=gen).to(torch.bfloat16) +base = torch.tensor([sl * ring for sl, _, _ in rows]) +pz = torch.tensor([p for _, p, _ in rows]) +lo = torch.tensor([l for _, _, l in rows]) +out = torch.zeros(N, H * hd, dtype=torch.bfloat16) +args = [dev_t(t) for t in (qa, K, V, base, pz, lo, out)] +launch(fn("tf_eagle_attn_kernel"), (N, H), 128, [("p", args[0].value), ("ll", H * hd), ("p", args[1].value), ("p", args[2].value), ("i", 1), ("p", args[3].value), ("p", args[4].value), ("p", args[5].value), ("i", ring), ("i", W), ("i", H), ("f", hd ** -0.5), ("p", args[6].value), ("ll", H * hd)]) +got = back(args[6], out).float() +err = 0.0 +for i, (sl, p, l) in enumerate(rows): + idx = torch.tensor([sl * ring + t % ring for t in range(max(l, p - W), p + 1)]) + sc = (qa[i].float().view(H, hd) @ K[idx].float().T) * hd ** -0.5 + want = sc.softmax(-1) @ V[idx].float() + err = max(err, float((got[i].view(H, hd) - r(want)).abs().max())) +res["attn"] = err +ring2, W2 = 527, 511 +K2 = torch.randn(2 * ring2 + 1, hd, generator=gen).to(torch.bfloat16) +V2 = torch.randn(2 * ring2 + 1, hd, generator=gen).to(torch.bfloat16) +rows2 = [(0, 600, 0), (1, 700, 400), (0, 511, 0), (1, 5, 0)] +N2 = len(rows2) +qb = torch.randn(N2, H * hd, generator=gen).to(torch.bfloat16) +out2 = torch.zeros(N2, H * hd, dtype=torch.bfloat16) +args2 = [dev_t(t) for t in (qb, K2, V2, torch.tensor([sl * ring2 for sl, _, _ in rows2]), torch.tensor([p for _, p, _ in rows2]), + torch.tensor([l for _, _, l in rows2]), out2)] +launch(fn("tf_eagle_attn_kernel"), (N2, H), 128, [("p", args2[0].value), ("ll", H * hd), ("p", args2[1].value), ("p", args2[2].value), ("i", 1), ("p", args2[3].value), ("p", args2[4].value), ("p", args2[5].value), ("i", ring2), ("i", W2), ("i", H), ("f", hd ** -0.5), ("p", args2[6].value), ("ll", H * hd)]) +got = back(args2[6], out2).float() +err = 0.0 +for i, (sl, p, l) in enumerate(rows2): + idx = torch.tensor([sl * ring2 + t % ring2 for t in range(max(l, p - W2), p + 1)]) + sc = (qb[i].float().view(H, hd) @ K2[idx].float().T) * hd ** -0.5 + want = sc.softmax(-1) @ V2[idx].float() + err = max(err, float((got[i].view(H, hd) - r(want)).abs().max())) +res["attn511"] = err + +# silu * up +I = 4096 +gu = torch.randn(N, 2 * I, generator=gen).to(torch.bfloat16) +y = torch.zeros(N, I, dtype=torch.bfloat16) +pg, py = dev_t(gu), dev_t(y) +launch(fn("tf_eagle_silu_mul_kernel"), (min((N * I + 255) // 256, 65535),), 256, [("p", pg.value), ("i", I), ("p", py.value), ("ll", N * I)]) +got = back(py, y).float() +want = r(r(F.silu(gu[:, :I].float())) * gu[:, I:].float()) +res["silu_mul"] = float((got - want).abs().max()) + +# top-1 and its confidence +Vv = 129280 +lg = torch.randn(N, Vv, generator=gen) * 4 +lg[2, 777] = 40.0 +lg[3, 5] = lg[3].max() + 1e-3 +ids = torch.zeros(N, dtype=torch.long) +conf = torch.zeros(N) +pl, pi, pc = dev_t(lg), dev_t(ids), dev_t(conf) +launch(fn("tf_eagle_top1_kernel"), (N,), 1024, [("p", pl.value), ("ll", Vv), ("i", Vv), ("p", 0), ("ll", 0), ("p", pi.value), ("p", pc.value), ("f", 1.0), ("f", 0.0)]) +gi, gc = back(pi, ids), back(pc, conf) +p1 = lg.softmax(-1).max(-1) +pr = p1.values.clamp(1e-7, 1 - 1e-7) +res["top1_ids_equal"] = bool((gi == p1.indices).all()) +res["top1_conf"] = float((gc - torch.log(pr / (1 - pr))).abs().max()) +parts = 64 +pm, pz, pq = torch.zeros(N, parts), torch.zeros(N, parts), torch.zeros(N, parts, dtype=torch.int32) +ppm, ppz, ppq = dev_t(pm), dev_t(pz), dev_t(pq) +ids3, conf3 = torch.zeros(N, dtype=torch.long), torch.zeros(N) +pi3, pc3 = dev_t(ids3), dev_t(conf3) +launch(fn("tf_eagle_top1_part_kernel"), (N, parts), 256, [("p", pl.value), ("ll", Vv), ("i", Vv), ("p", ppm.value), ("p", ppz.value), ("p", ppq.value)]) +launch(fn("tf_eagle_top1_join_kernel"), (N,), 32, [("p", ppm.value), ("p", ppz.value), ("p", ppq.value), ("i", parts), ("p", pi3.value), ("p", pc3.value), ("f", 1.0), ("f", 0.0)]) +res["top1_split_ids_equal"] = bool((back(pi3, ids3) == p1.indices).all()) +res["top1_split_conf"] = float((back(pc3, conf3) - torch.log(pr / (1 - pr))).abs().max()) + +# two_stage: top-k candidates (ties at the k-th value by the lower id), their exact logits, top-1 through the map +for k in (128, 512): + lk = (torch.randn(N, Vv, generator=gen) * 3).to(torch.bfloat16).float() # bf16 values: many ties + lk[1, :] = 0.5 # a row of one value + cand = torch.full((N, k), -1, dtype=torch.long) + pk, pcand = dev_t(lk), dev_t(cand) + launch(fn("tf_eagle_topk_kernel"), (N,), 1024, [("p", pk.value), ("ll", Vv), ("i", Vv), ("i", k), ("p", pcand.value), ("ll", k)]) + got = back(pcand, cand).sort(-1).values + order = torch.sort(lk, stable=True, dim=-1, descending=True).indices[:, :k] # stable: equal values by lower id + res[f"topk{k}_sets_equal"] = bool((got == order.sort(-1).values).all()) +Dk, k = 5120, 128 +Wt = (torch.randn(Vv, Dk, generator=gen) * 0.02).to(torch.bfloat16) +xo = torch.randn(N, Dk, generator=gen).to(torch.bfloat16) +cand = torch.randint(0, Vv, (N, k), generator=gen) +ex = torch.zeros(N, k) +pw, px, pcd, pex = dev_t(Wt), dev_t(xo), dev_t(cand), dev_t(ex) +f = fn("tf_eagle_rescore_kernel") +vals = [ctypes.c_uint64(pw.value), ctypes.c_int(Dk), ctypes.c_uint64(px.value), ctypes.c_longlong(Dk), ctypes.c_uint64(pcd.value), + ctypes.c_longlong(k), ctypes.c_int(k), ctypes.c_uint64(pex.value), ctypes.c_longlong(k)] +arr = (ctypes.c_void_p * len(vals))(*[ctypes.cast(ctypes.byref(v), ctypes.c_void_p) for v in vals]) +ok(cu.cuLaunchKernel(f, N, (k + 7) // 8, 1, 256, 1, 1, Dk * 4, None, arr, None), "launch") +ok(cu.cuCtxSynchronize(), "sync") +gx = back(pex, ex) +want = torch.einsum("nkd,nd->nk", Wt.float()[cand], xo.float()) +res["rescore"] = float((gx - want).abs().max() / want.abs().max()) +ids2, conf2 = torch.zeros(N, dtype=torch.long), torch.zeros(N) +pi2, pc2 = dev_t(ids2), dev_t(conf2) +launch(fn("tf_eagle_top1_kernel"), (N,), 1024, [("p", pex.value), ("ll", k), ("i", k), ("p", pcd.value), ("ll", k), ("p", pi2.value), ("p", pc2.value), ("f", 1.0), ("f", 0.0)]) +res["top1_map_ids_equal"] = bool((back(pi2, ids2) == cand.gather(1, gx.argmax(-1, keepdim=True))[:, 0]).all()) + + + +def fp8_rows(Wb: torch.Tensor): + """The converter's rule: s = amax(|W|) / 448 a row (1 where 0), q = e4m3 RNE of W / s.""" + + wf = Wb.float() + sc = wf.abs().amax(1) / 448.0 + sc = torch.where(sc == 0, torch.ones_like(sc), sc) + return (wf / sc[:, None]).to(torch.float8_e4m3fn), sc + + +gw = 0.0 +for (n_, k_) in ((5120, 4096), (4096, 5120), (512, 15360), (129280, 256), (256, 5120), (136, 10240)): + Wb = (torch.randn(n_, k_, generator=gen) * 0.02).to(torch.bfloat16) + Wb[3] = 0 + q8, s8 = fp8_rows(Wb) + for rows_ in (1, 4, 16): + xg = torch.randn(rows_, k_ + 8, generator=gen).to(torch.bfloat16) # rows k + 8 apart + for f8 in (False, True): + wd = q8.float() * s8[:, None] if f8 else Wb.float() + want = xg[:, :k_].float() @ wd.T + base_ = torch.randn(rows_, n_ + 8, generator=gen).to(torch.bfloat16) + for mode in (0, 1, 2): + o_ = torch.zeros(rows_, n_ + 8, dtype=torch.float32) if mode == 1 else base_.clone() + pw_ = dev_t(q8.view(torch.uint8) if f8 else Wb) + ps_ = dev_t(s8) + px_, po_ = dev_t(xg), dev_t(o_) + eb = 1 if f8 else 2 + G = min(32, k_ * eb // 16) + per = 32 // G + grid = (n_ + 8 * per - 1) // (8 * per) + if f8: + launch(fn("tf_eagle_gemv_fp8_kernel"), (grid,), 256, [("p", pw_.value), ("p", ps_.value), ("i", n_), ("i", k_), ("p", px_.value), ("ll", k_ + 8), ("i", rows_), ("p", po_.value), ("ll", n_ + 8), ("i", mode), ("i", G)]) + else: + launch(fn("tf_eagle_gemv_bf16_kernel"), (grid,), 256, [("p", pw_.value), ("i", n_), ("i", k_), ("p", px_.value), ("ll", k_ + 8), ("i", rows_), ("p", po_.value), ("ll", n_ + 8), ("i", mode), ("i", G)]) + got = back(po_, o_).float() + w_ = want if mode != 2 else base_[:, :n_].float() + want + if mode == 1: + e_ = float((got[:, :n_] - w_).abs().max() / w_.abs().max()) + else: + e_ = float((got[:, :n_] - r(w_)).abs().max() / w_.abs().max()) + e_ = max(e_, float((got[:, n_:] - o_[:, n_:].float()).abs().max())) # the gap columns stay + gw = max(gw, e_) + for t_ in (pw_, ps_, px_, po_): + cu.cuMemFree_v2(t_) +res["gemv"] = gw + +print(res) +bad = res["rmsnorm"] > 1e-2 or res["rope"] > 0.05 or res["attn"] > 0.02 or res["silu_mul"] > 0.05 or \ + not res["top1_ids_equal"] or res["top1_conf"] > 1e-3 or not res["topk128_sets_equal"] or not res["topk512_sets_equal"] or res["rescore"] > 1e-4 or not res["top1_map_ids_equal"] or \ + res["attn511"] > 0.02 or not res["top1_split_ids_equal"] or res["top1_split_conf"] > 1e-3 or res["gemv"] > 1e-2 +print("FAIL" if bad else "PASS") diff --git a/tools/dsv41-zig/eagle/eagle_kernels_time.py b/tools/dsv41-zig/eagle/eagle_kernels_time.py new file mode 100644 index 0000000000..36f5c5b1a5 --- /dev/null +++ b/tools/dsv41-zig/eagle/eagle_kernels_time.py @@ -0,0 +1,88 @@ +"""The TFE head's few-row linears, attention and top-1 timed on any CUDA GPU (CUDA events, the median of 50 runs): +the GB/s of each weight read. Driver API through ctypes, as eagle_kernels_test.py. + + python3 eagle_kernels_time.py EAGLE.fatbin +""" + +import ctypes +import json +import statistics +import sys + +import torch + +import eagle_kernels_test as kt # (runs its checks first: a broken kernel is not timed) + +cu = kt.cu + + +def timed(f, grid, block, args, reps=50): + vals = [] + for kind, v in args: + vals.append({"p": ctypes.c_uint64, "ll": ctypes.c_longlong, "i": ctypes.c_int, "f": ctypes.c_float}[kind](v)) + arr = (ctypes.c_void_p * len(vals))(*[ctypes.cast(ctypes.byref(v), ctypes.c_void_p) for v in vals]) + g = grid + (1,) * (3 - len(grid)) + e0, e1 = ctypes.c_void_p(), ctypes.c_void_p() + kt.ok(cu.cuEventCreate(ctypes.byref(e0), 0), "event") + kt.ok(cu.cuEventCreate(ctypes.byref(e1), 0), "event") + out = [] + for _ in range(reps + 3): + kt.ok(cu.cuEventRecord(e0, None), "record") + kt.ok(cu.cuLaunchKernel(f, g[0], g[1], g[2], block, 1, 1, 0, None, arr, None), "launch") + kt.ok(cu.cuEventRecord(e1, None), "record") + kt.ok(cu.cuEventSynchronize(e1), "sync") + ms = ctypes.c_float() + kt.ok(cu.cuEventElapsedTime(ctypes.byref(ms), e0, e1), "elapsed") + out.append(ms.value * 1e3) + return statistics.median(out[3:]) + + +gen = torch.Generator().manual_seed(1) +res = {} +for name, (n, k) in {"w_in": (5120, 10240), "wq": (4096, 5120), "wo": (5120, 4096), "w13": (8192, 5120), + "up": (129280, 256), "fc": (5120, 15360)}.items(): + Wb = (torch.randn(n, k, generator=gen) * 0.02).to(torch.bfloat16) + q8 = Wb.float().to(torch.float8_e4m3fn) + s8 = torch.ones(n) + for rows in (1, 4): + x = torch.randn(rows, k, generator=gen).to(torch.bfloat16) + o = torch.zeros(rows, n, dtype=torch.bfloat16) + px, po, ps = kt.dev_t(x), kt.dev_t(o), kt.dev_t(s8) + for f8 in (False, True): + pw = kt.dev_t(q8.view(torch.uint8) if f8 else Wb) + eb = 1 if f8 else 2 + G = min(32, k * eb // 16) + grid = (n + 8 * (32 // G) - 1) // (8 * (32 // G)) + if f8: + us = timed(kt.fn("tf_eagle_gemv_fp8_kernel"), (grid,), 256, [("p", pw.value), ("p", ps.value), ("i", n), ("i", k), ("p", px.value), ("ll", k), ("i", rows), ("p", po.value), ("ll", n), ("i", 0), ("i", G)]) + else: + us = timed(kt.fn("tf_eagle_gemv_bf16_kernel"), (grid,), 256, [("p", pw.value), ("i", n), ("i", k), ("p", px.value), ("ll", k), ("i", rows), ("p", po.value), ("ll", n), ("i", 0), ("i", G)]) + res[f"{name}_r{rows}_{'fp8' if f8 else 'bf16'}"] = f"{us:.1f} us {n * k * eb / us / 1e3:.0f} GB/s" + cu.cuMemFree_v2(pw) + for t in (px, po, ps): + cu.cuMemFree_v2(t) + +# attention: 1 row and 4 rows, window 128 and 511 +H, hd = 32, 128 +for W in (128, 511): + ring = W + 16 + K = torch.randn(4 * ring + 1, hd, generator=gen).to(torch.bfloat16) + V = torch.randn(4 * ring + 1, hd, generator=gen).to(torch.bfloat16) + for N in (1, 4): + q = torch.randn(N, H * hd, generator=gen).to(torch.bfloat16) + args = [kt.dev_t(t) for t in (q, K, V, torch.arange(N) * ring, torch.full((N,), 5000), torch.zeros(N, dtype=torch.long), + torch.zeros(N, H * hd, dtype=torch.bfloat16))] + us = timed(kt.fn("tf_eagle_attn_kernel"), (N, H), 128, [("p", args[0].value), ("ll", H * hd), ("p", args[1].value), ("p", args[2].value), ("i", 1), ("p", args[3].value), ("p", args[4].value), ("p", args[5].value), ("i", ring), ("i", W), ("i", H), ("f", hd ** -0.5), ("p", args[6].value), ("ll", H * hd)]) + res[f"attn_w{W}_r{N}"] = f"{us:.1f} us" + +# top-1 of 129280 logits: one block against the split +Vv = 129280 +for N in (1, 4): + lg = torch.randn(N, Vv, generator=gen) + pl, pi, pc = kt.dev_t(lg), kt.dev_t(torch.zeros(N, dtype=torch.long)), kt.dev_t(torch.zeros(N)) + res[f"top1_one_r{N}"] = f"{timed(kt.fn('tf_eagle_top1_kernel'), (N,), 1024, [('p', pl.value), ('ll', Vv), ('i', Vv), ('p', 0), ('ll', 0), ('p', pi.value), ('p', pc.value), ('f', 1.0), ('f', 0.0)]):.1f} us" + pm, pz, pq = kt.dev_t(torch.zeros(N, 64)), kt.dev_t(torch.zeros(N, 64)), kt.dev_t(torch.zeros(N, 64, dtype=torch.int32)) + a = timed(kt.fn("tf_eagle_top1_part_kernel"), (N, 64), 256, [("p", pl.value), ("ll", Vv), ("i", Vv), ("p", pm.value), ("p", pz.value), ("p", pq.value)]) + b = timed(kt.fn("tf_eagle_top1_join_kernel"), (N,), 32, [("p", pm.value), ("p", pz.value), ("p", pq.value), ("i", 64), ("p", pi.value), ("p", pc.value), ("f", 1.0), ("f", 0.0)]) + res[f"top1_split_r{N}"] = f"{a:.1f} + {b:.1f} us" +print(json.dumps(res, indent=1)) diff --git a/tools/dsv41-zig/eagle/eagle_rand_head.py b/tools/dsv41-zig/eagle/eagle_rand_head.py new file mode 100644 index 0000000000..931bb554a0 --- /dev/null +++ b/tools/dsv41-zig/eagle/eagle_rand_head.py @@ -0,0 +1,50 @@ +"""A random TFE head checkpoint in eagle_train.py's format (HEAD.pt: {"cfg", "state"}), to test the converter, the +reference and the engine on a head shape before a trained head exists. Norm weights are drawn from 0.5 to 1.5, so a +norm in the wrong place shows. + + python3 eagle_rand_head.py PATH/eagle_train.py OUT.pt --taps 26 2 37 1000 [--tap-norm] [--qk-norm] [--window 511] + [--cand-k 128] + +--taps: the training order (the converter puts them in the engine's order); 1000 is the final hidden state. +Then: eagle_ref.py cross OUT.pt --train-py PATH/eagle_train.py (the reference against the training code), and +eagle_convert.py + eagle_ref.py vectors for the engine's self-test (TF_DS_EAGLE_CHECK). +""" + +import argparse +import importlib.util + +import torch + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("train_py") + p.add_argument("out") + p.add_argument("--taps", type=int, nargs="+", required=True) + p.add_argument("--tap-norm", action="store_true") + p.add_argument("--qk-norm", action="store_true") + p.add_argument("--window", type=int, default=None) + p.add_argument("--cand-k", type=int, default=0) + p.add_argument("--seed", type=int, default=5) + a = p.parse_args() + spec = importlib.util.spec_from_file_location("eagle_train", a.train_py) + tr = importlib.util.module_from_spec(spec) + spec.loader.exec_module(tr) + torch.manual_seed(a.seed) + cfg = {"layers": len(a.taps), "tap_norm": a.tap_norm, "heads": 32, "head_dim": 128, "inter": 4096, "rank": 256, + "dim": 5120, "vocab": 129280, "qk_norm": a.qk_norm, "cand_k": a.cand_k} + if a.window is not None: + cfg["window"] = a.window + h = tr.make_head(cfg) + with torch.no_grad(): + for _, w in h.named_parameters(): + if w.dim() == 1 or (w.dim() == 2 and w.shape[0] == len(a.taps)): # the norms (tap_norms is [taps, dim]) + w.copy_(torch.rand_like(w) + 0.5) + out = dict(h.cfg) + out["tap_layers"] = list(a.taps) + print(out) + torch.save({"cfg": out, "state": h.state_dict()}, a.out) + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/eagle_ref.py b/tools/dsv41-zig/eagle/eagle_ref.py new file mode 100644 index 0000000000..d32ce8e696 --- /dev/null +++ b/tools/dsv41-zig/eagle/eagle_ref.py @@ -0,0 +1,323 @@ +"""The TFE draft head as the Zig engine runs it (zig/src/families/dsv41/eagle.zig), in torch on the CPU: the reference +for the engine's self-test vectors, and a cross-check against the training code's chain (eagle_train.py chain_forward). + +The engine's order of work for one stream (pending token at position q0, real rows below q0): + real row p (p < q0): g_p = RMSNorm(fc(taps_p)); x = w_in [RMSNorm(embed(x_{p+1})) ; RMSNorm(g_p)]; the layer with + RoPE position p; its key and value go into the stream's ring at p. + draft step 1: the real row q0 - 1's output o_1 -> logits -> draft d_1 (top-1). + draft step j > 1: position q0 + j - 2, token d_{j-1}, g = o_{j-1}; its key and value go into the ring too. + attention of a row at position P: the ring rows at positions max(lo, P - window) .. P (lo: the stream's first row). + qk_norm: q and k RMSNorm per head (q_norm, k_norm over head_dim) before RoPE; v as it is. +Activations are bf16 between the operations (fp32 inside each), as the engine's kernels and cuBLAS calls give them. + + python3 eagle_ref.py cross HEAD.pt --train-py PATH/eagle_train.py the chain equals chain_forward's + python3 eagle_ref.py vectors HEAD.tfe OUT.tfev --embed EMBED.bf16 [--T 160] [--steps 5] [--fp8] self-test vectors + python3 eagle_ref.py fp8cmp HEAD.tfe --embed EMBED.bf16 the FP8 form's chain against the bf16 one's + +--fp8: the linears are the file's FP8 forms (fp32(q) * s, eagle_convert.py), as the engine's TF_DS_EAGLE_FP8=1 runs. +""" + +import argparse +import importlib.util +import json +import math +import struct + +import torch +import torch.nn.functional as F + +from eagle_convert import LINEARS, read + +class Ref: + """The converted head's tensors (bf16) and its math; `rnd`: round to bf16 between the operations.""" + + def __init__(self, head: dict, t: dict, embed: torch.Tensor, rnd: bool = True, fp8: bool = False): + self.h, self.embed, self.rnd = head, embed.float(), rnd + self.t = {k: v.float() for k, v in t.items() if not k.endswith((".fp8", ".scale"))} + if fp8: + for n in LINEARS: + self.t[n + ".weight"] = t[n + ".fp8"].float() * t[n + ".scale"].float()[:, None] + self.eps = head["eps"] + + def r(self, x): + return x.to(torch.bfloat16).float() if self.rnd else x + + def rms(self, x, w): + return self.r(x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * w) + + def lin(self, x, name): + return self.r(x @ self.t[name].T) + + def rope(self, x, pos: torch.Tensor): + """RoPE on the first rope_dim features of each head (halves rotated), positions pos [rows].""" + + rd, theta = self.h["rope_dim"], self.h["rope_theta"] + half = rd // 2 + f = theta ** (-torch.arange(half, dtype=torch.float32) / half) + a = pos.float()[:, None] * f[None] + cos, sin = a.cos()[:, None], a.sin()[:, None] + x1, x2 = x[..., :half], x[..., half:rd] + out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos, x[..., rd:]], -1) + return self.r(out) + + def fuse(self, taps): + """taps [rows, L, D] in the engine's order -> g [rows, D].""" + + if self.h["tap_norm"]: + taps = self.rms(taps, self.t["tap_norms"]) + return self.rms(self.lin(taps.flatten(-2), "fc.weight"), self.t["fc_norm"]) + + def feedback(self, o): + return self.rms(o, self.t["fb_norm"]) if self.h["tap_norm"] else o + + def layer_in(self, ids, g): + e = self.r(self.embed[ids]) + x = self.lin(torch.cat([self.rms(e, self.t["e_norm"]), self.rms(g, self.t["g_norm"])], -1), "w_in.weight") + hn = self.rms(x, self.t["attn_norm"]) + H, hd = self.h["heads"], self.h["head_dim"] + q = self.lin(hn, "wq.weight").view(-1, H, hd) + kv = self.lin(hn, "wkv.weight") + k, v = kv[:, :hd], kv[:, hd:] + if self.h.get("qk_norm"): + q, k = self.rms(q, self.t["q_norm"]), self.rms(k, self.t["k_norm"]) + return x, q, k, v + + def add_lin(self, x, y, name): + """x + y @ W^T, the sum in fp32 and rounded once (the engine's linearDenseAdd and own GEMV mode 2).""" + + return self.r(x + y @ self.t[name].T) + + def layer_out(self, x, att): + x = self.add_lin(x, att.flatten(-2), "wo.weight") + gu = self.lin(self.rms(x, self.t["mlp_norm"]), "w13.weight") + g, u = gu.chunk(2, -1) + y = self.r(self.r(F.silu(g)) * u) # (tf_eagle_silu_mul_kernel rounds SiLU(g) to bf16 before the multiply) + return self.add_lin(x, y, "w2.weight") + + def logits(self, o): + """fp32 logits (the engine's up projection writes fp32): the low-rank head's over the vocabulary.""" + + d = self.lin(self.rms(o, self.t["out_norm"]), "down.weight") + return d @ self.t["up.weight"].T + + def draft(self, o, k: int | None = None): + """(token, p_top1, the logits it was picked from) of one row o [1, D]. lowrank: the argmax of logits(o) (ties: + the lower id). two_stage: the top k (default cand_k) of the low-rank logits (ties at the k-th: the lower ids), + their exact logits from the target head (fp32 sums), the argmax there; p over those k.""" + + lg = self.logits(o)[0] + if self.h.get("head_type", "lowrank") != "two_stage": + return int(lg.argmax()), float(lg.softmax(-1).max()), lg + k = k or self.h["cand_k"] + cand = torch.sort(lg, stable=True, dim=-1, descending=True).indices[:k] + on = self.rms(o, self.t["out_norm"])[0] + ex = self.t["target_head"][cand] @ on + best = int(ex.argmax()) + return int(cand[best]), float(ex.softmax(-1)[best]), ex + + def attend(self, q, keys, vals): + """q [H, hd]; keys / vals [n, hd] -> [H, hd] (fp32 scores, softmax, fp32 sum, bf16 out).""" + + s = (q @ keys.T) * self.h["head_dim"] ** -0.5 + return self.r(s.softmax(-1) @ vals) + + +def chain(ref: Ref, taps, ids, q0: int, steps: int, lo: int = 0): + """The engine's chain for one stream: taps [>= q0, L, D], ids [>= q0 + 1] (ids[q0]: the pending token), real rows + lo .. q0 - 1. -> (drafts [steps], o [steps, D], p_top1 [steps], logits [steps, V]).""" + + W = ref.h["window"] + P = torch.arange(lo, q0) + g = ref.fuse(taps[lo:q0]) + x, q, k, v = ref.layer_in(ids[lo + 1:q0 + 1], g) + q, k = ref.rope(q, P), ref.rope(k[:, None], P)[:, 0] + ring_k = {int(p): k[i] for i, p in enumerate(P)} + ring_v = {int(p): v[i] for i, p in enumerate(P)} + att = [] + for i, p in enumerate(P.tolist()): + win = list(range(max(lo, p - W), p + 1)) + att.append(ref.attend(q[i], torch.stack([ring_k[w] for w in win]), torch.stack([ring_v[w] for w in win]))) + o_all = ref.layer_out(x, torch.stack(att)) + o = o_all[-1:] + drafts, outs, probs, lgs = [], [], [], [] + for j in range(1, steps + 1): + if j > 1: + pos = q0 + j - 2 + x, qj, kj, vj = ref.layer_in(torch.tensor([drafts[-1]]), ref.feedback(o)) + qj, kj = ref.rope(qj, torch.tensor([pos])), ref.rope(kj[:, None], torch.tensor([pos]))[:, 0] + ring_k[pos], ring_v[pos] = kj[0], vj[0] + win = list(range(max(lo, pos - W), pos + 1)) + att = ref.attend(qj[0], torch.stack([ring_k[w] for w in win]), torch.stack([ring_v[w] for w in win])) + o = ref.layer_out(x, att[None]) + d, p, lg = ref.draft(o) + drafts.append(d) + outs.append(o[0]) + probs.append(p) + lgs.append(lg) + return drafts, torch.stack(outs), probs, torch.stack(lgs) + + +def load_train(path: str): + spec = importlib.util.spec_from_file_location("eagle_train", path) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def fake_inputs(L: int, D: int, V: int, T: int, seed: int): + """Random taps with row scales like the target's (a spread of 0.05 to 20) and random token ids.""" + + g = torch.Generator().manual_seed(seed) + scale = torch.exp(torch.empty(T, L, 1).uniform_(math.log(0.05), math.log(20.0), generator=g)) + taps = (torch.randn(T, L, D, generator=g) * scale).to(torch.bfloat16) + ids = torch.randint(0, V, (T,), generator=g) + return taps, ids + + +def cross(a) -> None: + """chain() without rounding against eagle_train.chain_forward (fp32, the same bf16-valued weights and inputs).""" + + import tempfile + + import eagle_convert + + tr = load_train(a.train_py) + ck = torch.load(a.head, map_location="cpu", weights_only=True) + cfg = ck["cfg"] + hm = tr.make_head({**cfg, "layers": len(cfg["tap_layers"])}) + hm.load_state_dict(ck["state"]) + with torch.no_grad(): + for p in hm.parameters(): + p.copy_(p.to(torch.bfloat16).float()) + hm.eval() + head, tensors = eagle_convert.convert(ck, a.target_head) + two = head["head_type"] == "two_stage" + if two: + hm.lm = tensors["target_head"].float() + D, V, L = head["dim"], head["vocab"], len(head["taps"]) + g = torch.Generator().manual_seed(7) + embed = (torch.randn(V, D, generator=g) * 0.02).to(torch.bfloat16) + ref = Ref(head, tensors, embed, rnd=False) + perm = [head["train_taps"].index(t) for t in head["taps"]] # engine order <- training order + T, S, q0s = cases_of(a, head["window"]) + taps, ids = fake_inputs(L, D, V, T + S, 3) + worst = 0.0 + out = [] + drafts_equal = True + for q0 in q0s: + d, o, _, lgs = chain(ref, taps.float(), ids, q0, S) + ids_tf = ids.clone() + ids_tf[q0 + 1:q0 + S] = torch.tensor(d[:S - 1]) # chain_forward's tokens = our drafts + taps_train = torch.empty_like(taps) + taps_train[:, perm] = taps # back to the training order + n = q0 + S + with torch.no_grad(): + outs = tr.chain_forward(hm, embed.float(), taps_train[:n].float(), ids_tf[:n], torch.arange(n), + torch.zeros(n, dtype=torch.long), S, hidden=True) + t = q0 - 1 + for j in range(S): + err = float((outs[j][t] - o[j]).abs().max() / o[j].abs().max()) + worst = max(worst, err) + case = {"q0": q0, "step": j + 1, "rel_err": round(err, 7)} + if two: # the trainer's draft from its own o: the argmax of its two_stage logits + with torch.no_grad(): + case["draft_equal"] = int(hm.logits(outs[j][t:t + 1]).argmax()) == d[j] + drafts_equal = drafts_equal and case["draft_equal"] + out.append(case) + print(json.dumps({"cross": "chain vs chain_forward", "worst_rel_err": worst, "drafts_equal": drafts_equal, + "cases": out})) + if worst > 1e-4 or not drafts_equal: + raise SystemExit("MISMATCH") + + +MAGIC_V = b"TFEVECS1" + + +def cases_of(a, window: int): + """(T, steps, q0 list): by default T = window + 32 rows, and q0 at 1, 2, 64, window + 1, window + 2 and T.""" + + T = a.T or window + 32 + return T, a.steps, a.q0 or [1, 2, 64, window + 1, window + 2, T] + + +def vectors(a) -> None: + """Self-test vectors for the engine: cases of one stream with the pending token at q0 (real rows 0 .. q0 - 1) and + the expected drafts, outputs and top-1 probabilities of `steps` steps. --embed: the target's embedding (raw bf16 + [vocab, dim], the checkpoint's embed.weight bytes), the table the engine's head reads.""" + + head, tensors = read(a.head) + D, V, L = head["dim"], head["vocab"], len(head["taps"]) + T, S, q0s = cases_of(a, head["window"]) + taps, ids = fake_inputs(L, D, V, T + S, 11) + embed = torch.frombuffer(bytearray(open(a.embed, "rb").read()), dtype=torch.int16).view(torch.bfloat16).view(V, D) + ref = Ref(head, tensors, embed, rnd=True, fp8=a.fp8) + cases = [] + for q0 in q0s: + d, o, p, lg = chain(ref, taps.float(), ids, q0, S) + top2 = lg.topk(2, -1).values + cases.append({"q0": q0, "drafts": d, "p_top1": p, "gap": (top2[:, 0] - top2[:, 1]).tolist(), "o": o}) + print(json.dumps({"q0": q0, "drafts": d, "p_top1": [round(x, 5) for x in p]})) + meta = {"T": T, "steps": S, "dim": D, "layers": L, "taps": head["taps"], "vocab": V, "fp8": a.fp8, + "q0": [c["q0"] for c in cases], "drafts": [c["drafts"] for c in cases], + "p_top1": [c["p_top1"] for c in cases], "gap": [c["gap"] for c in cases]} + blobs = [taps.view(torch.int16).numpy().tobytes(), ids.to(torch.int32).numpy().tobytes(), + torch.cat([c["o"] for c in cases]).float().numpy().tobytes()] + meta["blobs"] = ["taps bf16 [T + steps, L, D]", "ids int32 [T + steps]", "o fp32 [cases, steps, D]"] + meta["sizes"] = [len(b) for b in blobs] + text = json.dumps(meta).encode() + with open(a.out, "wb") as f: + f.write(MAGIC_V + struct.pack(" None: + """The FP8 form against the bf16 one on the self-test cases: the step-1 low-rank logits (the same inputs), the + step-1 drafts and the free-running drafts' agreement (each chain on its own drafts).""" + + head, tensors = read(a.head) + D, V, L = head["dim"], head["vocab"], len(head["taps"]) + T, S, q0s = cases_of(a, head["window"]) + taps, ids = fake_inputs(L, D, V, T + S, 11) + embed = torch.frombuffer(bytearray(open(a.embed, "rb").read()), dtype=torch.int16).view(torch.bfloat16).view(V, D) + rb = Ref(head, tensors, embed, rnd=True) + r8 = Ref(head, tensors, embed, rnd=True, fp8=True) + worst, same1, n1, agree, steps = 0.0, 0, 0, 0, 0 + out = [] + for q0 in q0s: + db, ob, _, _ = chain(rb, taps.float(), ids, q0, S) + d8, o8, _, _ = chain(r8, taps.float(), ids, q0, S) + lb, l8 = rb.logits(ob[:1])[0], r8.logits(o8[:1])[0] # the low-rank logits (a two_stage head's first stage) + err = float((l8 - lb).abs().max() / lb.abs().max()) + worst = max(worst, err) + same1 += int(db[0] == d8[0]) + n1 += 1 + k = 0 + while k < S and db[k] == d8[k]: + k += 1 + agree += k + steps += S + out.append({"q0": q0, "step1_logit_rel_err": round(err, 5), "drafts_bf16": db, "drafts_fp8": d8}) + print(json.dumps({"fp8cmp": a.head, "cases": out, "worst_step1_logit_rel_err": worst, + "step1_top1_equal": f"{same1}/{n1}", "free_run_steps_equal": f"{agree}/{steps}"})) + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("mode", choices=["cross", "vectors", "fp8cmp"]) + p.add_argument("head") + p.add_argument("out", nargs="?") + p.add_argument("--train-py", default=None) + p.add_argument("--target-head", default=None, help="cross, a two_stage head: the head dump prefix") + p.add_argument("--embed", default=None, help="vectors: the target embedding, raw bf16 [vocab, dim]") + p.add_argument("--T", type=int, default=None, help="rows of taps (default: the window + 32)") + p.add_argument("--steps", type=int, default=5) + p.add_argument("--q0", type=int, nargs="+", default=None, help="default: 1 2 64 window+1 window+2 T") + p.add_argument("--fp8", action="store_true", help="vectors: the FP8 forms of the linears") + a = p.parse_args() + {"cross": cross, "vectors": vectors, "fp8cmp": fp8cmp}[a.mode](a) + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/fatbins_from_elf.py b/tools/dsv41-zig/eagle/fatbins_from_elf.py new file mode 100644 index 0000000000..6a040ab2c7 --- /dev/null +++ b/tools/dsv41-zig/eagle/fatbins_from_elf.py @@ -0,0 +1,54 @@ +"""The kernel fatbins a Zig CUDA binary embeds (zig/src/cuda/kernels.zig: one `kernels.Blob("fatbin_").bytes` +symbol each), written out as .fatbin: the input of `zig build -Dfatbins=DIR`, so a build reuses another build's +kernels byte for byte (no nvcc). The binary must keep its symbols (no -Dstrip). + + python3 fatbins_from_elf.py BINARY OUT_DIR +""" + +import os +import re +import struct +import sys + + +def main() -> None: + if len(sys.argv) != 3: + raise SystemExit(__doc__) + b = open(sys.argv[1], "rb").read() + out = sys.argv[2] + os.makedirs(out, exist_ok=True) + if b[:4] != b"\x7fELF" or b[4] != 2 or b[5] != 1: + raise SystemExit("not a 64-bit little-endian ELF") + shoff, = struct.unpack_from("= len(secs): + continue + sec = secs[st_shndx] + data = b[sec[3] + (st_value - sec[2]):][:st_size] + if struct.unpack_from("= 0: + try: + last = json.loads(line[k:].replace("{ ", "[").replace(" }", "]"))["accept_by_pos"] + except ValueError: + pass + return last + + +def one(cat, i, msgs, draft, temp, seed): + body = {"model": MODEL, "messages": msgs, "max_tokens": MAXTOK[cat], "seed": seed, "draft": draft} + if temp is not None: + body["temperature"] = temp + if cat == "tools": + body["tools"] = TOOLS_SPEC + t0 = time.time() + rec = {"cat": cat, "i": i, "draft": draft, "temp": temp, "seed": seed} + try: + req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"}) + r = json.load(urllib.request.urlopen(req, timeout=600)) + rt = r.get("tensorfold") or r.get("runtime") or {} + sp = r.get("speculative") or {} + ch = r["choices"][0] + rec.update(ok=True, wall=time.time() - t0, tokens=r["usage"]["completion_tokens"], + tps=rt.get("tokens_per_second"), ttft=rt.get("time_to_first_token"), token_sha=rt.get("token_sha"), + drafted=sp.get("drafted"), accepted=sp.get("accepted"), rounds=sp.get("rounds"), + finish=ch.get("finish_reason"), tool_calls=len(ch["message"].get("tool_calls") or []), + head=(ch["message"].get("content") or "")[:80], keys=sorted(r.keys())) + except Exception as e: # a failed request is a line too + rec.update(ok=False, wall=time.time() - t0, error=repr(e)[:300]) + return rec + + +def batch(cat, streams, draft, temp, out, phase): + msgs = CATS[cat] + recs = [None] * len(msgs) + a0 = accept_line() + t0 = time.time() + if streams == 1: + for i, m in enumerate(msgs): + recs[i] = one(cat, i, m, draft, temp, 1000 + i) + else: + th = [] + for i, m in enumerate(msgs): + def run(i=i, m=m): + recs[i] = one(cat, i, m, draft, temp, 1000 + i) + th.append(threading.Thread(target=run)) + for t in th: + t.start() + for t in th: + t.join() + wall = time.time() - t0 + time.sleep(1.5) + a1 = accept_line() + for r in recs: + out.write(json.dumps({"phase": phase, "kind": "req", "streams": streams, **r}) + "\n") + toks = sum(r.get("tokens") or 0 for r in recs) + acc = None + if a0 and a1: + acc = {"rounds": a1["rounds"] - a0["rounds"], "windows": a1["windows"] - a0["windows"], + "asked": [x - y for x, y in zip(a1["asked"], a0["asked"])], + "kept": [x - y for x, y in zip(a1["kept"], a0["kept"])]} + elif a1: + acc = {"cumulative": a1} + tps = sorted(r["tps"] for r in recs if r.get("tps")) + line = {"phase": phase, "kind": "batch", "cat": cat, "streams": streams, "draft": draft, "temp": temp, + "ok": sum(1 for r in recs if r.get("ok")), "n": len(recs), "tokens": toks, "wall": wall, + "agg_tps": toks / wall if wall > 0 else 0, "median_tps": tps[len(tps) // 2] if tps else None, + "drafted": sum(r.get("drafted") or 0 for r in recs), "accepted": sum(r.get("accepted") or 0 for r in recs), + "accept_by_pos": acc} + out.write(json.dumps(line) + "\n") + out.flush() + print(json.dumps({k: line[k] for k in ("cat", "streams", "draft", "temp", "ok", "tokens", "agg_tps", "median_tps", "drafted", "accepted")}), flush=True) + return recs + + +def main(): + phase, path = sys.argv[1], sys.argv[2] + quick = "--quick" in sys.argv + with open(path, "a") as out: + shas = {} + for cat in ("prose", "code", "tools"): + for streams in ((1,) if quick else (1, 4)): + for draft in (True, False): + for r in batch(cat, streams, draft, None, out, phase): + shas.setdefault((cat, r["i"], "t1"), {})[(streams, draft)] = r.get("token_sha") + for draft in (True, False): + for r in batch(cat, 4, draft, 0, out, phase): + shas.setdefault((cat, r["i"], "t0"), {})[(4, draft)] = r.get("token_sha") + same = sum(1 for v in shas.values() if len(set(v.values())) == 1 and None not in v.values()) + verdict = {"phase": phase, "kind": "sha", "groups": len(shas), "all_equal": same, + "diff": [f"{k[0]}{k[1]}{k[2]}:" + ",".join(f"{s}{'d' if d else 'n'}={h}" for (s, d), h in v.items()) + for k, v in shas.items() if len(set(v.values())) != 1 or None in v.values()]} + out.write(json.dumps(verdict) + "\n") + print(json.dumps(verdict), flush=True) + print("PHASE_DONE " + phase, flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/tfe_seg.py b/tools/dsv41-zig/eagle/tfe_seg.py new file mode 100644 index 0000000000..ee6375ef3c --- /dev/null +++ b/tools/dsv41-zig/eagle/tfe_seg.py @@ -0,0 +1,57 @@ +"""Acceptance by position for each drafted batch of a phase, from the engine's TF_DS_ACCEPT_LOG lines (cumulative). + + python3 tfe_seg.py ENGINE_HEAD.log + +The drafted batches run in a fixed order (tfe_bench.py: per category 1 stream, 4 streams t1, 4 streams t0; draft +false between them). A drafted batch is a run of lines where the windows count rises; a draft-false batch between two +drafted ones leaves it flat for many rounds. Runs split at a flat stretch of 60 rounds or more. +""" + +import json +import sys + +ORDER = [(c, s, t) for c in ("prose", "code", "tools") for s, t in ((1, "t1"), (4, "t1"), (4, "t0"))] + + +def main(): + rows = [] + for line in open(sys.argv[1], errors="replace"): + k = line.find('{"accept_by_pos"') + if k >= 0: + try: + rows.append(json.loads(line[k:].replace("{ ", "[").replace(" }", "]"))["accept_by_pos"]) + except ValueError: + pass + segs = [] + start = prev = None + flat_from = None + for r in rows: + if prev is None: + prev = r + continue + if r["windows"] > prev["windows"]: + if start is None: + start = prev + flat_from = None + last = r + elif start is not None: + flat_from = flat_from or prev + if r["rounds"] - flat_from["rounds"] >= 60: + segs.append((start, last)) + start = None + flat_from = None + prev = r + if start is not None: + segs.append((start, last)) + print(f"{len(segs)} drafted runs (want {len(ORDER)} for a full phase, 3 for --quick)") + for i, (a, b) in enumerate(segs): + asked = [x - y for x, y in zip(b["asked"], a["asked"])] + kept = [x - y for x, y in zip(b["kept"], a["kept"])] + pos = [round(k / n, 2) if n else None for k, n in zip(kept, asked)][:5] + lab = ORDER[i] if len(segs) == len(ORDER) else (i,) + print(json.dumps({"batch": lab, "rounds": b["rounds"] - a["rounds"], "windows": b["windows"] - a["windows"], "asked": asked[:5], "kept": kept[:5], + "accept_by_pos": pos, "kept_per_window": round(sum(kept) / max(1, b["windows"] - a["windows"]), 3)})) + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/tfe_sum.py b/tools/dsv41-zig/eagle/tfe_sum.py new file mode 100644 index 0000000000..8396783992 --- /dev/null +++ b/tools/dsv41-zig/eagle/tfe_sum.py @@ -0,0 +1,39 @@ +"""Summary of a TFE window's bench.jsonl: per phase and batch the tokens, aggregate and median decode tok/s, drafted +and accepted; and token_sha of every (category, prompt, sampling) across the phases and draft settings. + + python3 tfe_sum.py BENCH.jsonl +""" + +import collections +import json +import sys + + +def main(): + reqs, batches = [], [] + for line in open(sys.argv[1]): + r = json.loads(line) + if r.get("kind") == "req": + reqs.append(r) + elif r.get("kind") == "batch": + batches.append(r) + print("phase | cat | streams | sampling | drafts | tokens | agg tok/s | median tok/s | drafted | accepted | acc rate") + for b in batches: + t = "t0" if b["temp"] == 0 else "t1" + acc = round(b["accepted"] / b["drafted"], 3) if b["drafted"] else None + print(f'{b["phase"]} | {b["cat"]} | {b["streams"]} | {t} | {"on" if b["draft"] else "off"} | {b["tokens"]} | ' + f'{b["agg_tps"]:.1f} | {b["median_tps"]:.1f} | {b["drafted"]} | {b["accepted"]} | {acc}') + shas = collections.defaultdict(set) + for r in reqs: + t = "t0" if r["temp"] == 0 else "t1" + shas[(r["cat"], r["i"], t)].add(r.get("token_sha")) + bad = {k: v for k, v in shas.items() if len(v) != 1 or None in v} + print(f"token_sha across all phases and draft settings: {len(shas) - len(bad)}/{len(shas)} groups one value") + for k, v in bad.items(): + print(" differs:", k, sorted(map(str, v))) + fails = [r for r in reqs if not r.get("ok")] + print(f"failed requests: {len(fails)}") + + +if __name__ == "__main__": + main() diff --git a/tools/dsv41-zig/eagle/tfe_tapcheck.py b/tools/dsv41-zig/eagle/tfe_tapcheck.py new file mode 100644 index 0000000000..850ad88f23 --- /dev/null +++ b/tools/dsv41-zig/eagle/tfe_tapcheck.py @@ -0,0 +1,102 @@ +"""The prompt path's TFE taps against themselves at other chunk cuts (stdlib only). The engine runs with +TF_DS_EAGLE_TAPDUMP=FILE (rank 0), so every absorb writes its tap rows (eagle.zig tapDump). + + python3 -I tfe_tapcheck.py send [URL] three prompts that share a prefix (lengths about N, N + 40, N - 40) + python3 -I tfe_tapcheck.py check DUMP [--gate 0,1] compare the "prompt" records at the positions they share + +A prompt chunk that reaches the head drops its rows below the replay floor (the prompt's length - the model window) at +the middle layer; a tap of a layer before that keeps the chunk's rows from its first, so the absorb reads it at the +cut. The three prompts cut at three places, so a wrong offset gives one position three different rows (a difference +near 1). Pass: each gated tap's max abs difference over max abs at the shared positions is at most 3e-2 (bf16 +rounding with other prefill chunk shapes), and the records have at least two different nonzero cuts. + +--gate: the tap indices (engine order) that must pass; the default is every tap. Gate only the taps of layers below +layers / 2: a tap at or past the middle layer depends on the replay floor (rows below it run no upper layer, so the +first replayed rows see less history), so it differs between prompts at the same position. 2026-10-24, taps +[2, 26, 37, final]: [0.0195, 0.84, 0.54, 0.26]. +""" + +import json +import struct +import sys +import urllib.request + +MODEL = "DeepSeek-v4.1-Flash-1017-TF" +WORDS = ("harbor lantern copper meadow signal river orchard basalt willow ember canyon ledger quartz monsoon " + "glacier sparrow tundra cobalt prairie lattice").split() + + +def text(n_words: int) -> str: + """Plain prose-like text of n_words words (a fixed sequence, so prompts share a prefix).""" + + out = [] + for i in range(n_words): + w = WORDS[(i * 7 + i // 13) % len(WORDS)] + out.append(w + ("." if i % 11 == 10 else ","if i % 5 == 4 else "")) + return " ".join(out) + + +def send(url: str) -> None: + base = text(2600) # about 3,300 tokens: two prompt chunks of 2,048 + prompts = [base, base + " " + text(30), base[: int(len(base) * 0.985)]] + for i, p in enumerate(prompts): + body = {"model": MODEL, "messages": [{"role": "user", "content": p}], "max_tokens": 2, "temperature": 0} + r = json.load(urllib.request.urlopen(urllib.request.Request(url, json.dumps(body).encode(), + {"Content-Type": "application/json"}), timeout=300)) + print(json.dumps({"prompt": i, "prompt_tokens": r["usage"]["prompt_tokens"]}), flush=True) + + +def records(path: str): + b = open(path, "rb").read() + at = 0 + while at < len(b): + n = struct.unpack_from(" None: + pr = [h for h in records(path) if h["kind"] == "prompt"] + print(json.dumps({"prompt_records": [{k: h[k] for k in ("slot", "start", "n", "cut")} for h in pr], + "round_records": sum(1 for h in records(path) if h["kind"] == "round")})) + if len(pr) < 2: + raise SystemExit("FAIL: fewer than 2 prompt records") + taps, dim = pr[0]["taps"], pr[0]["dim"] + worst = [0.0] * taps + pairs = 0 + for a in range(len(pr)): + for b in range(a + 1, len(pr)): + A, B = pr[a], pr[b] + lo = max(A["start"], B["start"]) + hi = min(A["start"] + A["n"], B["start"] + B["n"]) - 16 # (the chat template's last tokens differ) + for p in range(lo, hi, 4): + pairs += 1 + for j in range(taps): + x = bf16_row(A["rows"], p - A["start"], j, taps, dim) + y = bf16_row(B["rows"], p - B["start"], j, taps, dim) + m = max(max(abs(v) for v in x), 1e-12) + worst[j] = max(worst[j], max(abs(u - v) for u, v in zip(x, y)) / m) + cuts = sorted({h["cut"] for h in pr if h["cut"] > 0}) + gate = list(range(taps)) if gate is None else gate + ok = pairs > 0 and len(cuts) >= 2 and max(worst[j] for j in gate) <= 3e-2 + print(json.dumps({"tapcheck": "pass" if ok else "fail", "positions_compared": pairs, "cuts": cuts, "gate": gate, + "worst_rel_by_tap": [round(w, 6) for w in worst]})) + if not ok: + raise SystemExit(1) + + +if __name__ == "__main__": + if sys.argv[1] == "send": + send(sys.argv[2] if len(sys.argv) > 2 else "http://127.0.0.1:8899/v1/chat/completions") + else: + g = sys.argv[sys.argv.index("--gate") + 1] if "--gate" in sys.argv else None + check(sys.argv[2], [int(x) for x in g.split(",")] if g else None) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 885ab97d45..5be0ebda37 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -62,6 +62,7 @@ const kernels = [_]Kernel{ .{ .name = "dsv41_torch", .flags = torch_ops }, // ours: DeepSeek-V4.1 torch arithmetic at torch's rounding points .{ .name = "dsv41_rdma", .src = "dsv41/rdma_gather", .flags = &.{"-O3"} }, // ours: the decode gathers' RDMA kernel pair .{ .name = "dsv41_experts_par", .src = "dsv41/experts_par", .flags = &.{ "-O3", "-lineinfo" } }, // ours: the exl3 extension's gate / up with the parity split's rest in its grid (tensorfold_exl3_experts_v19's flags) + .{ .name = "dsv41_eagle", .src = "dsv41/eagle", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the TFE draft head's row kernels (eagle.zig) }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/dsv41/eagle.cu b/zig/kernels/cuda/dsv41/eagle.cu new file mode 100644 index 0000000000..a7c738ddd8 --- /dev/null +++ b/zig/kernels/cuda/dsv41/eagle.cu @@ -0,0 +1,524 @@ +// The TFE draft head's kernels (eagle.zig): fp32 math, bf16 rounds as eagle_ref.py; tools/dsv41-zig/eagle/README.md. +#include +#include +#include +#include + +namespace { + +__device__ __forceinline__ float bf(const __nv_bfloat16 x) { return __bfloat162float(x); } + +__device__ __forceinline__ float ninf() { return -__int_as_float(0x7f800000); } + +__device__ __forceinline__ float block_sum(float v, float* sh) { + for (int off = 16; off > 0; off >>= 1) v += __shfl_xor_sync(0xffffffffu, v, off); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5; + __syncthreads(); + if (lane == 0) sh[warp] = v; + __syncthreads(); + v = lane < nw ? sh[lane] : 0.f; + if (warp == 0) + for (int off = 16; off > 0; off >>= 1) v += __shfl_xor_sync(0xffffffffu, v, off); + if (threadIdx.x == 0) sh[0] = v; + __syncthreads(); + return sh[0]; +} + +__device__ __forceinline__ float block_max(float v, float* sh) { + for (int off = 16; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5; + __syncthreads(); + if (lane == 0) sh[warp] = v; + __syncthreads(); + v = lane < nw ? sh[lane] : ninf(); + if (warp == 0) + for (int off = 16; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); + if (threadIdx.x == 0) sh[0] = v; + __syncthreads(); + return sh[0]; +} + +} // namespace + +// out[r] = bf16(x[r] * rsqrt(mean(x[r]^2) + eps) * w), a row a block; x rows ldx apart, out rows ldo apart (elements). +extern "C" __global__ void tf_eagle_rmsnorm_kernel(const __nv_bfloat16* x, long long ldx, const __nv_bfloat16* w, + __nv_bfloat16* out, long long ldo, int d, float eps) { + __shared__ float sh[32]; + const __nv_bfloat16* row = x + (long long)blockIdx.x * ldx; + float ss = 0.f; + for (int i = threadIdx.x; i < d; i += blockDim.x) { + const float v = bf(row[i]); + ss += v * v; + } + const float r = rsqrtf(block_sum(ss, sh) / (float)d + eps); + __nv_bfloat16* o = out + (long long)blockIdx.x * ldo; + for (int i = threadIdx.x; i < d; i += blockDim.x) o[i] = __float2bfloat16_rn(bf(row[i]) * r * bf(w[i])); +} + +// RoPE in place on the first rd features of each head of row r at pos[r] (halves rotated); a row a block. +extern "C" __global__ void tf_eagle_rope_kernel(__nv_bfloat16* x, long long ld, int heads, int hd, int rd, + const long long* pos, float theta) { + const int half = rd / 2; + const float p = (float)pos[blockIdx.x]; + __nv_bfloat16* row = x + (long long)blockIdx.x * ld; + for (int t = threadIdx.x; t < heads * half; t += blockDim.x) { + const int h = t / half, i = t % half; + const float f = powf(theta, -((float)i / (float)half)); + const float a = p * f; + const float c = cosf(a), s = sinf(a); + __nv_bfloat16* q = row + h * hd; + const float x1 = bf(q[i]), x2 = bf(q[i + half]); + q[i] = __float2bfloat16_rn(x1 * c - x2 * s); + q[i + half] = __float2bfloat16_rn(x1 * s + x2 * c); + } +} + +// Window attention, a row and a query head a block (128 threads): ring keys max(lo, pos - window) .. pos; README. +extern "C" __global__ void __launch_bounds__(128) tf_eagle_attn_kernel( + const __nv_bfloat16* q, long long ldq, const __nv_bfloat16* K, const __nv_bfloat16* V, int kv_heads, + const long long* base, const long long* pos, const long long* lo, int ring, int window, int heads, float scale, + __nv_bfloat16* out, long long ldo) { + constexpr int HD = 128; + __shared__ float qs[HD]; + __shared__ float sc[512]; + __shared__ int rw[512]; + __shared__ float sh[32]; + __shared__ float4 part[4][32]; + const int r = blockIdx.x, h = blockIdx.y, t = threadIdx.x; + const int kvh = h / (heads / kv_heads); + const long long ldk = (long long)kv_heads * HD; + const long long P = pos[r]; + const long long first = max(lo[r], P - (long long)window); + const int n = (int)(P - first + 1); + if (n <= 0 || n > 512) { // no key (a caller's fault): zeros, not garbage + out[(long long)r * ldo + h * HD + t] = __float2bfloat16_rn(0.f); + return; + } + qs[t] = bf(q[(long long)r * ldq + h * HD + t]); + const long long b = base[r]; + for (int j = t; j < n; j += blockDim.x) rw[j] = (int)(b + (first + j) % ring); + __syncthreads(); + for (int j = t; j < n; j += blockDim.x) { + const uint4* k4 = reinterpret_cast(K + (long long)rw[j] * ldk + kvh * HD); + float s = 0.f; +#pragma unroll + for (int c = 0; c < HD / 8; ++c) { + const uint4 u = k4[c]; + const __nv_bfloat162* p2 = reinterpret_cast(&u); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float2 f = __bfloat1622float2(p2[e]); + s += qs[c * 8 + e * 2] * f.x + qs[c * 8 + e * 2 + 1] * f.y; + } + } + sc[j] = s * scale; + } + __syncthreads(); + float m = ninf(); + for (int j = t; j < n; j += blockDim.x) m = fmaxf(m, sc[j]); + m = block_max(m, sh); + float z = 0.f; + for (int j = t; j < n; j += blockDim.x) { + const float e = expf(sc[j] - m); + sc[j] = e; + z += e; + } + z = block_sum(z, sh); + const int lane = t & 31, warp = t >> 5; + float4 acc = make_float4(0.f, 0.f, 0.f, 0.f); +#pragma unroll 4 + for (int j = warp; j < n; j += 4) { + const uint2 u = *reinterpret_cast(V + (long long)rw[j] * ldk + kvh * HD + lane * 4); + const __nv_bfloat162* p2 = reinterpret_cast(&u); + const float2 f0 = __bfloat1622float2(p2[0]), f1 = __bfloat1622float2(p2[1]); + const float w = sc[j]; + acc.x += w * f0.x; + acc.y += w * f0.y; + acc.z += w * f1.x; + acc.w += w * f1.y; + } + part[warp][lane] = acc; + __syncthreads(); + const float* pf = reinterpret_cast(part); // part[w] as floats: feature lane * 4 + c = t + const float v = pf[t] + pf[HD + t] + pf[2 * HD + t] + pf[3 * HD + t]; + out[(long long)r * ldo + h * HD + t] = __float2bfloat16_rn(v / z); +} + +// SwiGLU: out [rows, inter] = bf16(bf16(silu(g)) * u) of gu [rows, 2 * inter] ([g ; u] a row), as torch's bf16 ops. +extern "C" __global__ void tf_eagle_silu_mul_kernel(const __nv_bfloat16* gu, int inter, __nv_bfloat16* out, + long long count) { + for (long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x; t < count; t += (long long)gridDim.x * blockDim.x) { + const long long r = t / inter, i = t % inter; + const float g = bf(gu[r * 2 * inter + i]), u = bf(gu[r * 2 * inter + inter + i]); + const float sg = bf(__float2bfloat16_rn(g / (1.f + expf(-g)))); + out[t] = __float2bfloat16_rn(sg * u); + } +} + +// A step's draft and confidence, a row a block: top-1 of [n] fp32 (tie: lower id), conf = a * logit(p) + b; README. +extern "C" __global__ void __launch_bounds__(1024) tf_eagle_top1_kernel(const float* logits, long long ld, int n, + const long long* map, long long ldm, + long long* ids, float* conf, float a, float b) { + __shared__ float sv[32]; + __shared__ long long si[32]; + __shared__ float sh[32]; + constexpr long long none = 0x7fffffffffffffffLL; + const float* row = logits + (long long)blockIdx.x * ld; + const long long* mrow = map ? map + (long long)blockIdx.x * ldm : nullptr; + float bv = ninf(); + long long bi = none; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const float v = row[i]; + const long long id = mrow ? mrow[i] : i; + if (v > bv || (v == bv && id < bi)) { + bv = v; + bi = id; + } + } + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const long long oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5; + if (lane == 0) { + sv[warp] = bv; + si[warp] = bi; + } + __syncthreads(); + if (warp == 0) { + bv = lane < nw ? sv[lane] : ninf(); + bi = lane < nw ? si[lane] : none; + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const long long oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + if (lane == 0) { + sv[0] = bv; + si[0] = bi; + } + } + __syncthreads(); + const float m = sv[0]; + const long long best = si[0]; + float z = 0.f; + for (int i = threadIdx.x; i < n; i += blockDim.x) z += expf(row[i] - m); + z = block_sum(z, sh); + if (threadIdx.x == 0) { + const float p = 1.f / z; + const float q = fminf(fmaxf(p, 1e-7f), 1.f - 1e-7f); + ids[blockIdx.x] = best == none ? 0 : best; + conf[blockIdx.x] = a * logf(q / (1.f - q)) + b; + } +} + +// The top-1 of long rows in two kernels: part (grid rows x parts) then join (grid rows, 32 threads); README. +extern "C" __global__ void __launch_bounds__(256) tf_eagle_top1_part_kernel(const float* logits, long long ld, int n, + float* pm, float* pz, int* pi) { + __shared__ float sv[32]; + __shared__ int si[32]; + __shared__ float sh[32]; + constexpr int none = 0x7fffffff; + const int parts = gridDim.y, per = (n + parts - 1) / parts; + const int lo = blockIdx.y * per, hi = min(n, lo + per); + const float* row = logits + (long long)blockIdx.x * ld; + float bv = ninf(); + int bi = none; + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) { + const float v = row[i]; + if (v > bv) { // (ids rise in a thread's loop: a tie keeps the lower id) + bv = v; + bi = i; + } + } + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5; + if (lane == 0) { + sv[warp] = bv; + si[warp] = bi; + } + __syncthreads(); + if (warp == 0) { + bv = lane < nw ? sv[lane] : ninf(); + bi = lane < nw ? si[lane] : none; + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + if (lane == 0) { + sv[0] = bv; + si[0] = bi; + } + } + __syncthreads(); + const float m = sv[0]; + const int best = si[0]; + float z = 0.f; + if (best != none) + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) z += expf(row[i] - m); + z = block_sum(z, sh); + if (threadIdx.x == 0) { + const long long at = (long long)blockIdx.x * parts + blockIdx.y; + pm[at] = m; + pz[at] = z; + pi[at] = best; + } +} + +extern "C" __global__ void __launch_bounds__(32) tf_eagle_top1_join_kernel(const float* pm, const float* pz, + const int* pi, int parts, long long* ids, + float* conf, float a, float b) { + constexpr int none = 0x7fffffff; + const long long at = (long long)blockIdx.x * parts; + const int lane = threadIdx.x; + float bv = ninf(); + int bi = none; + for (int j = lane; j < parts; j += 32) { + const float v = pm[at + j]; + const int id = pi[at + j]; + if (id != none && (v > bv || (v == bv && id < bi))) { + bv = v; + bi = id; + } + } + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > bv || (ov == bv && oi < bi)) { + bv = ov; + bi = oi; + } + } + float z = 0.f; + for (int j = lane; j < parts; j += 32) + if (pi[at + j] != none) z += pz[at + j] * expf(pm[at + j] - bv); + for (int off = 16; off > 0; off >>= 1) z += __shfl_xor_sync(0xffffffffu, z, off); + if (lane == 0) { + const float p = 1.f / z; + const float q = fminf(fmaxf(p, 1e-7f), 1.f - 1e-7f); + ids[blockIdx.x] = bi == none ? 0 : bi; + conf[blockIdx.x] = a * logf(q / (1.f - q)) + b; + } +} + +namespace { + +// 16 weights of a 16-byte chunk as fp32: 8 bf16 (F8 false, the first 8 of w set) or 16 e4m3. +template +__device__ __forceinline__ void chunk_w(const uint4 u, float* w) { + if constexpr (F8) { + const __nv_fp8x2_storage_t* p = reinterpret_cast(&u); +#pragma unroll + for (int e = 0; e < 8; ++e) { + const __half2 h = __half2(__nv_cvt_fp8x2_to_halfraw2(p[e], __NV_E4M3)); + const float2 f = __half22float2(h); + w[2 * e] = f.x; + w[2 * e + 1] = f.y; + } + } else { + const __nv_bfloat162* p = reinterpret_cast(&u); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float2 f = __bfloat1622float2(p[e]); + w[2 * e] = f.x; + w[2 * e + 1] = f.y; + } + } +} + +// y[r, o] = scale[o] * sum_i x[r, i] * W[o, i], rows <= 16, G lanes an output; mode 0 bf16, 1 fp32, 2 out + y. +template +__device__ __forceinline__ void gemv_rows(const uint8_t* W, const float* scale, int n, int k, const __nv_bfloat16* x, + long long ldx, int rows, void* out, long long ldo, int mode, int G) { + constexpr int E = F8 ? 16 : 8; // the weights of a chunk + const int chunks = k / E; + const long long ldw = (long long)k * (F8 ? 1 : 2); + const int lane = threadIdx.x & 31, sub = lane % G; + const int per_warp = 32 / G; + const long long warps = (long long)gridDim.x * (blockDim.x >> 5); + for (long long o = ((long long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5)) * per_warp + lane / G; ; + o += warps * per_warp) { + const bool live = o < n; + if (__all_sync(0xffffffffu, !live)) break; + float acc[16]; +#pragma unroll + for (int r = 0; r < 16; ++r) acc[r] = 0.f; + if (live) { + const uint4* wr = reinterpret_cast(W + o * ldw); +#pragma unroll 2 + for (int c = sub; c < chunks; c += G) { + float w[16]; + chunk_w(__ldg(wr + c), w); +#pragma unroll + for (int r = 0; r < 16; ++r) { + if (r < rows) { + const uint4* xr = reinterpret_cast(x + r * ldx + (long long)c * E); +#pragma unroll + for (int h = 0; h < E / 8; ++h) { + const uint4 u = __ldg(xr + h); + const __nv_bfloat162* p = reinterpret_cast(&u); +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float2 f = __bfloat1622float2(p[e]); + acc[r] = __fmaf_rn(w[h * 8 + 2 * e], f.x, acc[r]); // (FMAs: the build has + acc[r] = __fmaf_rn(w[h * 8 + 2 * e + 1], f.y, acc[r]); // --fmad=false) + } + } + } + } + } + } +#pragma unroll + for (int r = 0; r < 16; ++r) { + if (r < rows) + for (int off = G / 2; off > 0; off >>= 1) acc[r] += __shfl_xor_sync(0xffffffffu, acc[r], off); + } + if (live && sub == 0) { + const float sc = scale ? scale[o] : 1.f; +#pragma unroll + for (int r = 0; r < 16; ++r) { + if (r < rows) { + const float y = acc[r] * sc; + if (mode == 1) { + reinterpret_cast(out)[r * ldo + o] = y; + } else { + __nv_bfloat16* ob = reinterpret_cast<__nv_bfloat16*>(out) + r * ldo + o; + *ob = __float2bfloat16_rn(mode == 2 ? bf(*ob) + y : y); + } + } + } + } + } +} + +} // namespace + +// The head's few-row linears: x bf16 [rows <= 16, k] (16-byte aligned) @ W [n, k]^T, bf16 or e4m3 + row scale. +extern "C" __global__ void __launch_bounds__(256) tf_eagle_gemv_bf16_kernel(const uint8_t* W, int n, int k, + const __nv_bfloat16* x, long long ldx, + int rows, void* out, long long ldo, + int mode, int G) { + gemv_rows(W, nullptr, n, k, x, ldx, rows, out, ldo, mode, G); +} + +extern "C" __global__ void __launch_bounds__(256) tf_eagle_gemv_fp8_kernel(const uint8_t* W, const float* scale, int n, + int k, const __nv_bfloat16* x, long long ldx, + int rows, void* out, long long ldo, int mode, + int G) { + gemv_rows(W, scale, n, k, x, ldx, rows, out, ldo, mode, G); +} + +namespace { + +// An fp32's order as an unsigned key: a larger float, a larger key. +__device__ __forceinline__ unsigned int fkey(const float v) { + const unsigned int u = __float_as_uint(v); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +} // namespace + +// A two_stage head's candidates, a row a block: the k largest of [n] fp32 logits (radix select; tie: lower ids). +extern "C" __global__ void __launch_bounds__(1024) tf_eagle_topk_kernel(const float* logits, long long ld, int n, int k, + long long* cand, long long ldc) { + __shared__ unsigned int hist[256]; + __shared__ unsigned int s_prefix, s_left, s_out, s_eq; + __shared__ int wsum[32]; + const float* row = logits + (long long)blockIdx.x * ld; + long long* out = cand + (long long)blockIdx.x * ldc; + if (threadIdx.x == 0) { + s_prefix = 0u; + s_left = (unsigned int)k; + s_out = 0u; + s_eq = 0u; + } + unsigned int mask = 0u; + for (int shift = 24; shift >= 0; shift -= 8) { + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0u; + __syncthreads(); + const unsigned int prefix = s_prefix; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const unsigned int key = fkey(row[i]); + if ((key & mask) == prefix) atomicAdd(&hist[(key >> shift) & 255u], 1u); + } + __syncthreads(); + if (threadIdx.x == 0) { + unsigned int left = s_left, b = 255u; + for (;; --b) { + if (hist[b] >= left || b == 0u) break; + left -= hist[b]; + } + s_prefix = prefix | (b << shift); + s_left = left; + } + mask |= 255u << shift; + __syncthreads(); + } + const unsigned int kth = s_prefix; + const unsigned int take_eq = s_left; // the equal keys to take (the others: above kth) + for (int i = threadIdx.x; i < n; i += blockDim.x) { + if (fkey(row[i]) > kth) out[atomicAdd(&s_out, 1u)] = i; + } + // the equal keys in id order: a chunk of blockDim ids at a time, each thread's rank among the equal ones before it + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5; + for (int base = 0; base < n; base += blockDim.x) { + const int i = base + threadIdx.x; + const bool eq = i < n && fkey(row[i]) == kth; + const unsigned int bal = __ballot_sync(0xffffffffu, eq); + if (lane == 0) wsum[warp] = __popc(bal); + __syncthreads(); + int before = 0; + for (int w = 0; w < warp; ++w) before += wsum[w]; + before += __popc(bal & ((1u << lane) - 1u)); + const unsigned int taken = s_eq; + if (eq && taken + (unsigned int)before < take_eq) out[(unsigned int)k - take_eq + taken + before] = i; + __syncthreads(); + if (threadIdx.x == 0) { + int all = 0; + for (int w = 0; w < nw; ++w) all += wsum[w]; + s_eq = taken + (unsigned int)all; + } + __syncthreads(); + if (s_eq >= take_eq) break; + } +} + +// A two_stage head's exact logits: out[r, j] = W[cand[r, j]] . x[r] (fp32 sums), a warp a candidate. +extern "C" __global__ void __launch_bounds__(256) tf_eagle_rescore_kernel(const __nv_bfloat16* W, int d, + const __nv_bfloat16* x, long long ldx, + const long long* cand, long long ldc, int k, + float* out, long long ldo) { + extern __shared__ float xs[]; + const int r = blockIdx.x; + for (int i = threadIdx.x; i < d; i += blockDim.x) xs[i] = bf(x[(long long)r * ldx + i]); + __syncthreads(); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int j = blockIdx.y * 8 + warp; + if (j >= k) return; + const __nv_bfloat16* w = W + cand[(long long)r * ldc + j] * (long long)d; + float s = 0.f; + for (int i = lane * 2; i < d; i += 64) { + const __nv_bfloat162 p = *reinterpret_cast(w + i); + s += xs[i] * __low2float(p) + xs[i + 1] * __high2float(p); + } + for (int off = 16; off > 0; off >>= 1) s += __shfl_xor_sync(0xffffffffu, s, off); + if (lane == 0) out[(long long)r * ldo + j] = s; +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 1b355a7197..2c436af079 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -50,6 +50,7 @@ pub const dsv41_ops: []const u8 = if (available) &Blob("fatbin_dsv41_ops").bytes pub const dsv41_torch: []const u8 = if (available) &Blob("fatbin_dsv41_torch").bytes else &.{}; pub const dsv41_rdma: []const u8 = if (available) &Blob("fatbin_dsv41_rdma").bytes else &.{}; pub const dsv41_experts_par: []const u8 = if (available) &Blob("fatbin_dsv41_experts_par").bytes else &.{}; +pub const dsv41_eagle: []const u8 = if (available) &Blob("fatbin_dsv41_eagle").bytes else &.{}; pub const nemotron_norms: []const u8 = if (available) &Blob("fatbin_nemotron_norms").bytes else &.{}; pub const nemotron_route: []const u8 = if (available) &Blob("fatbin_nemotron_route").bytes else &.{}; pub const nemotron_mamba: []const u8 = if (available) &Blob("fatbin_nemotron_mamba").bytes else &.{}; diff --git a/zig/src/families/dsv41/cublas.zig b/zig/src/families/dsv41/cublas.zig index f5e47e31be..8912788c8a 100644 --- a/zig/src/families/dsv41/cublas.zig +++ b/zig/src/families/dsv41/cublas.zig @@ -82,6 +82,13 @@ pub const Blas = struct { try b.check(b.api.cublasGemmEx(b.handle, .t, .n, @intCast(n), @intCast(rows), @intCast(k), &one, w, bf16_type, @intCast(k), x, bf16_type, @intCast(ldx), &zero, out, out_type, @intCast(ldo), 68, 99)); } + /// linearDense with beta 1, bf16 out: out += x @ w^T in fp32, rounded once (the TFE head's residual adds). + pub fn linearDenseAdd(b: *const Blas, x: u64, ldx: usize, w: u64, out: u64, ldo: usize, rows: usize, k: usize, n: usize) Error!void { + const one: f32 = 1; + const bf16_type: c_int = 14; // CUDA_R_16BF + try b.check(b.api.cublasGemmEx(b.handle, .t, .n, @intCast(n), @intCast(rows), @intCast(k), &one, w, bf16_type, @intCast(k), x, bf16_type, @intCast(ldx), &one, out, bf16_type, @intCast(ldo), 68, 99)); + } + /// torch.matmul of x [rows, k] fp32 (row-major) and w [n, k] fp32 transposed: out [rows, n] fp32, as torch issues /// it (column-major C^T = W x^T: transa T, transb N, m = n, n = rows, lda = ldb = k, ldc = n, alpha 1, beta 0). pub fn xwT(b: *const Blas, x: u64, w: u64, out: u64, rows: usize, k: usize, n: usize) Error!void { diff --git a/zig/src/families/dsv41/dsv41.zig b/zig/src/families/dsv41/dsv41.zig index d1439eba28..672d567d52 100644 --- a/zig/src/families/dsv41/dsv41.zig +++ b/zig/src/families/dsv41/dsv41.zig @@ -48,6 +48,7 @@ pub const round2d = @import("round2d.zig"); pub const draft2d = @import("draft2d.zig"); pub const markov_tokens = @import("markov_tokens.zig"); pub const draft = @import("draft.zig"); +pub const eagle = @import("eagle.zig"); pub const sampling = @import("sampling.zig"); pub const model = @import("model.zig"); pub const lanes = @import("lanes.zig"); @@ -97,6 +98,7 @@ test { refAll(draft2d); refAll(markov_tokens); refAll(draft); + refAll(eagle); refAll(sampling); refAll(model); refAll(lanes); diff --git a/zig/src/families/dsv41/eagle.zig b/zig/src/families/dsv41/eagle.zig new file mode 100644 index 0000000000..1af0325584 --- /dev/null +++ b/zig/src/families/dsv41/eagle.zig @@ -0,0 +1,593 @@ +//! The TFE draft head (TF_DS_DRAFTER=eagle) on rank 0, an EAGLE-style head: tools/dsv41-zig/eagle/README.md. +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("prompt.zig"); + +const file = @import("eagle_file.zig"); + +const Engine = prompt.Engine; +pub const Shape = file.Shape; +pub const Header = file.Header; +const Lin = file.Lin; +const Out = file.Out; +/// TF_DS_EAGLE_CHECK=FILE: the head against eagle_ref.py's vectors (eagle_check.zig). +pub const selfTest = @import("eagle_check.zig").selfTest; + +/// A pass's streams at most (the pool's). +pub const max_streams = 16; +/// A stream's drafts a pass at most (a verify window holds the pending row and 5 drafts: lanes.window_rows). +pub const max_steps = 5; +/// The rows an extend or an absorb runs at once. +pub const chunk_rows = 160; +/// The taps a head reads at most. +pub const max_taps = 8; +/// The final hidden state's tap name (eagle_collect.py's FINAL): the target head's normed input. +pub const final_tap: u16 = 1000; +/// The keys an attention row reads at most (eagle.cu's score buffer): a window of 511. +pub const max_keys = 512; +/// A two_stage head's candidates a row at most. +pub const max_cand = 512; +/// The rows of one own GEMV launch (eagle.cu gemv_rows); more rows run in launches of this many. +pub const gemv_rows = 16; +/// The parts of a low-rank row's top-1 (eagle.cu tf_eagle_top1_part_kernel). +pub const top1_parts = 64; +pub const ring_extra = 16; + +/// Absorbed rows: `n` rows of positions start .. in `slot` (input row order); keep false: not drafting (no write). +pub const Span = struct { slot: usize, start: usize, n: usize, keep: bool = true }; + +/// A stream of a pass: its slot, its pending token at position q0, its ids by position (seq[p] for p < q0). +pub const Req = struct { slot: usize, q0: usize, pending: i64, seq: []const i32 }; + +/// What a slot's rings hold: g rows of positions [g_lo, g_hi), keys and values of the real rows below `ext` from `lo`. +const Slot = struct { live: bool = false, lo: usize = 0, ext: usize = 0, g_lo: usize = 0, g_hi: usize = 0 }; + +pub const Fns = struct { + rmsnorm: cuda.Function, + rope: cuda.Function, + attn: cuda.Function, + silu_mul: cuda.Function, + top1: cuda.Function, + topk: cuda.Function, + rescore: cuda.Function, + gemv_bf16: cuda.Function, + gemv_fp8: cuda.Function, + top1_part: cuda.Function, + top1_join: cuda.Function, +}; + +pub const Eagle = struct { + s: Shape, + buf: cuda.DeviceBuffer, // the weights + // the weights (row-major [out, in]; the norms bf16) + fc_norm: u64, + e_norm: u64, + g_norm: u64, + attn_norm: u64, + mlp_norm: u64, + out_norm: u64, + fc: Lin, + w_in: Lin, + wq: Lin, + wkv: Lin, + wo: Lin, + w13: Lin, + w2: Lin, + down: Lin, + up: Lin, + tap_norms: u64 = 0, + fb_norm: u64 = 0, + q_norm: u64 = 0, + k_norm: u64 = 0, + target: u64 = 0, // two_stage: the target head bf16 [vocab, dim] + topk: usize = 0, // two_stage: the candidates a row (cand_k, or TF_DS_EAGLE_TOPK) + mod: cuda.Module, + f: Fns, + slots: usize, + ring: usize, + // the rings: k, v bf16 [slots * ring + 1, kv_heads * 128], g bf16 [.., dim]; each one's last row is a sink row + kring: u64, + vring: u64, + gring: u64, + st: [max_streams]Slot = @splat(.{}), + // scratch (chunk_rows rows unless named) + ids: u64, // i64 + rrow: u64, // i64: ring rows + pos: u64, // i64 + tin: u64, // bf16 [rows, taps * dim]: an absorb's taps + fin: u64, // bf16 [rows, dim]: the final state's rows of a prompt chunk + ev: u64, // bf16 [rows, dim]: embedding rows + gv: u64, // bf16 [rows, dim]: g rows + cat: u64, // bf16 [rows, 2 dim] + x: u64, // bf16 [rows, dim] + h: u64, // bf16 [rows, dim] + kv: u64, // bf16 [rows, 2 kv_heads 128] + xs: [2]u64, // bf16 [max_streams, dim]: the steps' residual rows (alternate: the step before's is the feedback) + q: u64, // bf16 [max_streams, heads 128] + att: u64, + gu: u64, // bf16 [max_streams, 2 inter] + y: u64, // bf16 [max_streams, inter] + lr: u64, // bf16 [max_streams, rank] + lg: u64, // f32 [max_streams, vocab] + cand: u64, // i64 [max_streams, max_cand]: two_stage candidates + ex: u64, // f32 [max_streams, max_cand]: their exact logits + tpm: u64, // f32 [max_streams, top1_parts]: a low-rank top-1's parts (max, sum, id) + tpz: u64, + tpi: u64, // i32 + sp: u64, // i64 [max_steps, max_streams]: each step's positions, ring rows, then [max_streams] bases and lows + srow: u64, + sbase: u64, + slo: u64, + out: u64, // i64 [max_steps + 1, max_streams]: the drafts by step + conf: u64, // f32 [max_steps, max_streams] + invalid: u64, // u32: a gather's index outside its table + dbg: u64, // bf16 [max_steps, max_streams, dim]: each step's o (the self-test reads it) + conf_a: f32 = 1.0, + conf_b: f32 = 0.0, + // the last pass on the host + streams: usize = 0, + steps: usize = 0, + drafts: [max_streams][max_steps]i64 = undefined, + confs: [max_streams][max_steps]f32 = undefined, + // where the passes spend their time (ns, the host's clock around each pass) + passes: u64 = 0, + // TF_DS_EAGLE_GRAPH=1: the steps of a pass as one CUDA graph a (streams, steps) + graphs: ?std.AutoHashMapUnmanaged(GraphKey, cuda.graph.Exec) = null, + gpa: std.mem.Allocator = undefined, + io: std.Io = undefined, + // TF_DS_EAGLE_TAPDUMP: the file and its end + dump: ?std.Io.File = null, + dump_at: u64 = 0, + + const GraphKey = struct { streams: u32, steps: u32 }; + + /// The head of `path` (eagle_file.zig). + pub const open = file.open; + + pub fn close(eg: *Eagle, gpa: std.mem.Allocator) void { + if (eg.graphs) |*g| { + var it = g.valueIterator(); + while (it.next()) |x| x.deinit(); + g.deinit(gpa); + } + if (eg.dump) |f| f.close(eg.io); + eg.mod.unload(); + eg.buf.free(); + gpa.destroy(eg); + } + + /// A new prompt in `slot`: its rings hold nothing yet. + pub fn reset(eg: *Eagle, slot: usize) void { + if (slot < max_streams) eg.st[slot] = .{}; + } + + // -- the row kernels ------------------------------------------------------------------------------------------ + + fn rms(eg: *const Eagle, e: *const Engine, x: u64, ldx: usize, w: u64, out: u64, ldo: usize, rows: usize) !void { + return eg.rmsOf(e, x, ldx, w, out, ldo, rows, eg.s.dim); + } + + /// RMSNorm of `rows` rows of `d` features (rows ldx / ldo elements apart). + fn rmsOf(eg: *const Eagle, e: *const Engine, x: u64, ldx: usize, w: u64, out: u64, ldo: usize, rows: usize, d: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(c_longlong, @intCast(ldx))); + a.add(w); + a.add(out); + a.add(@as(c_longlong, @intCast(ldo))); + a.add(@as(c_int, @intCast(d))); + a.add(eg.s.eps); + try cuda.launch.launch(eg.f.rmsnorm, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 512 } }, e.s, &a); + } + + fn rope(eg: *const Eagle, e: *const Engine, x: u64, ld: usize, heads: usize, pos: u64, rows: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(c_longlong, @intCast(ld))); + a.add(@as(c_int, @intCast(heads))); + a.add(@as(c_int, @intCast(eg.s.head_dim))); + a.add(@as(c_int, @intCast(eg.s.rope_dim))); + a.add(pos); + a.add(eg.s.rope_theta); + try cuda.launch.launch(eg.f.rope, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, e.s, &a); + } + + /// x [rows, k] bf16 @ l^T into out (strides ldx, ldo) as `mode`: own GEMV (fp8; own bf16 to 16 rows), else cuBLAS. + fn lin(eg: *const Eagle, e: *const Engine, x: u64, ldx: usize, l: Lin, out: u64, ldo: usize, mode: Out, rows: usize) !void { + if (rows == 0) return; + const eb: usize = if (l.fp8) 1 else 2; + const fits = x % 16 == 0 and ldx % 8 == 0 and (l.k * eb) % 16 == 0 and l.w % 16 == 0; + if (l.fp8 or (l.own and rows <= gemv_rows and fits)) { + if (!fits) return error.EagleUnaligned; + const chunks = l.k * eb / 16; + var G: usize = 32; + while (G > chunks) G /= 2; + const per = 32 / G; + const blocks = (l.n + 8 * per - 1) / (8 * per); + const osz: usize = if (mode == .f32) 4 else 2; + var r0: usize = 0; + while (r0 < rows) : (r0 += gemv_rows) { + var a: cuda.Args = .{}; + a.add(l.w); + if (l.fp8) a.add(l.scale); + a.add(@as(c_int, @intCast(l.n))); + a.add(@as(c_int, @intCast(l.k))); + a.add(x + r0 * ldx * 2); + a.add(@as(c_longlong, @intCast(ldx))); + a.add(@as(c_int, @intCast(@min(gemv_rows, rows - r0)))); + a.add(out + r0 * ldo * osz); + a.add(@as(c_longlong, @intCast(ldo))); + a.add(@backingInt(mode)); + a.add(@as(c_int, @intCast(G))); + try cuda.launch.launch(if (l.fp8) eg.f.gemv_fp8 else eg.f.gemv_bf16, .{ .grid = .{ .x = @intCast(blocks) }, .block = .{ .x = 256 } }, e.s, &a); + } + return; + } + switch (mode) { + .bf16, .f32 => try e.blas.linearDense(x, ldx, l.w, out, ldo, mode == .f32, rows, l.k, l.n), + .add => try e.blas.linearDenseAdd(x, ldx, l.w, out, ldo, rows, l.k, l.n), + } + } + + /// Keys and values of `rows` normed rows h: wkv, k_norm (qk_norm), RoPE at `pos`, into ring rows `rrow` (device). + fn keysOut(eg: *const Eagle, e: *const Engine, h: u64, pos: u64, rrow: u64, rows: usize) !void { + const hd = eg.s.head_dim; + const hkv = eg.s.kv_heads * hd; + try eg.lin(e, h, eg.s.dim, eg.wkv, eg.kv, 2 * hkv, .bf16, rows); + if (eg.s.qk_norm) for (0..eg.s.kv_heads) |kh| try eg.rmsOf(e, eg.kv + kh * hd * 2, 2 * hkv, eg.k_norm, eg.kv + kh * hd * 2, 2 * hkv, rows, hd); + try eg.rope(e, eg.kv, 2 * hkv, eg.s.kv_heads, pos, rows); + try e.ops.scatterRows(e.s, eg.kv, 2 * hkv * 2, rrow, eg.kring, hkv * 2, hkv * 2, rows); + try e.ops.scatterRows(e.s, eg.kv + hkv * 2, 2 * hkv * 2, rrow, eg.vring, hkv * 2, hkv * 2, rows); + } + + /// The layer input of `rows` rows: x = w_in [RMSNorm(embed rows ev) ; RMSNorm(g rows, stride dim)] into `out`. + fn input(eg: *const Eagle, e: *const Engine, g: u64, out: u64, rows: usize) !void { + const D = eg.s.dim; + try eg.rms(e, eg.ev, D, eg.e_norm, eg.cat, 2 * D, rows); + try eg.rms(e, g, D, eg.g_norm, eg.cat + D * 2, 2 * D, rows); + try eg.lin(e, eg.cat, 2 * D, eg.w_in, out, D, .bf16, rows); + } + + // -- absorbs -------------------------------------------------------------------------------------------------- + + /// g = RMSNorm(fc(taps)) of `spans` (taps bf16 [rows, taps * dim], stride ldt) into each kept span's ring rows. + pub fn absorb(eg: *Eagle, e: *const Engine, taps: u64, ldt: usize, spans: []const Span) !void { + var total: usize = 0; + for (spans) |sp| total += sp.n; + if (total == 0) return; + if (total > chunk_rows) return error.EagleAbsorbTooLarge; + const D = eg.s.dim; + const L = eg.s.n_taps; + var rows: [chunk_rows]i64 = undefined; + var r: usize = 0; + for (spans) |sp| { + if (sp.slot >= eg.slots) return error.BadSlot; + // (more rows than the ring: two rows of one scatter go to one ring row, and either one can stay) + if (sp.keep and sp.n > eg.ring) return error.EagleAbsorbTooLarge; + for (0..sp.n) |j| { + rows[r] = @intCast(if (sp.keep) sp.slot * eg.ring + (sp.start + j) % eg.ring else eg.slots * eg.ring); + r += 1; + } + } + try prompt.upload(e, eg.rrow, &rows, total * 8); + var src = taps; + var ld = ldt; + if (eg.s.tap_norm) { + for (0..L) |j| try eg.rms(e, taps + j * D * 2, ldt, eg.tap_norms + j * D * 2, eg.tin + j * D * 2, L * D, total); + src = eg.tin; + ld = L * D; + } + try eg.lin(e, src, ld, eg.fc, eg.x, D, .bf16, total); + try eg.rms(e, eg.x, D, eg.fc_norm, eg.h, D, total); + try e.ops.scatterRows(e.s, eg.h, D * 2, eg.rrow, eg.gring, D * 2, D * 2, total); + for (spans) |sp| if (sp.keep) note(&eg.st[sp.slot], sp.start, sp.n, eg.ring); + } + + /// TF_DS_EAGLE_TAPDUMP: one record (u32 length, JSON header, bf16 tap rows); the format: eagle/README.md. + pub fn tapDump(eg: *Eagle, e: *const Engine, kind: []const u8, slot: usize, start: usize, n: usize, cut: usize, src: u64, ld: usize) !void { + const f = eg.dump orelse return; + if (n == 0) return; + const row = eg.s.n_taps * eg.s.dim * 2; + try e.s.synchronize(); + const host = try eg.gpa.alloc(u8, n * row); + defer eg.gpa.free(host); + for (0..n) |r| try e.d.check(e.d.api.cuMemcpyDtoH_v2(host[r * row ..].ptr, src + r * ld * 2, row), "cuMemcpyDtoH"); + var hb: [256]u8 = undefined; + const h = try std.fmt.bufPrint(&hb, "{{\"kind\": \"{s}\", \"slot\": {d}, \"start\": {d}, \"n\": {d}, \"cut\": {d}, \"taps\": {d}, \"dim\": {d}}}", .{ kind, slot, start, n, cut, eg.s.n_taps, eg.s.dim }); + var len: [4]u8 = undefined; + std.mem.writeInt(u32, &len, @intCast(h.len), .little); + try f.writePositionalAll(eg.io, &len, eg.dump_at); + try f.writePositionalAll(eg.io, h, eg.dump_at + 4); + try f.writePositionalAll(eg.io, host, eg.dump_at + 4 + h.len); + eg.dump_at += 4 + h.len + host.len; + } + + /// A round's windows (lanes.zig verify): layer taps at `taps` and, for a final-tap head, final states at `xc`. + pub fn absorbRound(eg: *Eagle, e: *const Engine, taps: u64, xc: u64, rows: usize, spans: []const Span) !void { + const D = eg.s.dim; + const L = eg.s.n_taps; + var src = taps; + if (eg.s.final()) { + if (rows > chunk_rows) return error.EagleAbsorbTooLarge; + const lt = L - 1; + if (lt > 0) try e.ops.copyRows(e.s, taps, lt * D * 2, eg.tin, L * D * 2, lt * D * 2, rows); + try e.ops.copyRows(e.s, xc, D * 2, eg.tin + lt * D * 2, L * D * 2, D * 2, rows); + src = eg.tin; + } + if (eg.dump != null) { + var r: usize = 0; + for (spans) |sp| { + if (sp.keep) try eg.tapDump(e, "round", sp.slot, sp.start, sp.n, 0, src + r * L * D * 2, L * D); + r += sp.n; + } + } + return eg.absorb(e, src, L * D, spans); + } + + // -- the pass ------------------------------------------------------------------------------------------------- + + /// Each stream's `steps` drafts and confidences; a stream without the g row of q0 - 1 gets none (confidence -30). + pub fn pass(eg: *Eagle, e: *const Engine, reqs: []const Req, steps: usize) !void { + const N = reqs.len; + if (N == 0 or N > max_streams or steps == 0 or steps > max_steps) return error.BadDraftPass; + const W = eg.s.window; + var ok: [max_streams]bool = @splat(false); + var from: [max_streams]usize = undefined; + // the extend's rows, every stream's in order, in chunks + var c_ids: [chunk_rows]i64 = undefined; + var c_row: [chunk_rows]i64 = undefined; + var c_pos: [chunk_rows]i64 = undefined; + var c_last: [max_streams]usize = undefined; // stream i's last row in the chunk (its x row) + var c_who: [max_streams]usize = undefined; + var nc: usize = 0; + var nl: usize = 0; + const cur: usize = 0; + for (reqs, 0..) |rq, i| { + if (rq.slot >= eg.slots or rq.q0 == 0) return error.BadSlot; + const st = &eg.st[rq.slot]; + if (!st.live or rq.q0 > st.g_hi or rq.q0 <= st.g_lo) continue; + var f = @min(st.ext, rq.q0 - 1); + if (rq.q0 - f > W + 1) f = rq.q0 - (W + 1); + if (f < st.g_lo) f = st.g_lo; + if (f > st.ext) st.lo = @max(st.lo, f); // rows below f never got their keys: out of every window + from[i] = f; + ok[i] = true; + for (f..rq.q0) |p| { + if (nc == chunk_rows) { + try eg.extend(e, c_ids[0..nc], c_row[0..nc], c_pos[0..nc], c_last[0..nl], c_who[0..nl], cur); + nc = 0; + nl = 0; + } + const tok: i64 = if (p + 1 < rq.q0) rq.seq[p + 1] else rq.pending; + c_ids[nc] = if (tok < 0) 0 else tok; // (an image position: no token embedding of its own) + c_row[nc] = @intCast(rq.slot * eg.ring + p % eg.ring); + c_pos[nc] = @intCast(p); + nc += 1; + } + c_last[nl] = nc - 1; + c_who[nl] = i; + nl += 1; + } + if (nc > 0) try eg.extend(e, c_ids[0..nc], c_row[0..nc], c_pos[0..nc], c_last[0..nl], c_who[0..nl], cur); + // the steps' tables: positions q0 - 1 + j - 1 (j from 1), ring rows, each stream's ring base and low + var hp: [max_steps * max_streams]i64 = @splat(0); + var hr: [max_steps * max_streams]i64 = @splat(0); + var hb: [max_streams]i64 = @splat(0); + var hl: [max_streams]i64 = @splat(0); + for (reqs, 0..) |rq, i| { + hb[i] = @intCast(rq.slot * eg.ring); + hl[i] = @intCast(if (ok[i]) eg.st[rq.slot].lo else rq.q0 - 1); + for (0..steps) |j| { + const p = rq.q0 - 1 + j; + hp[j * N + i] = @intCast(p); + hr[j * N + i] = @intCast(if (ok[i]) rq.slot * eg.ring + p % eg.ring else eg.slots * eg.ring); + } + } + try prompt.upload(e, eg.sp, &hp, steps * N * 8); + try prompt.upload(e, eg.srow, &hr, steps * N * 8); + try prompt.upload(e, eg.sbase, &hb, N * 8); + try prompt.upload(e, eg.slo, &hl, N * 8); + try eg.chain(e, N, steps, cur); + try e.s.synchronize(); + var ho: [(max_steps + 1) * max_streams]i64 = undefined; + var hc: [max_steps * max_streams]f32 = undefined; + try e.d.check(e.d.api.cuMemcpyDtoH_v2(&ho, eg.out, (steps + 1) * N * 8), "cuMemcpyDtoH"); + try e.d.check(e.d.api.cuMemcpyDtoH_v2(&hc, eg.conf, steps * N * 4), "cuMemcpyDtoH"); + for (reqs, 0..) |rq, i| { + for (0..steps) |j| { + eg.drafts[i][j] = if (ok[i]) ho[(j + 1) * N + i] else 0; + eg.confs[i][j] = if (ok[i]) hc[j * N + i] else -30.0; + } + if (ok[i]) eg.st[rq.slot].ext = rq.q0; + } + eg.streams = N; + eg.steps = steps; + eg.passes += 1; + } + + /// A pass's steps for N streams: eager, or with TF_DS_EAGLE_GRAPH one CUDA graph a (N, steps) (fixed buffers). + fn chain(eg: *Eagle, e: *const Engine, N: usize, steps: usize, cur: usize) !void { + const g = if (eg.graphs) |*m| m else return eg.chainBody(e, N, steps, cur); + const key: GraphKey = .{ .streams = @intCast(N), .steps = @intCast(steps) }; + if (g.getPtr(key)) |x| return x.launchOn(e.s); + try cuda.graph.beginCapture(e.s, .thread_local); + eg.chainBody(e, N, steps, cur) catch |err| { + if (cuda.graph.endCapture(e.s)) |got| { + var x = got; + x.deinit(); + } else |_| {} + return err; + }; + var graph = try cuda.graph.endCapture(e.s); + defer graph.deinit(); + var exec = try graph.instantiate(); + errdefer exec.deinit(); + try g.put(eg.gpa, key, exec); + try exec.launchOn(e.s); + } + + fn chainBody(eg: *Eagle, e: *const Engine, N: usize, steps: usize, cur: usize) !void { + const D = eg.s.dim; + var x = eg.xs[cur]; + var prev = eg.xs[1 - cur]; + for (1..steps + 1) |j| { + const pos = eg.sp + (j - 1) * N * 8; + const rrow = eg.srow + (j - 1) * N * 8; + if (j > 1) { + std.mem.swap(u64, &x, &prev); + try e.ops.gatherRows(e.s, e.w.embed, eg.s.vocab, eg.out + (j - 1) * N * 8, eg.ev, D * 2, N, eg.invalid); + var fb = prev; + if (eg.s.tap_norm) { + try eg.rms(e, prev, D, eg.fb_norm, eg.gv, D, N); + fb = eg.gv; + } + try eg.input(e, fb, x, N); + } + try eg.step(e, x, pos, rrow, N, j > 1); + try e.ops.copyRows(e.s, x, D * 2, eg.dbg + (j - 1) * N * D * 2, D * 2, D * 2, N); + try eg.rms(e, x, D, eg.out_norm, eg.h, D, N); + try eg.lin(e, eg.h, D, eg.down, eg.lr, eg.s.rank, .bf16, N); + try eg.lin(e, eg.lr, eg.s.rank, eg.up, eg.lg, eg.s.vocab, .f32, N); + if (eg.topk > 0) { + try eg.rescoreTop(e, N); + var a: cuda.Args = .{}; + a.add(eg.ex); + a.add(@as(c_longlong, max_cand)); + a.add(@as(c_int, @intCast(eg.topk))); + a.add(eg.cand); + a.add(@as(c_longlong, max_cand)); + a.add(eg.out + j * N * 8); + a.add(eg.conf + (j - 1) * N * 4); + a.add(eg.conf_a); + a.add(eg.conf_b); + try cuda.launch.launch(eg.f.top1, .{ .grid = .{ .x = @intCast(N) }, .block = .{ .x = 1024 } }, e.s, &a); + } else { + // the vocabulary's top-1 in top1_parts parts a row, then joined + var a: cuda.Args = .{}; + a.add(eg.lg); + a.add(@as(c_longlong, @intCast(eg.s.vocab))); + a.add(@as(c_int, @intCast(eg.s.vocab))); + a.add(eg.tpm); + a.add(eg.tpz); + a.add(eg.tpi); + try cuda.launch.launch(eg.f.top1_part, .{ .grid = .{ .x = @intCast(N), .y = top1_parts }, .block = .{ .x = 256 } }, e.s, &a); + var b: cuda.Args = .{}; + b.add(eg.tpm); + b.add(eg.tpz); + b.add(eg.tpi); + b.add(@as(c_int, top1_parts)); + b.add(eg.out + j * N * 8); + b.add(eg.conf + (j - 1) * N * 4); + b.add(eg.conf_a); + b.add(eg.conf_b); + try cuda.launch.launch(eg.f.top1_join, .{ .grid = .{ .x = @intCast(N) }, .block = .{ .x = 32 } }, e.s, &b); + } + } + } + + /// two_stage: each row's top `topk` low-rank logits (lg) as candidates, then their exact logits (ex) on rows h. + fn rescoreTop(eg: *Eagle, e: *const Engine, N: usize) !void { + var a: cuda.Args = .{}; + a.add(eg.lg); + a.add(@as(c_longlong, @intCast(eg.s.vocab))); + a.add(@as(c_int, @intCast(eg.s.vocab))); + a.add(@as(c_int, @intCast(eg.topk))); + a.add(eg.cand); + a.add(@as(c_longlong, max_cand)); + try cuda.launch.launch(eg.f.topk, .{ .grid = .{ .x = @intCast(N) }, .block = .{ .x = 1024 } }, e.s, &a); + var b: cuda.Args = .{}; + b.add(eg.target); + b.add(@as(c_int, @intCast(eg.s.dim))); + b.add(eg.h); + b.add(@as(c_longlong, @intCast(eg.s.dim))); + b.add(eg.cand); + b.add(@as(c_longlong, max_cand)); + b.add(@as(c_int, @intCast(eg.topk))); + b.add(eg.ex); + b.add(@as(c_longlong, max_cand)); + try cuda.launch.launch(eg.f.rescore, .{ .grid = .{ .x = @intCast(N), .y = @intCast((eg.topk + 7) / 8) }, .block = .{ .x = 256 }, .shared = @intCast(eg.s.dim * 4) }, e.s, &b); + } + + /// An extend chunk: the rows' layer inputs, their keys and values into the rings, last rows' x into xs[cur]. + fn extend(eg: *Eagle, e: *const Engine, ids: []const i64, rows: []const i64, pos: []const i64, last: []const usize, who: []const usize, cur: usize) !void { + const n = ids.len; + const D = eg.s.dim; + try prompt.upload(e, eg.ids, ids.ptr, n * 8); + try prompt.upload(e, eg.rrow, rows.ptr, n * 8); + try prompt.upload(e, eg.pos, pos.ptr, n * 8); + try e.ops.gatherRows(e.s, e.w.embed, eg.s.vocab, eg.ids, eg.ev, D * 2, n, eg.invalid); + try e.ops.gatherRows(e.s, eg.gring, eg.slots * eg.ring + 1, eg.rrow, eg.gv, D * 2, n, eg.invalid); + try eg.input(e, eg.gv, eg.x, n); + try eg.rms(e, eg.x, D, eg.attn_norm, eg.h, D, n); + try eg.keysOut(e, eg.h, eg.pos, eg.rrow, n); + for (last, who) |r, i| try e.ops.copyRows(e.s, eg.x + r * D * 2, D * 2, eg.xs[cur] + i * D * 2, D * 2, D * 2, 1); + } + + /// One step of N residual rows x (x becomes o): attention (keys and values first if `keys`), then the MLP. + fn step(eg: *Eagle, e: *const Engine, x: u64, pos: u64, rrow: u64, N: usize, keys: bool) !void { + const s = eg.s; + const D = s.dim; + const hq = s.heads * s.head_dim; + try eg.rms(e, x, D, eg.attn_norm, eg.h, D, N); + try eg.lin(e, eg.h, D, eg.wq, eg.q, hq, .bf16, N); + if (s.qk_norm) try eg.rmsOf(e, eg.q, s.head_dim, eg.q_norm, eg.q, s.head_dim, N * s.heads, s.head_dim); + try eg.rope(e, eg.q, hq, s.heads, pos, N); + if (keys) try eg.keysOut(e, eg.h, pos, rrow, N); + var a: cuda.Args = .{}; + a.add(eg.q); + a.add(@as(c_longlong, @intCast(hq))); + a.add(eg.kring); + a.add(eg.vring); + a.add(@as(c_int, @intCast(s.kv_heads))); + a.add(eg.sbase); + a.add(pos); + a.add(eg.slo); + a.add(@as(c_int, @intCast(eg.ring))); + a.add(@as(c_int, @intCast(s.window))); + a.add(@as(c_int, @intCast(s.heads))); + a.add(1.0 / @sqrt(@as(f32, @floatFromInt(s.head_dim)))); + a.add(eg.att); + a.add(@as(c_longlong, @intCast(hq))); + try cuda.launch.launch(eg.f.attn, .{ .grid = .{ .x = @intCast(N), .y = @intCast(s.heads) }, .block = .{ .x = 128 } }, e.s, &a); + try eg.lin(e, eg.att, hq, eg.wo, x, D, .add, N); + try eg.rms(e, x, D, eg.mlp_norm, eg.h, D, N); + try eg.lin(e, eg.h, D, eg.w13, eg.gu, 2 * s.inter, .bf16, N); + var b: cuda.Args = .{}; + b.add(eg.gu); + b.add(@as(c_int, @intCast(s.inter))); + b.add(eg.y); + b.add(@as(c_longlong, @intCast(N * s.inter))); + const count = N * s.inter; + try cuda.launch.launch(eg.f.silu_mul, .{ .grid = .{ .x = @intCast(@min((count + 255) / 256, 65535)) }, .block = .{ .x = 256 } }, e.s, &b); + try eg.lin(e, eg.y, s.inter, eg.w2, x, D, .add, N); + } +}; + +/// A slot after g rows start .. start + n: a gap (or a first window) starts it over; rewritten rows extend again. +fn note(st: *Slot, start: usize, n: usize, ring: usize) void { + if (!st.live or start > st.g_hi or start < st.g_lo) { + st.* = .{ .live = true, .lo = start, .ext = start, .g_lo = start, .g_hi = start }; + } else if (start < st.ext) st.ext = start; + st.g_hi = start + n; + if (st.g_hi - st.g_lo > ring) st.g_lo = st.g_hi - ring; +} + +test "a slot's rows: a first window, the next ones, a gap, a lag past the ring" { + var st: Slot = .{}; + note(&st, 100, 28, 144); // a prompt's last 28 rows + try std.testing.expectEqual(Slot{ .live = true, .lo = 100, .ext = 100, .g_lo = 100, .g_hi = 128 }, st); + st.ext = 128; // a pass at q0 128 + note(&st, 128, 4, 144); // its round: the pending row and 3 drafts + try std.testing.expectEqual(@as(usize, 128), st.ext); + try std.testing.expectEqual(@as(usize, 132), st.g_hi); + note(&st, 130, 6, 144); // 1 draft kept: the next window starts at 130, rows 130.. rewritten + try std.testing.expectEqual(@as(usize, 128), st.ext); // (still 128: the pass there extends from 128) + try std.testing.expectEqual(@as(usize, 136), st.g_hi); + note(&st, 300, 2, 144); // a gap: starts over + try std.testing.expectEqual(Slot{ .live = true, .lo = 300, .ext = 300, .g_lo = 300, .g_hi = 302 }, st); + for (0..100) |k| note(&st, 302 + 2 * k, 2, 144); + try std.testing.expectEqual(@as(usize, 502), st.g_hi); + try std.testing.expectEqual(@as(usize, 502 - 144), st.g_lo); +} diff --git a/zig/src/families/dsv41/eagle_check.zig b/zig/src/families/dsv41/eagle_check.zig new file mode 100644 index 0000000000..8061d36b64 --- /dev/null +++ b/zig/src/families/dsv41/eagle_check.zig @@ -0,0 +1,84 @@ +//! The TFE head's self-test (TF_DS_EAGLE_CHECK=FILE): the head against eagle_ref.py's vectors before serving. +const std = @import("std"); +const prompt = @import("prompt.zig"); +const eagle = @import("eagle.zig"); + +const Engine = prompt.Engine; +const Eagle = eagle.Eagle; +const max_steps = eagle.max_steps; +const chunk_rows = eagle.chunk_rows; + +const VecJson = struct { + T: usize, + steps: usize, + dim: usize, + layers: usize, + taps: []const u16, + q0: []const usize, + drafts: []const []const i64, + p_top1: []const []const f64, + gap: []const []const f64, + sizes: []const usize, +}; + +/// The head against eagle_ref.py's vectors in slot 0; pass: step-1 o within 3%, drafts equal until a gap < 0.05. +pub fn selfTest(eg: *Eagle, gpa: std.mem.Allocator, io: std.Io, e: *const Engine, path: []const u8) !bool { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const all = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 30)); + if (all.len < 16 or !std.mem.eql(u8, all[0..8], "TFEVECS1")) return error.NotEagleVectors; + const n: usize = @intCast(std.mem.readInt(u64, all[8..16], .little)); + const v = try std.json.parseFromSliceLeaky(VecJson, a, all[16 .. 16 + n], .{ .ignore_unknown_fields = true }); + const D = eg.s.dim; + const L = eg.s.n_taps; + if (v.dim != D or v.layers != L or !std.mem.eql(u16, v.taps, eg.s.tapList()) or v.steps > max_steps) return error.EagleVectorsMismatch; + const rows = v.T + v.steps; + var off: usize = 16 + n; + const taps = all[off..][0 .. rows * L * D * 2]; + off += taps.len; + const ids_b = all[off..][0 .. rows * 4]; + off += ids_b.len; + const o_ref = std.mem.bytesAsSlice(f32, @as([]align(1) const u8, all[off..][0 .. v.q0.len * v.steps * D * 4])); + const seq = try a.alloc(i32, rows); + for (seq, 0..) |*q, i| q.* = std.mem.readInt(i32, ids_b[i * 4 ..][0..4], .little); + const host_o = try a.alloc(u16, v.steps * D); + var pass_ok = true; + for (v.q0, 0..) |q0, ci| { + eg.reset(0); + var at: usize = 0; + while (at < q0) { + const k = @min(chunk_rows, eg.ring, q0 - at); + try prompt.upload(e, eg.tin, taps[at * L * D * 2 ..].ptr, k * L * D * 2); + try eg.absorb(e, eg.tin, L * D, &.{.{ .slot = 0, .start = at, .n = k }}); + at += k; + } + try eg.pass(e, &.{.{ .slot = 0, .q0 = q0, .pending = seq[q0], .seq = seq[0..q0] }}, v.steps); + try e.d.check(e.d.api.cuMemcpyDtoH_v2(host_o.ptr, eg.dbg, v.steps * D * 2), "cuMemcpyDtoH"); + var errs: [max_steps]f32 = @splat(0); + var p_eng: [max_steps]f32 = @splat(0); + var same = true; + for (0..v.steps) |j| { + var mx: f32 = 0; + var md: f32 = 0; + for (0..D) |i| { + const got: f32 = @bitCast(@as(u32, host_o[j * D + i]) << 16); + const want = o_ref[(ci * v.steps + j) * D + i]; + mx = @max(mx, @abs(want)); + md = @max(md, @abs(got - want)); + } + errs[j] = md / @max(mx, 1e-12); + p_eng[j] = 1.0 / (1.0 + @exp(-(eg.confs[0][j] - eg.conf_b) / eg.conf_a)); + if (same and eg.drafts[0][j] != v.drafts[ci][j]) { + if (v.gap[ci][j] >= 0.05) pass_ok = false; + same = false; // the chains differ from here: later steps compare nothing + } + if (!same) break; + } + if (errs[0] > 0.03) pass_ok = false; + std.debug.print("{{\"eagle_check\": {d}, \"rel_err\": {any}, \"drafts\": {any}, \"ref_drafts\": {any}, \"p\": {any}, \"ref_p\": {any}}}\n", .{ q0, errs[0..v.steps], eg.drafts[0][0..v.steps], v.drafts[ci], p_eng[0..v.steps], v.p_top1[ci] }); + } + eg.reset(0); + std.debug.print("{{\"eagle_check\": \"{s}\", \"cases\": {d}}}\n", .{ if (pass_ok) "pass" else "fail", v.q0.len }); + return pass_ok; +} diff --git a/zig/src/families/dsv41/eagle_file.zig b/zig/src/families/dsv41/eagle_file.zig new file mode 100644 index 0000000000..b14f87a284 --- /dev/null +++ b/zig/src/families/dsv41/eagle_file.zig @@ -0,0 +1,337 @@ +//! The TFE head file (eagle_convert.py): its header, the tensors a head runs, and its load (Eagle.open). +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("prompt.zig"); +const eagle = @import("eagle.zig"); + +const Engine = prompt.Engine; +const Eagle = eagle.Eagle; +const max_streams = eagle.max_streams; +const max_steps = eagle.max_steps; +const chunk_rows = eagle.chunk_rows; +const max_taps = eagle.max_taps; +const final_tap = eagle.final_tap; +const max_keys = eagle.max_keys; +const max_cand = eagle.max_cand; +const top1_parts = eagle.top1_parts; +const ring_extra = eagle.ring_extra; +const magic = "TFEHEAD1"; + +pub const Shape = struct { + dim: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + inter: usize, + rank: usize, + vocab: usize, + window: usize, + rope_dim: usize, + rope_theta: f32, + eps: f32, + tap_norm: bool, + qk_norm: bool = false, + cand_k: usize = 0, // a two_stage head's candidates (0: lowrank) + taps: [max_taps]u16 = @splat(0), + n_taps: usize = 0, + + pub fn tapList(s: *const Shape) []const u16 { + return s.taps[0..s.n_taps]; + } + + /// Whether the last tap is the final state. + pub fn final(s: *const Shape) bool { + return s.n_taps > 0 and s.taps[s.n_taps - 1] == final_tap; + } + + /// The layer taps (the final state left out), in the engine's order. + pub fn layerTaps(s: *const Shape) []const u16 { + return s.taps[0 .. s.n_taps - @intFromBool(s.final())]; + } +}; + +const TensorJson = struct { offset: u64, shape: []const u64, dtype: []const u8 }; +const HeaderJson = struct { + format: u32, + kind: []const u8, + dim: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + inter: usize, + rank: usize, + vocab: usize, + window: usize, + rope_dim: usize, + rope_theta: f32, + eps: f32, + tap_norm: bool, + qk_norm: bool = false, + head_type: []const u8 = "lowrank", + cand_k: usize = 0, + taps: []const u16, + tensors: std.json.ArrayHashMap(TensorJson), + data_bytes: u64, +}; + +/// A head file's header (eagle_convert.py): its shape and the data's first byte. +pub const Header = struct { + shape: Shape, + json: HeaderJson, + data_at: u64, + + /// Reads and checks the header of `path` (allocations in `a`, an arena the caller frees). + pub fn read(a: std.mem.Allocator, io: std.Io, path: []const u8) !Header { + var f = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer f.close(io); + var pre: [16]u8 = undefined; + if (try f.readPositionalAll(io, &pre, 0) != 16 or !std.mem.eql(u8, pre[0..8], magic)) return error.NotEagleHead; + const n = std.mem.readInt(u64, pre[8..16], .little); + if (n > 1 << 24) return error.NotEagleHead; + const text = try a.alloc(u8, @intCast(n)); + if (try f.readPositionalAll(io, text, 16) != text.len) return error.ShortEagleHead; + const j = try std.json.parseFromSliceLeaky(HeaderJson, a, text, .{ .ignore_unknown_fields = true }); + if (j.format != 1 or !std.mem.eql(u8, j.kind, "base")) return error.UnsupportedEagleHead; + if (j.taps.len == 0 or j.taps.len > max_taps or j.kv_heads == 0 or j.heads % j.kv_heads != 0) return error.UnsupportedEagleHead; + if (j.head_dim != 128 or j.rope_dim > j.head_dim or j.rope_dim % 2 != 0 or j.window + 1 > max_keys) return error.UnsupportedEagleHead; + const two = std.mem.eql(u8, j.head_type, "two_stage"); + if (!two and !std.mem.eql(u8, j.head_type, "lowrank")) return error.UnsupportedEagleHead; + if (two != (j.cand_k > 0) or j.cand_k > max_cand) return error.UnsupportedEagleHead; + var s: Shape = .{ .dim = j.dim, .heads = j.heads, .kv_heads = j.kv_heads, .head_dim = j.head_dim, .inter = j.inter, .rank = j.rank, .vocab = j.vocab, .window = j.window, .rope_dim = j.rope_dim, .rope_theta = j.rope_theta, .eps = j.eps, .tap_norm = j.tap_norm, .qk_norm = j.qk_norm, .cand_k = j.cand_k }; + for (j.taps, 0..) |t, i| { + if (t == final_tap and i + 1 != j.taps.len) return error.UnsupportedEagleHead; + if (i > 0 and t <= j.taps[i - 1]) return error.UnsupportedEagleHead; // the engine's order: ascending, final last + s.taps[i] = t; + } + s.n_taps = j.taps.len; + return .{ .shape = s, .json = j, .data_at = 16 + n }; + } +}; + +/// A head linear: weights [n, k] row-major, bf16 or e4m3 with an fp32 scale a row; own: the own GEMV (always fp8). +pub const Lin = struct { w: u64 = 0, scale: u64 = 0, n: usize = 0, k: usize = 0, fp8: bool = false, own: bool = false }; + +/// A linear's output: bf16, fp32, or added to the bf16 out (one rounding). +pub const Out = enum(c_int) { bf16 = 0, f32 = 1, add = 2 }; + +/// The tensors a head needs from its file: each one's place in the file and on the device (a multiple of 256 bytes). +const Plan = struct { + items: [40]Item = undefined, + n: usize = 0, + + const Item = struct { name: []const u8, dtype: []const u8, shape: [2]u64, dims: usize, dest: *u64, src: u64 = 0, bytes: u64 = 0, at: u64 = 0 }; + + fn add(p: *Plan, name: []const u8, dtype: []const u8, shape: []const u64, dest: *u64) void { + var sh: [2]u64 = .{ 0, 0 }; + for (shape, 0..) |x, i| sh[i] = x; + p.items[p.n] = .{ .name = name, .dtype = dtype, .shape = sh, .dims = shape.len, .dest = dest }; + p.n += 1; + } + + /// A linear `base` [n, k]: "BASE.weight" bf16, or "BASE.fp8" e4m3 and "BASE.scale" fp32 [n]. + fn lin(p: *Plan, l: *Lin, comptime base: []const u8, n: usize, k: usize, fp8: bool, own: bool) void { + l.* = .{ .n = n, .k = k, .fp8 = fp8, .own = own or fp8 }; + if (fp8) { + p.add(base ++ ".fp8", "f8e4m3", &.{ @intCast(n), @intCast(k) }, &l.w); + p.add(base ++ ".scale", "f32", &.{@intCast(n)}, &l.scale); + } else p.add(base ++ ".weight", "bf16", &.{ @intCast(n), @intCast(k) }, &l.w); + } + + /// The bytes on the device; each item's file range checked against the header. + fn resolve(p: *Plan, h: Header) !u64 { + var total: u64 = 0; + for (p.items[0..p.n]) |*it| { + const t = h.json.tensors.map.get(it.name) orelse { + std.log.err("eagle head: no tensor {s} (an FP8 head needs eagle_convert.py's FP8 forms)", .{it.name}); + return error.MissingEagleTensor; + }; + if (!std.mem.eql(u8, t.dtype, it.dtype) or !std.mem.eql(u64, t.shape, it.shape[0..it.dims])) { + std.log.err("eagle head: {s} is {s} {any}, want {s} {any}", .{ it.name, t.dtype, t.shape, it.dtype, it.shape[0..it.dims] }); + return error.EagleShapeMismatch; + } + const es: u64 = if (std.mem.eql(u8, it.dtype, "f32")) 4 else if (std.mem.eql(u8, it.dtype, "bf16")) 2 else 1; + var n: u64 = es; + for (it.shape[0..it.dims]) |x| n *= x; + if (t.offset + n > h.json.data_bytes) return error.ShortEagleHead; + it.src = h.data_at + t.offset; + it.bytes = n; + it.at = total; + total = (total + n + 255) / 256 * 256; + } + return total; + } +}; + +/// The head of `path` (header `hdr`) on the device: the module, rings for `slots` streams, scratch from the arena. +pub fn open(gpa: std.mem.Allocator, io: std.Io, e: *const Engine, a: *prompt.Arena, path: []const u8, hdr: Header, slots: usize) !*Eagle { + const s = hdr.shape; + const c = e.c; + if (s.dim != c.hidden or s.vocab != c.vocab) return error.EagleShapeMismatch; + if (slots == 0 or slots > max_streams) return error.BadStreams; + // a prompt gives the head its last c.window rows (the replay floor): a longer head window starts short + if (s.window > c.window) std.log.warn("eagle head: window {d}, but a prompt gives the head only its last {d} rows", .{ s.window, c.window }); + if (!cuda.kernels.available) return error.NoKernelImages; + const eg = try gpa.create(Eagle); + errdefer gpa.destroy(eg); + eg.* = undefined; + eg.st = @splat(.{}); + eg.s = s; + eg.conf_a = 1.0; + eg.conf_b = 0.0; + eg.streams = 0; + eg.steps = 0; + eg.passes = 0; + eg.gpa = gpa; + eg.io = io; + eg.dump = null; + eg.dump_at = 0; + eg.graphs = null; + if (std.c.getenv("TF_DS_EAGLE_GRAPH")) |v| { + if (std.mem.eql(u8, std.mem.span(v), "1")) eg.graphs = .empty; + } + eg.tap_norms = 0; + eg.fb_norm = 0; + eg.q_norm = 0; + eg.k_norm = 0; + eg.target = 0; + eg.topk = s.cand_k; + if (s.cand_k > 0) if (std.c.getenv("TF_DS_EAGLE_TOPK")) |v| { + const k = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadEagleTopK; + if (k == 0 or k > max_cand) return error.BadEagleTopK; + eg.topk = k; + }; + if (std.c.getenv("TF_DS_EAGLE_CONF")) |v| { + var it = std.mem.splitScalar(u8, std.mem.span(v), ','); + eg.conf_a = try std.fmt.parseFloat(f32, it.next() orelse return error.BadEagleConf); + eg.conf_b = try std.fmt.parseFloat(f32, it.next() orelse "0"); + } + // FP8: on when the file has the forms (TF_DS_EAGLE_FP8 unset), "1" requires them, "0" runs bf16 + const fp8_env = if (std.c.getenv("TF_DS_EAGLE_FP8")) |v| std.mem.span(v) else ""; + if (fp8_env.len > 0 and !std.mem.eql(u8, fp8_env, "0") and !std.mem.eql(u8, fp8_env, "1")) return error.BadEagleFp8; + const has_fp8 = hdr.json.tensors.map.get("fc.fp8") != null; + const fp8 = if (std.mem.eql(u8, fp8_env, "1")) true else if (std.mem.eql(u8, fp8_env, "0")) false else has_fp8; + if (fp8_env.len == 0 and !has_fp8) std.log.info("eagle head: no FP8 forms in {s}: the bf16 linears", .{path}); + var own_all = false; + var own_wo = true; + if (std.c.getenv("TF_DS_EAGLE_GEMV")) |v| { + const m = std.mem.span(v); + if (std.mem.eql(u8, m, "all")) { + own_all = true; + } else if (std.mem.eql(u8, m, "cublas")) { + own_wo = false; + } else if (!std.mem.eql(u8, m, "wo")) return error.BadEagleGemv; + } + // the tensors this head runs (one form of each linear), each read from the file into one buffer + const D = s.dim; + const hq = s.heads * s.head_dim; + const hkv = s.kv_heads * s.head_dim; + const L = s.n_taps; + const d: u64 = @intCast(D); + var plan: Plan = .{}; + plan.add("fc_norm", "bf16", &.{d}, &eg.fc_norm); + plan.add("e_norm", "bf16", &.{d}, &eg.e_norm); + plan.add("g_norm", "bf16", &.{d}, &eg.g_norm); + plan.add("attn_norm", "bf16", &.{d}, &eg.attn_norm); + plan.add("mlp_norm", "bf16", &.{d}, &eg.mlp_norm); + plan.add("out_norm", "bf16", &.{d}, &eg.out_norm); + plan.lin(&eg.fc, "fc", D, L * D, fp8, own_all); + plan.lin(&eg.w_in, "w_in", D, 2 * D, fp8, own_all); + plan.lin(&eg.wq, "wq", hq, D, fp8, own_all); + plan.lin(&eg.wkv, "wkv", 2 * hkv, D, fp8, own_all); + plan.lin(&eg.wo, "wo", D, hq, fp8, own_all or own_wo); + plan.lin(&eg.w13, "w13", 2 * s.inter, D, fp8, own_all); + plan.lin(&eg.w2, "w2", D, s.inter, fp8, own_all or own_wo); + plan.lin(&eg.down, "down", s.rank, D, fp8, own_all); + plan.lin(&eg.up, "up", s.vocab, s.rank, fp8, own_all); + if (s.tap_norm) { + plan.add("tap_norms", "bf16", &.{ @intCast(L), d }, &eg.tap_norms); + plan.add("fb_norm", "bf16", &.{d}, &eg.fb_norm); + } + if (s.qk_norm) { + plan.add("q_norm", "bf16", &.{@intCast(s.head_dim)}, &eg.q_norm); + plan.add("k_norm", "bf16", &.{@intCast(s.head_dim)}, &eg.k_norm); + } + if (s.cand_k > 0) plan.add("target_head", "bf16", &.{ @intCast(s.vocab), d }, &eg.target); + const bytes: usize = @intCast(try plan.resolve(hdr)); + eg.buf = try cuda.DeviceBuffer.alloc(e.d, bytes); + errdefer eg.buf.free(); + { + var f = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer f.close(io); + const piece = try gpa.alloc(u8, 64 << 20); + defer gpa.free(piece); + for (plan.items[0..plan.n]) |it| { + var done: u64 = 0; + while (done < it.bytes) { + const n: usize = @intCast(@min(piece.len, it.bytes - done)); + if (try f.readPositionalAll(io, piece[0..n], it.src + done) != n) return error.ShortEagleHead; + try eg.buf.upload(@intCast(it.at + done), piece[0..n]); + done += n; + } + it.dest.* = eg.buf.ptr + it.at; + } + } + if (std.c.getenv("TF_DS_EAGLE_TAPDUMP")) |v| eg.dump = try std.Io.Dir.cwd().createFile(io, std.mem.span(v), .{ .truncate = true }); + errdefer if (eg.dump) |f| f.close(io); + eg.mod = try cuda.Module.load(e.d, cuda.kernels.dsv41_eagle); + errdefer eg.mod.unload(); + eg.f = .{ + .rmsnorm = try eg.mod.function("tf_eagle_rmsnorm_kernel"), + .rope = try eg.mod.function("tf_eagle_rope_kernel"), + .attn = try eg.mod.function("tf_eagle_attn_kernel"), + .silu_mul = try eg.mod.function("tf_eagle_silu_mul_kernel"), + .top1 = try eg.mod.function("tf_eagle_top1_kernel"), + .topk = try eg.mod.function("tf_eagle_topk_kernel"), + .rescore = try eg.mod.function("tf_eagle_rescore_kernel"), + .gemv_bf16 = try eg.mod.function("tf_eagle_gemv_bf16_kernel"), + .gemv_fp8 = try eg.mod.function("tf_eagle_gemv_fp8_kernel"), + .top1_part = try eg.mod.function("tf_eagle_top1_part_kernel"), + .top1_join = try eg.mod.function("tf_eagle_top1_join_kernel"), + }; + eg.slots = slots; + eg.ring = s.window + ring_extra; + const rows = slots * eg.ring; + eg.kring = try a.take((rows + 1) * hkv * 2); + eg.vring = try a.take((rows + 1) * hkv * 2); + eg.gring = try a.take((rows + 1) * D * 2); + try prompt.fill(e, eg.kring, 0, (rows + 1) * hkv * 2); + try prompt.fill(e, eg.vring, 0, (rows + 1) * hkv * 2); + try prompt.fill(e, eg.gring, 0, (rows + 1) * D * 2); + const R = chunk_rows; + const M = max_streams; + eg.ids = try a.take(R * 8); + eg.rrow = try a.take(R * 8); + eg.pos = try a.take(R * 8); + eg.tin = try a.take(R * L * D * 2); + eg.fin = try a.take(R * D * 2); + eg.ev = try a.take(R * D * 2); + eg.gv = try a.take(R * D * 2); + eg.cat = try a.take(R * 2 * D * 2); + eg.x = try a.take(R * D * 2); + eg.h = try a.take(R * D * 2); + eg.kv = try a.take(R * 2 * hkv * 2); + eg.xs = .{ try a.take(M * D * 2), try a.take(M * D * 2) }; + eg.q = try a.take(M * hq * 2); + eg.att = try a.take(M * hq * 2); + eg.gu = try a.take(M * 2 * s.inter * 2); + eg.y = try a.take(M * s.inter * 2); + eg.lr = try a.take(M * s.rank * 2); + eg.lg = try a.take(M * s.vocab * 4); + eg.cand = try a.take(M * max_cand * 8); + eg.ex = try a.take(M * max_cand * 4); + eg.tpm = try a.take(M * top1_parts * 4); + eg.tpz = try a.take(M * top1_parts * 4); + eg.tpi = try a.take(M * top1_parts * 4); + eg.sp = try a.take(max_steps * M * 8); + eg.srow = try a.take(max_steps * M * 8); + eg.sbase = try a.take(M * 8); + eg.slo = try a.take(M * 8); + eg.out = try a.take((max_steps + 1) * M * 8); + eg.conf = try a.take(max_steps * M * 4); + eg.invalid = try a.take(4); + eg.dbg = try a.take(max_steps * M * D * 2); + try prompt.fill(e, eg.invalid, 0, 4); + try e.s.synchronize(); + std.log.info("eagle head: {s}, taps {any}, window {d}, qk_norm {}, candidates {d}, fp8 {}, own gemv {s}, {d} MiB weights, {d} slots", .{ path, s.tapList(), s.window, s.qk_norm, eg.topk, fp8, if (fp8 or own_all) "all" else if (own_wo) "wo, w2" else "none", bytes >> 20, slots }); + return eg; +} diff --git a/zig/src/families/dsv41/lanes.zig b/zig/src/families/dsv41/lanes.zig index f2de54524b..912d863750 100644 --- a/zig/src/families/dsv41/lanes.zig +++ b/zig/src/families/dsv41/lanes.zig @@ -18,6 +18,7 @@ const lanes = @import("lanes"); const model = @import("model.zig"); const round = @import("round.zig"); const draft = @import("draft.zig"); +const eagle = @import("eagle.zig"); const sampling = @import("sampling.zig"); const link = @import("link.zig"); const peer_watch = @import("peer_watch.zig"); @@ -249,6 +250,32 @@ pub const Profile = struct { prefills: u64 = 0, }; +/// TF_DS_ACCEPT_LOG=N: kept held drafts by position (a draft kept: it and all before equal the draws), on stderr. +pub const AcceptLog = struct { + every: u64, + rounds: u64 = 0, + windows: u64 = 0, + asked: [window_rows]u64 = @splat(0), + kept: [window_rows]u64 = @splat(0), + + fn note(x: *AcceptLog, drafts: []const u32, sampled: []const u32) void { + if (drafts.len == 0) return; + x.windows += 1; + var ok = true; + for (drafts, 0..) |d, j| { + x.asked[j] += 1; + ok = ok and j < sampled.len and sampled[j] == d; + if (ok) x.kept[j] += 1; + } + } + + fn endRound(x: *AcceptLog) void { + x.rounds += 1; + if (x.rounds % x.every != 0) return; + std.debug.print("{{\"accept_by_pos\": {{\"rounds\": {d}, \"windows\": {d}, \"asked\": {any}, \"kept\": {any}}}}}\n", .{ x.rounds, x.windows, x.asked, x.kept }); + } +}; + /// A round-cost table "30.0,35.2,..." (TF_DS_ROUND_MS's form): ms by rows from 1. pub fn parseMs(a: std.mem.Allocator, text: []const u8) ![]const f64 { var out: std.ArrayList(f64) = .empty; @@ -318,9 +345,21 @@ pub const Lanes = struct { next: u64 = 0, snaps: std.ArrayList(*Snap) = .empty, // kept prompt states, oldest first next_snap: u64 = 1, + eagle_drafts: usize = eagle.max_steps, // TF_DS_EAGLE_DRAFTS: a stream's TFE drafts a round at most (1 to 5) + accept_log: ?AcceptLog = null, // TF_DS_ACCEPT_LOG pub fn init(gpa: std.mem.Allocator, m: *Model) !Lanes { var ln: Lanes = .{ .gpa = gpa, .m = m, .peer = if (m.world > 1) m.peer() else null }; + if (std.c.getenv("TF_DS_ACCEPT_LOG")) |v| { + const every = std.fmt.parseInt(u64, std.mem.span(v), 10) catch return error.BadAcceptLog; + if (every == 0) return error.BadAcceptLog; + ln.accept_log = .{ .every = every }; + } + if (std.c.getenv("TF_DS_EAGLE_DRAFTS")) |v| { + const k = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadEagleDrafts; + if (k == 0 or k > eagle.max_steps) return error.BadEagleDrafts; + ln.eagle_drafts = k; + } if (m.world > 1) ln.watch = try watchLinks(gpa, m, m.followers()); // the round costs by rows: ROUND_MS[min(rows, len) - 1] (the served table, or TF_DS_ROUND_MS) const t = m.round_ms; @@ -380,7 +419,7 @@ pub const Lanes = struct { .mtp = drafting, .speculate = drafting, .speculate_early = false, - .drafts = if (self.m.dr) |dr| @intCast(dr.n) else 1, + .drafts = @intCast(self.blockDrafts()), .window_costs = self.wc[0..], .mtp_step_ms = draft.draft_ms, .streams_exact = true, @@ -393,6 +432,13 @@ pub const Lanes = struct { }; } + /// A stream's drafts a round at most: DSpark's block, the TFE head's setting, else 1. + fn blockDrafts(self: *const Lanes) usize { + if (self.m.dr) |dr| return dr.n; + if (self.m.eg != null) return self.eagle_drafts; + return 1; + } + fn of(ptr: *anyopaque) *Lanes { return @ptrCast(@alignCast(ptr)); } @@ -726,7 +772,9 @@ pub const Lanes = struct { @memcpy(o.drafts[0..w.held], l.held[0..w.held]); @memcpy(o.drafts[w.held..][0..w.tokens.len], w.tokens); l.nheld = 0; + if (self.accept_log) |*x| x.note(o.drafts[0..w.held], o.sampled[0..w.rows()]); } + if (self.accept_log) |*x| x.endRound(); } /// The caches are positional: the next window starts at the stream's new length and writes over the rows past it. @@ -743,8 +791,8 @@ pub const Lanes = struct { const self = of(ptr); const m = self.m; const gpa = self.gpa; - if (m.dr == null) return error.NoDraftHead; - const block = m.dr.?.n; + if (!m.drafting()) return error.NoDraftHead; + const block = self.blockDrafts(); var tokens: [max_streams]i64 = undefined; var q0: [max_streams]i64 = undefined; var slots: [max_streams]i64 = undefined; @@ -769,6 +817,29 @@ pub const Lanes = struct { n += 1; } if (n == 0) return; + if (m.eg) |eg| { + // the TFE head runs on rank 0 alone (rank 1 gets the drafts in the verify frame): no frame + var reqs: [max_streams]eagle.Req = undefined; + for (0..n) |i| { + const sl: usize = @intCast(slots[i]); + reqs[i] = .{ .slot = sl, .q0 = @intCast(q0[i]), .pending = tokens[i], .seq = self.seqs[sl].items }; + } + const t0 = m.now(); + try m.eaglePass(reqs[0..n], steps); + if (self.prof) |*p| { + p.pass += m.now() - t0; + p.passes += 1; + p.pass_streams += n; + } + for (asked[0..n], depths[0..n], 0..) |l, k, i| { + for (0..k) |j| { + l.held[j] = @intCast(eg.drafts[i][j]); + l.confs[j] = eg.confs[i][j]; + } + l.nheld = k; + } + return; + } if (try self.begin(.pass)) { try self.wire.int(gpa, u32, @intCast(n)); try self.wire.int(gpa, u32, @intCast(steps)); @@ -810,7 +881,7 @@ pub const Lanes = struct { survivals(l.confs[0..out.len], out); const live = @max(1, self.streams.count()); if (self.served_k == .all or (self.served_k == .solo and live == 1)) { - const drafts: usize = if (self.m.dr) |dr| dr.n else 1; + const drafts = self.blockDrafts(); const kmax = @min(out.len, draft.depth(live, drafts, self.m.eng.round_rows)); const k = draft.chooseK(l.confs[0..out.len], kmax, live, drafts, self.m.eng.round_rows); for (out[k..]) |*o| o.* = 0.0; diff --git a/zig/src/families/dsv41/model.zig b/zig/src/families/dsv41/model.zig index e977074a01..8aab96ece3 100644 --- a/zig/src/families/dsv41/model.zig +++ b/zig/src/families/dsv41/model.zig @@ -14,6 +14,8 @@ const prompt = @import("prompt.zig"); const prompt2d = @import("prompt2d.zig"); const round = @import("round.zig"); const draft = @import("draft.zig"); +const eagle = @import("eagle.zig"); +const tri_basic = @import("tri_basic.zig"); const engram = @import("engram.zig"); const engram_io = @import("engram_io.zig"); const engram_aio = @import("engram_aio.zig"); @@ -67,6 +69,12 @@ pub const Options = struct { vision_bias: ?[]const u8 = null, // gate_bias_vl.safetensors (null: the kit's vision/): the image-span routing bias }; +/// TF_DS_DRAFTER=eagle: the TFE head drafts on rank 0, no rank loads DSpark; unset or "dspark": DSpark (all ranks). +pub fn eagleDrafter() bool { + const v = std.c.getenv("TF_DS_DRAFTER") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "eagle"); +} + /// TF_DS_2D_GU=parity (the served switch): the 2D split's experts with half the gate / up blocks a pair and the ninth by /// expert parity (plan.Split.parity); the rank cache must be one written for it. pub fn parityGU() bool { @@ -122,6 +130,7 @@ pub const Model = struct { pgraphs: ?draft.PassGraphs = null, // the drafter pass as CUDA graphs (with graphs) dpool: ?draft.Pool = null, dr: ?draft.Drafter = null, + eg: ?*eagle.Eagle = null, // TF_DS_DRAFTER=eagle: the TFE draft head (rank 0 only; eagle.zig) tables: ?engram_io.Tables = null, epool: ?*engram_io.Pool = null, aio: ?*engram_aio.Aio = null, @@ -200,6 +209,7 @@ pub const Model = struct { m.t_absorb = 0; m.dpool = null; m.dr = null; + m.eg = null; m.tables = null; m.epool = null; m.aio = null; @@ -228,9 +238,25 @@ pub const Model = struct { errdefer m.cache_file.close(io); var sp = try prompt2d.split(o.rank, o.world); sp.parity = sp.pair != null and parityGU(); // TF_DS_2D_GU=parity: the balanced experts split - m.w = try weights.load(gpa, d, .{ .cache = .{ .file = m.cache_file, .index = &m.ix, .io = io } }, c.*, sp, o.drafts); + const use_eagle = eagleDrafter(); + if (use_eagle and o.world != 2) return error.NotPortedYet; + const dspark = o.drafts and !use_eagle; + m.w = try weights.load(gpa, d, .{ .cache = .{ .file = m.cache_file, .index = &m.ix, .io = io } }, c.*, sp, dspark); errdefer m.w.deinit(); - if (o.drafts and m.w.dspark == null) return error.NoDrafter; + if (dspark and m.w.dspark == null) return error.NoDrafter; + // the TFE head's taps are the ones the forward writes (rank 0; rank 1 runs no head and writes none) + var eagle_hdr: ?eagle.Header = null; + if (use_eagle and o.drafts) { + if (o.rank == 0) { + const path = std.mem.span(std.c.getenv("TF_DS_EAGLE_HEAD") orelse return error.NoEagleHead); + eagle_hdr = try eagle.Header.read(a, io, path); + const lt = eagle_hdr.?.shape.layerTaps(); + for (lt) |t| if (t >= c.layers) return error.UnsupportedEagleHead; + c.dspark_taps = .{}; + for (lt, 0..) |t, i| c.dspark_taps.items[i] = t; + c.dspark_taps.len = lt.len; + } else c.dspark_taps = .{}; + } const neighbor_rdma = try neighbor.configured(); if (neighbor_rdma and o.world != 4) return error.NotFourNodes; @@ -246,7 +272,7 @@ pub const Model = struct { .streams = @intCast(o.streams), .round_rows = @intCast(o.round_rows), .rails = if (o.rdma_devices) |list| @intCast(std.mem.count(u8, list, ",") + 1) else 0, - .drafts = @intFromBool(o.drafts), + .drafts = @intFromBool(dspark), .replay_floor = @intFromBool(prompt.replay_floor), .idx_split = @intCast(prompt.split_from), .neighbor_rdma = @intFromBool(neighbor_rdma), @@ -347,7 +373,7 @@ pub const Model = struct { if (o.side) try m.rd.useSide(&m.eng); if (o.prefetch) try m.rd.usePrefetch(&m.eng); if (m.two) |t| if (t.ir != 0 and o.side and m.par == null) try m.rd.usePar(&m.eng); // parity in two launches: the rest beside the main gate / up - if (o.drafts) { + if (dspark) { m.dpool = try draft.Pool.init(&m.eng, &m.arena, m.streams); m.dr = try draft.Drafter.init(&m.eng, &m.arena, sp); if (o.graphs) { @@ -355,6 +381,12 @@ pub const Model = struct { m.dr.?.graphs = &m.pgraphs.?; } } + if (eagle_hdr) |hdr| { + m.eg = try eagle.Eagle.open(gpa, io, &m.eng, &m.arena, std.mem.span(std.c.getenv("TF_DS_EAGLE_HEAD").?), hdr, m.streams); + // TF_DS_EAGLE_CHECK=FILE: the head against eagle_ref.py's vectors before serving (its verdict on stderr) + if (std.c.getenv("TF_DS_EAGLE_CHECK")) |v| _ = try eagle.selfTest(m.eg.?, gpa, io, &m.eng, std.mem.span(v)); + } + errdefer if (m.eg) |x| x.close(gpa); try m.synchronize(); try d.check(d.api.cuCtxSynchronize(), "cuCtxSynchronize"); // the setup's legacy-stream copies too @@ -402,6 +434,7 @@ pub const Model = struct { m.spans.deinit(m.gpa); if (m.graphs) |*g| g.deinit(); if (m.pgraphs) |*g| g.deinit(); + if (m.eg) |x| x.close(m.gpa); if (m.timer) |t| { t.deinit(); m.gpa.destroy(t); @@ -515,7 +548,7 @@ pub const Model = struct { } pub fn drafting(m: *const Model) bool { - return m.dr != null; + return m.dr != null or m.eg != null; } // -- the primitives (both ranks, in the same order) ---------------------------------------------------------- @@ -633,6 +666,7 @@ pub const Model = struct { m.view = try m.caches.view(&m.eng, slot, base); for (m.rings, m.ring_view) |r, *v| v.* = r + slot * m.eng.ringBytes(); m.fill_slot = slot; + if (m.eg) |x| x.reset(slot); } /// The prompt's chunk of n rows from position `start` (seq: the stream's ids through the chunk's end) through every @@ -654,7 +688,7 @@ pub const Model = struct { try prompt.begin(e, ch, m.ids64[0..n], start, m.host_pos); if (ch.spans.len > 0) try prompt.embedSpans(e, ch); if (pt) |t| try t.mark(m.stream, .embed); - const taps = m.dr != null; + const taps = m.dr != null or m.eg != null; var floor: usize = 0; var kv_done: ?usize = null; // Engram rows read ahead on a thread while the GPU runs the layers before them (the served reader pool's @@ -738,6 +772,7 @@ pub const Model = struct { for (0..c.dspark_taps.slice().len) |j| try e.ops.copyRows(e.s, prompt.tapRows(e, ch, j), dw * 2, dr.at + j * dw * 2, dr.taps_w * 2, dw * 2, ch.n); try dr.absorb(e, ch, &m.dpool.?, m.fill_slot, dr.at, ch.n, ch.start); } + if (m.eg) |eg| try m.eagleAbsorbChunk(eg, start); // no synchronize: the next chunk's host steps (its first Engram layer's hashing, table reads and decode, ~0.14 s // at 2,048 rows) run while this chunk's kernels do; readers of the results (the prompt's logits) synchronize. // The profile (--profile 2) times a chunk to its end, so it alone waits. @@ -792,6 +827,13 @@ pub const Model = struct { } if (n > 0) try dr.absorbMany(e, &m.ch, &m.dpool.?, items[0..n]); } + if (m.eg) |eg| { + // the drafting windows' g rows (the others' rows go to the ring's sink row) + const wins = rows.windows orelse return error.BadRound; + var spans: [max_streams]eagle.Span = undefined; + for (wins, absorb, 0..) |win, yes, i| spans[i] = .{ .slot = @intCast(rows.slots.?[win.row]), .start = @intCast(rows.pos[win.row]), .n = win.n, .keep = yes }; + try eg.absorbRound(e, m.rd.taps, m.rd.xc, rows.ids.len, spans[0..wins.len]); + } try m.synchronize(); if (m.prof) m.t_absorb = m.now() - t1; } @@ -850,6 +892,39 @@ pub const Model = struct { return m.logits[0..m.vocab]; } + /// A prompt chunk's last ring rows into the head: taps (below layers / 2 at ch.start - start0) and final states. + fn eagleAbsorbChunk(m: *Model, eg: *eagle.Eagle, start0: usize) !void { + const e = &m.eng; + const c = e.c; + const ch = &m.ch; + const D = c.hidden; + const L = eg.s.n_taps; + const total = @min(ch.n, eg.ring); + const cut = ch.start - start0; + var at = ch.n - total; // (the rows before: no window reaches them) + while (at < ch.n) { + const k = @min(eagle.chunk_rows, ch.n - at); + for (eg.s.layerTaps(), 0..) |t, j| { + const off = (if (t < c.layers / 2) cut else 0) + at; + try e.ops.copyRows(e.s, prompt.tapRows(e, ch, j) + off * D * 2, D * 2, eg.tin + j * D * 2, L * D * 2, D * 2, k); + } + if (eg.s.final()) { + try tri_basic.collapseNorm(e.t, ch.h + at * c.hc * D * 2, ch.pre + at * c.hc * 4, e.w.norm, eg.fin, c.eps, k, D); + try e.ops.copyRows(e.s, eg.fin, D * 2, eg.tin + (L - 1) * D * 2, L * D * 2, D * 2, k); + } + try eg.tapDump(e, "prompt", m.fill_slot, ch.start + at, k, cut, eg.tin, L * D); + try eg.absorb(e, eg.tin, L * D, &.{.{ .slot = m.fill_slot, .start = ch.start + at, .n = k }}); + at += k; + } + } + + /// The TFE head's pass (rank 0): each request's drafts and confidences into m.eg's drafts / confs. + pub fn eaglePass(m: *Model, reqs: []const eagle.Req, steps: usize) !void { + const eg = m.eg orelse return error.NoDrafter; + try eg.pass(&m.eng, reqs, steps); + try m.ringsOk(); + } + /// The drafter's batched pass for N streams: each one's pending token at position q0 in its slot (its rings hold /// its positions below q0), `steps` drafts each (the drafter's drafts and confidences on the host). pub fn pass(m: *Model, tokens: []const i64, q0: []const i64, slots: []const i64, steps: usize) !void {