// F5-TTS DiT custom WebGPU engine (fp16 weights, fp32 accumulators).
const PRELUDE = `enable f16;\n`;

const REDUCE8 = `
  if (tid < 128u) { red[tid] += red[tid + 128u]; } workgroupBarrier();
  if (tid < 64u) { red[tid] += red[tid + 64u]; } workgroupBarrier();
  if (tid < 32u) { red[tid] += red[tid + 32u]; } workgroupBarrier();
  if (tid < 16u) { red[tid] += red[tid + 16u]; } workgroupBarrier();
  if (tid < 8u) { red[tid] += red[tid + 8u]; } workgroupBarrier();
  if (tid < 4u) { red[tid] += red[tid + 4u]; } workgroupBarrier();
  if (tid < 2u) { red[tid] += red[tid + 2u]; } workgroupBarrier();
  if (tid < 1u) { red[tid] += red[tid + 1u]; } workgroupBarrier();
`;
const REDUCE8MAX = `
  if (tid < 128u) { red[tid] = max(red[tid], red[tid + 128u]); } workgroupBarrier();
  if (tid < 64u) { red[tid] = max(red[tid], red[tid + 64u]); } workgroupBarrier();
  if (tid < 32u) { red[tid] = max(red[tid], red[tid + 32u]); } workgroupBarrier();
  if (tid < 16u) { red[tid] = max(red[tid], red[tid + 16u]); } workgroupBarrier();
  if (tid < 8u) { red[tid] = max(red[tid], red[tid + 8u]); } workgroupBarrier();
  if (tid < 4u) { red[tid] = max(red[tid], red[tid + 4u]); } workgroupBarrier();
  if (tid < 2u) { red[tid] = max(red[tid], red[tid + 2u]); } workgroupBarrier();
  if (tid < 1u) { red[tid] = max(red[tid], red[tid + 1u]); } workgroupBarrier();
`;

const GEMM = `
enable chromium_experimental_subgroup_matrix;
diagnostic(off, chromium.subgroup_matrix_uniformity);
struct P { M: u32, N: u32, K: u32, pad: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> A: array<f16>;
@group(0) @binding(2) var<storage, read> B: array<f16>;
@group(0) @binding(3) var<storage, read> bias: array<f16>;
@group(0) @binding(4) var<storage, read_write> C: array<f16>;
const TR: u32 = 128u; const TC: u32 = 64u; const TK: u32 = 32u;
var<workgroup> tile_A: array<f16, TR * TK>;
var<workgroup> tile_B: array<f16, TC * TK>;
var<workgroup> scratch: array<array<f32, 64>, 8>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let bm = wg.x * TR; let bn = wg.y * TC;
  let sg = tid / 32u; let lane = tid % 32u;
  let sg_row = sg / 2u; let sg_col = sg % 2u;
  var acc: array<subgroup_matrix_result<f32, 8, 8>, 16>;
  for (var i = 0u; i < 16u; i++) { acc[i] = subgroup_matrix_result<f32, 8, 8>(); }
  let nk = min((p.K + TK - 1u) / TK, 2048u);
  for (var kt = 0u; kt < nk; kt++) {
    let k0 = kt * TK;
    for (var i = 0u; i < 16u; i++) {
      let idx = i * 256u + tid; let r = idx / TK; let c = idx % TK;
      let gr = bm + r; let gc = k0 + c;
      var v: f16 = f16(0.0);
      if (gr < p.M && gc < p.K) { v = A[gr * p.K + gc]; }
      tile_A[idx] = v;
    }
    for (var i = 0u; i < 8u; i++) {
      let idx = i * 256u + tid; let gn = idx % TC; let gk = idx / TC;
      let gr = bn + gn; let gc = k0 + gk;
      var v: f16 = f16(0.0);
      if (gr < p.N && gc < p.K) { v = B[gc * p.N + gr]; }
      tile_B[gn * TK + gk] = v;
    }
    workgroupBarrier();
    for (var step = 0u; step < TK; step += 8u) {
      let a_off = (sg_row * 32u) * TK + step;
      let m0 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off, false, TK);
      let m1 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 8u * TK, false, TK);
      let m2 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 16u * TK, false, TK);
      let m3 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 24u * TK, false, TK);
      let b_off = (sg_col * 32u) * TK + step;
      let n0 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off, true, TK);
      let n1 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 8u * TK, true, TK);
      let n2 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 16u * TK, true, TK);
      let n3 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 24u * TK, true, TK);
      acc[0] = subgroupMatrixMultiplyAccumulate(m0, n0, acc[0]);
      acc[1] = subgroupMatrixMultiplyAccumulate(m0, n1, acc[1]);
      acc[2] = subgroupMatrixMultiplyAccumulate(m0, n2, acc[2]);
      acc[3] = subgroupMatrixMultiplyAccumulate(m0, n3, acc[3]);
      acc[4] = subgroupMatrixMultiplyAccumulate(m1, n0, acc[4]);
      acc[5] = subgroupMatrixMultiplyAccumulate(m1, n1, acc[5]);
      acc[6] = subgroupMatrixMultiplyAccumulate(m1, n2, acc[6]);
      acc[7] = subgroupMatrixMultiplyAccumulate(m1, n3, acc[7]);
      acc[8] = subgroupMatrixMultiplyAccumulate(m2, n0, acc[8]);
      acc[9] = subgroupMatrixMultiplyAccumulate(m2, n1, acc[9]);
      acc[10] = subgroupMatrixMultiplyAccumulate(m2, n2, acc[10]);
      acc[11] = subgroupMatrixMultiplyAccumulate(m2, n3, acc[11]);
      acc[12] = subgroupMatrixMultiplyAccumulate(m3, n0, acc[12]);
      acc[13] = subgroupMatrixMultiplyAccumulate(m3, n1, acc[13]);
      acc[14] = subgroupMatrixMultiplyAccumulate(m3, n2, acc[14]);
      acc[15] = subgroupMatrixMultiplyAccumulate(m3, n3, acc[15]);
    }
    workgroupBarrier();
  }
  let row0 = bm + sg_row * 32u;
  let col0 = bn + sg_col * 32u;
  for (var f = 0u; f < 16u; f++) {
    subgroupMatrixStore(&scratch[sg], 0u, acc[f], false, 8u);
    workgroupBarrier();
    let fr = f / 4u; let fc = f % 4u;
    for (var e = lane; e < 64u; e += 32u) {
      let i = e / 8u; let j = e % 8u;
      let gr = row0 + fr * 8u + i; let gc = col0 + fc * 8u + j;
      if (gr < p.M && gc < p.N) { C[gr * p.N + gc] = f16(clamp(scratch[sg][e] + f32(bias[gc]), -60000.0, 60000.0)); }
    }
    workgroupBarrier();
  }
}
`;

