Release 2026-10-05: row graphs for 64, 128 and 256 tokens, shared-state pairs, final Gated DeltaNet kernel
90285c9 verified | """The r13_kernel forms on Kev-4B shapes, torch only (random inputs; no full-model forward). | |
| python r15_kernel_test.py --model 4b --out results/r15_kernel_test_4b.json | |
| The 4B GatedDeltaNet has 16 key heads and 32 value heads (ratio 2): the guarded forward copies every q / k head with | |
| kev_qwen35_patch._litert_interleave_heads (concat + rank-4 reshape, patch guard (6)) before the chunk kernel, so the | |
| kernel sees 32 heads of q / k / v with dk = dv = 128 (the 0.8B: 16 heads, no copy). Two levels: | |
| 1. kernel: q / k drawn with 16 heads and interleaved to 32 (every head pair identical, as in the 4B), v with 32 heads, | |
| in four regimes (R1 random keys + strong decay, R2 weak decay, R3 correlated keys + weak decay, R4 one key for every | |
| token + no decay), L 256 / 130 (tail pad) / 200, with and without an initial state. Per form, the kernel part of the | |
| form (A = R64+ec, B = R64+ec+dd+vs8; sp acts in the GatedDeltaNet forward, not in the kernel) vs the unmodified | |
| kernel (kev_qwen35_patch._rank4_chunk_gated_delta_rule = the loop kernel) and vs the float64 token-by-token | |
| recurrence; vs8 multiplies v by 2^8 at the kernel entry, so the output AND the recurrent state it carries are exactly | |
| 2^8 larger (the delta rule is linear in (v, state) jointly): an initial state must come in scaled by 2^8 too (= the | |
| shared-state pair's gdn_state contract under vs8: the state outputs are 2^8 x the stock ones), and the output / final | |
| state are compared after / 2^8. The unscaled-state run (the wrong contract) is recorded as a control. | |
| 2. layer: real 4B GatedDeltaNet layers (weights of merged/kev-4b-v1.0 read from the safetensors, nothing else loaded) | |
| in the patched class; input = random hidden states with a right pad (valid = 1 then 0: the pad guard); the guarded | |
| forward with the loop kernel vs with r13_kernel.apply(<holder>, form) (the kernel, the softplus rewrite in the | |
| guarded forward's F, vs8's gated norm). Reported: out max |diff| at the real positions (absolute and relative to max | |
| |out|), finiteness, and the interleave calls (2 per forward = q and k). | |
| Never overwrites --out.""" | |
| import argparse | |
| import json | |
| import time | |
| import torch | |
| import kev_qwen35_patch as P | |
| import r13_kernel as R | |
| from r2_common import CKPT, K, MODEL, add_model_arg, dump_json | |
| FORMS = {"A": "R64+sp+ec", "B": "R64+sp+ec+dd+vs8"} | |
| KERNELS = {"A": dict(inverse="R64", expclamp=True, decay="cumsum", vscale_log2=0), | |
| "B": dict(inverse="R64", expclamp=True, decay="direct", vscale_log2=8)} | |
| def regime(name, L, seed, hk, hv, D, init_state=False): | |
| """q / k with hk heads interleaved to hv (ratio hv / hk), v with hv heads.""" | |
| gen = torch.Generator().manual_seed(seed) | |
| r = hv // hk | |
| q = torch.randn(1, L, hk, D, generator=gen) | |
| if name in ("R1", "R2"): | |
| k = torch.randn(1, L, hk, D, generator=gen) | |
| elif name == "R3": | |
| u = torch.randn(1, 1, hk, D, generator=gen) | |
| u = u / u.norm(dim=-1, keepdim=True) * D ** 0.5 | |
| k = u + 0.3 * torch.randn(1, L, hk, D, generator=gen) | |
| else: | |
| k = torch.randn(1, 1, hk, D, generator=gen).expand(1, L, hk, D).contiguous() | |
| v = torch.randn(1, L, hv, D, generator=gen) | |
| if name == "R1": | |
| beta = 0.2 + 0.7 * torch.rand(1, L, hv, generator=gen) | |
| g = -(0.05 + 2.95 * torch.rand(1, L, hv, generator=gen)) | |
| elif name in ("R2", "R3"): | |
| beta = (0.2 if name == "R2" else 0.3) + (0.7 if name == "R2" else 0.7) * torch.rand(1, L, hv, generator=gen) | |
| g = -0.02 * torch.rand(1, L, hv, generator=gen) | |
| else: | |
| beta = torch.full((1, L, hv), 0.5) | |
| g = torch.zeros(1, L, hv) | |
| s0 = torch.randn(1, hv, D, D, generator=gen) * 0.1 if init_state else None | |
| if r > 1: | |
| q, k = P._litert_interleave_heads(q, r), P._litert_interleave_heads(k, r) | |
| return q, k, v, g, beta, s0 | |
| def truth_f64(q, k, v, g, beta, s0): | |
| """The gated delta recurrence in float64 (l2norm eps 1e-6, q scaled by D^-0.5), any head count.""" | |
| D = q.shape[-1] | |
| q, k, v, g, beta = [x.to(torch.float64) for x in (q, k, v, g, beta)] | |
| q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6) | |
| k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6) | |
| q, k, v = [x.transpose(1, 2) for x in (q, k, v)] # [1, H, L, D] | |
| g, beta = g.transpose(1, 2), beta.transpose(1, 2) # [1, H, L] | |
| q = q * D ** -0.5 | |
| H = q.shape[1] | |
| s = torch.zeros(1, H, D, v.shape[-1], dtype=torch.float64) if s0 is None else s0.to(torch.float64).clone() | |
| outs = [] | |
| for t in range(q.shape[2]): | |
| s = s * g[:, :, t].exp()[..., None, None] | |
| kv = (s * k[:, :, t, :, None]).sum(-2) | |
| delta = (v[:, :, t] - kv) * beta[:, :, t, None] | |
| s = s + k[:, :, t, :, None] * delta[:, :, None, :] | |
| outs.append((s * q[:, :, t, :, None]).sum(-2)) | |
| return torch.stack(outs, 2).transpose(1, 2), s | |
| def kernel_level(hk, hv, D): | |
| rows = [] | |
| for name, L, seed, init in (("R1", 256, 11, False), ("R1", 130, 12, True), ("R2", 256, 21, False), | |
| ("R2", 130, 22, True), ("R3", 256, 31, False), ("R3", 130, 32, True), | |
| ("R4", 200, 41, True), ("R4", 256, 42, False)): | |
| q, k, v, g, beta, s0 = regime(name, L, seed, hk, hv, D, init) | |
| o_t, s_t = truth_f64(q, k, v, g, beta, s0) | |
| o_ref, s_ref = P._rank4_chunk_gated_delta_rule(q, k, v, g, beta, initial_state=s0, output_final_state=True, | |
| use_qk_l2norm_in_kernel=True) | |
| row = {"regime": name, "L": L, "initial_state": init, "heads_q_k_drawn": hk, "heads_kernel": hv, "dim": D, | |
| "head_pairs_identical": bool(torch.equal(q[:, :, 0::2], q[:, :, 1::2])) if hv > hk else None, | |
| "max_abs_truth_out": float(o_t.abs().max()), "max_abs_truth_state": float(s_t.abs().max()), | |
| "stock_vs_truth_out": float((o_ref.double() - o_t).abs().max()), | |
| "stock_vs_truth_state": float((s_ref.double() - s_t).abs().max()), "forms": {}} | |
| for f, kw in KERNELS.items(): | |
| kern = R.make_kernel(**kw) | |
| scale = float(2 ** kw["vscale_log2"]) | |
| # vs<k>: the kernel scales v by 2^k at its entry; the recurrent state it carries (and outputs) is then 2^k x | |
| # the stock state, so a continued state must come in scaled too (the pair's gdn_state contract under vs<k>) | |
| s_in = None if s0 is None else s0 * scale | |
| o, s = kern(q, k, v, g, beta, initial_state=s_in, output_final_state=True, use_qk_l2norm_in_kernel=True) | |
| o, s = o / scale, s / scale | |
| if s0 is not None and scale != 1.0: # the wrong contract (unscaled state) for the record | |
| o_u, _ = kern(q, k, v, g, beta, initial_state=s0, output_final_state=True, use_qk_l2norm_in_kernel=True) | |
| unscaled_state_vs_stock = float((o_u / scale - o_ref).abs().max()) | |
| else: | |
| unscaled_state_vs_stock = None | |
| row["forms"][f] = {"kernel": kern.__name__, "finite": bool(torch.isfinite(o).all() and torch.isfinite(s).all()), | |
| "vs_stock_out": float((o - o_ref).abs().max()), | |
| "vs_stock_state": float((s - s_ref).abs().max()), | |
| "vs_truth_out": float((o.double() - o_t).abs().max()), | |
| "vs_truth_state": float((s.double() - s_t).abs().max()), | |
| "initial_state_scaled_by": scale if s0 is not None else None, | |
| "unscaled_initial_state_vs_stock_out": unscaled_state_vs_stock} | |
| rows.append(row) | |
| print(json.dumps({"regime": name, "L": L, "init": init, "stock_vs_truth": row["stock_vs_truth_out"], | |
| **{f: (round(x["vs_stock_out"], 10), round(x["vs_truth_out"], 10)) for f, x in row["forms"].items()}}), | |
| flush=True) | |
| return rows | |
| class Holder(torch.nn.Module): | |
| """r13_kernel.apply / reset walk model.modules() and model.layers.""" | |
| def __init__(self, gdn): | |
| super().__init__() | |
| self.layers = torch.nn.ModuleList([gdn]) | |
| def load_gdn(layer): | |
| """One PatchedQwen3_5GatedDeltaNet with the checkpoint's weights for decoder layer `layer`.""" | |
| from safetensors import safe_open | |
| from transformers import AutoConfig | |
| cfg = AutoConfig.from_pretrained(str(CKPT)) | |
| cfg._attn_implementation = "eager" | |
| m = P.PatchedQwen3_5GatedDeltaNet(cfg, layer).eval() | |
| assert m.chunk_gated_delta_rule is P._rank4_chunk_gated_delta_rule | |
| idx = json.loads((CKPT / "model.safetensors.index.json").read_text())["weight_map"] | |
| prefix = f"layers.{layer}.linear_attn." | |
| names = [n for n in idx if n.startswith(prefix)] | |
| sd = {} | |
| for shard in sorted({idx[n] for n in names}): | |
| with safe_open(str(CKPT / shard), framework="pt") as f: | |
| for n in names: | |
| if idx[n] == shard: | |
| sd[n[len(prefix):]] = f.get_tensor(n) | |
| m.load_state_dict(sd, strict=True) | |
| return m, {"layer": layer, "tensors": sorted(sd), "num_k_heads": m.num_k_heads, "num_v_heads": m.num_v_heads, | |
| "head_k_dim": m.head_k_dim, "head_v_dim": m.head_v_dim, "conv_dim": m.conv_dim, "norm": type(m.norm).__name__} | |
| def layer_level(layers, L=200, n_real=173, seeds=(5, 6), scales=(1.0, 3.0)): | |
| import r11_fp16_safe as S11 | |
| g = P.PatchedQwen3_5GatedDeltaNet.forward.__globals__ | |
| inner = g["_litert_interleave_heads"] | |
| calls = {"n": 0} | |
| def counted(x, r): | |
| calls["n"] += 1 | |
| return inner(x, r) | |
| g["_litert_interleave_heads"] = counted | |
| out = [] | |
| try: | |
| for layer in layers: | |
| m, info = load_gdn(layer) | |
| holder = Holder(m) | |
| valid = torch.zeros(1, L) | |
| valid[0, :n_real] = 1.0 | |
| for seed in seeds: | |
| for sc in scales: | |
| x = torch.randn(1, L, m.hidden_size, generator=torch.Generator().manual_seed(seed)) * sc | |
| m._litert_valid = valid | |
| c0 = calls["n"] | |
| with torch.no_grad(): | |
| ref = m(x) | |
| stock_calls = calls["n"] - c0 | |
| row = {**info, "seed": seed, "input_scale": sc, "L": L, "n_real": n_real, | |
| "interleave_calls_stock": stock_calls, "max_abs_out_stock": float(ref[0, :n_real].abs().max()), | |
| "forms": {}} | |
| for f, spec in FORMS.items(): | |
| applied = R.apply(holder, spec) | |
| try: | |
| c1 = calls["n"] | |
| with torch.no_grad(): | |
| y = m(x) | |
| d = (y - ref)[0, :n_real].abs().max() | |
| row["forms"][f] = {"spec": spec, "kernel": applied["kernel"], "r11": applied["r11_applied"], | |
| "finite": bool(torch.isfinite(y).all()), | |
| "interleave_calls": calls["n"] - c1, | |
| "vs_stock_out_max_abs": float(d), | |
| "vs_stock_out_rel": float(d / ref[0, :n_real].abs().max())} | |
| finally: | |
| R.reset(holder) | |
| assert S11.APPLIED == [], S11.APPLIED | |
| out.append(row) | |
| print(json.dumps({"layer": layer, "seed": seed, "scale": sc, "interleave_stock": stock_calls, | |
| **{f: (x["vs_stock_out_max_abs"], x["vs_stock_out_rel"], x["interleave_calls"]) | |
| for f, x in row["forms"].items()}}), flush=True) | |
| finally: | |
| g["_litert_interleave_heads"] = inner | |
| return out | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--out", required=True) | |
| add_model_arg(ap) | |
| a = ap.parse_args() | |
| out = K / a.out | |
| assert not out.exists(), out | |
| torch.set_num_threads(4) | |
| t0 = time.perf_counter() | |
| cfg = json.loads((CKPT / "config.json").read_text()) | |
| hk, hv, D = cfg["linear_num_key_heads"], cfg["linear_num_value_heads"], cfg["linear_key_head_dim"] | |
| doc = {"step": "r13_kernel forms on the model's GatedDeltaNet shapes (random inputs)", | |
| "model": MODEL, "checkpoint": str(CKPT.relative_to(K)), "forms": FORMS, "kernel_parts": KERNELS, | |
| "heads_k": hk, "heads_v": hv, "dim": D} | |
| doc["kernel"] = kernel_level(hk, hv, D) | |
| gdn_layers = [i for i, t in enumerate(cfg["layer_types"]) if t == "linear_attention"] | |
| doc["layer"] = layer_level([gdn_layers[0], gdn_layers[3]]) | |
| kmax = {f: max(r["forms"][f]["vs_stock_out"] / max(r["max_abs_truth_out"], 1e-30) for r in doc["kernel"]) for f in FORMS} | |
| smax = {f: max(r["forms"][f]["vs_stock_state"] / max(r["max_abs_truth_state"], 1e-30) for r in doc["kernel"]) | |
| for f in FORMS} | |
| lmax = {f: max(r["forms"][f]["vs_stock_out_rel"] for r in doc["layer"]) for f in FORMS} | |
| doc["summary"] = { | |
| "kernel_vs_stock_out_rel_max": kmax, "kernel_vs_stock_state_rel_max": smax, | |
| "stock_vs_truth_out_max": max(r["stock_vs_truth_out"] for r in doc["kernel"]), | |
| "forms_vs_truth_out_max": {f: max(r["forms"][f]["vs_truth_out"] for r in doc["kernel"]) for f in FORMS}, | |
| "layer_vs_stock_out_rel_max": lmax, | |
| "all_finite": all(x["finite"] for r in doc["kernel"] for x in r["forms"].values()) | |
| and all(x["finite"] for r in doc["layer"] for x in r["forms"].values()), | |
| "head_pairs_identical": all(r["head_pairs_identical"] for r in doc["kernel"]) if hv > hk else None, | |
| "interleave_calls_per_forward": sorted({r["interleave_calls_stock"] for r in doc["layer"]} | |
| | {x["interleave_calls"] for r in doc["layer"] for x in r["forms"].values()}), | |
| } | |
| doc["seconds"] = round(time.perf_counter() - t0, 1) | |
| dump_json(out, doc) | |
| print(json.dumps(doc["summary"], indent=1)) | |
| if __name__ == "__main__": | |
| main() | |