# Tensor Roll — Recursive CUDA-Q Model Quantizer # Copyright (C) 2026 SnapKitty Collective # SPDX-License-Identifier: AGPL-3.0-or-later """The recursive TensorRoll operator, per-partition metrics, searched compression policy, and the explicit classical/quantum dispatch boundary. Honesty contract (enforced here, not just documented): * Classical execution on this machine is NumPy on CPU. Every classical report carries backend="cpu-numpy", device="cpu". * Quantum-candidate execution runs on the classical state-vector simulator in ``qsim.py``. Every such report carries backend="qsim-classical". Nothing here claims real QPU/GPU execution. """ from __future__ import annotations import time from dataclasses import dataclass, field from typing import Callable, Dict, List, Optional import numpy as np from .util import fmt_bytes, rss_mb DECISIONS = ( "PRESERVE", "QUANTIZE", "FACTORIZE", "MERGED", "ROUTED", "RECONSTRUCTED", "PRUNED", "QUANTUM-ENCODED", ) #: Max elements a block may have to be eligible for the quantum path #: (amplitude encoding needs 2**n amplitudes; we cap at 2**6). QUANTUM_ELIGIBLE_MAX = 256 # 8 qubits; honest near-term register size CLASSICAL_BACKEND = "cpu-numpy" QSIM_BACKEND = "qsim-classical" # --------------------------------------------------------------------------- # Metrics # --------------------------------------------------------------------------- def tensor_metrics(W: np.ndarray) -> Dict: W = np.asarray(W) flat = W.ravel().astype(np.float64) m = { "shape": list(W.shape), "ndim": int(W.ndim), "params": int(W.size), "dtype": str(W.dtype), "norm": float(np.linalg.norm(flat)), "variance": float(flat.var()) if flat.size else 0.0, } if W.ndim == 2 and min(W.shape) > 1: s = np.linalg.svd(W.astype(np.float64), compute_uv=False) total = float((s ** 2).sum()) p = (s ** 2) / total if total > 0 else np.zeros_like(s) p = p[p > 0] m["rank"] = int(np.sum(s > max(W.shape) * np.finfo(float).eps * (s[0] if s.size else 0.0))) m["entropy"] = float(-(p * np.log(p)).sum()) if p.size else 0.0 m["spectral"] = float((s[0] ** 2) / total) if total > 0 else 0.0 m["sing_vals"] = [float(x) for x in s[:8]] else: m["rank"] = int(1 if m["norm"] > 0 else 0) m["entropy"] = 0.0 m["spectral"] = 1.0 if m["norm"] > 0 else 0.0 m["sing_vals"] = [] return m def relative_error(A: np.ndarray, B: np.ndarray) -> float: A = A.astype(np.float64) denom = float(np.linalg.norm(A)) if denom == 0: return 0.0 if float(np.linalg.norm(B)) == 0 else float("inf") return float(np.linalg.norm(A - B) / denom) def cosine_sim(A: np.ndarray, B: np.ndarray) -> float: a = A.astype(np.float64).ravel() b = B.astype(np.float64).ravel() na, nb = float(np.linalg.norm(a)), float(np.linalg.norm(b)) if na == 0 or nb == 0: return 1.0 if na == nb else 0.0 return float(np.dot(a, b) / (na * nb)) # --------------------------------------------------------------------------- # Invocation reports — every kernel call, classical or quantum, reports these. # --------------------------------------------------------------------------- def new_report( *, backend: str, device: str, kernel: str, tensor_dims, precision: str, exec_time_s: float, memory_bytes: float, recursion_depth: int, shots=None, measurement_stats=None, reconstruction_error=None, extra=None, ) -> Dict: rep = { "backend": backend, "device": device, "kernel": kernel, "tensor_dims": [tuple(d) for d in tensor_dims], "precision": precision, "exec_time_s": float(exec_time_s), "memory_bytes": float(memory_bytes), "recursion_depth": int(recursion_depth), "shots": shots, "measurement_stats": measurement_stats, "reconstruction_error": reconstruction_error, } if extra: rep.update(extra) return rep def dispatch_matmul(A: np.ndarray, B: np.ndarray, *, depth: int = 0, label: str = "gemm") -> tuple: """Classical GEMM baseline. Backend is honestly cpu-numpy.""" A = np.asarray(A) B = np.asarray(B) mem = float(A.nbytes + B.nbytes) t0 = time.perf_counter() C = A @ B dt = time.perf_counter() - t0 mem += float(C.nbytes) rep = new_report( backend=CLASSICAL_BACKEND, device="cpu", kernel=label, tensor_dims=[A.shape, B.shape, C.shape], precision=str(A.dtype), exec_time_s=dt, memory_bytes=mem, recursion_depth=depth, ) return C, rep # --------------------------------------------------------------------------- # Compression option evaluation (measured, not hard-coded) # --------------------------------------------------------------------------- def _svd_options(W: np.ndarray): U, s, Vt = np.linalg.svd(W.astype(np.float64), full_matrices=False) return U, s, Vt def _vq_codebook(W: np.ndarray, k: int = 16, iters: int = 12, seed: int = 0): """From-scratch k-means vector quantization of W's elements. Returns (reconstructed_W, cost_params). Deterministic.""" flat = W.ravel().astype(np.float64) rng = np.random.default_rng(seed) cents = np.quantile(flat, np.linspace(0, 1, k)).astype(np.float64) for _ in range(iters): assign = np.abs(flat[:, None] - cents[None, :]).argmin(axis=1) for j in range(k): m = assign == j if np.any(m): cents[j] = flat[m].mean() assign = np.abs(flat[:, None] - cents[None, :]).argmin(axis=1) Wd = cents[assign].reshape(W.shape) bits_per_idx = float(np.ceil(np.log2(k))) cost = float(k + flat.size * bits_per_idx / 32.0) return Wd, cost def evaluate_options(name: str, W: np.ndarray, *, quantum_policy: str = "off", qsim_encode=None, rank: int = 0) -> List[Dict]: """Measure candidate representations of W. Each option carries a real (cost_params, retained_energy, reconstruction_error) triple computed from actual transforms — the policy search below picks winners, so no decision is hard-coded here. Decision classes searched: PRESERVE, QUANTIZE (incl. codebook VQ), FACTORIZE, MERGED, ROUTED, RECONSTRUCTED, PRUNED, QUANTUM-ENCODED.""" W = np.asarray(W, dtype=np.float64) params = int(W.size) opts: List[Dict] = [] def add(decision, cost, energy, recon, detail=""): opts.append({ "tensor": name, "decision": decision, "cost_params": float(cost), "retained_energy": float(energy), "reconstruction_error": float(recon), "detail": detail, }) add("PRESERVE", params, 1.0, 0.0, "bit-exact") if W.ndim == 2 and min(W.shape) > 1: U, s, Vt = _svd_options(W) total_e = float((s ** 2).sum()) # FACTORIZE at several ranks — measured energies fracs = [0.75, 0.5, 0.25] if rank and 0 < rank < min(W.shape): fracs.append(rank / min(W.shape)) for frac in fracs: r = max(1, int(len(s) * frac)) cost = r * (W.shape[0] + W.shape[1]) energy = float((s[:r] ** 2).sum() / total_e) if total_e else 1.0 Wr = (U[:, :r] * s[:r]) @ Vt[:r, :] add("FACTORIZE", cost, energy, relative_error(W, Wr), f"svd-rank-{r}") # QUANTIZE int8 — actually quantize and measure qmin, qmax = W.min(), W.max() if qmax > qmin: scale = (qmax - qmin) / 255.0 Wq = np.round((W - qmin) / scale).astype(np.int64) Wd = Wq.astype(np.float64) * scale + qmin add("QUANTIZE", params * 8 / 32, 1.0 - relative_error(W, Wd) ** 2, relative_error(W, Wd), "int8-symmetric-range") # QUANTIZE via codebook vector quantization (k-means, k=16) Wv, vq_cost = _vq_codebook(W, k=16) vq_err = relative_error(W, Wv) add("QUANTIZE", vq_cost, max(0.0, 1.0 - vq_err ** 2), vq_err, "codebook-vq-k16") # PRUNED — magnitude prune 25% / 50%, sparse cost model flat = np.abs(W.ravel()) for frac in (0.25, 0.5): k = int(flat.size * frac) if k > 0: thr = np.partition(flat, k)[k] Wp = np.where(np.abs(W) >= thr, W, 0.0) kept = int(np.count_nonzero(Wp)) # sparse CSR-ish cost: values + indices cost = kept * 2 energy = 1.0 - relative_error(W, Wp) ** 2 add("PRUNED", cost, energy, relative_error(W, Wp), f"magnitude-{int(frac*100)}pct-sparse") # QUANTUM-ENCODED — only tiny blocks, routed through the real qsim # RECONSTRUCTED — recursive-measurement classical reconstruction if W.size <= QUANTUM_ELIGIBLE_MAX and qsim_encode is not None and quantum_policy != "off": t0 = time.perf_counter() rec, qrep = qsim_encode(W) dt = time.perf_counter() - t0 err = relative_error(W, rec) # Honest cost: classical reconstruction must still be stored. add("QUANTUM-ENCODED", params, max(0.0, 1.0 - err ** 2), err, f"qsim shots={qrep['shots']} t={dt:.4f}s") # RECONSTRUCTED: deeper recursive measurement -> classical # amplitude reconstruction stored as fp32 (measured, distinct # from the single-level quantum-encoded path above). from . import qsim as _Q rec2, rep2 = _Q.tensor_roll_recursive( W, depth=2, shots=qrep.get("shots", 2048), seed=0) err2 = relative_error(W, rec2) add("RECONSTRUCTED", params, max(0.0, 1.0 - err2 ** 2), err2, f"recursive-measurement depth=2 shots={rep2['shots']}") if quantum_policy == "explore": add("ROUTED", params, max(0.0, 1.0 - err ** 2), err, "routed-to-qsim-for-measurement") else: # 1-D tensors (biases, norms): merge into parent op — MERGED. add("MERGED", params, 1.0, 0.0, "folded-into-parent-op") qmin, qmax = W.min(), W.max() if qmax > qmin: scale = (qmax - qmin) / 255.0 Wd = np.round((W - qmin) / scale) * scale + qmin add("QUANTIZE", params * 8 / 32, 1.0 - relative_error(W, Wd) ** 2, relative_error(W, Wd), "int8-1d") return opts # --------------------------------------------------------------------------- # Recursive TensorRoll operator # --------------------------------------------------------------------------- @dataclass class RollNode: name: str shape: tuple depth: int metrics: Dict = field(default_factory=dict) children: List["RollNode"] = field(default_factory=list) options: List[Dict] = field(default_factory=list) decision: Optional[Dict] = None def roll_tree(name: str, W: np.ndarray, *, axis: int = 0, depth: int = 3, quantum_policy: str = "off", qsim_encode=None, rank: int = 0, _level: int = 0, _counter=None) -> RollNode: """Recursive partitioner backing TensorRoll(T, axis, depth, rank, quantum_policy). Recursively partitions T along `axis`; every node measures norm, rank, entropy, variance, spectral contribution. Leaves evaluate real compression options (rank-aware when `rank` > 0). Internal nodes aggregate child energy. """ if _counter is None: _counter = [0] W = np.asarray(W) node = RollNode(name=name, shape=tuple(W.shape), depth=_level, metrics=tensor_metrics(W)) _counter[0] += 1 node.metrics["roll_id"] = _counter[0] splittable = (_level < depth and W.ndim >= 1 and W.shape[axis] >= 4) if splittable: halves = np.array_split(W, 2, axis=axis) for i, h in enumerate(halves): node.children.append( roll_tree(f"{name}/{i}", h, axis=axis, depth=depth, quantum_policy=quantum_policy, qsim_encode=qsim_encode, rank=rank, _level=_level + 1, _counter=_counter)) else: node.options = evaluate_options( name, W, quantum_policy=quantum_policy, qsim_encode=qsim_encode, rank=rank) return node def select_plan(leaves: List[RollNode], budget_params: float) -> Dict[str, Dict]: """Global representation search: greedy knapsack over measured (cost, retained_energy) options subject to the hard parameter budget. Returns {tensor_name: winning_option}.""" # Start every leaf at its cheapest option, then upgrade greedily. srt = [sorted(leaf.options, key=lambda o: o["cost_params"]) for leaf in leaves] choice = [0] * len(leaves) total = sum(s[0]["cost_params"] for s in srt) improved = True while improved: improved = False best_gain, best_move = 0.0, None for i, opts in enumerate(srt): cur = choice[i] for j in range(cur + 1, len(opts)): dcost = opts[j]["cost_params"] - opts[cur]["cost_params"] denergy = opts[j]["retained_energy"] - opts[cur]["retained_energy"] if total + dcost <= budget_params and dcost > 0 and denergy > 0: gain = denergy / dcost if gain > best_gain: best_gain, best_move = gain, (i, j, dcost) if best_move: i, j, dcost = best_move choice[i] = j total += dcost improved = True plan = {} for i, leaf in enumerate(leaves): plan[leaf.name] = srt[i][choice[i]] plan["_total_cost_params"] = total plan["_budget_params"] = float(budget_params) return plan def iter_leaves(node: RollNode): if node.children: for c in node.children: yield from iter_leaves(c) else: yield node def TensorRoll(T, axis: int = -1, depth: int = 2, rank: int = 8, quantum_policy: str = "balanced", qsim_encode=None, budget_params: float = 0.0, name: str = "T") -> Dict: """First-class recursive operator. TensorRoll(T, axis, depth, rank, quantum_policy): 1. recursively partitions T along `axis` to `depth`, measuring every node (norm, rank, entropy, variance, spectral contribution); 2. evaluates candidate representations per leaf — PRESERVE, QUANTIZE (incl. codebook VQ), FACTORIZE (rank-aware via `rank`), MERGED, ROUTED, RECONSTRUCTED, PRUNED, QUANTUM-ENCODED (the last gated by `quantum_policy`); 3. searches the global representation plan under `budget_params` (0 = unconstrained: keep the best-energy option per leaf); 4. returns the plan plus measured totals. No decision is hard-coded; every option carries measured (cost, retained_energy, reconstruction_error). Backend labels: ordinary options execute on cpu-numpy; quantum-gated options execute on qsim-classical (statevector simulator, NOT a QPU). """ T = np.asarray(T, dtype=np.float64) ax = axis % T.ndim if T.ndim else 0 tree = roll_tree(name, T, axis=ax, depth=depth, quantum_policy=quantum_policy, qsim_encode=qsim_encode, rank=rank) leaves = list(iter_leaves(tree)) full = float(sum(l.metrics["params"] for l in leaves)) budget = float(budget_params) if budget_params > 0 else full plan = select_plan(leaves, budget) for leaf in leaves: leaf.decision = plan[leaf.name] decisions = {} for leaf in leaves: d = leaf.decision["decision"] decisions[d] = decisions.get(d, 0) + 1 energies = [leaf.decision["retained_energy"] for leaf in leaves] return { "tree": tree, "plan": plan, "n_leaves": len(leaves), "decisions": decisions, "total_cost_params": plan["_total_cost_params"], "budget_params": plan["_budget_params"], "full_params": full, "mean_retained_energy": float(np.mean(energies)) if energies else 1.0, "quantum_policy": quantum_policy, "rank": rank, } def render_roll_line(node: RollNode, idx: int) -> str: m = node.metrics dec = node.decision["decision"] if node.decision else "—" backend = "qsim-classical" if node.decision and "QUANTUM" in dec else CLASSICAL_BACKEND return ( f"[ROLL {idx:04d}] source: {node.name} shape: {list(node.shape)} " f"depth: {node.depth} params: {m['params']} " f"rank: {m['rank']} entropy: {m['entropy']:.3f} " f"spectral: {m['spectral']:.3f} " f"cost: {node.decision['cost_params']:.0f} " f"retained-energy: {node.decision['retained_energy']:.4f} " f"recon-err: {node.decision['reconstruction_error']:.2e} " f"backend: {backend} decision: {dec}" )