const GEMM_SIMT_I4 = `
struct P { M: u32, N: u32, K: u32, group: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> A: array<f16>;
@group(0) @binding(2) var<storage, read> Pk: array<u32>;
@group(0) @binding(3) var<storage, read> Sc: array<f16>;
@group(0) @binding(4) var<storage, read> bias: array<f16>;
@group(0) @binding(5) var<storage, read_write> C: array<f16>;
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 32u;
var<workgroup> As: array<f32, 2048>;
var<workgroup> Bs: array<f32, 2048>;
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let bm = wg.x * BM; let bn = wg.y * BN;
  let ty = tid / 16u; let tx = tid % 16u;
  var acc: array<vec4<f32>, 4>;
  for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); }
  let nk = min((p.K + BK - 1u) / BK, 2048u);
  for (var kt = 0u; kt < nk; kt++) {
    let k0 = kt * BK;
    for (var i = 0u; i < 8u; i++) {
      let idx = i * 256u + tid; let r = idx / BK; let c = idx % BK;
      let gr = bm + r; let gc = k0 + c;
      var v: f32 = 0.0;
      if (gr < p.M && gc < p.K) { v = f32(A[gr * p.K + gc]); }
      As[r * BK + c] = v;
    }
    // dequant B tile: 256 words (16 kp x 16 words of 4 columns)
    let kp = tid / 16u; let n0 = (tid % 16u) * 4u;
    let gk2 = k0 / 2u + kp; let gn0 = bn + n0;
    var word: u32 = 0u;
    if (gk2 < p.K / 2u && gn0 < p.N) { word = Pk[(gk2 * p.N + gn0) / 4u]; }
    let grp = (k0 + 2u * kp) / p.group;
    for (var j = 0u; j < 4u; j++) {
      let n = n0 + j;
      if (bn + n < p.N) {
        let byte = (word >> (8u * j)) & 0xFFu;
        let s = f32(Sc[grp * p.N + bn + n]);
        Bs[(2u * kp) * BN + n] = (f32(byte & 0xFu) - 8.0) * s;
        Bs[(2u * kp + 1u) * BN + n] = (f32(byte >> 4u) - 8.0) * s;
      } else {
        Bs[(2u * kp) * BN + n] = 0.0;
        Bs[(2u * kp + 1u) * BN + n] = 0.0;
      }
    }
    workgroupBarrier();
    for (var k = 0u; k < BK; k++) {
      let a0 = As[(ty * 4u + 0u) * BK + k];
      let a1 = As[(ty * 4u + 1u) * BK + k];
      let a2 = As[(ty * 4u + 2u) * BK + k];
      let a3 = As[(ty * 4u + 3u) * BK + k];
      let b0 = Bs[k * BN + tx * 4u + 0u];
      let b1 = Bs[k * BN + tx * 4u + 1u];
      let b2 = Bs[k * BN + tx * 4u + 2u];
      let b3 = Bs[k * BN + tx * 4u + 3u];
      let bv = vec4<f32>(b0, b1, b2, b3);
      acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv;
    }
    workgroupBarrier();
  }
  for (var r = 0u; r < 4u; r++) {
    let gr = bm + ty * 4u + r;
    if (gr < p.M) {
      for (var c = 0u; c < 4u; c++) {
        let gc = bn + tx * 4u + c;
        if (gc < p.N) { C[gr * p.N + gc] = f16(clamp(acc[r][c] + f32(bias[gc]), -60000.0, 60000.0)); }
      }
    }
  }
}
`;

const GEMM_SIMT = `
struct P { M: u32, N: u32, K: u32, pad: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> A: array<f16>;
@group(0) @binding(2) var<storage, read> B: array<f16>;
@group(0) @binding(3) var<storage, read> bias: array<f16>;
@group(0) @binding(4) var<storage, read_write> C: array<f16>;
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 32u;
var<workgroup> As: array<f32, 2048>;
var<workgroup> Bs: array<f32, 2048>;
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let bm = wg.x * BM; let bn = wg.y * BN;
  let ty = tid / 16u; let tx = tid % 16u;
  var acc: array<vec4<f32>, 4>;
  for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); }
  let nk = min((p.K + BK - 1u) / BK, 2048u);
  for (var kt = 0u; kt < nk; kt++) {
    let k0 = kt * BK;
    for (var i = 0u; i < 8u; i++) {
      let idx = i * 256u + tid; let r = idx / BK; let c = idx % BK;
      let gr = bm + r; let gc = k0 + c;
      var v: f32 = 0.0;
      if (gr < p.M && gc < p.K) { v = f32(A[gr * p.K + gc]); }
      As[r * BK + c] = v;
    }
    for (var i = 0u; i < 8u; i++) {
      let idx = i * 256u + tid; let r = idx / BN; let c = idx % BN;
      let gr = k0 + r; let gc = bn + c;
      var v: f32 = 0.0;
      if (gr < p.K && gc < p.N) { v = f32(B[gr * p.N + gc]); }
      Bs[r * BN + c] = v;
    }
    workgroupBarrier();
    for (var k = 0u; k < BK; k++) {
      let a0 = As[(ty * 4u + 0u) * BK + k];
      let a1 = As[(ty * 4u + 1u) * BK + k];
      let a2 = As[(ty * 4u + 2u) * BK + k];
      let a3 = As[(ty * 4u + 3u) * BK + k];
      let b0 = Bs[k * BN + tx * 4u + 0u];
      let b1 = Bs[k * BN + tx * 4u + 1u];
      let b2 = Bs[k * BN + tx * 4u + 2u];
      let b3 = Bs[k * BN + tx * 4u + 3u];
      let bv = vec4<f32>(b0, b1, b2, b3);
      acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv;
    }
    workgroupBarrier();
  }
  for (var r = 0u; r < 4u; r++) {
    let gr = bm + ty * 4u + r;
    if (gr < p.M) {
      for (var c = 0u; c < 4u; c++) {
        let gc = bn + tx * 4u + c;
        if (gc < p.N) { C[gr * p.N + gc] = f16(clamp(acc[r][c] + f32(bias[gc]), -60000.0, 60000.0)); }
      }
    }
  }
}
`;

const GEMV = `
struct P { N: u32, K: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> A: array<f16>;
@group(0) @binding(2) var<storage, read> B: array<f16>;
@group(0) @binding(3) var<storage, read> bias: array<f16>;
@group(0) @binding(4) var<storage, read_write> C: array<f16>;
var<workgroup> shA: array<f32, 1024>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  for (var i = tid; i < p.K; i += 256u) { shA[i] = f32(A[i]); }
  workgroupBarrier();
  let base = wg.x * 1024u + tid * 4u;
  var acc = vec4<f32>(0.0);
  let KK = min(p.K, 65536u);
  for (var k = 0u; k < KK; k++) {
    let a = shA[k];
    let o = k * p.N + base;
    acc += a * vec4<f32>(f32(B[o]), f32(B[o + 1u]), f32(B[o + 2u]), f32(B[o + 3u]));
  }
  for (var j = 0u; j < 4u; j++) {
    let c = base + j;
    if (c < p.N) { C[c] = f16(acc[j] + f32(bias[c])); }
  }
}
`;

const GEMV_I4 = `
struct P { N: u32, K: u32, group: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> A: array<f16>;
@group(0) @binding(2) var<storage, read> Pk: array<u32>;
@group(0) @binding(3) var<storage, read> Sc: array<f16>;
@group(0) @binding(4) var<storage, read> bias: array<f16>;
@group(0) @binding(5) var<storage, read_write> C: array<f16>;
var<workgroup> shA: array<f32, 1024>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  for (var i = tid; i < p.K; i += 256u) { shA[i] = f32(A[i]); }
  workgroupBarrier();
  let base = wg.x * 1024u + tid * 4u;
  var acc = vec4<f32>(0.0);
  let KK = min(p.K, 65536u);
  for (var k = 0u; k < KK; k++) {
    let a = shA[k];
    let word = Pk[((k / 2u) * p.N + base) / 4u];
    let sh = (k & 1u) * 4u;
    let s = vec4<f32>(f32(Sc[(k / p.group) * p.N + base]), f32(Sc[(k / p.group) * p.N + base + 1u]),
                      f32(Sc[(k / p.group) * p.N + base + 2u]), f32(Sc[(k / p.group) * p.N + base + 3u]));
    let q = vec4<f32>(f32((word >> sh) & 0xFu), f32((word >> (8u + sh)) & 0xFu),
                      f32((word >> (16u + sh)) & 0xFu), f32((word >> (24u + sh)) & 0xFu));
    acc += a * (q - vec4<f32>(8.0)) * s;
  }
  let bb = vec4<f32>(f32(bias[base]), f32(bias[base + 1u]), f32(bias[base + 2u]), f32(bias[base + 3u]));
  for (var j = 0u; j < 4u; j++) { if (base + j < p.N) { C[base + j] = f16(acc[j] + bb[j]); } }
}
`;

