Custom WGSL DiT engine (int4, CFG-free 4-step)

Add an engine backend: the 22-block DiT runs on hand-written WebGPU
kernels (tiled/sgmat GEMM, fused LN+adaLN, RoPE, tiled QK/softmax/AV,
grouped convpos, GRN) with 4-bit packed weights and per-group scales
dequantized in-kernel. The worker loads only the ORT vocoder for this
path; ONNX Runtime remains the fallback.

- 229.7 MB int4 bundle, ~1.2 s per clip on WebGPU (was 2.8 s int4 ORT)
- flow: 4 steps, sway -1, cfg 0 (B=1, no null pass)
- tuning: tts.backend / tts.gemm switch kernels without touching the worker
- fixes found on the way: u32 vs f32 uniforms, second SiLU on the time MLP,
RoPE inv_freq exponent, AV score-tile column, f32 attention scores,
safe tanh in Mish

engine/f5engine.js

@@ -0,0 +1,1324 @@
+// 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;
+  }
+}

tts.worker.js

@@ -1,4 +1,5 @@
 import * as ort from './ort/ort.mjs';
+import { F5Engine } from './engine/f5engine.js';
 
 const SAMPLE_RATE = 24000;
 const MEL_BINS = 100;
@@ -14,6 +15,8 @@ const MIN_GEN_FRAMES = 16;
 const MAX_GEN_FRAMES = 640;
 const TARGET_RMS = 0.1;
 let refTextEncoder = null;
+let engine = null;
+let engineDevice = null;
 
 ort.env.wasm.wasmPaths = '/ort/';
 ort.env.wasm.numThreads = Math.min(8, Math.max(1, (self.navigator?.hardwareConcurrency || 4) - 1));
@@ -121,7 +124,49 @@ async function createSession(url, label, executionProviders) {
   return session;
 }
 
+async function initEngine() {
+  const want = (tuning?.tts?.backend ?? 'engine') === 'engine';
+  if (!want) return;
+  const info = await probeCaps();
+  if (!info.webgpu || !info.shaderF16) return;
+  try {
+    const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
+    if (!adapter) return;
+    const feats = ['shader-f16'];
+    let sgmat = false;
+    try { sgmat = adapter.features.has('chromium-experimental-subgroup-matrix'); } catch {}
+    if (sgmat) feats.push('chromium-experimental-subgroup-matrix');
+    const device = await adapter.requestDevice({
+      requiredFeatures: feats,
+      requiredLimits: {
+        maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
+        maxBufferSize: adapter.limits.maxBufferSize,
+        maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize,
+      },
+    });
+    engineDevice = device;
+    engine = await F5Engine.create(device, {
+      gemm: tuning?.tts?.gemm || (sgmat ? 'sgmat' : 'simt'),
+      branches: 1,
+      metaUrl: `${base}/f5w_q4.json`,
+      binUrl: `${base}/f5w_q4.bin`,
+    });
+    progress('session', { message: 'DiT kernels ready (WebGPU)' });
+  } catch (err) {
+    engine = null;
+    engineDevice = null;
+    progress('fallback', { message: `engine unavailable, using ORT: ${err?.message || err}` });
+  }
+}
+
 async function loadGraphs() {
+  if (engine) {
+    if (!vocoderSession) {
+      const g = manifest.graphs || {};
+      vocoderSession = await createSession(`${base}/${g.vocoder_fp16 || g.vocoder}`, 'vocoder · WebGPU', ['webgpu', 'wasm']);
+    }
+    return;
+  }
   if (ditSession && vocoderSession) return;
   const graphs = manifest.graphs || {};
   const fp16Available = !!(graphs.dit_step_fp16 && graphs.vocoder_fp16);
@@ -151,7 +196,7 @@ async function loadGraphs() {
 
 async function warmup() {
   const frames = tuning?.tts?.warmupFrames ?? manifest.vocoder_chunk_frames ?? 375;
-  await runAndDispose(ditSession, () => ditInputs(null, frames, 16, null, null, null));
+  if (!engine) await runAndDispose(ditSession, () => ditInputs(null, frames, 16, null, null, null));
   await runAndDispose(vocoderSession, () => ({ mel: f16Tensor(new Float32Array(MEL_BINS * frames), [1, MEL_BINS, frames]) }));
 }
 
@@ -495,6 +540,7 @@ async function init(message) {
   if (message.base) base = message.base;
   await loadTuning();
   manifest = await fetchJson(`${base}/manifest.json`);
+  await initEngine();
   await loadGraphs();
   progress('warmup', { message: 'warming up kernels' });
   await warmup();
@@ -514,7 +560,7 @@ async function getRefIds(voice) {
 }
 
 async function synthesize(message) {
-  if (!ditSession || !vocoderSession) throw new Error('worker not initialized');
+  if ((!engine && !ditSession) || !vocoderSession) throw new Error('worker not initialized');
   const started = performance.now();
   const { ids: genIds, unknown } = await tokenize(message.text);
   if (!genIds.length) throw new Error(`no vocab tokens for: ${String(message.text).slice(0, 20)}`);
@@ -532,17 +578,31 @@ async function synthesize(message) {
   cfgStrength = tuning?.tts?.cfgStrength ?? manifest.sampler?.cfg_strength ?? CFG_STRENGTH;
   const steps = tuning?.tts?.flowSteps ?? manifest.sampler?.steps ?? DEFAULT_STEPS;
   progress('synth', { message: 'flow matching', step: 0, steps });
-  const mel = await flowSample(
-    (x, time) => ditStep(x, time, voice, ids, textLen, frames, refFrames),
-    frames,
-    MEL_BINS,
-    {
-      steps,
-      sway: manifest.sampler?.sway_sampling_coef ?? -1,
-      seed: message.seed ?? 1234,
-      onStep: (step, total) => progress('synth', { step, steps: total, message: `flow step ${step}/${total}` }),
-    },
-  );
+  let mel;
+  if (engine) {
+    engine.setup(frames);
+    await engine.setCond(padReference(voice, refFrames, frames), ids);
+    engine.setX(gaussianNoise(frames, MEL_BINS, message.seed ?? 1234));
+    await engineDevice.queue.onSubmittedWorkDone();
+    const sched = flowSchedule(steps, manifest.sampler?.sway_sampling_coef ?? -1);
+    for (let i = 0; i < steps; i++) {
+      progress('synth', { step: i + 1, steps, message: `flow step ${i + 1}/${steps}` });
+      await engine.step(sched[i], sched[i + 1] - sched[i]);
+    }
+    mel = await engine.readX();
+  } else {
+    mel = await flowSample(
+      (x, time) => ditStep(x, time, voice, ids, textLen, frames, refFrames),
+      frames,
+      MEL_BINS,
+      {
+        steps,
+        sway: manifest.sampler?.sway_sampling_coef ?? -1,
+        seed: message.seed ?? 1234,
+        onStep: (step, total) => progress('synth', { step, steps: total, message: `flow step ${step}/${total}` }),
+      },
+    );
+  }
   progress('synth', { message: 'vocoding', step: steps, steps });
   const pcm = await decodeMel(mel.subarray(refFrames * MEL_BINS), genFrames);
   let sum = 0;

tuning.json

@@ -9,7 +9,9 @@
  "tts": {
   "flowSteps": 4,
   "cfgStrength": 0,
-  "warmupFrames": 64
+  "warmupFrames": 64,
+  "backend": "engine",
+  "gemm": "simt"
  },
  "session": {
   "preferPrecision": "auto"