import json import os import subprocess import sys import time import traceback os.environ.setdefault("CHIPFORMER_WROM_HEADS", "320") R = {"sections": {}} def sh(cmd, tail=600): t = time.time() p = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True) out = ((p.stdout or "") + (p.stderr or ""))[-tail:] print(f"[{time.time() - t:.0f}s rc={p.returncode}] {out}", flush=True) return p.returncode, out def safe(name): def deco(fn): try: R["sections"][name] = fn() print(f"OK {name}: {json.dumps(R['sections'][name])[:300]}", flush=True) except Exception as e: R["sections"][name] = {"error": f"{type(e).__name__}: {e}", "tb": traceback.format_exc()[-800:]} print(f"FAIL {name}: {type(e).__name__}: {e}", flush=True) return fn return deco @safe("hardware") def hardware(): import torch g = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"], capture_output=True, text=True).stdout.strip() name, mem, drv = [s.strip() for s in g.split(",")] return {"gpu": name, "vram_gb": round(int(mem.split()[0]) / 1024, 1), "driver": drv, "torch": torch.__version__, "cuda": torch.version.cuda, "capability": ".".join(map(str, torch.cuda.get_device_capability())), "cpus": os.cpu_count(), "ram_gb": round(int(open("/proc/meminfo").readline().split()[1]) / 1048576, 1), "python": sys.version.split()[0]} A = R["sections"] @safe("setup") def setup(): sh("apt-get -qq install -y clang lld >/dev/null 2>&1; clang --version | head -1") rc1, _ = sh("cd /content && rm -rf qwenputer && git clone -q https://github.com/divagr18/qwenputer " "&& cd qwenputer && git checkout -q 2717232") rc2, _ = sh("cd /content && rm -rf bundle && mkdir bundle && unzip -q colab_bundle.zip -d bundle " "&& cp -r bundle/colab/patches/* qwenputer/ && cp bundle/colab/tl/* tl/") rc3, out = sh("cd /content/tl && python3 export_full.py 2>&1 | tail -1") rc4, _ = sh("cd /content/tl && clang --target=riscv32-unknown-elf -march=rv32im_zba -mabi=ilp32 " "-O2 -DNTOK=40 -nostdlib -nostartfiles -static -fuse-ld=lld " "-T /content/qwenputer/experimental/riscv/guests/stories260k/vd.ld " "-I /content/qwenputer/experimental/riscv/chipformer/rv/env -I . " "-o tl40.elf /content/qwenputer/experimental/riscv/chipformer/rv/env/crt0.S " "tl.c vd2.S wrom_tl.c minic.c 2>&1 | grep ' error' | head -2") return {"export": out.strip(), "rc": [rc1, rc2, rc3, rc4]} @safe("model") def model(): from huggingface_hub import snapshot_download p = snapshot_download("Maykeye/TinyLLama-v0", allow_patterns=["model.safetensors", "config.json", "tokenizer.json"]) for f in ("model.safetensors", "config.json", "tokenizer.json"): subprocess.run(["cp", f"{p}/{f}", f"/content/tl/{f}"], check=True) from safetensors import safe_open n = 0 with safe_open("/content/tl/model.safetensors", "pt") as f: for k in f.keys(): m = 1 for d in f.get_slice(k).get_shape(): m *= d n += m return {"params": n, "bf16_mb": round(n * 2 / 2 ** 20, 2), "int4_mb": round(n * 0.5 / 2 ** 20, 2), "fp32_mb": round(n * 4 / 2 ** 20, 2)} @safe("trace") def trace(): import collections from chipformer.rv import elf, ref from chipformer.rv import spec as S m = ref.Machine.from_program(elf.program("/content/tl/tl40.elf")) mix = collections.Counter() toks = 0 t0 = time.time() while not m.halted and m.ram[0xE000 // 4 + 100] == 0: w = m.rom[(m.pc - S.ROM_BASE) >> 2] if m.n != S.N_MAX else None mix[ref.decode(w).op if w is not None else "R"] += 1 toks += len(m.record()) tot = sum(mix.values()) return {"instructions_one_token": tot, "trace_tokens_one_token": toks, "cpu_seconds": round(time.time() - t0, 1), "tokens_per_instruction": round(toks / tot, 3), "mix": {k: v for k, v in mix.most_common(12)}, "vdot": mix["VDOT"], "vdot_pct": round(100 * mix["VDOT"] / tot, 1), "overhead_pct": round(100 * (tot - mix["VDOT"]) / tot, 1)} @safe("executor") def executor(): import numpy as np import torch from chipformer.rv import core, elf, fastexec, ref from chipformer.rv import spec as S c = core.build() prog = elf.program("/content/tl/tl40.elf") want = ref.Machine.from_program(prog).run(10 ** 8) out = {"circuit": c.stats(), "reference_tokens": len(want)} core.snapshot(c, ref.Machine.from_program(prog)) N = 200000 fr = fastexec.FreeRun(c, 1, "cuda", False, v2=True) fr.load(0, S.tok("GO")) t0 = time.time() got = fr.run(N)[0] torch.cuda.synchronize() dt = time.time() - t0 exp = np.array(want[1:N + 1]) bad = np.nonzero(got != exp)[0] out.update(tokens=N, seconds=round(dt, 1), tokens_per_second=round(N / dt, 1), mismatches=int(len(bad)), first_mismatch=None if not len(bad) else int(bad[0]), minutes_per_model_token=round(2470000 / (N / dt) / 60, 1)) return out @safe("merge") def merge(): from huggingface_hub import snapshot_download p = snapshot_download("Qwen/Qwen3.5-4B", allow_patterns=["*.safetensors", "*.json", "*.jinja"]) rc, out = sh(f"cd /content/qwenputer/experimental/riscv && CHIPFORMER_WROM_HEADS=320 " f"python3 -m chipformer.rv.merge {p} --out /content/merged 2>&1 | tail -4", 1500) size = subprocess.run(["bash", "-lc", "du -sm /content/merged | cut -f1"], capture_output=True, text=True).stdout.strip() import re layers = re.findall(r"'layers': (\d+)", out) return {"rc": rc, "size_mb": int(size or 0), "circuit_layers": int(layers[0]) if layers else None, "tail": out.strip()[-300:]} @safe("degradation") def degradation(): import torch from transformers import AutoModelForCausalLM, AutoTokenizer tok = AutoTokenizer.from_pretrained("/content/merged") base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16, device_map="cuda").eval() merged = AutoModelForCausalLM.from_pretrained("/content/merged", dtype=torch.bfloat16, device_map="cuda").eval() def logsm(z): z = z - z.max(-1, keepdim=True).values return z - torch.log(torch.exp(z).sum(-1, keepdim=True)) texts = ["The capital of France is Paris, and the river that runs through it is the Seine.", "In 1969, humans first landed on the Moon, a milestone of the space race.", "A neural network learns by adjusting its weights to reduce a loss function.", "The recipe calls for two cups of flour, one egg, and a pinch of salt.", "Photosynthesis converts light energy into chemical energy inside chloroplasts.", "The committee approved the budget after a long debate about public transport.", "Mount Everest rises 8,849 metres above sea level on the border of Nepal and China.", "She opened the book, read the first page twice, and put it down again."] rows, kls, pos = [], [], 0 with torch.no_grad(): for t in texts: ids = tok(t, return_tensors="pt").to("cuda") a, b = base(**ids).logits[0].float(), merged(**ids).logits[0].float() n = min(a.shape[-1], b.shape[-1]) a, b = a[..., :n], b[..., :n] k = (torch.exp(logsm(a)) * (logsm(a) - logsm(b))).sum(-1).mean().item() agree = (a.argmax(-1) == b.argmax(-1)).float().mean().item() rows.append({"kw": round(k, 6), "top1": round(agree, 3), "positions": int(a.shape[0])}) kls.append(k) pos += int(a.shape[0]) ids = tok("The capital of France is", return_tensors="pt").to("cuda") with torch.no_grad(): out = merged.generate(**ids, max_new_tokens=20, do_sample=False) text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True) merged.generate(**ids, max_new_tokens=9, do_sample=False) torch.cuda.synchronize() t0 = time.time() merged.generate(**ids, max_new_tokens=32, do_sample=False) torch.cuda.synchronize() dt = time.time() - t0 return {"per_text": rows, "kl_mean": round(sum(kls) / len(kls), 6), "positions": pos, "sample_continuation": text, "decode_tok_s": round(32 / dt, 2), "peak_vram_gb": round(torch.cuda.max_memory_allocated() / 2 ** 30, 1), "base_layers": len(base.model.layers), "merged_layers": len(merged.model.layers)} json.dump(R, open("/content/results.json", "w"), indent=1) print("RESULTS " + json.dumps(R), flush=True)