const LNM = `
struct P { M: u32, D: u32, offS: u32, offT: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read> MOD: array<f16>;
@group(0) @binding(3) var<storage, read_write> Y: array<f16>;
var<workgroup> red: array<f32, 256>;
var<workgroup> stats: array<f32, 2>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let row = wg.x; let D = p.D; let off = row * D;
  var s = 0.0;
  for (var i = tid; i < D; i += 256u) { s += f32(X[off + i]); }
  red[tid] = s; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { stats[0] = red[0] / f32(D); }
  workgroupBarrier();
  let mu = stats[0];
  var vs = 0.0;
  for (var i = tid; i < D; i += 256u) { let d = f32(X[off + i]) - mu; vs += d * d; }
  red[tid] = vs; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { stats[1] = inverseSqrt(red[0] / f32(D) + 1e-6); }
  workgroupBarrier();
  let inv = stats[1];
  for (var i = tid; i < D; i += 256u) {
    let sc = 1.0 + f32(MOD[p.offS + i]);
    let sh = f32(MOD[p.offT + i]);
    Y[off + i] = f16((f32(X[off + i]) - mu) * inv * sc + sh);
  }
}
`;

const LNA = `
struct P { M: u32, D: u32, pad0: u32, pad1: u32, eps: f32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read> Wt: array<f16>;
@group(0) @binding(3) var<storage, read> Bs: array<f16>;
@group(0) @binding(4) var<storage, read_write> Y: array<f16>;
var<workgroup> red: array<f32, 256>;
var<workgroup> stats: array<f32, 2>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let row = wg.x; let D = p.D; let off = row * D;
  var s = 0.0;
  for (var i = tid; i < D; i += 256u) { s += f32(X[off + i]); }
  red[tid] = s; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { stats[0] = red[0] / f32(D); }
  workgroupBarrier();
  let mu = stats[0];
  var vs = 0.0;
  for (var i = tid; i < D; i += 256u) { let d = f32(X[off + i]) - mu; vs += d * d; }
  red[tid] = vs; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { stats[1] = inverseSqrt(red[0] / f32(D) + p.eps); }
  workgroupBarrier();
  let inv = stats[1];
  for (var i = tid; i < D; i += 256u) {
    Y[off + i] = f16((f32(X[off + i]) - mu) * inv * f32(Wt[i]) + f32(Bs[i]));
  }
}
`;

const ROPE = `
struct P { M: u32, N: u32, total: u32, pad: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> FREQ: array<f32>;
@group(0) @binding(2) var<storage, read_write> X: array<f16>;
const SC: f32 = 0.35355339059327379;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.total) { return; }
  let pair = i % 32u; let t = i / 32u; let h = t % 16u; let row = t / 16u;
  let pos = row % p.N;
  let theta = f32(pos) * FREQ[pair];
  let c = cos(theta); let s = sin(theta);
  let base = row * 1024u + h * 64u + pair * 2u;
  let x0 = f32(X[base]); let x1 = f32(X[base + 1u]);
  X[base] = f16((x0 * c - x1 * s) * SC);
  X[base + 1u] = f16((x0 * s + x1 * c) * SC);
}
`;

const QK = `
struct P { N: u32, B: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> Q: array<f16>;
@group(0) @binding(2) var<storage, read> K: array<f16>;
@group(0) @binding(3) var<storage, read_write> S: array<f32>;
var<workgroup> Qs: array<f16, 4096>;
var<workgroup> Ks: array<f16, 4096>;
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let b = wg.z / 16u; let h = wg.z % 16u;
  let qi = wg.x * 64u; let kj = wg.y * 64u;
  let N = min(p.N, 8192u);
  for (var i = 0u; i < 16u; i++) {
    let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u;
    let qr = qi + r; var v: f16 = f16(0.0);
    if (qr < N) { v = Q[((b * N + qr) * 16u + h) * 64u + c]; }
    Qs[idx] = v;
  }
  for (var i = 0u; i < 16u; i++) {
    let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u;
    let kr = kj + r; var v: f16 = f16(0.0);
    if (kr < N) { v = K[((b * N + kr) * 16u + h) * 64u + c]; }
    Ks[idx] = v;
  }
  workgroupBarrier();
  let ty = tid / 16u; let tx = tid % 16u;
  var acc: array<vec4<f32>, 4>;
  for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); }
  for (var k = 0u; k < 64u; k++) {
    let a0 = f32(Qs[(ty * 4u + 0u) * 64u + k]);
    let a1 = f32(Qs[(ty * 4u + 1u) * 64u + k]);
    let a2 = f32(Qs[(ty * 4u + 2u) * 64u + k]);
    let a3 = f32(Qs[(ty * 4u + 3u) * 64u + k]);
    let b0 = f32(Ks[(tx * 4u + 0u) * 64u + k]);
    let b1 = f32(Ks[(tx * 4u + 1u) * 64u + k]);
    let b2 = f32(Ks[(tx * 4u + 2u) * 64u + k]);
    let b3 = f32(Ks[(tx * 4u + 3u) * 64u + k]);
    let bv = vec4<f32>(b0, b1, b2, b3);
    acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv;
  }
  let bidx = wg.z;
  for (var r = 0u; r < 4u; r++) {
    let gr = qi + ty * 4u + r;
    if (gr < N) {
      for (var c = 0u; c < 4u; c++) {
        let gc = kj + tx * 4u + c;
        if (gc < N) { S[(bidx * N + gr) * N + gc] = acc[r][c]; }
      }
    }
  }
}
`;

const SOFTMAX = `
struct P { N: u32, rows: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> S: array<f32>;
var<workgroup> red: array<f32, 256>;
var<workgroup> st: array<f32, 2>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let off = wg.x * p.N;
  var m = -3.4e38;
  for (var i = tid; i < p.N; i += 256u) { m = max(m, S[off + i]); }
  red[tid] = m; workgroupBarrier();
  ${REDUCE8MAX}
  if (tid == 0u) { st[0] = red[0]; }
  workgroupBarrier();
  let mx = st[0];
  var sum = 0.0;
  for (var i = tid; i < p.N; i += 256u) { sum += exp(S[off + i] - mx); }
  red[tid] = sum; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { st[1] = 1.0 / red[0]; }
  workgroupBarrier();
  let inv = st[1];
  for (var i = tid; i < p.N; i += 256u) { S[off + i] = exp(S[off + i] - mx) * inv; }
}
`;

const AV = `
struct P { N: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> S: array<f32>;
@group(0) @binding(2) var<storage, read> V: array<f16>;
@group(0) @binding(3) var<storage, read_write> O: array<f16>;
var<workgroup> Ss: array<f32, 4096>;
var<workgroup> Vs: array<f16, 4096>;
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let b = wg.z / 16u; let h = wg.z % 16u;
  let qi = wg.x * 64u;
  let N = min(p.N, 8192u); let bidx = wg.z;
  let ty = tid / 16u; let tx = tid % 16u;
  var acc: array<vec4<f32>, 4>;
  for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); }
  let nk = min((N + 63u) / 64u, 256u);
  for (var kt = 0u; kt < nk; kt++) {
    let k0 = kt * 64u;
    for (var i = 0u; i < 16u; i++) {
      let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u;
      let gr = qi + r; let skey = k0 + c; let vkey = k0 + r;
      var v0: f32 = 0.0;
      if (gr < N && skey < N) { v0 = S[(bidx * N + gr) * N + skey]; }
      Ss[r * 64u + c] = v0;
      var v1: f16 = f16(0.0);
      if (vkey < N) { v1 = V[((b * N + vkey) * 16u + h) * 64u + c]; }
      Vs[r * 64u + c] = v1;
    }
    workgroupBarrier();
    for (var k = 0u; k < 64u; k++) {
      let a0 = Ss[(ty * 4u + 0u) * 64u + k];
      let a1 = Ss[(ty * 4u + 1u) * 64u + k];
      let a2 = Ss[(ty * 4u + 2u) * 64u + k];
      let a3 = Ss[(ty * 4u + 3u) * 64u + k];
      let b0 = f32(Vs[k * 64u + tx * 4u + 0u]);
      let b1 = f32(Vs[k * 64u + tx * 4u + 1u]);
      let b2 = f32(Vs[k * 64u + tx * 4u + 2u]);
      let b3 = f32(Vs[k * 64u + tx * 4u + 3u]);
      let bv = vec4<f32>(b0, b1, b2, b3);
      acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv;
    }
    workgroupBarrier();
  }
  for (var r = 0u; r < 4u; r++) {
    let gr = qi + ty * 4u + r;
    if (gr < N) {
      for (var c = 0u; c < 4u; c++) {
        O[((b * N + gr) * 16u + h) * 64u + tx * 4u + c] = f16(acc[r][c]);
      }
    }
  }
}
`;

