| """Standalone in-weights runner for `Qwen3.5-9B-Math-GraphCalc` (and `0.8B`). |
| |
| Requires `qwenputer` (`https://github.com/divagr18/qwenputer`) on `PYTHONPATH` or pointed to by |
| `QWENPUTER_ROOT` (defaults to `./qwenputer` or `../qwenputer`): |
| |
| git clone https://github.com/divagr18/qwenputer |
| QWENPUTER_ROOT=./qwenputer python runner.py --ask "Calculate 37429 * 8613 + 194820 and explain your steps." |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| import gc |
| import json |
| import os |
| import re |
| import sys |
| import time |
| import zlib |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| from tokenizers import Tokenizer |
| from transformers import AutoConfig, AutoTokenizer, DynamicCache, Qwen3_5ForCausalLM |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as q35_mod |
|
|
| MODEL_DIR = Path(__file__).resolve().parent |
|
|
|
|
| def _locate_qwenputer_rv() -> Path: |
| candidates = [] |
| if os.environ.get("QWENPUTER_ROOT"): |
| candidates.append(Path(os.environ["QWENPUTER_ROOT"])) |
| candidates.extend([ |
| Path.cwd() / "qwenputer", |
| Path.cwd().parent / "qwenputer", |
| MODEL_DIR / "qwenputer", |
| MODEL_DIR.parent / "qwenputer", |
| MODEL_DIR.parent.parent / "qwenputer", |
| MODEL_DIR.parent.parent.parent / "qwenputer", |
| ]) |
| for root in candidates: |
| rv = root.resolve() / "experimental" / "riscv" |
| if (rv / "chipformer").is_dir(): |
| return rv |
| raise RuntimeError( |
| "Could not locate `chipformer` from `qwenputer` (`experimental/riscv/chipformer`). " |
| "Clone https://github.com/divagr18/qwenputer and set QWENPUTER_ROOT=/path/to/qwenputer." |
| ) |
|
|
|
|
| try: |
| import chipformer |
| except ImportError: |
| sys.path.insert(0, str(_locate_qwenputer_rv())) |
|
|
| from chipformer.q35 import fast as F_chip |
| from chipformer.q35 import ir |
| from chipformer.q35.ir import Target |
| from chipformer.rv import bridge as BR |
| from chipformer.rv import core, elf, hf, ref |
| from chipformer.rv import spec as S |
|
|
| CARBON_SPECIALS = [ |
| "<tool_response>", |
| "</tool_response>", |
| "<think>", |
| "</think>", |
| "<|im_end|>", |
| "<|im_start|>", |
| "user", |
| "assistant", |
| "\n\n", |
| ] |
|
|
| GRAPHCALC_TOOL = { |
| "type": "function", |
| "function": { |
| "name": "graphcalc", |
| "description": ( |
| "Exact in-weights graphing & math calculator (62-digit integers/rationals, scientific " |
| "functions, equation solver, and coordinate plotter). Pass an arithmetic expression, " |
| "an equation in x (LHS = RHS), or a function f(x); optionally pass xmin and xmax to " |
| "plot y = f(x) and compute roots, vertex/extrema, definite integral, and value table." |
| ), |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "expression": { |
| "type": "string", |
| "description": "Math expression, equation (LHS = RHS), or function of x (e.g. 'x^2 - 4*x + 3')", |
| }, |
| "xmin": { |
| "type": "number", |
| "description": "Optional domain/plot minimum x (default -5 for plots, -20 for nonlinear solve)", |
| }, |
| "xmax": { |
| "type": "number", |
| "description": "Optional domain/plot maximum x (default 5 for plots, 20 for nonlinear solve)", |
| }, |
| }, |
| "required": ["expression"], |
| }, |
| }, |
| } |
|
|
| GRAPHCALC_SYSTEM_PROMPT = ( |
| "You are a clear, encouraging step-by-step math tutor equipped with an embedded graphing " |
| "calculator (`graphcalc`). Always call `graphcalc` first to compute exact arithmetic, roots, " |
| "coordinate tables, or function plots, then explain every step clearly so a student can learn." |
| ) |
|
|
| _orig_chunk_gdn = q35_mod.torch_chunk_gated_delta_rule |
| _orig_rec_gdn = q35_mod.torch_recurrent_gated_delta_rule |
|
|
|
|
| def _fp32_chunk_gdn(*args, **kwargs): |
| with torch.autocast("cuda", enabled=False): |
| return _orig_chunk_gdn(*args, **kwargs) |
|
|
|
|
| def _fp32_rec_gdn(*args, **kwargs): |
| with torch.autocast("cuda", enabled=False): |
| return _orig_rec_gdn(*args, **kwargs) |
|
|
|
|
| q35_mod.torch_chunk_gated_delta_rule = _fp32_chunk_gdn |
| q35_mod.torch_recurrent_gated_delta_rule = _fp32_rec_gdn |
| F_chip.exact_matmuls() |
|
|
|
|
| class CarbonText(BR.Text): |
| """Qwen3.5 bridge vocabulary with native ChatML special tokens (256..264) and space HALT echo.""" |
|
|
| def __init__(self, tokenizer_path: str | Path) -> None: |
| p = Path(tokenizer_path) |
| tj = p if p.is_file() else (p / "tokenizer.json") |
| self.tok = Tokenizer.from_file(str(tj)) |
| vocab = self.tok.get_vocab(with_added_tokens=True) |
| base_vocab = {s: i for s, i in vocab.items() if not s.startswith("<|chip_")} |
| self.n = max(base_vocab.values()) + 1 |
| dec = {c: b for b, c in BR.bytes_to_unicode().items()} |
| added = ( |
| {a.content for a in self.tok.get_added_tokens_decoder().values()} |
| if hasattr(self.tok, "get_added_tokens_decoder") |
| else set() |
| ) |
| self.bytes = [b""] * self.n |
| for s, i in base_vocab.items(): |
| if s in added or not all(ch in dec for ch in s): |
| self.bytes[i] = s.encode("utf-8") |
| else: |
| self.bytes[i] = bytes(dec[ch] for ch in s) |
| self.TCO = base_vocab["<tool_call>"] |
| self.TCC = base_vocab["</tool_call>"] |
| u = BR.bytes_to_unicode() |
| special_ids = [] |
| for sp in CARBON_SPECIALS: |
| if sp in base_vocab: |
| special_ids.append(base_vocab[sp]) |
| else: |
| enc = self.tok.encode(sp, add_special_tokens=False).ids |
| assert len(enc) == 1, f"Special {sp!r} is not a single Qwen token: {enc}" |
| special_ids.append(enc[0]) |
| self.printable = [base_vocab[u[b]] for b in range(256)] + special_ids |
| self.pindex = {q: p for p, q in enumerate(self.printable)} |
| self.NL = self.printable[32] |
|
|
|
|
| def _compact_layer_inplace(layer: nn.Module, is_chip: bool) -> tuple[int, int]: |
| before = after = 0 |
| if is_chip and hasattr(layer, "linear_attn") and isinstance(layer.linear_attn.out_proj, nn.Linear): |
| F_chip.prune_heads(layer.linear_attn) |
| for parent in list(layer.modules()): |
| for name, child in list(parent.named_children()): |
| if isinstance(child, nn.Linear) and child.bias is None: |
| sp = F_chip.SparseLinear(child.weight) |
| before += child.weight.numel() |
| after += sp.weight.numel() |
| setattr(parent, name, sp) |
| return before, after |
|
|
|
|
| def _load_and_compact_qwen35(path: Path, n_chip: int, dev: str) -> tuple[Qwen3_5ForCausalLM, dict]: |
| layout_file = path / "chip_layout.json" |
| if not layout_file.is_file(): |
| model = Qwen3_5ForCausalLM.from_pretrained( |
| str(path), |
| dtype=torch.bfloat16, |
| attn_implementation="sdpa", |
| ) |
| chip_b = chip_a = qwen_b = qwen_a = 0 |
| for i, layer in enumerate(model.model.layers): |
| b, a = _compact_layer_inplace(layer, is_chip=(i < n_chip)) |
| if i < n_chip: |
| chip_b += b |
| chip_a += a |
| else: |
| qwen_b += b |
| qwen_a += a |
| return model.to(dev).eval(), { |
| "chip_dense": chip_b, |
| "chip_compact": chip_a, |
| "qwen_dense": qwen_b, |
| "qwen_compact": qwen_a, |
| } |
|
|
| layout = json.loads(layout_file.read_text(encoding="utf-8")) |
| blob = (path / "chip_indices.bin.zlib").read_bytes() |
| cfg = AutoConfig.from_pretrained(str(path)) |
| cfg._attn_implementation = "sdpa" |
| with torch.device("meta"): |
| model = Qwen3_5ForCausalLM(cfg) |
|
|
| tensors_spec: dict[str, dict] = layout["tensors"] |
| by_layer: list[list[str]] = [[] for _ in range(cfg.num_hidden_layers)] |
| for k in tensors_spec: |
| if k.startswith("model.layers."): |
| lid = int(k.split(".", 3)[2]) |
| by_layer[lid].append(k) |
|
|
| chip_before = chip_after = qwen_before = qwen_after = head_before = head_after = 0 |
| with safe_open(str(path / "model.safetensors"), framework="pt", device="cpu") as f: |
| def _unpack_one(k: str) -> torch.Tensor: |
| spec = tensors_spec[k] |
| vals = f.get_tensor(k) |
| if spec["mode"] == "dense": |
| return vals |
| shape = tuple(spec["shape"]) |
| raw = zlib.decompress(blob[spec["offset"] : spec["offset"] + spec["length"]]) |
| idx = np.cumsum(np.frombuffer(raw, dtype=np.uint32), dtype=np.int64) |
| out = torch.zeros(int(np.prod(shape)), dtype=vals.dtype) |
| out[torch.from_numpy(idx)] = vals |
| return out.view(shape) |
|
|
| for lid in range(cfg.num_hidden_layers): |
| prefix = f"model.layers.{lid}." |
| sub_sd = {k[len(prefix) :]: _unpack_one(k) for k in by_layer[lid]} |
| layer = model.model.layers[lid] |
| layer.load_state_dict(sub_sd, assign=True, strict=True) |
| del sub_sd |
| b, a = _compact_layer_inplace(layer, is_chip=(lid < n_chip)) |
| if lid < n_chip: |
| chip_before += b |
| chip_after += a |
| else: |
| qwen_before += b |
| qwen_after += a |
| layer.to(dev) |
|
|
| norm_w = _unpack_one("model.norm.weight").to(dev) |
| model.model.norm.load_state_dict({"weight": norm_w}, assign=True, strict=True) |
|
|
| emb_w = _unpack_one("model.embed_tokens.weight").to(dev) |
| model.model.embed_tokens.load_state_dict({"weight": emb_w}, assign=True, strict=True) |
|
|
| if "lm_head.weight" in tensors_spec and not getattr(cfg, "tie_word_embeddings", True): |
| head_w_cpu = _unpack_one("lm_head.weight") |
| head_before = head_w_cpu.numel() |
| hc0, hc1 = F_chip._hull(F_chip.nonzero_cols(head_w_cpu)) |
| hblock = nn.Parameter(head_w_cpu[:, hc0:hc1].contiguous().to(dev), requires_grad=False) |
| head_after = hblock.numel() |
| model.lm_head = F_chip.SparseLinear( |
| head_w_cpu, rows=(0, head_w_cpu.shape[0]), cols=(hc0, hc1), block=hblock |
| ) |
| del head_w_cpu |
| else: |
| model.lm_head.load_state_dict({"weight": emb_w}, assign=True, strict=True) |
| model.tie_weights() |
| head_before = head_after = emb_w.numel() |
|
|
| model.model.rotary_emb = q35_mod.Qwen3_5TextRotaryEmbedding(config=cfg).to(dev) |
| model.eval() |
| if dev == "cuda": |
| with torch.no_grad(): |
| for _ in range(2): |
| _wc = DynamicCache(config=cfg) |
| model.model( |
| input_ids=torch.zeros((1, 2), dtype=torch.long, device=dev), |
| past_key_values=_wc, |
| use_cache=True, |
| ) |
| model.model( |
| input_ids=torch.zeros((1, 1), dtype=torch.long, device=dev), |
| past_key_values=_wc, |
| use_cache=True, |
| ) |
| del _wc |
| gc.collect() |
| torch.cuda.empty_cache() |
| else: |
| gc.collect() |
| return model, { |
| "chip_dense": chip_before, |
| "chip_compact": chip_after, |
| "qwen_dense": qwen_before, |
| "qwen_compact": qwen_after, |
| "head_dense": head_before, |
| "head_compact": head_after, |
| } |
|
|
|
|
| class CarbonGraphCalcRunner: |
| """Self-contained runner for `Qwen3.5-Math-GraphCalc` checkpoints.""" |
|
|
| def __init__( |
| self, |
| model_dir: str | Path = MODEL_DIR, |
| firmware_path: str | Path | None = None, |
| device: str | None = None, |
| fast: bool = True, |
| executor: bool | str = True, |
| ) -> None: |
| torch.set_grad_enabled(False) |
| F_chip.exact_matmuls() |
| self.path = Path(model_dir) |
| self.meta = json.loads((self.path / "chipformer.json").read_text(encoding="utf-8")) |
| self.dev = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.n_chip = int(self.meta["chip_layers"]) |
| self.v0 = int(self.meta["chip_token_offset"]) |
| self.d0 = int(self.meta["chip_dims_offset"]) |
|
|
| self.model, self.compact_stats = _load_and_compact_qwen35(self.path, self.n_chip, self.dev) |
| self.lm = self.model.model |
| if isinstance(self.model.lm_head, F_chip.SparseLinear): |
| self.head_c0 = self.model.lm_head.c0 |
| self.head_c1 = self.model.lm_head.c1 |
| self.qwen_head_w = self.model.lm_head.weight[: self.v0] |
| else: |
| self.head_c0 = 0 |
| self.head_c1 = self.model.config.hidden_size |
| self.qwen_head_w = self.model.lm_head.weight[: self.v0] |
| self._p4_buf = torch.zeros((4, 1, 1), dtype=torch.long, device=self.dev) |
|
|
| self.tok = AutoTokenizer.from_pretrained(str(self.path)) |
| self.T = CarbonText(self.path / "tokenizer.json") |
|
|
| cfg = self.model.config |
| self.c = core.build( |
| Target(cfg.linear_num_key_heads, cfg.linear_num_value_heads, cfg.intermediate_size), |
| self.T, |
| ) |
| self.code_dims = torch.tensor(self.meta["code_dims"], device=self.dev, dtype=torch.long) |
| self.eos = {248046, 248044} |
| self.go_model_id = self.v0 + S.tok("GO") |
|
|
| fw = firmware_path or (self.path / self.meta.get("firmware", "graphcalc.elf")) |
| self.firmware_path = Path(fw) |
| self.prog = elf.program(self.firmware_path) |
| core.snapshot(self.c, self.fresh_machine()) |
|
|
| self.executor = executor if self.dev == "cuda" else False |
| self.verify = self.executor == "verify" |
| if self.executor: |
| try: |
| self._build_exec() |
| except Exception as exc: |
| if self.verify: |
| raise |
| print(f"[runner] Triton Executor2 unavailable ({type(exc).__name__}: {exc}); using CUDA-graphed HF chip layers.", flush=True) |
| self.executor = False |
| self.verify = False |
|
|
| self.fast = bool(fast and self.dev == "cuda" and (not self.executor or self.verify)) |
| if self.fast: |
| self._build_fast() |
|
|
| def _build_exec(self) -> None: |
| from chipformer.rv.chat import ExecRow |
|
|
| self.keep_ir = torch.tensor(self.c.keep, device=self.dev, dtype=torch.long) |
| cols = self.meta["code_dims"] + [self.meta["tau_dim"], self.meta["mu_dim"]] |
| self.keep_cols = torch.tensor(cols, device=self.dev, dtype=torch.long) |
| self.emb_cache: dict[int, torch.Tensor] = {} |
| self.erow = ExecRow(self) |
| self.erow._capture() |
|
|
| def _exec_emb(self, ids: list[int]) -> torch.Tensor: |
| new = sorted({i for i in ids if i not in self.emb_cache}) |
| if new: |
| rows = self.erow.ex.embed( |
| [i - self.v0 if i >= self.v0 else BR.text_token(i) for i in new] |
| ) |
| for i, r in zip(new, rows): |
| self.emb_cache[i] = r |
| return torch.stack([self.emb_cache[i] for i in ids]) |
|
|
| @torch.no_grad() |
| def _exec_step_one(self, tok_id: int) -> torch.Tensor: |
| row = self.emb_cache.get(tok_id) |
| if row is None: |
| st = tok_id - self.v0 if tok_id >= self.v0 else BR.text_token(tok_id) |
| row = self.erow.ex.embed([st])[0] |
| self.emb_cache[tok_id] = row |
| self.erow.xbuf[0].copy_(row) |
| self.erow.graph.replay() |
| return self.erow.kbuf |
|
|
| @torch.no_grad() |
| def _exec_start_collect(self, ids: list[int]) -> torch.Tensor: |
| if self.erow.graph is None: |
| self.erow._capture() |
| for st in self.erow.heads: |
| st["S"].zero_() |
| st["hist"].zero_() |
| emb = self._exec_emb(ids) |
| keeps = torch.empty((len(ids), len(self.keep_cols)), dtype=torch.float32, device=self.dev) |
| for t in range(len(ids)): |
| self.erow.xbuf.copy_(emb[t : t + 1]) |
| self.erow.graph.replay() |
| keeps[t].copy_(self.erow.kbuf) |
| self.erow.ex.load_init() |
| return keeps |
|
|
| def _exec_hc(self, tok_id: int, keep: torch.Tensor) -> torch.Tensor: |
| hc = torch.zeros(self.lm.embed_tokens.weight.shape[1], dtype=torch.bfloat16, device=self.dev) |
| hc[: self.d0] = self.lm.embed_tokens.weight[tok_id, : self.d0] |
| hc[self.keep_cols] = keep.to(torch.bfloat16) |
| return hc |
|
|
| def _build_fast(self) -> None: |
| cfg = copy.deepcopy(self.model.config) |
| cfg.num_hidden_layers = self.n_chip + 1 |
| cfg.layer_types = ["linear_attention"] * self.n_chip + ["full_attention"] |
| self.ccfg = cfg |
| dummy = q35_mod.Qwen3_5DecoderLayer(cfg, self.n_chip).to(self.dev, torch.bfloat16) |
| for p in dummy.parameters(): |
| p.data.zero_() |
| with torch.device("meta"): |
| cm = q35_mod.Qwen3_5TextModel(cfg) |
| cm.embed_tokens = self.lm.embed_tokens |
| cm.layers = nn.ModuleList([*self.lm.layers[: self.n_chip], dummy]) |
| cm.norm = nn.Identity() |
| cm.rotary_emb = q35_mod.Qwen3_5TextRotaryEmbedding(config=cfg).to(self.dev) |
| self.cm = cm.eval() |
| self.graph = None |
|
|
| def _chip_cache(self): |
| if getattr(self, "ccache", None) is None: |
| self.ccache = DynamicCache(config=self.ccfg) |
| else: |
| for L in range(self.n_chip): |
| lay = self.ccache.layers[L] |
| lay.conv_states[0].zero_() |
| lay.recurrent_states[0].zero_() |
| lay.has_previous_state[0] = False |
| return self.ccache |
|
|
| @torch.no_grad() |
| def _chip_body(self): |
| h = self.cm.embed_tokens(self.g_ids) |
| for L in range(self.n_chip): |
| h = self.cm.layers[L]( |
| h, |
| position_embeddings=None, |
| attention_mask=None, |
| position_ids=None, |
| past_key_values=self.ccache, |
| ) |
| self.g_out.copy_(h) |
|
|
| @torch.no_grad() |
| def _capture(self): |
| self.g_ids = torch.zeros((1, 1), dtype=torch.long, device=self.dev) |
| self.g_out = torch.zeros( |
| (1, 1, self.model.config.hidden_size), dtype=torch.bfloat16, device=self.dev |
| ) |
| lays = [self.ccache.layers[L] for L in range(self.n_chip)] |
| saved = [(lay.conv_states[0].clone(), lay.recurrent_states[0].clone()) for lay in lays] |
| s = torch.cuda.Stream() |
| s.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(s): |
| for _ in range(3): |
| self._chip_body() |
| torch.cuda.current_stream().wait_stream(s) |
| self.graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(self.graph): |
| self._chip_body() |
| for lay, (cs, rs) in zip(lays, saved): |
| lay.conv_states[0].copy_(cs) |
| lay.recurrent_states[0].copy_(rs) |
|
|
| @torch.no_grad() |
| def _chip_step(self, tok_id: int, g: int, cache, compiled: bool = True) -> torch.Tensor: |
| if not compiled or self.graph is None: |
| x = torch.tensor([[tok_id]], device=self.dev) |
| out = self.cm( |
| input_ids=x, |
| past_key_values=cache, |
| use_cache=True, |
| cache_position=torch.tensor([g], device=self.dev), |
| ) |
| return out.last_hidden_state[0, -1].clone() |
| self.g_ids.fill_(tok_id) |
| self.graph.replay() |
| return self.g_out[0, -1].clone() |
|
|
| @torch.no_grad() |
| def _qwen_step(self, hc: torch.Tensor, cache, pos: int, full_head: bool = False) -> torch.Tensor: |
| h = hc.view(1, 1, -1) |
| p4 = self._p4_buf |
| p4.fill_(pos) |
| pe = self.lm.rotary_emb(h, p4[1:]) |
| for i in range(self.n_chip, len(self.lm.layers)): |
| h = self.lm.layers[i]( |
| h, |
| position_embeddings=pe, |
| attention_mask=None, |
| position_ids=p4[0], |
| past_key_values=cache, |
| ) |
| nh = self.lm.norm(h)[0, -1] |
| if full_head: |
| return self.model.lm_head(nh) |
| return F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w) |
|
|
| @torch.no_grad() |
| def _qwen_prefill( |
| self, x_ids: torch.Tensor, keeps: torch.Tensor, full_head: bool = False |
| ) -> tuple[DynamicCache, torch.Tensor]: |
| p_len = int(x_ids.shape[1]) |
| h = torch.zeros( |
| (1, p_len, self.model.config.hidden_size), dtype=torch.bfloat16, device=self.dev |
| ) |
| h[0, :, : self.d0] = self.lm.embed_tokens.weight[x_ids[0], : self.d0] |
| h[0, :, self.keep_cols] = keeps.to(torch.bfloat16) |
| cache = DynamicCache(config=self.model.config) |
| pos_vec = torch.arange(0, p_len, device=self.dev, dtype=torch.long).unsqueeze(0) |
| p4 = pos_vec.unsqueeze(0).expand(4, -1, -1) |
| pe = self.lm.rotary_emb(h, p4[1:]) |
| for i in range(self.n_chip, len(self.lm.layers)): |
| h = self.lm.layers[i]( |
| h, |
| position_embeddings=pe, |
| attention_mask=None, |
| position_ids=p4[0], |
| past_key_values=cache, |
| ) |
| nh = self.lm.norm(h)[0, -1] |
| logits = ( |
| self.model.lm_head(nh) |
| if full_head |
| else F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w) |
| ) |
| return cache, logits |
|
|
| @torch.no_grad() |
| def _qwen_chunk( |
| self, hcs: list[torch.Tensor], cache, pos0: int, full_head: bool = False |
| ) -> torch.Tensor: |
| if len(hcs) == 1: |
| return self._qwen_step(hcs[0], cache, pos0, full_head=full_head) |
| k_len = len(hcs) |
| h = torch.stack(hcs, dim=0).unsqueeze(0) |
| pos_vec = torch.arange(pos0, pos0 + k_len, device=self.dev, dtype=torch.long).unsqueeze(0) |
| p4 = pos_vec.unsqueeze(0).expand(4, -1, -1) |
| pe = self.lm.rotary_emb(h, p4[1:]) |
| attn_mask = torch.zeros((1, 1, k_len, pos0 + k_len), dtype=h.dtype, device=self.dev) |
| future = torch.triu(torch.ones((k_len, k_len), dtype=torch.bool, device=self.dev), diagonal=1) |
| attn_mask[0, 0, :, pos0:][future] = torch.finfo(h.dtype).min |
| for i in range(self.n_chip, len(self.lm.layers)): |
| lmask = attn_mask if self.model.config.layer_types[i] == "full_attention" else None |
| h = self.lm.layers[i]( |
| h, |
| position_embeddings=pe, |
| attention_mask=lmask, |
| position_ids=p4[0], |
| past_key_values=cache, |
| ) |
| nh = self.lm.norm(h)[0, -1] |
| if full_head: |
| return self.model.lm_head(nh) |
| return F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w) |
|
|
| @torch.no_grad() |
| def _layers_slow(self, tok_id: int, cache, pos: int | None): |
| h = self.lm.embed_tokens(torch.tensor([[tok_id]], device=self.dev)) |
| pe = tp = None |
| if pos is not None: |
| p4 = torch.full((4, 1, 1), pos, device=self.dev, dtype=torch.long) |
| tp = p4[0] |
| pe = self.lm.rotary_emb(h, p4[1:]) |
| hc = None |
| n = self.n_chip if pos is None else len(self.lm.layers) |
| for i in range(n): |
| h = self.lm.layers[i]( |
| h, |
| position_embeddings=pe, |
| attention_mask=None, |
| position_ids=tp, |
| past_key_values=cache, |
| ) |
| if i == self.n_chip - 1: |
| hc = h |
| logits = None if pos is None else self.model.lm_head(self.lm.norm(h))[0, -1] |
| return hc[0, -1], logits |
|
|
| def _decode_code_list(self, vals: list[float]) -> int | None: |
| if not any(v != 0.0 for v in vals): |
| return None |
| nf = core.NFAM |
| fid = 0 |
| for i in range(nf): |
| if vals[i] > 0.0: |
| fid |= 1 << i |
| p = 0 |
| for i in range(len(vals) - nf): |
| if vals[nf + i] > 0.0: |
| p |= 1 << i |
| t = core.decode_code(fid, p, self.T) |
| if t < 0: |
| raise RuntimeError(f"Invalid circuit token prediction (family={fid}, payload={p})") |
| return t - S.TEXT_TOKEN if S.is_text(t) else self.v0 + t |
|
|
| def code_token(self, code: torch.Tensor) -> int | None: |
| return self._decode_code_list(code.tolist()) |
|
|
| def circuit_token(self, hc: torch.Tensor) -> int | None: |
| return self.code_token(hc[self.code_dims]) |
|
|
| def fresh_machine(self) -> ref.Machine: |
| m = ref.Machine.from_program(self.prog) |
| BR.park(m) |
| return m |
|
|
| def prompt_ids( |
| self, |
| user_prompt: str | list[dict], |
| system_prompt: str | None = GRAPHCALC_SYSTEM_PROMPT, |
| use_tools: bool = True, |
| ) -> list[int]: |
| if isinstance(user_prompt, str): |
| msgs = [] |
| if system_prompt: |
| msgs.append({"role": "system", "content": system_prompt}) |
| msgs.append({"role": "user", "content": user_prompt}) |
| else: |
| msgs = list(user_prompt) |
| tools = [GRAPHCALC_TOOL] if use_tools else None |
| out = self.tok.apply_chat_template( |
| msgs, |
| tools=tools, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| tokenize=True, |
| ) |
| if not isinstance(out, list): |
| out = out["input_ids"] |
| return [int(x) for x in out] |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| ids: list[int], |
| max_new_text: int = 512, |
| max_chip_tokens: int = 500_000, |
| max_calls: int = 1, |
| check_head: bool = False, |
| check_lock: bool = False, |
| ) -> dict: |
| m = self.fresh_machine() |
| lock = None |
| if check_lock: |
| lock = BR.Session(self.T, self.fresh_machine()) |
| lock.history(ids) |
|
|
| x = torch.tensor([ids], device=self.dev, dtype=torch.long) |
| core.snapshot(self.c, m) |
| n_code = len(self.code_dims) |
|
|
| if self.executor and not self.verify: |
| keeps = self._exec_start_collect(ids) |
| cache, logits = self._qwen_prefill(x, keeps, full_head=check_head) |
| else: |
| lm_out = self.lm(input_ids=x, use_cache=True) |
| cache = lm_out.past_key_values |
| if self.executor: |
| self.erow.start(ids) |
| if self.fast: |
| ccache = self._chip_cache() |
| h = self.cm.embed_tokens(x) |
| for L in range(self.n_chip): |
| h = self.cm.layers[L]( |
| h, |
| position_embeddings=None, |
| attention_mask=None, |
| position_ids=None, |
| past_key_values=ccache, |
| ) |
| hf.inject(self.c, ccache) |
| if self.graph is None: |
| self._capture() |
| g = len(ids) |
| elif not self.executor: |
| hf.inject(self.c, cache) |
| logits = self.model.lm_head(lm_out.last_hidden_state[0, -1]) |
| pos = len(ids) |
| gen: list[int] = [] |
| n_qwen_text = n_echo_text = n_chip_tok = n_calls = head_mismatch = n_verify_bad = 0 |
| nxt = int(torch.argmax(logits[: self.v0])) |
| prev_owned = False |
| in_tool_call = False |
| pending_hcs: list[torch.Tensor] = [] |
| pending_pos0: int | None = None |
| t0 = time.time() |
| expect: list[int] = [] |
| lock_error = None |
|
|
| while n_chip_tok < max_chip_tokens: |
| if n_qwen_text >= max_new_text and (not prev_owned) and (not in_tool_call): |
| break |
|
|
| gen.append(nxt) |
| if nxt == self.T.TCO and not prev_owned: |
| in_tool_call = True |
| elif nxt == self.T.TCC: |
| in_tool_call = False |
| if nxt == self.go_model_id: |
| n_calls += 1 |
|
|
| if lock is not None: |
| st = nxt - self.v0 if nxt >= self.v0 else BR.text_token(nxt) |
| if expect: |
| e = expect.pop(0) |
| if st != e: |
| lock_error = {"at": len(gen) - 1, "got": st, "want": e} |
| break |
| elif st < S.VOCAB: |
| lock_error = {"at": len(gen) - 1, "got": S.name(st), "want": "text"} |
| break |
| else: |
| k = len(lock.toks) |
| try: |
| lock.qwen([nxt]) |
| except RuntimeError as exc: |
| lock_error = {"at": len(gen) - 1, "got": str(exc), "want": "halt"} |
| break |
| expect = list(lock.toks[k + 1 :]) |
|
|
| if nxt in self.eos and not prev_owned: |
| break |
|
|
| chip_in = nxt >= self.v0 |
| hc = keep = None |
| if self.executor: |
| keep = self._exec_step_one(nxt) |
| if not chip_in: |
| hc = self._exec_hc(nxt, keep) |
| if self.fast and (self.verify or not self.executor): |
| hf_hc = self._chip_step(nxt, g, ccache) |
| g += 1 |
| if self.verify: |
| same = ( |
| bool((hf_hc[self.keep_cols] == keep.to(torch.bfloat16)).all()) |
| if chip_in |
| else bool((hf_hc == hc).all()) |
| ) |
| if not same: |
| n_verify_bad += 1 |
| hc = hf_hc |
| elif not self.executor: |
| hc, logits = self._layers_slow(nxt, cache, None if chip_in else pos) |
|
|
| if chip_in: |
| n_chip_tok += 1 |
| nxt = ( |
| self._decode_code_list(keep[:n_code].tolist()) |
| if hc is None |
| else self.circuit_token(hc) |
| ) |
| prev_owned = True |
| if nxt is None: |
| raise RuntimeError("Circuit relinquished ownership mid-chip-token") |
| else: |
| if prev_owned: |
| n_echo_text += 1 |
| else: |
| n_qwen_text += 1 |
| ct = ( |
| self._decode_code_list(keep[:n_code].tolist()) |
| if (keep is not None and not self.verify) |
| else self.circuit_token(hc) |
| ) |
| if ct is not None and (self.fast or self.executor) and not check_head: |
| if pending_pos0 is None: |
| pending_pos0 = pos |
| pending_hcs.append(hc) |
| pos += 1 |
| nxt = ct |
| prev_owned = True |
| else: |
| if self.fast or self.executor: |
| if pending_hcs: |
| pending_hcs.append(hc) |
| logits = self._qwen_chunk( |
| pending_hcs, cache, pending_pos0, full_head=check_head |
| ) |
| pending_hcs.clear() |
| pending_pos0 = None |
| else: |
| logits = self._qwen_step(hc, cache, pos, full_head=check_head) |
| pos += 1 |
| if ct is not None: |
| if check_head and int(torch.argmax(logits)) != ct: |
| head_mismatch += 1 |
| nxt = ct |
| prev_owned = True |
| else: |
| q_logits = logits[: self.v0] |
| if n_calls >= max_calls or n_qwen_text >= max(16, max_new_text - 48): |
| q_logits = q_logits.clone() |
| q_logits[self.T.TCO] = float("-inf") |
| nxt = int(torch.argmax(q_logits)) |
| prev_owned = False |
|
|
| dt = time.time() - t0 |
| parsed = self.parse_turn(gen) |
| return { |
| "tokens": gen, |
| "text_tokens": n_qwen_text, |
| "echo_text_tokens": n_echo_text, |
| "chip_tokens": n_chip_tok, |
| "tool_calls_executed": n_calls, |
| "seconds": round(dt, 3), |
| "head_mismatch": head_mismatch, |
| "lock_error": lock_error, |
| "verify_mismatch": n_verify_bad if self.verify else None, |
| **parsed, |
| } |
|
|
| def text_of(self, ids: list[int]) -> str: |
| return self.tok.decode([i for i in ids if i < self.v0], skip_special_tokens=False) |
|
|
| def parse_turn(self, gen_ids: list[int]) -> dict: |
| raw = self.text_of(gen_ids) |
| calls = re.findall(r"<tool_call>\s*(.*?)\s*</tool_call>", raw, flags=re.DOTALL) |
| responses = re.findall(r"<tool_response>\n?(.*?)\n?</tool_response>", raw, flags=re.DOTALL) |
| cleaned = re.sub( |
| r"<tool_call>.*?</tool_response>\s*(?:<\|im_end\|>\s*<\|im_start\|>assistant\s*<think>\s*</think>\s*)?", |
| "\n\n", |
| raw, |
| flags=re.DOTALL, |
| ) |
| cleaned = re.sub(r"<tool_call>.*$", "", cleaned, flags=re.DOTALL) |
| cleaned = re.sub(r"</?think>", "", cleaned) |
| for stop_s in ("<|im_end|>", "<|endoftext|>", "<|im_start|>"): |
| if stop_s in cleaned: |
| cleaned = cleaned.split(stop_s, 1)[0] |
| cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip() |
| return { |
| "raw_text": raw, |
| "tool_calls": calls, |
| "tool_responses": responses, |
| "final_explanation": cleaned, |
| } |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description="Run Qwen3.5-Math-GraphCalc in-weights") |
| ap.add_argument( |
| "--ask", |
| default="Graph y = x^3 - 3*x from x = -2 to 2, identify its roots and turning points, and state the local maximum y-value.", |
| help="Math or graphing question to solve", |
| ) |
| ap.add_argument("--pure-hf", action="store_true", help="Run un-fused HF Qwen3_5DecoderLayer weights (executor=False)") |
| ap.add_argument("--max-new-text", type=int, default=512) |
| args = ap.parse_args() |
|
|
| runner = CarbonGraphCalcRunner(MODEL_DIR, fast=True, executor=not args.pure_hf) |
| res = runner.generate(runner.prompt_ids(args.ask), max_new_text=args.max_new_text, check_lock=True) |
| if res["tool_responses"]: |
| print("=== IN-WEIGHTS RV32 <tool_response> ===") |
| print(res["tool_responses"][0]) |
| print() |
| print("=== STEP-BY-STEP TUTOR EXPLANATION ===") |
| print(res["final_explanation"]) |
| print(f"\n[chip_tokens={res['chip_tokens']} | lock_error={res['lock_error']} | {res['seconds']}s]") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
