Spaces:

magenta-community
/

magenta-rt-stack

Running

App Files Files Community

magenta-rt-stack / depth-timing.html

multimodalart HF Staff

Upload depth-timing.html with huggingface_hub

e7304d3 verified about 2 months ago

Raw

History Blame Contribute Delete

51.9 kB

	<!DOCTYPE html>
	<html lang="en">
	<head>
	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width, initial-scale=1">
	<title>DEPTH WGSL timing (grid-parallel)</title>
	<style>
	html, body {
	margin: 0;
	padding: 0;
	background: #0b0f14;
	color: #e6edf3;
	font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", system-ui, sans-serif;
	}
	#wrap {
	max-width: 880px;
	margin: 0 auto;
	padding: 28px 24px 60px;
	}
	h1 {
	font-size: 26px;
	font-weight: 700;
	margin: 0 0 6px;
	}
	.sub { color: #9aa7b4; font-size: 14px; margin-bottom: 22px; }
	#status {
	font-size: 16px;
	padding: 14px 16px;
	border-radius: 10px;
	background: #16202b;
	border: 1px solid #243240;
	margin-bottom: 22px;
	white-space: pre-wrap;
	}
	.running { color: #ffd479; }
	.ok { color: #7ee787; }
	.bad { color: #ff7b72; }
	#result { display: none; }
	.card {
	background: #111a23;
	border: 1px solid #243240;
	border-radius: 12px;
	padding: 18px 20px;
	margin-bottom: 18px;
	}
	.row {
	display: flex;
	justify-content: space-between;
	align-items: baseline;
	padding: 6px 0;
	border-bottom: 1px solid #1c2733;
	gap: 16px;
	}
	.row:last-child { border-bottom: none; }
	.row .k { color: #9aa7b4; font-size: 14px; }
	.row .v { font-size: 16px; font-weight: 600; text-align: right; word-break: break-word; }
	.hero {
	text-align: center;
	padding: 26px 18px;
	}
	.hero .label { color: #9aa7b4; font-size: 14px; letter-spacing: .08em; text-transform: uppercase; }
	.hero .big {
	font-size: 64px;
	font-weight: 800;
	line-height: 1.05;
	margin: 8px 0 2px;
	font-variant-numeric: tabular-nums;
	}
	.hero .unit { font-size: 26px; color: #9aa7b4; font-weight: 600; }
	.badge {
	display: inline-block;
	padding: 6px 16px;
	border-radius: 999px;
	font-size: 20px;
	font-weight: 800;
	letter-spacing: .05em;
	}
	.badge.pass { background: #133a1e; color: #7ee787; border: 1px solid #2ea043; }
	.badge.fail { background: #3a1414; color: #ff7b72; border: 1px solid #da3633; }
	pre.codes {
	font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
	font-size: 13px;
	color: #cdd9e5;
	background: #0d141b;
	padding: 10px 12px;
	border-radius: 8px;
	overflow-x: auto;
	margin: 6px 0 0;
	}
	.err {
	color: #ff7b72;
	font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
	font-size: 13px;
	white-space: pre-wrap;
	}
	</style>
	</head>
	<body>
	<div id="wrap">
	<h1>Magenta-RT2 small — DEPTH (WGSL) timing</h1>
	<div class="sub">Grid-parallel WebGPU depth transformer · W0/W1 split storage buffers · bit-exact correctness + real device timing</div>
	<div id="status" class="running">starting…</div>
	<div id="result"></div>
	</div>

	<script type="module">
	"use strict";

	const BASE = "https://huggingface.co/magenta-community/magenta-rt-onnx-small/resolve/main/depth_wgsl/";
	const CACHE_NAME = "depth-wgsl-v1";

	const statusEl = document.getElementById("status");
	const resultEl = document.getElementById("result");

	function setStatus(msg, cls) {
	statusEl.textContent = msg;
	statusEl.className = cls \|\| "";
	}
	function showError(where, e) {
	const msg = (e && (e.stack \|\| e.message)) ? (e.stack \|\| e.message) : String(e);
	statusEl.className = "bad";
	statusEl.innerHTML = '<b>Error (' + where + ')</b>\n<span class="err">' + escapeHtml(msg) + '</span>';
	console.error(where, e);
	}
	function escapeHtml(s) {
	return String(s).replace(/[&<>]/g, c => ({ "&": "&", "<": "<", ">": ">" }[c]));
	}

	// ---------------- constants (mirror depth_numpy.py / depth_gpu.ts) ----------------
	const D_IN = 1024;
	const D_MODEL = 768;
	const D_FF = 3072;
	const VOCAB = 12294;
	const NUM_CODEBOOKS = 12;
	const CODEBOOK_SIZE = 1024;
	const NUM_RESERVED = 6;
	const NUM_LAYERS = 2;
	const NUM_HEADS = 6;
	const UPH = 128; // per head dim
	const HD = NUM_HEADS * UPH; // 768
	const QKV = 3 * HD; // 2304
	const SOFT_CAP = 30.0;
	const EPS = 1e-6;
	const R_SOFTPLUS_0 = 1.442695041;

	// ---- grid matmul tiling (matches depth_gpu.ts) -------------------------------
	// LEGACY orient kernel (used only by Wo, which is already k-contiguous):
	// workgroup_size = TILE_ROWS * LANES threads.
	const TILE_ROWS = 64;
	const LANES = 4;
	const WG_SIZE = TILE_ROWS * LANES; // 256
	function nWG(N) { return Math.ceil(N / TILE_ROWS); }

	// ---- BANDWIDTH matmul tiling (the big single-token mat-VECs) -----------------
	// Weights are transposed ONCE at upload to k-CONTIGUOUS (W[n*K+k]) in the WT
	// buffer so consecutive K-reduction lanes read consecutive addresses (coalesced).
	// Each output is reduced by RED lanes; ROWS outputs per workgroup; ceil(N/ROWS)
	// workgroups fill the GPU. RED partials reduced by subgroupAdd when available,
	// else a shared-memory tree. All accumulation f32.
	const RED = 32; // K-reduction width per output (warp / subgroup width)
	const ROWS = 8; // output columns per workgroup
	const T_WG = ROWS * RED; // 256
	function nWGT(N) { return Math.ceil(N / ROWS); }
	let USE_SUBGROUPS = false; // set after device creation

	// ---------------- dispatch counter (intercept) ----------------
	let DISPATCH_COUNT = 0;
	const _origDispatch = GPUComputePassEncoder.prototype.dispatchWorkgroups;
	GPUComputePassEncoder.prototype.dispatchWorkgroups = function (x, y, z) {
	DISPATCH_COUNT++;
	return _origDispatch.call(this, x, y, z);
	};

	// ---------------- cached fetch (Cache API) ----------------
	async function cachedFetch(url, onProgress) {
	let cache = null;
	try { cache = await caches.open(CACHE_NAME); } catch (e) { cache = null; }
	if (cache) {
	const hit = await cache.match(url);
	if (hit) return hit;
	}
	const resp = await fetch(url);
	if (!resp.ok) throw new Error("fetch " + url + " -> HTTP " + resp.status);

	// stream so we can report progress on the big file
	const total = Number(resp.headers.get("content-length")) \|\| 0;
	if (resp.body && onProgress) {
	const reader = resp.body.getReader();
	const chunks = [];
	let received = 0;
	for (;;) {
	const { done, value } = await reader.read();
	if (done) break;
	chunks.push(value);
	received += value.byteLength;
	onProgress(received, total);
	}
	const blob = new Blob(chunks);
	const buffered = new Response(blob, {
	headers: { "content-type": resp.headers.get("content-type") \|\| "application/octet-stream" },
	});
	if (cache) { try { await cache.put(url, buffered.clone()); } catch (e) {} }
	return buffered;
	}
	if (cache) { try { await cache.put(url, resp.clone()); } catch (e) {} }
	return resp;
	}

	function fmtBytes(n) {
	if (n >= 1 << 30) return (n / (1 << 30)).toFixed(2) + " GB";
	if (n >= 1 << 20) return (n / (1 << 20)).toFixed(1) + " MB";
	if (n >= 1 << 10) return (n / (1 << 10)).toFixed(1) + " KB";
	return n + " B";
	}

	// ======================================================================
	// GRID-PARALLEL WGSL (ported verbatim from depth_gpu.ts).
	//
	// Weights are split across TWO storage buffers (W0/W1) so each binding stays
	// under the device's maxStorageBufferBindingSize. The full ~148MB blob exceeds
	// the 128MB cap on some devices/drivers (incl. browser/Dawn), which silently
	// zero an over-sized binding. wv(i) routes a global float index to W0 or W1.
	// W_SPLIT is filled in at module-build time below (depends on total_floats).
	// ======================================================================
	const CONSTS = `
	const D_IN : u32 = ${D_IN}u;
	const D_MODEL : u32 = ${D_MODEL}u;
	const D_FF : u32 = ${D_FF}u;
	const VOCAB : u32 = ${VOCAB}u;
	const NUM_HEADS : u32 = ${NUM_HEADS}u;
	const UPH : u32 = ${UPH}u;
	const HD : u32 = ${HD}u; // 768
	const QKV : u32 = ${QKV}u; // 2304
	const SOFT_CAP : f32 = ${SOFT_CAP};
	const EPS : f32 = ${EPS};
	const R_SOFTPLUS_0 : f32 = ${R_SOFTPLUS_0};
	const NUM_RESERVED : u32 = ${NUM_RESERVED}u;
	const CODEBOOK_SIZE : u32 = ${CODEBOOK_SIZE}u;
	const TILE_ROWS : u32 = ${TILE_ROWS}u;
	const LANES : u32 = ${LANES}u;
	const RED : u32 = ${RED}u;
	const ROWS : u32 = ${ROWS}u;
	`;

	// Built once W_SPLIT is known. Kernels that don't read weights (the sampler)
	// use CONSTS only, so layout:"auto" doesn't prune+misbind the W0/W1 buffers.
	function buildHeaders(W_SPLIT) {
	const WBUFS = `
	@group(0) @binding(0) var<storage, read> W0 : array<f32>;
	@group(0) @binding(6) var<storage, read> W1 : array<f32>;
	const W_SPLIT : u32 = ${W_SPLIT}u;
	fn wv(i : u32) -> f32 {
	if (i < W_SPLIT) { return W0[i]; }
	return W1[i - W_SPLIT];
	}

	fn softplus(x: f32) -> f32 {
	return log(1.0 + exp(-abs(x))) + max(x, 0.0);
	}
	fn gelu_tanh(x: f32) -> f32 {
	let c : f32 = 0.7978845608028654; // sqrt(2/pi)
	return 0.5 * x * (1.0 + tanh(c * (x + 0.044715 * x * x * x)));
	}
	`;
	return { HEADER: CONSTS + WBUFS };
	}

	// ----- grid matmul #1 (orient-1, W[in*N + out]); optional pre-RMSNorm + epilogue
	function gridMatmul1(HEADER, opts) {
	const K = opts.K;
	const inDecl =
	`@group(0) @binding(${opts.inBinding}) var<storage, read> vin : array<f32>;`;
	const outDecl =
	`@group(0) @binding(${opts.outBinding}) var<storage, read_write> vout : array<f32>;`;
	const uniformDecl =
	`struct U { wOff:u32, colBase:u32, normOff:u32, biasOff:u32, };
	@group(0) @binding(5) var<uniform> u : U;`;

	const normReduce = opts.rmsnorm
	? `
	// ---- fold pre-RMSNorm: reduce sum of squares over K (==D_MODEL) ----
	var lss : f32 = 0.0;
	for (var i : u32 = tid; i < ${K}u; i = i + ${WG_SIZE}u) {
	let xv = vin[i]; lss = lss + xv * xv;
	}
	red[tid] = lss;
	workgroupBarrier();
	for (var s : u32 = ${WG_SIZE / 2}u; s > 0u; s = s >> 1u) {
	if (tid < s) { red[tid] = red[tid] + red[tid + s]; }
	workgroupBarrier();
	}
	let rms = 1.0 / sqrt(red[0] / f32(${K}u) + EPS);
	workgroupBarrier();
	// stage normed input
	for (var i : u32 = tid; i < ${K}u; i = i + ${WG_SIZE}u) {
	inS[i] = (vin[i] * rms) * wv(u.normOff + i);
	}
	workgroupBarrier();`
	: `
	// stage raw input
	for (var i : u32 = tid; i < ${K}u; i = i + ${WG_SIZE}u) {
	inS[i] = vin[i];
	}
	workgroupBarrier();`;

	let epi = "";
	if (opts.epilogue === "geluBias") {
	epi = `acc = gelu_tanh(acc + wv(u.biasOff + col));`;
	} else if (opts.epilogue === "softcap") {
	const biasAdd = opts.hasBias ? `acc = acc + wv(u.biasOff + col);` : ``;
	epi = `${biasAdd}
	acc = SOFT_CAP * tanh(acc / SOFT_CAP);`;
	} else {
	epi = opts.hasBias ? `acc = acc + wv(u.biasOff + col);` : ``;
	}

	return HEADER + `
	${inDecl}
	${outDecl}
	${uniformDecl}

	var<workgroup> inS : array<f32, ${K}>;
	var<workgroup> red : array<f32, ${WG_SIZE}>;
	var<workgroup> part : array<f32, ${WG_SIZE}>;

	@compute @workgroup_size(${WG_SIZE})
	fn main(@builtin(local_invocation_id) lid : vec3<u32>,
	@builtin(workgroup_id) wid : vec3<u32>) {
	let tid = lid.x;
	let row = tid / LANES; // 0..TILE_ROWS-1
	let lane = tid % LANES; // 0..LANES-1
	${normReduce}

	let outIdx = wid.x * TILE_ROWS + row; // logical output column index
	var acc : f32 = 0.0;
	let col = u.colBase + outIdx; // physical column in W
	let wbase = u.wOff + col;
	for (var d : u32 = lane; d < ${K}u; d = d + LANES) {
	acc = acc + inS[d] * wv(wbase + d * ${opts.NCOL}u);
	}
	part[tid] = acc;
	workgroupBarrier();
	if (lane == 0u) {
	var s : f32 = part[row * LANES];
	for (var l : u32 = 1u; l < LANES; l = l + 1u) {
	s = s + part[row * LANES + l];
	}
	var acc : f32 = s;
	${epi}
	vout[outIdx] = acc;
	}
	}
	`;
	}

	// ----- BANDWIDTH matmul (k-contiguous WT @ binding 7) — see depth_gpu.ts -----
	function gridMatmulT(HEADER, opts) {
	const K = opts.K;
	const sg = USE_SUBGROUPS;
	const inDecl =
	`@group(0) @binding(${opts.inBinding}) var<storage, read> vin : array<f32>;`;
	const outDecl =
	`@group(0) @binding(${opts.outBinding}) var<storage, read_write> vout : array<f32>;`;
	const uniformDecl =
	`struct U { wOff:u32, colBase:u32, normOff:u32, biasOff:u32, };
	@group(0) @binding(5) var<uniform> u : U;`;
	const wtDecl = `@group(0) @binding(7) var<storage, read> WT : array<f32>;`;

	const normReduce = opts.rmsnorm
	? `
	var lss : f32 = 0.0;
	for (var i : u32 = tid; i < ${K}u; i = i + ${T_WG}u) {
	let xv = vin[i]; lss = lss + xv * xv;
	}
	red[tid] = lss;
	workgroupBarrier();
	for (var s : u32 = ${T_WG / 2}u; s > 0u; s = s >> 1u) {
	if (tid < s) { red[tid] = red[tid] + red[tid + s]; }
	workgroupBarrier();
	}
	let rms = 1.0 / sqrt(red[0] / f32(${K}u) + EPS);
	workgroupBarrier();
	for (var i : u32 = tid; i < ${K}u; i = i + ${T_WG}u) {
	inS[i] = (vin[i] * rms) * wv(u.normOff + i);
	}
	workgroupBarrier();`
	: `
	for (var i : u32 = tid; i < ${K}u; i = i + ${T_WG}u) { inS[i] = vin[i]; }
	workgroupBarrier();`;

	let epi = "";
	if (opts.epilogue === "geluBias") {
	epi = `acc = gelu_tanh(acc + wv(u.biasOff + col));`;
	} else if (opts.epilogue === "softcap") {
	const biasAdd = opts.hasBias ? `acc = acc + wv(u.biasOff + col);` : ``;
	epi = `${biasAdd}
	acc = SOFT_CAP * tanh(acc / SOFT_CAP);`;
	} else {
	epi = opts.hasBias ? `acc = acc + wv(u.biasOff + col);` : ``;
	}

	const reduceWrite = sg
	? `
	let total = subgroupAdd(acc);
	if (lane == 0u) {
	var accf : f32 = total;
	${epi.replace(/\bacc\b/g, "accf")}
	vout[outIdx] = accf;
	}`
	: `
	part[tid] = acc;
	workgroupBarrier();
	for (var s : u32 = ${RED / 2}u; s > 0u; s = s >> 1u) {
	if (lane < s) { part[tid] = part[tid] + part[tid + s]; }
	workgroupBarrier();
	}
	if (lane == 0u) {
	var accf : f32 = part[row * RED];
	${epi.replace(/\bacc\b/g, "accf")}
	vout[outIdx] = accf;
	}`;

	const enableSg = sg ? "enable subgroups;\n" : "";

	return enableSg + HEADER + `
	${inDecl}
	${outDecl}
	${uniformDecl}
	${wtDecl}

	var<workgroup> inS : array<f32, ${K}>;
	${opts.rmsnorm ? `var<workgroup> red : array<f32, ${T_WG}>;` : ``}
	${sg ? "" : `var<workgroup> part : array<f32, ${T_WG}>;`}

	@compute @workgroup_size(${T_WG})
	fn main(@builtin(local_invocation_id) lid : vec3<u32>,
	@builtin(workgroup_id) wid : vec3<u32>) {
	let tid = lid.x;
	let row = tid / RED;
	let lane = tid % RED;
	${normReduce}

	let outIdx = wid.x * ROWS + row;
	let col = u.colBase + outIdx;
	let wbase = u.wOff + col * ${K}u;
	var acc : f32 = 0.0;
	for (var k : u32 = lane; k < ${K}u; k = k + RED) {
	acc = acc + inS[k] * WT[wbase + k];
	}
	${reduceWrite}
	}
	`;
	}

	// ----- grid matmul #2 (orient-2, W[out*K + in]) — attention Wo
	function gridMatmul2(HEADER, opts) {
	const K = opts.K;
	return HEADER + `
	@group(0) @binding(${opts.inBinding}) var<storage, read> vin : array<f32>;
	@group(0) @binding(${opts.outBinding}) var<storage, read_write> vout : array<f32>;
	struct U { wOff:u32, colBase:u32, normOff:u32, biasOff:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> inS : array<f32, ${K}>;
	var<workgroup> part : array<f32, ${WG_SIZE}>;

	@compute @workgroup_size(${WG_SIZE})
	fn main(@builtin(local_invocation_id) lid : vec3<u32>,
	@builtin(workgroup_id) wid : vec3<u32>) {
	let tid = lid.x;
	let row = tid / LANES;
	let lane = tid % LANES;
	for (var i : u32 = tid; i < ${K}u; i = i + ${WG_SIZE}u) { inS[i] = vin[i]; }
	workgroupBarrier();
	let outIdx = wid.x * TILE_ROWS + row;
	let col = u.colBase + outIdx;
	let wbase = u.wOff + col * ${K}u;
	var acc : f32 = 0.0;
	for (var d : u32 = lane; d < ${K}u; d = d + LANES) {
	acc = acc + inS[d] * wv(wbase + d);
	}
	part[tid] = acc;
	workgroupBarrier();
	if (lane == 0u) {
	var s : f32 = part[row * LANES];
	for (var l : u32 = 1u; l < LANES; l = l + 1u) { s = s + part[row * LANES + l]; }
	vout[outIdx] = s;
	}
	}
	`;
	}

	// ----- QKV combined matmul (folds attn-pre RMSNorm) -> q, k[slot], v[slot]
	// Bandwidth version: Wq/Wk/Wv transposed k-contiguous in WT (binding 7); for
	// output column oi the K-stripe is WT[Wq + oi*768 + k] (coalesced across lanes).
	function buildQKV(HEADER) {
	const sg = USE_SUBGROUPS;
	const enableSg = sg ? "enable subgroups;\n" : "";
	const reduceWrite = sg
	? `
	let sq = subgroupAdd(qa);
	let sk = subgroupAdd(ka);
	let sv = subgroupAdd(va);
	if (lane == 0u) {
	qbuf[oi] = sq;
	let cbase = u.slot * HD;
	kcache[cbase + oi] = sk;
	vcache[cbase + oi] = sv;
	}`
	: `
	pq[tid] = qa; pk[tid] = ka; pv[tid] = va;
	workgroupBarrier();
	for (var s : u32 = ${RED / 2}u; s > 0u; s = s >> 1u) {
	if (lane < s) {
	pq[tid] = pq[tid] + pq[tid + s];
	pk[tid] = pk[tid] + pk[tid + s];
	pv[tid] = pv[tid] + pv[tid + s];
	}
	workgroupBarrier();
	}
	if (lane == 0u) {
	qbuf[oi] = pq[row * RED];
	let cbase = u.slot * HD;
	kcache[cbase + oi] = pk[row * RED];
	vcache[cbase + oi] = pv[row * RED];
	}`;
	return enableSg + HEADER + `
	@group(0) @binding(1) var<storage, read> vin : array<f32>; // x768 residual
	@group(0) @binding(2) var<storage, read_write> qbuf : array<f32>; // [768]
	@group(0) @binding(3) var<storage, read_write> kcache : array<f32>; // [T,768]
	@group(0) @binding(4) var<storage, read_write> vcache : array<f32>; // [T,768]
	@group(0) @binding(7) var<storage, read> WT : array<f32>;
	struct U { normOff:u32, Wq:u32, Wk:u32, Wv:u32, slot:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> inS : array<f32, ${D_MODEL}>;
	var<workgroup> red : array<f32, ${T_WG}>;
	${sg ? "" : `var<workgroup> pq : array<f32, ${T_WG}>;
	var<workgroup> pk : array<f32, ${T_WG}>;
	var<workgroup> pv : array<f32, ${T_WG}>;`}

	@compute @workgroup_size(${T_WG})
	fn main(@builtin(local_invocation_id) lid : vec3<u32>,
	@builtin(workgroup_id) wid : vec3<u32>) {
	let tid = lid.x;
	let row = tid / RED;
	let lane = tid % RED;
	// fold attn-pre RMSNorm over D_MODEL
	var lss : f32 = 0.0;
	for (var i : u32 = tid; i < D_MODEL; i = i + ${T_WG}u) {
	let xv = vin[i]; lss = lss + xv * xv;
	}
	red[tid] = lss;
	workgroupBarrier();
	for (var s : u32 = ${T_WG / 2}u; s > 0u; s = s >> 1u) {
	if (tid < s) { red[tid] = red[tid] + red[tid + s]; }
	workgroupBarrier();
	}
	let rms = 1.0 / sqrt(red[0] / f32(D_MODEL) + EPS);
	workgroupBarrier();
	for (var i : u32 = tid; i < D_MODEL; i = i + ${T_WG}u) {
	inS[i] = (vin[i] * rms) * wv(u.normOff + i);
	}
	workgroupBarrier();

	let oi = wid.x * ROWS + row; // output column in [0,768)
	let qb = u.Wq + oi * HD;
	let kb = u.Wk + oi * HD;
	let vb = u.Wv + oi * HD;
	var qa : f32 = 0.0; var ka : f32 = 0.0; var va : f32 = 0.0;
	for (var k : u32 = lane; k < D_MODEL; k = k + RED) {
	let hv = inS[k];
	qa = qa + hv * WT[qb + k];
	ka = ka + hv * WT[kb + k];
	va = va + hv * WT[vb + k];
	}
	${reduceWrite}
	}
	`;
	}

	// ----- attention core (small): per-head softmax over T<=12 keys -> ctx[768]
	function buildAttnCore(HEADER) {
	return HEADER + `
	@group(0) @binding(1) var<storage, read> qbuf : array<f32>; // [768]
	@group(0) @binding(2) var<storage, read> kcache : array<f32>; // [T,768]
	@group(0) @binding(3) var<storage, read> vcache : array<f32>; // [T,768]
	@group(0) @binding(4) var<storage, read_write> ctxbuf : array<f32>;// [768]
	struct U { per_dim_scale:u32, cacheLen:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> logitsW : array<f32, ${NUM_HEADS * NUM_CODEBOOKS}>;
	var<workgroup> red : array<f32, 2>;

	@compute @workgroup_size(256)
	fn main(@builtin(local_invocation_id) lid : vec3<u32>) {
	let t = lid.x;
	let T = u.cacheLen;
	let qscale = R_SOFTPLUS_0 * (1.0 / sqrt(f32(UPH)));
	for (var h : u32 = 0u; h < NUM_HEADS; h = h + 1u) {
	for (var k : u32 = t; k < T; k = k + 256u) {
	var acc : f32 = 0.0;
	let kbase = k * HD + h * UPH;
	let qbase = h * UPH;
	for (var d : u32 = 0u; d < UPH; d = d + 1u) {
	let sv = qscale * softplus(wv(u.per_dim_scale + d));
	acc = acc + (qbuf[qbase + d] * sv) * kcache[kbase + d];
	}
	logitsW[h * T + k] = acc;
	}
	}
	workgroupBarrier();
	for (var h : u32 = 0u; h < NUM_HEADS; h = h + 1u) {
	if (t == 0u) {
	var m : f32 = logitsW[h * T + 0u];
	for (var k : u32 = 1u; k < T; k = k + 1u) { m = max(m, logitsW[h * T + k]); }
	red[0] = m;
	}
	workgroupBarrier();
	let m = red[0];
	if (t == 0u) {
	var s : f32 = 0.0;
	for (var k : u32 = 0u; k < T; k = k + 1u) {
	let e = exp(logitsW[h * T + k] - m);
	logitsW[h * T + k] = e; s = s + e;
	}
	red[1] = s;
	}
	workgroupBarrier();
	let denom = red[1];
	for (var d : u32 = t; d < UPH; d = d + 256u) {
	var acc : f32 = 0.0;
	for (var k : u32 = 0u; k < T; k = k + 1u) {
	let w = logitsW[h * T + k] / denom;
	acc = acc + w * vcache[k * HD + h * UPH + d];
	}
	ctxbuf[h * UPH + d] = acc;
	}
	workgroupBarrier();
	}
	}
	`;
	}

	// ----- post-norm + residual (small, 768): x = x + rmsnorm(y) * scale
	function buildResid(HEADER) {
	return HEADER + `
	@group(0) @binding(1) var<storage, read_write> x : array<f32>; // [768] residual
	@group(0) @binding(2) var<storage, read> y : array<f32>; // [768] block out
	struct U { post:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> red : array<f32, 256>;
	@compute @workgroup_size(256)
	fn main(@builtin(local_invocation_id) lid : vec3<u32>) {
	let t = lid.x;
	var ss : f32 = 0.0;
	for (var i : u32 = t; i < D_MODEL; i = i + 256u) { let v = y[i]; ss = ss + v*v; }
	red[t] = ss;
	workgroupBarrier();
	for (var s : u32 = 128u; s > 0u; s = s >> 1u) {
	if (t < s) { red[t] = red[t] + red[t + s]; }
	workgroupBarrier();
	}
	let rms = 1.0 / sqrt(red[0] / f32(D_MODEL) + EPS);
	workgroupBarrier();
	for (var i : u32 = t; i < D_MODEL; i = i + 256u) {
	x[i] = x[i] + (y[i] * rms) * wv(u.post + i);
	}
	}
	`;
	}

	// ----- final LayerNorm (small, 768)
	function buildFinalLN(HEADER) {
	return HEADER + `
	@group(0) @binding(1) var<storage, read> x : array<f32>; // [768]
	@group(0) @binding(2) var<storage, read_write> xn : array<f32>; // [768]
	struct U { final_scale:u32, final_bias:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> red : array<f32, 256>;
	@compute @workgroup_size(256)
	fn main(@builtin(local_invocation_id) lid : vec3<u32>) {
	let t = lid.x;
	var sm : f32 = 0.0;
	for (var i : u32 = t; i < D_MODEL; i = i + 256u) { sm = sm + x[i]; }
	red[t] = sm;
	workgroupBarrier();
	for (var s : u32 = 128u; s > 0u; s = s >> 1u) {
	if (t < s) { red[t] = red[t] + red[t + s]; }
	workgroupBarrier();
	}
	let mean = red[0] / f32(D_MODEL);
	workgroupBarrier();
	var vs : f32 = 0.0;
	for (var i : u32 = t; i < D_MODEL; i = i + 256u) { let dd = x[i] - mean; vs = vs + dd*dd; }
	red[t] = vs;
	workgroupBarrier();
	for (var s : u32 = 128u; s > 0u; s = s >> 1u) {
	if (t < s) { red[t] = red[t] + red[t + s]; }
	workgroupBarrier();
	}
	let rstd = 1.0 / sqrt(red[0] / f32(D_MODEL) + EPS);
	workgroupBarrier();
	for (var i : u32 = t; i < D_MODEL; i = i + 256u) {
	xn[i] = ((x[i] - mean) * rstd) * wv(u.final_scale + i) + wv(u.final_bias + i);
	}
	}
	`;
	}

	// ----- sample (small): over the 1024-wide sliced+softcapped logits, gumbel+argmax.
	// WEIGHT-FREE: uses CONSTS only (no W0/W1), so layout:"auto" doesn't prune+misbind
	// the weight buffers and silently read zeros.
	function buildSample() {
	return CONSTS + `
	@group(0) @binding(1) var<storage, read> logits1024 : array<f32>; // [1024]
	@group(0) @binding(2) var<storage, read> noise : array<f32>; // [Q*VOCAB]
	@group(0) @binding(3) var<storage, read_write> outTok : array<u32>;// [NUM_CODEBOOKS]
	struct U { lo:u32, temperature:f32, slot:u32, noiseBase:u32, };
	@group(0) @binding(5) var<uniform> u : U;

	var<workgroup> bestScore : array<f32, 256>;
	var<workgroup> bestIdx : array<u32, 256>;
	@compute @workgroup_size(256)
	fn main(@builtin(local_invocation_id) lid : vec3<u32>) {
	let t = lid.x;
	let lo = u.lo;
	var bScore : f32 = -3.0e38;
	var bIdx : u32 = 0xffffffffu;
	for (var j : u32 = t; j < CODEBOOK_SIZE; j = j + 256u) {
	var score : f32 = logits1024[j]; // already soft-capped, in-slice
	let gidx = lo + j; // global token index
	if (u.temperature > 0.0) {
	var un = noise[u.noiseBase + gidx];
	un = clamp(un, 1.0e-10, 1.0 - 1.0e-7);
	let g = -log(-log(un));
	score = score + g * u.temperature;
	}
	if (score > bScore \|\| (score == bScore && gidx < bIdx)) {
	bScore = score; bIdx = gidx;
	}
	}
	bestScore[t] = bScore; bestIdx[t] = bIdx;
	workgroupBarrier();
	for (var s : u32 = 128u; s > 0u; s = s >> 1u) {
	if (t < s) {
	let a = bestScore[t]; let ai = bestIdx[t];
	let b = bestScore[t + s]; let bi = bestIdx[t + s];
	if (b > a \|\| (b == a && bi < ai)) { bestScore[t] = b; bestIdx[t] = bi; }
	}
	workgroupBarrier();
	}
	// Write sampled token into this step's slot of the GPU tokens buffer.
	// No CPU readback mid-loop: the next step's embed kernel reads outTok[slot].
	if (t == 0u) { outTok[u.slot] = bestIdx[0]; }
	}
	`;
	}

	// ----- embed_next: x1024 = embedding[tok] * embed_scale
	function buildEmbed(HEADER) {
	return HEADER + `
	@group(0) @binding(1) var<storage, read> tokBuf : array<u32>; // [NUM_CODEBOOKS]
	@group(0) @binding(2) var<storage, read_write> xout : array<f32>; // [1024]
	struct EU { embeddingOff : u32, embed_scale : f32, slot : u32 };
	@group(0) @binding(3) var<uniform> u : EU;

	@compute @workgroup_size(256)
	fn main(@builtin(local_invocation_id) lid : vec3<u32>) {
	let t = lid.x;
	// Gather row of the token sampled at the PREVIOUS step (slot), GPU-side.
	let tok = tokBuf[u.slot];
	let base = u.embeddingOff + tok * D_IN;
	for (var i : u32 = t; i < D_IN; i = i + 256u) {
	xout[i] = wv(base + i) * u.embed_scale;
	}
	}
	`;
	}

	// ---------------- main ----------------
	async function main() {
	if (!navigator.gpu) {
	throw new Error("navigator.gpu unavailable — this browser/origin has no WebGPU.");
	}

	// -------- download (cache via Cache API) --------
	setStatus("downloading weights + test data (depth.bin is 148MB; cached after first run)…", "running");

	const [layoutResp, testsResp, binResp] = await Promise.all([
	cachedFetch(BASE + "depth_layout.json"),
	cachedFetch(BASE + "testcases.json"),
	cachedFetch(BASE + "depth.bin", (recv, total) => {
	const pct = total ? " (" + (100 * recv / total).toFixed(0) + "%)" : "";
	setStatus("downloading depth.bin: " + fmtBytes(recv) +
	(total ? " / " + fmtBytes(total) : "") + pct + "…", "running");
	}),
	]);

	const layout = await layoutResp.json();
	const tests = await testsResp.json();
	const binBytes = new Uint8Array(await binResp.arrayBuffer());
	const EMBED_SCALE = layout.embed_scale;
	const TOTAL_FLOATS = layout.total_floats;
	// float-aligned split point for W0/W1 (matches depth_gpu.ts)
	const W_SPLIT = Math.ceil(TOTAL_FLOATS / 2);

	function off(name) { return layout.tensors[name].offset; }
	function layerOffsets(li) {
	const p = "l" + li + "_";
	return {
	attn_pre: off(p + "attn_pre"), attn_post: off(p + "attn_post"),
	Wq: off(p + "Wq"), Wk: off(p + "Wk"), Wv: off(p + "Wv"),
	per_dim_scale: off(p + "per_dim_scale"), Wo: off(p + "Wo"),
	ffn_pre: off(p + "ffn_pre"), ffn_post: off(p + "ffn_post"),
	Wi: off(p + "Wi"), bi: off(p + "bi"),
	Wo_ffn: off(p + "Wo_ffn"), bo: off(p + "bo"),
	};
	}

	// -------- GPU setup --------
	setStatus("initializing WebGPU device and uploading weights…", "running");

	const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
	if (!adapter) throw new Error("no WebGPU adapter");

	// adapter info (best-effort across implementations)
	let adapterInfo = { vendor: "", architecture: "", device: "", description: "" };
	try {
	if (adapter.info) adapterInfo = adapter.info;
	else if (adapter.requestAdapterInfo) adapterInfo = await adapter.requestAdapterInfo();
	} catch (e) {}

	// Use subgroups (Metal / A100-native) only when exposed AND the hardware
	// subgroup width is exactly RED, so subgroupAdd reduces one output's RED
	// partials and nothing else. Otherwise use the shared-memory tree fallback.
	const subgroupsOk = adapter.features.has("subgroups") &&
	adapter.limits.minSubgroupSize === RED &&
	adapter.limits.maxSubgroupSize === RED;
	USE_SUBGROUPS = subgroupsOk;

	const device = await adapter.requestDevice({
	requiredFeatures: subgroupsOk ? ["subgroups"] : [],
	requiredLimits: {
	maxStorageBufferBindingSize: Math.min(
	Math.max(256 * 1024 * 1024, binBytes.byteLength),
	adapter.limits.maxStorageBufferBindingSize),
	maxBufferSize: Math.min(
	Math.max(256 * 1024 * 1024, binBytes.byteLength),
	adapter.limits.maxBufferSize),
	maxComputeWorkgroupStorageSize: Math.min(
	32768, adapter.limits.maxComputeWorkgroupStorageSize),
	},
	});
	let deviceLost = null;
	device.lost.then((info) => { deviceLost = info; });
	console.log("subgroups: " + USE_SUBGROUPS + " (RED=" + RED + " ROWS=" + ROWS + " T_WG=" + T_WG + ")");

	function makeStorage(floats, usageExtra = 0) {
	return device.createBuffer({
	size: Math.max(4, floats * 4),
	usage: GPUBufferUsage.STORAGE \| GPUBufferUsage.COPY_DST \|
	GPUBufferUsage.COPY_SRC \| usageExtra,
	});
	}

	// ---- TWO weight buffers (W0/W1 split) ----
	// W0 = floats [0, W_SPLIT), W1 = floats [W_SPLIT, TOTAL_FLOATS). Splitting keeps
	// each binding under maxStorageBufferBindingSize so Dawn doesn't zero an
	// over-sized binding.
	const W0_BYTES = W_SPLIT * 4;
	const W1_BYTES = binBytes.byteLength - W0_BYTES;
	const w0Buf = device.createBuffer({
	size: W0_BYTES,
	usage: GPUBufferUsage.STORAGE \| GPUBufferUsage.COPY_DST,
	});
	const w1Buf = device.createBuffer({
	size: W1_BYTES,
	usage: GPUBufferUsage.STORAGE \| GPUBufferUsage.COPY_DST,
	});
	device.queue.writeBuffer(w0Buf, 0, binBytes, 0, W0_BYTES);
	device.queue.writeBuffer(w1Buf, 0, binBytes, W0_BYTES, W1_BYTES);

	// ---- transposed (k-contiguous) weight buffer WT ----
	// The big mat-VECs are packed in-major (W[k*N+n]) in depth.bin; transpose each
	// ONCE here into WT laid out k-contiguous (WT[n*K+k]) so the bandwidth kernels
	// read consecutive k on consecutive lanes -> fully coalesced. depth.bin is
	// unchanged. `Wo` is already k-contiguous and stays in W0/W1.
	const binF32 = new Float32Array(
	binBytes.buffer, binBytes.byteOffset, binBytes.byteLength / 4);
	const WT_TENSORS = [{ src: "adapter", K: D_IN, N: D_MODEL }];
	for (let li = 0; li < NUM_LAYERS; li++) {
	const p = "l" + li + "_";
	WT_TENSORS.push(
	{ src: p + "Wq", K: D_MODEL, N: HD },
	{ src: p + "Wk", K: D_MODEL, N: HD },
	{ src: p + "Wv", K: D_MODEL, N: HD },
	{ src: p + "Wi", K: D_MODEL, N: D_FF },
	{ src: p + "Wo_ffn", K: D_FF, N: D_MODEL });
	}
	WT_TENSORS.push({ src: "to_logits_w", K: D_MODEL, N: VOCAB });
	const wtOffset = {};
	let wtFloats = 0;
	for (const t of WT_TENSORS) { wtOffset[t.src] = wtFloats; wtFloats += t.K * t.N; }
	const wtOff = (name) => wtOffset[name];
	const wtHost = new Float32Array(wtFloats);
	for (const t of WT_TENSORS) {
	const srcOff = off(t.src), dstOff = wtOffset[t.src], K = t.K, N = t.N;
	for (let k = 0; k < K; k++) {
	const srcRow = srcOff + k * N;
	for (let n = 0; n < N; n++) wtHost[dstOff + n * K + k] = binF32[srcRow + n];
	}
	}
	const wtBuf = device.createBuffer({
	size: wtHost.byteLength,
	usage: GPUBufferUsage.STORAGE \| GPUBufferUsage.COPY_DST,
	});
	device.queue.writeBuffer(wtBuf, 0, wtHost);

	// activation buffers (all GPU-resident)
	const xinBuf = makeStorage(D_IN); // x1024
	const x768Buf = makeStorage(D_MODEL); // residual stream
	const qBuf = makeStorage(HD);
	const ctxBuf = makeStorage(HD);
	const attnOutBuf = makeStorage(D_MODEL); // raw Wo output
	const ffBuf = makeStorage(D_FF); // ffn hidden
	const ffnOutBuf = makeStorage(D_MODEL); // raw Wo_ffn output
	const xnBuf = makeStorage(D_MODEL); // final layernormed
	const logits1024Buf = makeStorage(CODEBOOK_SIZE);
	const kBufs = [makeStorage(NUM_CODEBOOKS * D_MODEL), makeStorage(NUM_CODEBOOKS * D_MODEL)];
	const vBufs = [makeStorage(NUM_CODEBOOKS * D_MODEL), makeStorage(NUM_CODEBOOKS * D_MODEL)];
	// GPU-resident tokens buffer: one u32 slot per codebook step. The sampler at
	// step q writes tokens[q]; the embed kernel at step q reads tokens[q]. Read
	// back to the CPU ONCE, after all 12 steps.
	const tokBuf = device.createBuffer({
	size: NUM_CODEBOOKS * 4,
	usage: GPUBufferUsage.STORAGE \| GPUBufferUsage.COPY_DST \| GPUBufferUsage.COPY_SRC,
	});
	const tokReadBuf = device.createBuffer({
	size: NUM_CODEBOOKS * 4,
	usage: GPUBufferUsage.MAP_READ \| GPUBufferUsage.COPY_DST,
	});
	// Whole-noise buffer holds all Q rows ([Q*VOCAB]); uploaded once per depth.
	const noiseBuf = makeStorage(NUM_CODEBOOKS * VOCAB);

	function makeUniform(bytes) {
	return device.createBuffer({
	size: bytes,
	usage: GPUBufferUsage.UNIFORM \| GPUBufferUsage.COPY_DST,
	});
	}

	function pipeline(code) {
	const module = device.createShaderModule({ code });
	return device.createComputePipeline({
	layout: "auto",
	compute: { module, entryPoint: "main" },
	});
	}

	// ---- build WGSL with the W0/W1 header ----
	const { HEADER } = buildHeaders(W_SPLIT);

	// ---- grid matmul pipelines (bandwidth kernels read k-contiguous WT) ----
	const pAdapter = pipeline(gridMatmulT(HEADER, {
	K: D_IN, inBinding: 1, outBinding: 2,
	rmsnorm: false, epilogue: "none", hasBias: false,
	}));
	const pQKV = pipeline(buildQKV(HEADER));
	const pAttnCore = pipeline(buildAttnCore(HEADER));
	const pWo = pipeline(gridMatmul2(HEADER, { K: HD, inBinding: 1, outBinding: 2 }));
	const pResid = pipeline(buildResid(HEADER));
	const pWi = pipeline(gridMatmulT(HEADER, {
	K: D_MODEL, inBinding: 1, outBinding: 2,
	rmsnorm: true, epilogue: "geluBias", hasBias: true,
	}));
	const pWoFfn = pipeline(gridMatmulT(HEADER, {
	K: D_FF, inBinding: 1, outBinding: 2,
	rmsnorm: false, epilogue: "none", hasBias: true,
	}));
	const pFinalLN = pipeline(buildFinalLN(HEADER));
	const pToLogits = pipeline(gridMatmulT(HEADER, {
	K: D_MODEL, inBinding: 1, outBinding: 2,
	rmsnorm: false, epilogue: "softcap", hasBias: true,
	}));
	const pSample = pipeline(buildSample());
	const pEmbed = pipeline(buildEmbed(HEADER));

	// ---- uniform buffers ----
	// All per-step uniforms are pre-created (one buffer per (step,layer)) and
	// filled ONCE up front. Nothing in the depth loop calls writeBuffer, so no
	// per-step queue write can serialize the otherwise back-to-back submits.
	const uAdapter = makeUniform(16);
	const uWo = [makeUniform(16), makeUniform(16)];
	const uWi = [makeUniform(16), makeUniform(16)];
	const uWoFfn = [makeUniform(16), makeUniform(16)];
	// to_logits varies per step (colBase = lo): one buffer per step.
	const uToLogits = Array.from({ length: NUM_CODEBOOKS }, () => makeUniform(16));
	// { normOff,Wq,Wk,Wv,slot }: slot varies per step -> one per (step, layer).
	const uQKV = Array.from({ length: NUM_CODEBOOKS },
	() => [makeUniform(32), makeUniform(32)]);
	// { per_dim_scale, cacheLen }: cacheLen varies per step -> one per (step, layer).
	const uAttnCore = Array.from({ length: NUM_CODEBOOKS },
	() => [makeUniform(16), makeUniform(16)]);
	const uResidAttn = [makeUniform(16), makeUniform(16)];
	const uResidFfn = [makeUniform(16), makeUniform(16)];
	const uFinalLN = makeUniform(16);
	// { lo, temperature, slot, noiseBase }: per step.
	const uSample = Array.from({ length: NUM_CODEBOOKS }, () => makeUniform(16));
	// { embeddingOff, embed_scale, slot }: per step.
	const uEmbed = Array.from({ length: NUM_CODEBOOKS }, () => makeUniform(16));

	// ---- fill static uniforms ----
	// wOff for transposed bandwidth matmuls indexes WT (wtOff), not depth.bin.
	{
	const a = new Uint32Array(4);
	a[0] = wtOff("adapter");
	a[1] = 0; // colBase
	device.queue.writeBuffer(uAdapter, 0, a);
	}
	for (let li = 0; li < NUM_LAYERS; li++) {
	const L = layerOffsets(li);
	const p = "l" + li + "_";
	// QKV + attn core: slot / cacheLen vary per step, so fill all NUM_CODEBOOKS
	// buffers up front (one per (step, layer)). Wq/Wk/Wv index WT.
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const a = new Uint32Array(8);
	a[0] = L.attn_pre; a[1] = wtOff(p + "Wq"); a[2] = wtOff(p + "Wk"); a[3] = wtOff(p + "Wv");
	a[4] = q; // slot
	device.queue.writeBuffer(uQKV[q][li], 0, a);
	const c = new Uint32Array(4);
	c[0] = L.per_dim_scale;
	c[1] = q + 1; // cacheLen
	device.queue.writeBuffer(uAttnCore[q][li], 0, c);
	}
	{
	const a = new Uint32Array(4);
	a[0] = L.attn_post;
	device.queue.writeBuffer(uResidAttn[li], 0, a);
	}
	{
	const a = new Uint32Array(4);
	a[0] = L.Wo; a[1] = 0;
	device.queue.writeBuffer(uWo[li], 0, a);
	}
	{
	const a = new Uint32Array(4);
	a[0] = wtOff(p + "Wi"); a[1] = 0; a[2] = L.ffn_pre; a[3] = L.bi;
	device.queue.writeBuffer(uWi[li], 0, a);
	}
	{
	const a = new Uint32Array(4);
	a[0] = wtOff(p + "Wo_ffn"); a[1] = 0; a[2] = 0; a[3] = L.bo;
	device.queue.writeBuffer(uWoFfn[li], 0, a);
	}
	{
	const a = new Uint32Array(4);
	a[0] = L.ffn_post;
	device.queue.writeBuffer(uResidFfn[li], 0, a);
	}
	}
	{
	const a = new Uint32Array(4);
	a[0] = off("final_scale");
	a[1] = off("final_bias");
	device.queue.writeBuffer(uFinalLN, 0, a);
	}
	// embed (per step): { embeddingOff, embed_scale, slot }. The embed kernel at
	// step q reads tokens[q] (just sampled this step) for the next step's input.
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const e = new ArrayBuffer(16);
	new Uint32Array(e, 0, 1)[0] = off("embedding");
	new Float32Array(e, 4, 1)[0] = EMBED_SCALE;
	new Uint32Array(e, 8, 1)[0] = q; // slot
	device.queue.writeBuffer(uEmbed[q], 0, e);
	}
	// to_logits (per step): colBase = lo = NUM_RESERVED + q*CODEBOOK_SIZE.
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const lo = NUM_RESERVED + q * CODEBOOK_SIZE;
	const a = new Uint32Array(4);
	a[0] = wtOff("to_logits_w"); // wOff into WT (k-contiguous, all 12294 cols)
	a[1] = lo; // colBase -> global column = WT row
	a[2] = 0; // normOff (unused)
	a[3] = off("to_logits_b"); // biasOff (global column lo+j, in W0/W1)
	device.queue.writeBuffer(uToLogits[q], 0, a);
	}

	// ---- bind groups ----
	// WT@7 always; W0@0/W1@6 only when the kernel calls wv() (biases/norm scales).
	// The adapter (no rmsnorm/bias) never calls wv, so layout:"auto" prunes W0/W1
	// and binding them would fail validation -> must pass usesWv=false there.
	function bgMatmul1(p, inBuf, outBuf, uBuf, usesWv = true) {
	const entries = [
	{ binding: 7, resource: { buffer: wtBuf } },
	{ binding: 1, resource: { buffer: inBuf } },
	{ binding: 2, resource: { buffer: outBuf } },
	{ binding: 5, resource: { buffer: uBuf } },
	];
	if (usesWv) {
	entries.push(
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } });
	}
	return device.createBindGroup({ layout: p.getBindGroupLayout(0), entries });
	}
	const bgAdapter = bgMatmul1(pAdapter, xinBuf, x768Buf, uAdapter, false);
	const bgWo = [0, 1].map((li) => device.createBindGroup({
	layout: pWo.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: ctxBuf } },
	{ binding: 2, resource: { buffer: attnOutBuf } },
	{ binding: 5, resource: { buffer: uWo[li] } },
	],
	}));
	// Per-step (q) x per-layer (li) bind groups: only the per-step uniform buffer
	// (uQKV[q][li] / uAttnCore[q][li]) differs vs. a static bind group.
	const bgQKV = Array.from({ length: NUM_CODEBOOKS }, (_, q) =>
	[0, 1].map((li) => device.createBindGroup({
	layout: pQKV.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 7, resource: { buffer: wtBuf } },
	{ binding: 1, resource: { buffer: x768Buf } },
	{ binding: 2, resource: { buffer: qBuf } },
	{ binding: 3, resource: { buffer: kBufs[li] } },
	{ binding: 4, resource: { buffer: vBufs[li] } },
	{ binding: 5, resource: { buffer: uQKV[q][li] } },
	],
	})));
	const bgAttnCore = Array.from({ length: NUM_CODEBOOKS }, (_, q) =>
	[0, 1].map((li) => device.createBindGroup({
	layout: pAttnCore.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: qBuf } },
	{ binding: 2, resource: { buffer: kBufs[li] } },
	{ binding: 3, resource: { buffer: vBufs[li] } },
	{ binding: 4, resource: { buffer: ctxBuf } },
	{ binding: 5, resource: { buffer: uAttnCore[q][li] } },
	],
	})));
	const bgResidAttn = [0, 1].map((li) => device.createBindGroup({
	layout: pResid.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: x768Buf } },
	{ binding: 2, resource: { buffer: attnOutBuf } },
	{ binding: 5, resource: { buffer: uResidAttn[li] } },
	],
	}));
	const bgWi = [0, 1].map((li) => bgMatmul1(pWi, x768Buf, ffBuf, uWi[li]));
	const bgWoFfn = [0, 1].map((li) => bgMatmul1(pWoFfn, ffBuf, ffnOutBuf, uWoFfn[li]));
	const bgResidFfn = [0, 1].map((li) => device.createBindGroup({
	layout: pResid.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: x768Buf } },
	{ binding: 2, resource: { buffer: ffnOutBuf } },
	{ binding: 5, resource: { buffer: uResidFfn[li] } },
	],
	}));
	const bgFinalLN = device.createBindGroup({
	layout: pFinalLN.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: x768Buf } },
	{ binding: 2, resource: { buffer: xnBuf } },
	{ binding: 5, resource: { buffer: uFinalLN } },
	],
	});
	// Per-step bind groups for slices/uniforms that vary with q.
	const bgToLogits = Array.from({ length: NUM_CODEBOOKS },
	(_, q) => bgMatmul1(pToLogits, xnBuf, logits1024Buf, uToLogits[q]));
	const bgSample = Array.from({ length: NUM_CODEBOOKS }, (_, q) =>
	device.createBindGroup({
	layout: pSample.getBindGroupLayout(0),
	entries: [
	{ binding: 1, resource: { buffer: logits1024Buf } },
	{ binding: 2, resource: { buffer: noiseBuf } },
	{ binding: 3, resource: { buffer: tokBuf } },
	{ binding: 5, resource: { buffer: uSample[q] } },
	],
	}));
	const bgEmbed = Array.from({ length: NUM_CODEBOOKS }, (_, q) =>
	device.createBindGroup({
	layout: pEmbed.getBindGroupLayout(0),
	entries: [
	{ binding: 0, resource: { buffer: w0Buf } },
	{ binding: 6, resource: { buffer: w1Buf } },
	{ binding: 1, resource: { buffer: tokBuf } },
	{ binding: 2, resource: { buffer: xinBuf } },
	{ binding: 3, resource: { buffer: uEmbed[q] } },
	],
	}));

	// precomputed dispatch grid sizes (bandwidth kernels: ceil(N/ROWS) workgroups)
	const WG_ADAPTER = nWGT(D_MODEL); // 768/ROWS
	const WG_QKV = nWGT(HD); // 768/ROWS
	const WG_WO = nWG(D_MODEL); // 12 (legacy gridMatmul2)
	const WG_WI = nWGT(D_FF); // 3072/ROWS
	const WG_WOFFN = nWGT(D_MODEL); // 768/ROWS
	const WG_TOLOGITS = nWGT(CODEBOOK_SIZE);// 1024/ROWS

	// -------- one depth (12 steps), GPU-resident, one pass per step --------
	async function runDepth(temporalOut, noiseFlat, temperature) {
	// ---- one-time per-depth uploads (NO writes inside the step loop) ----
	device.queue.writeBuffer(xinBuf, 0, temporalOut);
	// Upload the whole [Q*VOCAB] noise blob once; the sampler indexes it per
	// step via uSample[q].noiseBase = q*VOCAB.
	device.queue.writeBuffer(noiseBuf, 0, noiseFlat);
	// Fill per-step sample uniforms (depend on temperature -> per-depth).
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const lo = NUM_RESERVED + q * CODEBOOK_SIZE;
	const buf = new ArrayBuffer(16);
	new Uint32Array(buf, 0, 1)[0] = lo;
	new Float32Array(buf, 4, 1)[0] = temperature;
	new Uint32Array(buf, 8, 1)[0] = q; // slot
	new Uint32Array(buf, 12, 1)[0] = q * VOCAB; // noiseBase
	device.queue.writeBuffer(uSample[q], 0, buf);
	}

	// ---- 12 GPU-resident steps: one compute pass each, back-to-back submits.
	// No awaits, maps, readbacks, or writeBuffer inside this loop. Each step's
	// sampler writes tokens[q] on the GPU; the next step's embed reads tokens[q]
	// on the GPU. The whole loop stays GPU-resident. ----
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const enc = device.createCommandEncoder();
	const pass = enc.beginComputePass();

	pass.setPipeline(pAdapter);
	pass.setBindGroup(0, bgAdapter);
	pass.dispatchWorkgroups(WG_ADAPTER);

	for (let li = 0; li < NUM_LAYERS; li++) {
	pass.setPipeline(pQKV);
	pass.setBindGroup(0, bgQKV[q][li]);
	pass.dispatchWorkgroups(WG_QKV);

	pass.setPipeline(pAttnCore);
	pass.setBindGroup(0, bgAttnCore[q][li]);
	pass.dispatchWorkgroups(1);

	pass.setPipeline(pWo);
	pass.setBindGroup(0, bgWo[li]);
	pass.dispatchWorkgroups(WG_WO);

	pass.setPipeline(pResid);
	pass.setBindGroup(0, bgResidAttn[li]);
	pass.dispatchWorkgroups(1);

	pass.setPipeline(pWi);
	pass.setBindGroup(0, bgWi[li]);
	pass.dispatchWorkgroups(WG_WI);

	pass.setPipeline(pWoFfn);
	pass.setBindGroup(0, bgWoFfn[li]);
	pass.dispatchWorkgroups(WG_WOFFN);

	pass.setPipeline(pResid);
	pass.setBindGroup(0, bgResidFfn[li]);
	pass.dispatchWorkgroups(1);
	}

	pass.setPipeline(pFinalLN);
	pass.setBindGroup(0, bgFinalLN);
	pass.dispatchWorkgroups(1);

	pass.setPipeline(pToLogits);
	pass.setBindGroup(0, bgToLogits[q]);
	pass.dispatchWorkgroups(WG_TOLOGITS);

	pass.setPipeline(pSample);
	pass.setBindGroup(0, bgSample[q]);
	pass.dispatchWorkgroups(1);

	if (q < NUM_CODEBOOKS - 1) {
	pass.setPipeline(pEmbed);
	pass.setBindGroup(0, bgEmbed[q]);
	pass.dispatchWorkgroups(1);
	}
	pass.end();
	device.queue.submit([enc.finish()]);
	}

	// ---- SINGLE readback of all 12 tokens, after the whole loop ----
	const enc2 = device.createCommandEncoder();
	enc2.copyBufferToBuffer(tokBuf, 0, tokReadBuf, 0, NUM_CODEBOOKS * 4);
	device.queue.submit([enc2.finish()]);
	await tokReadBuf.mapAsync(GPUMapMode.READ);
	const toks = new Uint32Array(tokReadBuf.getMappedRange().slice(0));
	tokReadBuf.unmap();

	const codes = [];
	for (let q = 0; q < NUM_CODEBOOKS; q++) {
	const tok = toks[q];
	const code = ((Number(tok) - NUM_RESERVED) % CODEBOOK_SIZE + CODEBOOK_SIZE) % CODEBOOK_SIZE;
	codes.push(code);
	}
	return codes;
	}

	// -------- correctness over the 3 testcases --------
	setStatus("running correctness over testcases…", "running");
	const allCases = tests.cases;
	const caseReports = [];
	let allMatch = true;
	let dispatchesPerDepth = 0;

	for (let ci = 0; ci < allCases.length; ci++) {
	const c = allCases[ci];
	DISPATCH_COUNT = 0;
	const gpuCodes = await runDepth(
	new Float32Array(c.temporal_out),
	new Float32Array(c.noise),
	c.temperature,
	);
	const ref = c.ref_codes;
	const match = gpuCodes.length === ref.length && gpuCodes.every((v, i) => v === ref[i]);
	allMatch = allMatch && match;
	if (ci === 0) dispatchesPerDepth = DISPATCH_COUNT;
	caseReports.push({ ci, gpuCodes, ref, match });
	console.log("case " + ci + ": GPU=[" + gpuCodes.join(",") +
	"] REF=[" + ref.join(",") + "] MATCH=" + match);
	}
	if (deviceLost) throw new Error("device lost: " + deviceLost.message);

	// -------- timing: median over 50 runs, fixed input (testcase 0) --------
	setStatus("timing depth (warmup + 50 runs)…", "running");
	const c0 = allCases[0];
	const t_temporal = new Float32Array(c0.temporal_out);
	const t_noise = new Float32Array(c0.noise);
	const t_temp = c0.temperature;

	const WARMUP = 5;
	const RUNS = 50;
	for (let i = 0; i < WARMUP; i++) {
	await runDepth(t_temporal, t_noise, t_temp);
	await device.queue.onSubmittedWorkDone();
	}
	const times = [];
	for (let i = 0; i < RUNS; i++) {
	const t0 = performance.now();
	await runDepth(t_temporal, t_noise, t_temp);
	await device.queue.onSubmittedWorkDone();
	const t1 = performance.now();
	times.push(t1 - t0);
	if (i % 10 === 0) {
	setStatus("timing depth (run " + (i + 1) + "/" + RUNS + ")…", "running");
	}
	}
	times.sort((a, b) => a - b);
	const median = times[Math.floor(times.length / 2)];
	const tmin = times[0];
	const tmax = times[times.length - 1];

	if (deviceLost) throw new Error("device lost: " + deviceLost.message);

	// -------- render --------
	renderResult({
	adapterInfo, dispatchesPerDepth, median, tmin, tmax, RUNS,
	allMatch, caseReports, binBytes,
	});
	setStatus("done.", "ok");
	}

	function renderResult(r) {
	const ai = r.adapterInfo \|\| {};
	const adapterStr = [ai.vendor, ai.architecture, ai.device, ai.description]
	.filter(Boolean).join(" · ") \|\| "(adapter info unavailable)";

	const casesHtml = r.caseReports.map((c) =>
	'<div class="row"><span class="k">case ' + c.ci +
	' — MATCH=' + c.match + '</span>' +
	'<span class="v ' + (c.match ? "ok" : "bad") + '">' +
	(c.match ? "MATCH" : "MISMATCH") + '</span></div>' +
	'<pre class="codes">GPU [' + c.gpuCodes.join(", ") + ']\nREF [' + c.ref.join(", ") + ']</pre>'
	).join("");

	resultEl.innerHTML =
	'<div class="card hero">' +
	'<div class="label">median depth time</div>' +
	'<div class="big">' + r.median.toFixed(2) + '<span class="unit"> ms</span></div>' +
	'<div class="sub">12-step depth, median of ' + r.RUNS + ' runs ' +
	'(min ' + r.tmin.toFixed(2) + ' / max ' + r.tmax.toFixed(2) + ' ms)</div>' +
	'<div style="margin-top:14px"><span class="badge ' +
	(r.allMatch ? "pass" : "fail") + '">' +
	(r.allMatch ? "CORRECTNESS PASS" : "CORRECTNESS FAIL") + '</span></div>' +
	'</div>' +
	'<div class="card">' +
	'<div class="row"><span class="k">GPU adapter</span><span class="v">' +
	escapeHtml(adapterStr) + '</span></div>' +
	'<div class="row"><span class="k">dispatchWorkgroups / depth</span><span class="v">' +
	r.dispatchesPerDepth + '</span></div>' +
	'<div class="row"><span class="k">weights uploaded (W0+W1)</span><span class="v">' +
	fmtBytes(r.binBytes.byteLength) + '</span></div>' +
	'<div class="row"><span class="k">median depth time</span><span class="v">' +
	r.median.toFixed(3) + ' ms</span></div>' +
	'</div>' +
	'<div class="card">' +
	'<div class="row"><span class="k">testcase correctness (bit-exact codes)</span>' +
	'<span class="v ' + (r.allMatch ? "ok" : "bad") + '">' +
	(r.allMatch ? "all MATCH" : "FAIL") + '</span></div>' +
	casesHtml +
	'</div>';
	resultEl.style.display = "block";
	}

	// auto-run on load
	main().catch((e) => showError("main", e));
	</script>
	</body>
	</html>