const ADDGATE = `
struct P { M: u32, D: u32, gOff: u32, pad: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> X: array<f16>;
@group(0) @binding(2) var<storage, read> Z: array<f16>;
@group(0) @binding(3) var<storage, read> MOD: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.M * p.D) { return; }
  let d = i % p.D;
  X[i] = f16(f32(X[i]) + f32(MOD[p.gOff + d]) * f32(Z[i]));
}
`;

const GELU = `
struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> X: array<f16>;
fn erf(x: f32) -> f32 {
  let ax = abs(x);
  let t = 1.0 / (1.0 + 0.3275911 * ax);
  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-ax * ax);
  return select(-y, y, x >= 0.0);
}
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  let x = f32(X[gid.x]);
  X[gid.x] = f16(0.5 * x * (1.0 + erf(x * 0.70710678118654752)));
}
`;

const SILU = `
struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> X: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  let x = f32(X[gid.x]);
  X[gid.x] = f16(x / (1.0 + exp(-x)));
}
`;

const MISH = `
struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> X: array<f16>;
fn tanh_safe(v: f32) -> f32 {
  if (v > 15.0) { return 1.0; }
  if (v < -15.0) { return -1.0; }
  return tanh(v);
}
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  let x = f32(X[gid.x]);
  var sp = x;
  if (x < 20.0) { sp = log(1.0 + exp(x)); }
  X[gid.x] = f16(x * tanh_safe(sp));
}
`;

const CONCAT3 = `
struct P { N: u32, B: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read> C: array<f16>;
@group(0) @binding(3) var<storage, read> T: array<f16>;
@group(0) @binding(4) var<storage, read_write> O: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.B * p.N * 712u) { return; }
  let col = i % 712u; let row = i / 712u;
  let b = row / p.N; let f = row % p.N;
  var v: f16 = f16(0.0);
  if (col < 100u) { v = X[f * 100u + col]; }
  else if (col < 200u) { if (b == 0u) { v = C[f * 100u + (col - 100u)]; } }
  else { if (b == 0u) { v = T[f * 512u + (col - 200u)]; } }
  O[i] = v;
}
`;

const GATHER = `
struct P { N: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> IDS: array<i32>;
@group(0) @binding(2) var<storage, read> EMB: array<f16>;
@group(0) @binding(3) var<storage, read> FRQ: array<f16>;
@group(0) @binding(4) var<storage, read_write> O: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.N * 512u) { return; }
  let f = i / 512u; let d = i % 512u;
  let id = IDS[f] + 1;
  if (id <= 0) { O[i] = f16(0.0); return; }
  O[i] = f16(f32(EMB[u32(id) * 512u + d]) + f32(FRQ[f * 512u + d]));
}
`;

const DWCONV = `
struct P { M: u32, D: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read> Wt: array<f16>;
@group(0) @binding(3) var<storage, read> Bs: array<f16>;
@group(0) @binding(4) var<storage, read_write> Y: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.M * p.D) { return; }
  let d = i32(i % p.D); let f = i32(i / p.D);
  var acc = f32(Bs[u32(d)]);
  for (var k = 0; k < 7; k++) {
    let ff = f - 3 + k;
    if (ff >= 0 && ff < i32(p.M)) { acc += f32(Wt[u32(d) * 7u + u32(k)]) * f32(X[u32(ff) * p.D + u32(d)]); }
  }
  Y[i] = f16(acc);
}
`;

const GRNRED = `
struct P { M: u32, D: u32, off: u32, pad: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read_write> NM: array<f32>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>) {
  let d = wg.x + p.off;
  var s = 0.0;
  for (var f = 0u; f < p.M; f++) { let v = f32(X[f * p.D + d]); s += v * v; }
  NM[d] = sqrt(s);
}
`;

const GRNMEAN = `
struct P { D: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> NM: array<f32>;
@group(0) @binding(2) var<storage, read_write> MEAN: array<f32>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_index) tid: u32) {
  var acc = 0.0;
  for (var c = tid; c < p.D; c += 256u) { acc += NM[c]; }
  red[tid] = acc; workgroupBarrier();
  ${REDUCE8}
  if (tid == 0u) { MEAN[0] = red[0] / f32(p.D); }
}
`;

const GRNAPPLY = `
struct P { M: u32, D: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read_write> Y: array<f16>;
@group(0) @binding(3) var<storage, read> G: array<f16>;
@group(0) @binding(4) var<storage, read> B: array<f16>;
@group(0) @binding(5) var<storage, read> NM: array<f32>;
@group(0) @binding(6) var<storage, read> MEAN: array<f32>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.M * p.D) { return; }
  let d = i % p.D;
  let v = f32(X[i]);
  Y[i] = f16(v * (NM[d] / (MEAN[0] + 1e-6)) * f32(G[d]) + f32(B[d]) + v);
}
`;

const ADDRES = `
struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> Y: array<f16>;
@group(0) @binding(2) var<storage, read> X: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  Y[gid.x] = f16(f32(Y[gid.x]) + f32(X[gid.x]));
}
`;

const ZEROFILL = `
struct P { M: u32, D: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> IDS: array<i32>;
@group(0) @binding(2) var<storage, read_write> X: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.M * p.D) { return; }
  if (IDS[i / p.D] < 0) { X[i] = f16(0.0); }
}
`;

