Spaces:

magenta-community
/

magenta-rt-stack

Running

File size: 51,893 Bytes

<!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 &mdash; DEPTH (WGSL) timing</h1>
  <div class="sub">Grid-parallel WebGPU depth transformer &middot; W0/W1 split storage buffers &middot; bit-exact correctness + real device timing</div>
  <div id="status" class="running">starting&hellip;</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 => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;" }[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 +
    ' &mdash; 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>