97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""
|
|
gen_wkv7_fixture.py -- Emit a C header of fixed-seed WKV7 fixtures for the
|
|
gradient-check test (test_wkv7_bwd.c).
|
|
|
|
It imports a *verified* WKV7 reference oracle (wkv7_reference.py, whose manual
|
|
reverse-scan backward matches autograd to rel<1e-9 in fp64; it is not shipped
|
|
here) and dumps, for a few fixed-seed shapes: the inputs (r,w,k,v,a,b), the
|
|
upstream grad dy, the forward y, and the reference grads (dr,dw,dk,dv,da,db).
|
|
|
|
Run: python3 rwkvx/tests/gen_wkv7_fixture.py /path/to/wkv7_reference.py \
|
|
> rwkvx/tests/wkv7_fixture.h
|
|
(or set WKV7_REFERENCE=/path/to/wkv7_reference.py instead of the argument)
|
|
(torch CPU is fine). The generated header is committed so the C test runs
|
|
without torch; regenerate only if the oracle changes.
|
|
"""
|
|
import os
|
|
import sys
|
|
import importlib.util
|
|
|
|
ORACLE = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("WKV7_REFERENCE")
|
|
if not ORACLE or not os.path.isfile(ORACLE):
|
|
sys.exit("usage: gen_wkv7_fixture.py <path/to/wkv7_reference.py> "
|
|
"(or set WKV7_REFERENCE)")
|
|
|
|
spec = importlib.util.spec_from_file_location("wkv7_reference", ORACLE)
|
|
oracle = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(oracle)
|
|
|
|
import torch # noqa: E402
|
|
|
|
|
|
def emit_array(name, t):
|
|
"""Emit a flat C float array from a tensor (row-major [B,T,H,N])."""
|
|
flat = t.detach().to(torch.float64).reshape(-1).tolist()
|
|
print(f"static const float {name}[] = {{")
|
|
line = " "
|
|
for i, x in enumerate(flat):
|
|
tok = f"{x:.9e}f,"
|
|
if len(line) + len(tok) > 96:
|
|
print(line)
|
|
line = " "
|
|
line += tok + " "
|
|
if line.strip():
|
|
print(line)
|
|
print("};")
|
|
|
|
|
|
def emit_case(idx, B, T, H, N, seed, dy_mode):
|
|
ins = oracle._rand_inputs(B, T, H, N, torch.float64, seed=seed,
|
|
requires_grad=True)
|
|
r, w, k, v, a, b = ins
|
|
y = oracle.wkv7_forward(r, w, k, v, a, b)
|
|
if dy_mode == "ones":
|
|
dy = torch.ones_like(y)
|
|
else:
|
|
g = torch.Generator().manual_seed(seed + 1000)
|
|
dy = torch.randn(*y.shape, generator=g, dtype=torch.float64)
|
|
y.backward(dy)
|
|
grads = {"dr": r.grad, "dw": w.grad, "dk": k.grad,
|
|
"dv": v.grad, "da": a.grad, "db": b.grad}
|
|
|
|
p = f"c{idx}_"
|
|
print(f"\n/* ---- case {idx}: B={B} T={T} H={H} N={N} seed={seed} "
|
|
f"dy={dy_mode} ---- */")
|
|
print(f"#define {p.upper()}B {B}")
|
|
print(f"#define {p.upper()}T {T}")
|
|
print(f"#define {p.upper()}H {H}")
|
|
print(f"#define {p.upper()}N {N}")
|
|
for nm, t in [("r", r), ("w", w), ("k", k), ("v", v), ("a", a), ("b", b),
|
|
("dy", dy), ("y", y)]:
|
|
emit_array(p + nm, t)
|
|
for nm, t in grads.items():
|
|
emit_array(p + nm, t)
|
|
|
|
|
|
def main():
|
|
print("/* AUTO-GENERATED by rwkvx/tests/gen_wkv7_fixture.py -- DO NOT EDIT.")
|
|
print(" * Source oracle: wkv7_reference.py (fp64 PyTorch WKV7 reference)")
|
|
print(" * (verified: manual backward == autograd, rel<1e-9 fp64).")
|
|
print(" * Reference grads are fp64, stored as float. */")
|
|
print("#ifndef WKV7_FIXTURE_H")
|
|
print("#define WKV7_FIXTURE_H")
|
|
# canonical fixture from the doc's gradient-check plan
|
|
emit_case(0, 1, 4, 1, 4, 42, "ones")
|
|
# larger, multi-head, random dy -> exercises the reverse scan harder
|
|
emit_case(1, 1, 16, 2, 16, 7, "randn")
|
|
# a third with B>1 and different N
|
|
emit_case(2, 2, 8, 3, 8, 123, "randn")
|
|
print("\n#endif /* WKV7_FIXTURE_H */")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|