const CONVPOS = `
struct P { N: u32, M: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read> Wt: array<f16>;
@group(0) @binding(3) var<storage, read> Bs: array<f16>;
@group(0) @binding(4) var<storage, read_write> Y: array<f16>;
var<workgroup> Ws: array<f16, 2 * 64 * 31>;
var<workgroup> Xs: array<f32, 2 * 94>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
  let g = wg.x / 16u; let pt = wg.x % 16u;
  let p0 = pt * 64u;
  let b = wg.y;
  let tx = tid % 16u; let ty = tid / 16u;
  let c0 = g * 64u + ty * 4u;
  var acc: array<f32, 16>;
  for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; }
  for (var cc = 0u; cc < 32u; cc++) {
    let ci0 = cc * 2u;
    for (var i = tid; i < 2u * 64u * 31u; i += 256u) {
      let k = i % 31u; let oc = (i / 31u) % 64u; let ic = i / (31u * 64u);
      Ws[i] = Wt[((g * 64u + oc) * 64u + (ci0 + ic)) * 31u + k];
    }
    for (var i = tid; i < 2u * 94u; i += 256u) {
      let ic = i / 94u; let pp = i % 94u;
      let pos = i32(p0) - 15 + i32(pp);
      var v: f32 = 0.0;
      if (pos >= 0 && pos < i32(p.N)) { v = f32(X[(b * p.N + u32(pos)) * 1024u + g * 64u + ci0 + ic]); }
      Xs[i] = v;
    }
    workgroupBarrier();
    for (var ic = 0u; ic < 2u; ic++) {
      for (var k = 0u; k < 31u; k++) {
        let x0 = Xs[ic * 94u + tx * 4u + k];
        let x1 = Xs[ic * 94u + tx * 4u + 1u + k];
        let x2 = Xs[ic * 94u + tx * 4u + 2u + k];
        let x3 = Xs[ic * 94u + tx * 4u + 3u + k];
        let w0 = f32(Ws[(ic * 64u + ty * 4u + 0u) * 31u + k]);
        let w1 = f32(Ws[(ic * 64u + ty * 4u + 1u) * 31u + k]);
        let w2 = f32(Ws[(ic * 64u + ty * 4u + 2u) * 31u + k]);
        let w3 = f32(Ws[(ic * 64u + ty * 4u + 3u) * 31u + k]);
        acc[0] += x0 * w0; acc[1] += x0 * w1; acc[2] += x0 * w2; acc[3] += x0 * w3;
        acc[4] += x1 * w0; acc[5] += x1 * w1; acc[6] += x1 * w2; acc[7] += x1 * w3;
        acc[8] += x2 * w0; acc[9] += x2 * w1; acc[10] += x2 * w2; acc[11] += x2 * w3;
        acc[12] += x3 * w0; acc[13] += x3 * w1; acc[14] += x3 * w2; acc[15] += x3 * w3;
      }
    }
    workgroupBarrier();
  }
  for (var oc = 0u; oc < 4u; oc++) {
    for (var j = 0u; j < 4u; j++) {
      let pos = p0 + tx * 4u + j;
      if (pos < p.N) {
        Y[(b * p.N + pos) * 1024u + c0 + oc] = f16(clamp(acc[j * 4u + oc] + f32(Bs[c0 + oc]), -60000.0, 60000.0));
      }
    }
  }
}
`;

const FLOW = `
struct P { total: u32, N: u32, B: u32, pad1: u32, dt: f32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read_write> X: array<f16>;
@group(0) @binding(2) var<storage, read> V: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  let i = gid.x;
  let d = i % 100u; let f = i / 100u;
  let v0 = f32(V[f * 100u + d]);
  var v = v0;
  if (p.B == 2u) { let v1 = f32(V[(p.N + f) * 100u + d]); v = v0 + (v0 - v1) * 2.0; }
  X[i] = f16(f32(X[i]) + p.dt * v);
}
`;

const TEB = `
struct P { N: u32, B: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> T: array<f16>;
@group(0) @binding(2) var<storage, read_write> O: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= p.B * p.N * 512u) { return; }
  let row = i / 512u; let d = i % 512u;
  if (row < p.N) { O[i] = T[row * 512u + d]; } else { O[i] = f16(0.0); }
}
`;

const COPY = `
struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 };
@group(0) @binding(0) var<uniform> p: P;
@group(0) @binding(1) var<storage, read> X: array<f16>;
@group(0) @binding(2) var<storage, read_write> Y: array<f16>;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= p.total) { return; }
  Y[gid.x] = X[gid.x];
}
`;

export const KERNELS = {
  gemm: GEMM, gemv: GEMV, lnm: LNM, lna: LNA, rope: ROPE, qk: QK, softmax: SOFTMAX,
  av: AV, addgate: ADDGATE, gelu: GELU, silu: SILU, mish: MISH, concat3: CONCAT3,
  gather: GATHER, dwconv: DWCONV, grnred: GRNRED, grnmean: GRNMEAN, grnapply: GRNAPPLY,
  addres: ADDRES, zerofill: ZEROFILL, convpos: CONVPOS, flow: FLOW, teb: TEB, copy: COPY,
};

function f16tof32(h) {
  const sign = (h & 0x8000) ? -1 : 1, exp = (h >> 10) & 0x1f, frac = h & 0x3ff;
  if (exp === 0) return sign * Math.pow(2, -14) * (frac / 1024);
  if (exp === 31) return frac ? NaN : sign * Infinity;
  return sign * Math.pow(2, exp - 15) * (1 + frac / 1024);
}

export class F5Engine {
  static async create(device, opts = {}) {
    const e = new F5Engine(device, opts);
    await e._init();
    return e;
  }
  constructor(device, opts) { this.device = device; this.opts = opts; this.f32 = opts.dtype === 'f32'; this.es = this.f32 ? 4 : 2; }

  async _init() {
    const dev = this.device;
    const [mr, br] = await Promise.all([
      fetch(this.opts.metaUrl || '/webgpu/f5w.json'),
      fetch(this.opts.binUrl || '/webgpu/f5w.bin'),
    ]);
    this.meta = await mr.json();
    this.binSize = Number(br.headers.get('content-length')) || 0;
    this.wsize = Math.ceil((this.f32 ? this.binSize * 2 : this.binSize) / 256) * 256;
    this._planChunks();
    this.wbufs = this.chunkStarts.map((st, i) => {
      const end = i + 1 < this.chunkStarts.length ? this.chunkStarts[i + 1] : this.wsize;
      return dev.createBuffer({ size: Math.ceil((end - st) / 256) * 256, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    });
    await this._streamWeights(br);
    this.ubufs = [];
    this.ubNext = 0;
    this.compileErrors = [];
    const cre = (name, code) => {
      if (this.f32) code = code.replace(/enable f16;\n/g, '').replace(/\bf16\b/g, 'f32');
      const mod = dev.createShaderModule({ code: (this.f32 ? '' : PRELUDE) + code });
      mod.getCompilationInfo().then((ci) => {
        for (const m of ci.messages) if (m.type === 'error') this.compileErrors.push(name + ': ' + m.message + ' @' + m.lineNum);
      });
      return dev.createComputePipelineAsync({ layout: 'auto', compute: { module: mod, entryPoint: 'main' } });
    };
    this.gemmKind = this.opts.gemm || (this.device.features.has('chromium-experimental-subgroup-matrix') ? 'sgmat' : 'simt');
    console.log('DBG gemmKind', this.gemmKind, 'opts', JSON.stringify(Object.keys(this.opts || {})));
    const defs = {
      gemm: this.gemmKind === 'sgmat' ? GEMM : GEMM_SIMT, gemm_i4: GEMM_SIMT_I4, gemv: GEMV, gemv_i4: GEMV_I4, lnm: LNM, lna: LNA, rope: ROPE,
      qk: QK, softmax: SOFTMAX, av: AV, addgate: ADDGATE, gelu: GELU,
      silu: SILU, mish: MISH, concat3: CONCAT3, gather: GATHER,
      dwconv: DWCONV, grnred: GRNRED, grnmean: GRNMEAN,
      grnapply: GRNAPPLY, addres: ADDRES, zerofill: ZEROFILL,
      convpos: CONVPOS, flow: FLOW, teb: TEB, copy: COPY,
    };
    this.pipe = {};
    for (const [name, code] of Object.entries(defs)) {
      this.pipe[name] = await cre(name, code);
    }
    this.pipeNames = new Map(Object.entries(this.pipe).map(([k, v]) => [v, k]));
  }

  _planChunks() {
    const mult = this.f32 ? 2 : 1;
    const ents = Object.entries(this.meta).map(([n, t]) => {
      if (t.q) return { n, off: t.packed * mult, len: (t.scales + t.scale_bytes) - t.packed };
      const len = Math.ceil(t.shape.reduce((a, b) => a * b, 1) * this.es / 4) * 4;
      return { n, off: t.offset * mult, len };
    }).sort((a, b) => a.off - b.off);
    const CHUNK = 640 << 20;
    const starts = [0];
    let cur = 0;
    this.bufOf = new Map();
    for (const e of ents) {
      if (cur > 0 && cur + e.len > CHUNK) { starts.push(e.off); cur = 0; }
      this.bufOf.set(e.n, starts.length - 1);
      cur = e.off + e.len - starts[starts.length - 1];
    }
    this.chunkStarts = starts;
  }

  _writeAt(globalOff, data, byteLen) {
    let off = globalOff, src = 0, remain = byteLen;
    while (remain > 0) {
      let bi = this.chunkStarts.length - 1;
      for (let i = 0; i < this.chunkStarts.length; i++) {
        const end = i + 1 < this.chunkStarts.length ? this.chunkStarts[i + 1] : this.wsize;
        if (off < end) { bi = i; break; }
      }
      const base = this.chunkStarts[bi];
      const bo = off - base;
      const end = bi + 1 < this.chunkStarts.length ? this.chunkStarts[bi + 1] : this.wsize;
      const n = Math.min(remain, end - off);
      if (data instanceof Float32Array) {
        this.device.queue.writeBuffer(this.wbufs[bi], bo, data.subarray(src / 4, (src + n) / 4));
      } else {
        this.device.queue.writeBuffer(this.wbufs[bi], bo, data.subarray(src, src + n));
      }
      off += n; src += n; remain -= n;
    }
  }

  async _streamWeights(res) {
    // Stream the bundle to the GPU in ~32 MB slices; convert f16 -> f32 if needed.
    const dev = this.device;
    const reader = res.body.getReader();
    const CH = 32 << 20;
    let stage = new Uint8Array(CH + 8);
    let staged = 0, gpuOff = 0;
    const flush = async (final) => {
      if (!staged) return;
      let n = staged;
      const extra = final ? (4 - (n % 4)) % 4 : 0;
      if (extra) { stage.fill(0, n, n + extra); n += extra; }
      if (this.f32) {
        const elems = n >> 1;
        const out = new Float32Array(elems);
        out.set(new Float16Array(stage.buffer, 0, elems));
        this._writeAt(gpuOff, out, elems * 4);
        gpuOff += elems * 4;
      } else {
        this._writeAt(gpuOff, stage.subarray(0, n), n);
        gpuOff += n;
      }
      staged = 0;
      if (this.opts.onProgress) this.opts.onProgress(gpuOff / (this.f32 ? this.binSize * 2 : this.binSize));
    };
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      let off = 0;
      while (off < value.length) {
        const take = Math.min(value.length - off, stage.length - staged);
        stage.set(value.subarray(off, off + take), staged);
        staged += take; off += take;
        if (staged >= CH) await flush(false);
      }
    }
    await flush(true);
    const pad = Math.ceil(gpuOff / 256) * 256 - gpuOff;
    if (pad > 0) this._writeAt(gpuOff, new Uint8Array(pad), pad);
  }

  _t(name) { const t = this.meta[name]; if (!t) throw new Error('missing ' + name); return t; }
  _wslice(name) {
    const t = this._t(name);
    const n = t.shape.reduce((a, b) => a * b, 1);
    const mult = this.f32 ? 2 : 1;
    const bi = this.bufOf.get(name);
    const base = this.chunkStarts[bi];
    return { buffer: this.wbufs[bi], offset: t.offset * mult - base, size: Math.ceil(n * this.es / 4) * 4 };
  }
  _alc(bytes) {
    const size = Math.ceil(bytes / 256) * 256;
    return this.device.createBuffer({ size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC });
  }
  _ub(vals, floats) {
    const ab = new ArrayBuffer(256);
    const dv = new DataView(ab);
    for (let i = 0; i < vals.length; i++) {
      if (floats && floats.indexOf(i) >= 0) dv.setFloat32(i * 4, vals[i], true);
      else dv.setUint32(i * 4, vals[i] >>> 0, true);
    }
    const buf = this.device.createBuffer({ size: 256, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    this.device.queue.writeBuffer(buf, 0, ab);
    this.ubNext++;
    this.ubufs.push(buf);
    return { buffer: buf, offset: 0, size: 256 };
  }
  _bg(pipe, u, bufs) {
    const wrap = (b) => (b && b.buffer) ? b : { buffer: b, offset: 0, size: b.size };
    return this.device.createBindGroup({
      layout: pipe.getBindGroupLayout(0),
      entries: [{ binding: 0, resource: wrap(u) }, ...bufs.map((b, i) => ({ binding: i + 1, resource: wrap(b) }))],
    });
  }

  setup(N) {
    const dev = this.device;
    this.N = N; this.B = this.opts.branches || 2; this.M2 = this.B * N;
    const A = (name, elems) => (this[name] = this._alc(elems * this.es));
    A('xBuf', N * 100); A('condBuf', N * 100);
    A('teBuf', N * 512); A('tebBuf', 2 * N * 512); A('catBuf', 2 * N * 712);
    A('hBuf', 2 * N * 1024); A('tmpBuf', 2 * N * 1024); A('qBuf', 2 * N * 1024);
    A('kBuf', 2 * N * 1024); A('vBuf', 2 * N * 1024); A('attnBuf', 2 * N * 1024);
    A('ffBuf', 2 * N * 2048); A('velBuf', 2 * N * 100);
    A('timeIn', 256); A('tembBuf', 1024); A('temb2Buf', 1024);
    A('modBuf', 137216); this.scoresBuf = this._alc(this.B * 16 * N * N * 4);
    A('tBuf', N * 1024); A('tBuf2', N * 1024); A('tBuf3', N * 2048);
    A('tnmBuf', 1024 * 4); A('tmeanBuf', 4);
    this.idsBuf = this._alc(N * 4);
    this.ubNext = 0;
    this.ropeFreq = dev.createBuffer({ size: 256, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    const f = new Float32Array(64);
    for (let j = 0; j < 32; j++) f[j] = Math.pow(10000, -j / 32);
    dev.queue.writeBuffer(this.ropeFreq, 0, f);
    this.calls = [];
    this._buildStepCalls();
    const stepCalls = this.calls;
    this.calls = [];
    this._buildTextCalls();
    this.textCalls = this.calls;
    this.calls = stepCalls;
  }

  _add(pipe, bg, x, y = 1, z = 1, meta) {
    const c = { pipe, bg, x: Math.ceil(x), y, z, name: this.pipeNames.get(pipe), meta };
    this.calls.push(c);
    if (this.debug) this.debug.push(c);
    return c;
  }

  _qdesc(name) {
    const t = this._t(name);
    const mult = this.f32 ? 2 : 1;
    const bi = this.bufOf.get(name);
    const base = this.chunkStarts[bi];
    return {
      packed: { buffer: this.wbufs[bi], offset: t.packed * mult - base, size: t.packed_bytes },
      scales: { buffer: this.wbufs[bi], offset: t.scales * mult - base, size: t.scale_bytes },
      group: t.group,
    };
  }

  _gemm(A, wn, bn, C, M, K, N) {
    const t = this._t(wn);
    if (t.q) {
      if (!(M > 0 && N > 0 && K > 0 && M < 1e5 && N < 1e6 && K < 1e6)) throw new Error(`bad gemm args ${M},${K},${N} (${wn})`);
      const q = this._qdesc(wn);
      if (q.packed.offset + q.packed.size > q.packed.buffer.size) throw new Error(`q4 packed OOB ${wn}`);
      const u = this._ub([M, N, K, q.group]);
      this._add(this.pipe.gemm_i4, this._bg(this.pipe.gemm_i4, u, [A, q.packed, q.scales, this._wslice(bn), C]),
        Math.ceil(M / 64), Math.ceil(N / 64));
      return;
    }
    const wa = this._wslice(wn), ba = this._wslice(bn);
    if (!(M > 0 && N > 0 && K > 0 && M < 1e5 && N < 1e6 && K < 1e6)) throw new Error(`bad gemm args ${M},${K},${N} (${wn})`);
    if (wa.offset + wa.size > wa.buffer.size) throw new Error(`gemm weight OOB ${wn}`);
    if (ba.offset + ba.size > ba.buffer.size) throw new Error(`gemm bias OOB ${bn}`);
    const u = this._ub([M, N, K, 0]);
    this._add(this.pipe.gemm, this._bg(this.pipe.gemm, u, [A, wa, ba, C]), Math.ceil(M / (this.gemmKind === 'sgmat' ? 128 : 64)), Math.ceil(N / 64), 1,
      { vals: [M, N, K], wa: [wa.offset, wa.size, wa.buffer.size], ba: [ba.offset, ba.size, ba.buffer.size] });
  }
  _gemv(A, wn, bn, C, K, N) {
    const t = this._t(wn);
    if (t.q) {
      const q = this._qdesc(wn);
      const u = this._ub([N, K, q.group, 0]);
      this._add(this.pipe.gemv_i4, this._bg(this.pipe.gemv_i4, u, [A, q.packed, q.scales, this._wslice(bn), C]), Math.ceil(N / 1024));
      return;
    }
    const u = this._ub([N, K, 0, 0]);
    this._add(this.pipe.gemv, this._bg(this.pipe.gemv, u, [A, this._wslice(wn), this._wslice(bn), C]), Math.ceil(N / 1024));
  }
  _elem(pipe, X, total) {
    const u = this._ub([total, 0, 0, 0]);
    this._add(pipe, this._bg(pipe, u, [X]), Math.ceil(total / 256));
  }

  _buildStepCalls() {
    const N = this.N, M2 = this.M2;
    this._gemm(this.timeIn, 'time.0', 'time.0.b', this.tembBuf, 1, 256, 1024);
    this._elem(this.pipe.silu, this.tembBuf, 1024);
    this._gemm(this.tembBuf, 'time.2', 'time.2.b', this.temb2Buf, 1, 1024, 1024);
    this._elem(this.pipe.silu, this.temb2Buf, 1024);
    this._gemv(this.temb2Buf, 'ada.w', 'ada.b', this.modBuf, 1024, 137216);
    this._add(this.pipe.concat3, this._bg(this.pipe.concat3, this._ub([N, this.B, 0, 0]), [this.xBuf, this.condBuf, this.tebBuf, this.catBuf]), Math.ceil(this.B * N * 712 / 256));
    this._gemm(this.catBuf, 'in.proj', 'in.proj.b', this.hBuf, M2, 712, 1024);
    const cp = (X, wn, bn, Y) => {
      const u = this._ub([N, M2, 0, 0]);
      this._add(this.pipe.convpos, this._bg(this.pipe.convpos, u, [X, this._wslice(wn), this._wslice(bn), Y]), 16 * 16, this.B);
    };
    cp(this.hBuf, 'convpos.0.w', 'convpos.0.b', this.tmpBuf);
    this._elem(this.pipe.mish, this.tmpBuf, M2 * 1024);
    cp(this.tmpBuf, 'convpos.2.w', 'convpos.2.b', this.qBuf);
    this._elem(this.pipe.mish, this.qBuf, M2 * 1024);
    {
      const u = this._ub([M2 * 1024, 0, 0, 0]);
      this._add(this.pipe.addres, this._bg(this.pipe.addres, u, [this.hBuf, this.qBuf]), Math.ceil(M2 * 1024 / 256));
    }
    for (let L = 0; L < 22; L++) {
      const base = L * 6144;
      {
        const u = this._ub([M2, 1024, base + 1024, base]);
        this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2);
      }
      this._gemm(this.tmpBuf, `b${L}.q`, `b${L}.q.b`, this.qBuf, M2, 1024, 1024);
      this._gemm(this.tmpBuf, `b${L}.k`, `b${L}.k.b`, this.kBuf, M2, 1024, 1024);
      this._gemm(this.tmpBuf, `b${L}.v`, `b${L}.v.b`, this.vBuf, M2, 1024, 1024);
      for (const buf of [this.qBuf, this.kBuf]) {
        const u = this._ub([M2, N, M2 * 16 * 32, 0]);
        this._add(this.pipe.rope, this._bg(this.pipe.rope, u, [this.ropeFreq, buf]), Math.ceil(M2 * 16 * 32 / 256));
      }
      {
        const u = this._ub([N, 0, 0, 0]);
        this._add(this.pipe.qk, this._bg(this.pipe.qk, u, [this.qBuf, this.kBuf, this.scoresBuf]), Math.ceil(N / 64), Math.ceil(N / 64), this.B * 16);
        const u2 = this._ub([N, this.B * 16 * N, 0, 0]);
        this._add(this.pipe.softmax, this._bg(this.pipe.softmax, u2, [this.scoresBuf]), this.B * 16 * N);
        const u3 = this._ub([N, 0, 0, 0]);
        this._add(this.pipe.av, this._bg(this.pipe.av, u3, [this.scoresBuf, this.vBuf, this.attnBuf]), Math.ceil(N / 64), 1, this.B * 16);
      }
      this._gemm(this.attnBuf, `b${L}.o`, `b${L}.o.b`, this.tmpBuf, M2, 1024, 1024);
      {
        const u = this._ub([M2, 1024, base + 2048, 0]);
        this._add(this.pipe.addgate, this._bg(this.pipe.addgate, u, [this.hBuf, this.tmpBuf, this.modBuf]), Math.ceil(M2 * 1024 / 256));
      }
      {
        const u = this._ub([M2, 1024, base + 4096, base + 3072]);
        this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2);
      }
      this._gemm(this.tmpBuf, `b${L}.up`, `b${L}.up.b`, this.ffBuf, M2, 1024, 2048);
      this._elem(this.pipe.gelu, this.ffBuf, M2 * 2048);
      this._gemm(this.ffBuf, `b${L}.down`, `b${L}.down.b`, this.tmpBuf, M2, 2048, 1024);
      {
        const u = this._ub([M2, 1024, base + 5120, 0]);
        this._add(this.pipe.addgate, this._bg(this.pipe.addgate, u, [this.hBuf, this.tmpBuf, this.modBuf]), Math.ceil(M2 * 1024 / 256));
      }
    }
    const fbase = 22 * 6144;
    {
      const u = this._ub([M2, 1024, fbase, fbase + 1024]);
      this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2);
    }
    this._gemm(this.tmpBuf, 'out.w', 'out.b', this.velBuf, M2, 1024, 100);
    this.flowUniform = this._ub([N * 100, N, this.B, 0]);
    this._add(this.pipe.flow, this._bg(this.pipe.flow, this.flowUniform, [this.xBuf, this.velBuf]), Math.ceil(N * 100 / 256));
  }

  _buildTextCalls() {
    const N = this.N;
    {
      const u = this._ub([N, 0, 0, 0]);
      this._add(this.pipe.gather, this._bg(this.pipe.gather, u, [this.idsBuf, this._wslice('txt.emb'), this._wslice('txt.freqs'), this.teBuf]), Math.ceil(N * 512 / 256));
    }
    for (let b = 0; b < 4; b++) {
      {
        const u = this._ub([N, 512, 0, 0]);
        this._add(this.pipe.dwconv, this._bg(this.pipe.dwconv, u, [this.teBuf, this._wslice(`tb${b}.dw.w`), this._wslice(`tb${b}.dw.b`), this.tBuf]), Math.ceil(N * 512 / 256));
      }
      {
        const u = this._ub([N, 512, 0, 0, 1e-5, 0, 0, 0], [4]);
        this._add(this.pipe.lna, this._bg(this.pipe.lna, u, [this.tBuf, this._wslice(`tb${b}.ln.w`), this._wslice(`tb${b}.ln.b`), this.tBuf3]), N);
      }
      this._gemm(this.tBuf3, `tb${b}.p1`, `tb${b}.p1.b`, this.tBuf2, N, 512, 1024);
      this._elem(this.pipe.gelu, this.tBuf2, N * 1024);
      {
        const u = this._ub([N, 1024, 0, 0]);
        this._add(this.pipe.grnred, this._bg(this.pipe.grnred, u, [this.tBuf2, this.tnmBuf]), 1024);
        const u2 = this._ub([1024, 0, 0, 0]);
        this._add(this.pipe.grnmean, this._bg(this.pipe.grnmean, u2, [this.tnmBuf, this.tmeanBuf]), 1);
        const u3 = this._ub([N, 1024, 0, 0]);
        this._add(this.pipe.grnapply, this._bg(this.pipe.grnapply, u3, [this.tBuf2, this.tBuf, this._wslice(`tb${b}.grn.g`), this._wslice(`tb${b}.grn.b`), this.tnmBuf, this.tmeanBuf]), Math.ceil(N * 1024 / 256));
      }
      this._gemm(this.tBuf, `tb${b}.p2`, `tb${b}.p2.b`, this.tBuf2, N, 1024, 512);
      {
        const u = this._ub([N * 512, 0, 0, 0]);
        this._add(this.pipe.addres, this._bg(this.pipe.addres, u, [this.tBuf2, this.teBuf]), Math.ceil(N * 512 / 256));
        const u2 = this._ub([N, 512, 0, 0]);
        this._add(this.pipe.zerofill, this._bg(this.pipe.zerofill, u2, [this.idsBuf, this.tBuf2]), Math.ceil(N * 512 / 256));
        const u3 = this._ub([N * 512, 0, 0, 0]);
        this._add(this.pipe.copy, this._bg(this.pipe.copy, u3, [this.tBuf2, this.teBuf]), Math.ceil(N * 512 / 256));
      }
    }
    {
      const u = this._ub([N, this.B, 0, 0]);
      this._add(this.pipe.teb, this._bg(this.pipe.teb, u, [this.teBuf, this.tebBuf]), Math.ceil(this.B * N * 512 / 256));
    }
  }

  textNames() { return this.textCalls.map((c) => c.name); }
  async testGemm2(M, K, N, mode) {
    const dev = this.device;
    const mk = (bytes) => dev.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC });
    const A = mode === 'dedicated' ? mk(M * K * 2) : this.tBuf3;
    const C = mode === 'dedicated' ? mk(M * N * 2) : this.tBuf2;
    const W = mode === 'dedicated' ? { buffer: mk(K * N * 2), offset: 0, size: K * N * 2 } : this._wslice('tb0.p1');
    const Bs = mode === 'dedicated' ? { buffer: mk(N * 2), offset: 0, size: N * 2 } : this._wslice('tb0.p1.b');
    const u = this._ub([M, N, K, 0]);
    const bg = this._bg(this.pipe.gemm, u, [A, W, Bs, C]);
    this.encodeCalls([{ pipe: this.pipe.gemm, bg, x: Math.ceil(M / 128), y: Math.ceil(N / 64), z: 1 }]);
    await this.device.queue.onSubmittedWorkDone();
    return mode + ' ok';
  }

  async testGemm(M, K, N) {
    const u = this._ub([M, N, K, 0]);
    const bg = this._bg(this.pipe.gemm, u, [this.tBuf3, this._wslice('tb0.p1'), this._wslice('tb0.p1.b'), this.tBuf2]);
    this.encodeCalls([{ pipe: this.pipe.gemm, bg, x: Math.ceil(M / 64), y: Math.ceil(N / 64), z: 1 }]);
    await this.device.queue.onSubmittedWorkDone();
    return M + 'x' + K + 'x' + N;
  }

  debugText() { return this.textCalls.map((c) => ({ name: c.name, x: c.x, y: c.y, meta: c.meta })); }

  async runTextPrefix(k) {
    this.encodeCalls(this.textCalls.slice(0, k));
    await this.device.queue.onSubmittedWorkDone();
    return k;
  }

  async runOne(list, k) {
    this.encodeCalls([list[k]]);
    await this.device.queue.onSubmittedWorkDone();
    return list[k].name;
  }

  async setCond(condMelF32, ids, limit) {
    const dev = this.device, N = this.N;
    const ca = this.f32 ? new Float32Array(N * 100) : new Float16Array(N * 100);
    ca.set(condMelF32.subarray(0, N * 100));
    dev.queue.writeBuffer(this.condBuf, 0, ca);
    const idarr = new Int32Array(N).fill(-1);
    for (let i = 0; i < Math.min(ids.length, N); i++) idarr[i] = ids[i];
    dev.queue.writeBuffer(this.idsBuf, 0, idarr);
    const list = (limit === undefined) ? this.textCalls : this.textCalls.slice(0, limit);
    if (list.length) this.encodeCalls(list);
    if (limit !== 0) await dev.queue.onSubmittedWorkDone();
  }

  setX(xF32) {
    const f = this.f32 ? new Float32Array(this.N * 100) : new Float16Array(this.N * 100);
    f.set(xF32.subarray(0, this.N * 100));
    this.device.queue.writeBuffer(this.xBuf, 0, f);
  }

  encodeCalls(list) {
    const enc = this.device.createCommandEncoder();
    const p = enc.beginComputePass();
    for (const c of list) {
      p.setPipeline(c.pipe);
      p.setBindGroup(0, c.bg);
      p.dispatchWorkgroups(c.x, c.y, c.z);
    }
    p.end();
    this.device.queue.submit([enc.finish()]);
  }

  writeTime(t) {
    const half = 128;
    const tb = this.f32 ? new Float32Array(256) : new Float16Array(256);
    for (let j = 0; j < 128; j++) {
      const fr = Math.exp(-Math.log(10000) / (half - 1) * j);
      const a = t * 1000 * fr;
      tb[j] = Math.sin(a); tb[128 + j] = Math.cos(a);
    }
    this.device.queue.writeBuffer(this.timeIn, 0, tb);
  }

  async step(t, dt) {
    this.writeTime(t);
    this.device.queue.writeBuffer(this.flowUniform.buffer, this.flowUniform.offset + 16, new Float32Array([dt]));
    this.encodeCalls(this.calls);
    await this.device.queue.onSubmittedWorkDone();
  }

  async runStepPrefix(k, t) {
    this.writeTime(t === undefined ? 0.5 : t);
    this.encodeCalls(this.calls.slice(0, k));
    await this.device.queue.onSubmittedWorkDone();
  }

  async readBuf(buf, elems) {
    const bytes = elems * this.es;
    const staging = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    const enc = this.device.createCommandEncoder();
    enc.copyBufferToBuffer(buf, 0, staging, 0, bytes);
    this.device.queue.submit([enc.finish()]);
    await staging.mapAsync(GPUMapMode.READ);
    const ab = staging.getMappedRange().slice(0);
    staging.unmap();
    return this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab));
  }

  async readVel() {
    const bytes = this.M2 * 100 * this.es;
    const buf = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    const enc = this.device.createCommandEncoder();
    enc.copyBufferToBuffer(this.velBuf, 0, buf, 0, bytes);
    this.device.queue.submit([enc.finish()]);
    await buf.mapAsync(GPUMapMode.READ);
    const ab = buf.getMappedRange().slice(0);
    const out = this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab));
    buf.unmap();
    return out;
  }

  async readX() {
    const bytes = this.N * 100 * this.es;
    const buf = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    const enc = this.device.createCommandEncoder();
    enc.copyBufferToBuffer(this.xBuf, 0, buf, 0, bytes);
    this.device.queue.submit([enc.finish()]);
    await buf.mapAsync(GPUMapMode.READ);
    const ab = buf.getMappedRange().slice(0);
    const out = this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab));
    buf.unmap();
    return out;
  }
}