engine/autotune.js

// Adapter-aware routing autotune.
//
// The 'auto' thresholds in decoder.js encode sweeps from ONE device (RTX
// 5070 Ti laptop). Re-running the sweep suite on the Intel xe-lpg iGPU
// (2026-07-06) flipped four of them outright: fuse_ln LOSES 51% at b1 (its
// dGPU home turf), all-5 tiled projections WIN from B≥32 (the dGPU starves
// below 128), the fused lm_head argmax loses ~5%, and the lm_head
// gemv→tiled crossover moves 16 → 32. Apple/Mali/Adreno have never been
// measured at all. Rather than grow a vendor table, this measures the four
// decisions on THE device at hand (~1s dGPU / ~8s iGPU, once per load) and
// returns a threshold table; callers resolve it per batch with
// tunedOptions(tuned, B) and pass the result through translateBatch.
//
// The tuned space is VALUE-PRESERVING routing only — every arm is a kernel/
// layout choice already quality-gated on its own (proj_sweep divergence
// gate, fuse_ln_equiv, argmax_fuse_equiv, q8 goldens; near-tie argmax flips
// between arms are the known, accepted f16 reality). ffn 'q8' is NOT in the
// space: it fails the golden bar (27/30) no matter how fast it is.
//
// Ties keep the shipped default (a candidate must beat it by WIN_MARGIN),
// so on the 5070 Ti the policy resolves to exactly today's behavior.

import { runEncoder } from './encoder.js';
import { createDecodeState, encodeDecodeStep } from './decoder.js';
import {
  createUniformParamPool, getDispatchStats, shouldUseUniformParamPool,
} from './pipelines.js';
import { maxBatchForLimits } from './shapes.js';
import { EOS, VOCAB } from './constants.js';

const ALL5 = ['fc1', 'fc2', 'self_out', 'cross_q', 'cross_out'];

// decoder.js 'auto' behavior, expressed as thresholds (see the consts there).
export const DEFAULT_ROUTING = Object.freeze({
  fuseLnMaxB: 4,
  lmHeadMinB: 16,
  lmHeadFuse: 'auto',
  projTiledMinB: 128,
  projTiledKinds: ['fc1', 'fc2'],
  // 'auto' = megakernel at B ≤ DECODE_MEGA_MAX_B; 'off' disables it. Won
  // −12–16% b1 e2e on the NVIDIA it was tuned on, but LOSES on Metal (Mac
  // mega_sweep 2026-07-06: +7% b1, and mega+fusedLn 'auto' is 3.1× off the
  // lean unfused path) — so it is a measured decision like the others.
  decodeMega: 'auto',
  // Subgroup reductions at the WT-GEMV + LN sites. Default OFF: a wash on
  // the NVIDIA reference (sg_sweep b8 −3.3%, e2e noise), but −12.7% b8 /
  // −7.4% b16 steady-state on M5 Max Metal (barriers are what a
  // one-workgroup-per-row kernel pays there) — measured per device below.
  sg: 'off',
  // FFN split-K at large batch: 'auto' keeps FFN_SPLITK_AUTO (sk8 — tuned on
  // the RTX, where sk0 measures +27% at b576), 'off' turns it off from
  // B ≥ FFN_SK_OFF_MIN_B. Measured because the constant does NOT transfer:
  // on the Intel Xe iGPU sk8 makes the b576 fc1+fc2 chain 1.43× SLOWER than
  // sk0 (xe_sweep 2026-07-13, RR×7: 27.25 → 18.99ms; b256 12.37 → 9.86) —
  // ~63 workgroups already saturate 4 Xe cores and the 8× f32 partials
  // traffic just burns shared bandwidth. At b128 split-K still wins on BOTH
  // devices (Xe 6.48 vs 7.62) — hence the ≥256 threshold, not a full off.
  ffnSkLarge: 'auto',
});

// Batch floor for a measured ffnSkLarge 'off' — below this split-K wins on
// every device measured so far (see the DEFAULT_ROUTING note).
export const FFN_SK_OFF_MIN_B = 256;

// The split-K probe's bucketed source length (kept short: the encoder run is
// setup cost, the decode steps are the measurement).
const SK_PROBE_S = 64;

// Batch for the split-K probe, or null to skip it. B=576 gives the strongest
// e2e signal (FFN is ~8% of a b576 step on Xe; at b256 it drops to ~5%, right
// at WIN_MARGIN — a probe there would 'keep' on noise on the very device the
// decision exists for). Two hard gates shrink or drop it (Codex review of
// 5a0b88c, 2026-07-13):
//   - maxB: the caller's production batch ceiling (app absCapForEnv — 192 on
//     mobile). A device that never decodes B ≥ FFN_SK_OFF_MIN_B gets zero
//     benefit from the policy, and the b576 probe state is exactly the shape
//     that jetsam'd an iPhone 17 PM (decode ~650MB wired, 2026-07-09) — skip.
//   - binding limit: maxBatchForLimits(S=64) floored to whole 64-row tiles.
//     Moxhi-30 at the 128MiB default passes 576 exactly (per-row ceiling is
//     encoder ffnTmp: 585 → 576); a Hachimi-60-class model (FFN 2304) caps at
//     455 → 448, which still clears the ≥256 floor — measured there instead
//     of failing the whole tune into an uncacheable partial table.
export function resolveSkProbeB(ctx, dtypeBytes, maxB) {
  const bindingB = Math.floor(maxBatchForLimits(ctx, SK_PROBE_S, dtypeBytes) / 64) * 64;
  const b = Math.min(576, maxB ?? 576, bindingB);
  return b >= FFN_SK_OFF_MIN_B ? b : null;
}

// A candidate wins only if it beats the default by >5%. The defaults come
// from 7-round × 50-step sweeps on the reference GPU; this quick 3-round ×
// 24-step probe must not overturn them on noise (measured: ±4% b1 arm drift
// on a busy dGPU flipped fuseLn between two back-to-back runs at a 2%
// margin, while every real iGPU win measures ≥5.5%).
const WIN_MARGIN = 0.95;

// Part of the app's autotune cache key (autotuneCacheKey): bump whenever
// WGSL kernel code, the tuned space, or routing semantics change, so tables
// measured against old kernels re-tune instead of steering new ones.
export const KERNEL_REV = 13;

// Bump when the probe method changes (points, rounds, S, or verdict rules).
// KERNEL_REV tracks kernel/routing-space compatibility; this tracks how the
// measurements that produced one tuned table should be interpreted.
export const AUTOTUNE_PROTOCOL_REV = 2;

const median = (xs) => {
  const s = [...xs].sort((a, b) => a - b);
  const m = s.length >> 1;
  return s.length % 2 ? s[m] : (s[m - 1] + s[m]) / 2;
};

export class ReductionSafetyError extends Error {
  constructor(message, verdicts) {
    super(message);
    this.name = 'ReductionSafetyError';
    this.code = 'WEBMT_NO_SAFE_REDUCTION';
    this.verdicts = verdicts;
  }
}

// Turn the two tri-state probes into one explicit safety policy. `false`
// means a route was measured clean, `true` means it was measured broken,
// and `null` means the probe itself could not produce a verdict. Performance
// tuning is allowed to choose between routes only when BOTH are proven safe.
export function resolveReductionSafety({ treeBug, sgMismap, sgFeature }) {
  const valid = (v) => v === true || v === false || v === null;
  if (!valid(treeBug) || !valid(sgMismap)) {
    throw new Error(`invalid reduction verdicts: treeBug=${treeBug}, sgMismap=${sgMismap}`);
  }
  const treeSafe = treeBug === false;
  const sgSafe = !!sgFeature && sgMismap === false;
  if (treeSafe && sgSafe) {
    return { status: 'both-safe', retry: false, sgOk: true, forceSg: false };
  }
  if (treeSafe) {
    return { status: 'tree-safe', retry: false, sgOk: false, forceSg: false };
  }
  if (sgSafe) {
    return { status: 'subgroup-safe', retry: false, sgOk: true, forceSg: true };
  }
  const retry = treeBug === null || (!!sgFeature && sgMismap === null);
  return { status: retry ? 'unresolved' : 'unsafe', retry, sgOk: false, forceSg: false };
}

// Rotate rather than merely reverse: with three or more arms, each one gets
// every position in the thermal ramp before the sequence repeats.
export function rotatedArmOrder(names, round) {
  if (!names.length) return [];
  const offset = ((round % names.length) + names.length) % names.length;
  return [...names.slice(offset), ...names.slice(0, offset)];
}

export function pairedRatios(runs, alt, base = 'def') {
  const a = runs[alt];
  const b = runs[base];
  if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length || !a.length) {
    throw new Error(`pairedRatios: ${alt}/${base} need equal non-empty run arrays`);
  }
  return a.map((v, i) => v / b[i]);
}

// Three rounds are a screening pass. Only a strong, consistent >=10% win
// settles early; a clean non-win keeps the shipped default; noisy or near-
// threshold results ask for four more paired rounds. At seven rounds the
// candidate needs both a >5% paired-median win and >=5 favorable pairs.
export function pairedPerfVerdict(runs, alt, base = 'def') {
  const ratios = pairedRatios(runs, alt, base);
  const med = median(ratios);
  const favorable = ratios.filter((r) => r < 1).length;
  if (ratios.length <= 3) {
    // A three-round early win must be unanimous. Two fast pairs plus one
    // loss is exactly the noisy pattern that flipped sg on the reference GPU.
    if (med <= 0.90 && favorable === ratios.length) return 'win';
    if (med >= 0.98 && Math.min(...ratios) >= 0.95) return 'keep';
    return 'more';
  }
  const need = Math.ceil(ratios.length * (2 / 3));
  return med < WIN_MARGIN && favorable >= need ? 'win' : 'keep';
}

// Performance-only subgroup routing changes reduction order and can flip a
// near-tie token even when both routes are numerically valid. Demand stronger
// and longer evidence than ordinary routing knobs: at least seven paired
// rounds, >=10% median gain, and every round favorable.
export function pairedStableSgWin(runs, alt = 'sgOn', base = 'def') {
  const ratios = pairedRatios(runs, alt, base);
  if (ratios.length < 7) return false;
  const favorable = ratios.filter((ratio) => ratio < 1).length;
  return median(ratios) <= 0.90 && favorable === ratios.length;
}

// Deterministic filler tokens (LCG) — timing only depends on shapes, but
// keep runs reproducible and inside the vocab (clear of the specials).
function syntheticBatch(B, S) {
  const ids = new Uint32Array(B * S);
  let x = 0x9e3779b9;
  for (let i = 0; i < ids.length; i++) {
    x = (Math.imul(x, 1664525) + 1013904223) >>> 0;
    ids[i] = 100 + (x % (VOCAB - 200));
  }
  for (let b = 0; b < B; b++) ids[b * S + S - 1] = EOS;
  return { ids, lens: new Uint32Array(B).fill(S), B, S };
}

// One lazy parameter pool per canary/full-tune session. The session owns the
// pool explicitly so cleanup stays in this module instead of hiding mutable
// state on ctx (callers often pass a shallow ctx copy).
function createAutotuneParamPoolSession(ctx, enabled) {
  let pool = null;
  let destroyed = false;
  let decodeTransientUniforms = 0;
  return {
    poolFor(state) {
      if (destroyed) throw new Error('autotune uniform-pool session is destroyed');
      if (!shouldUseUniformParamPool(enabled, state)) return null;
      pool ??= createUniformParamPool(ctx.device, { banks: 2 });
      return pool;
    },
    recordDecodeUniformDelta(delta) {
      decodeTransientUniforms += Math.max(0, delta);
    },
    invalidateBindings() {
      return pool?.invalidateBindings() ?? 0;
    },
    snapshot() {
      return { decodeTransientUniforms };
    },
    destroy() {
      if (destroyed) return;
      pool?.destroy();
      pool = null;
      destroyed = true;
    },
  };
}

// Wall-clock µs per decode step for one routing arm — same grouped
// record/submit shape as translateBatch (GROUP_STEPS=8), fresh state per
// call so KV growth is identical across arms.
async function timeArm(ctx, weights, encRun, B, steps, opts, paramPoolSession) {
  const { device } = ctx;
  const state = createDecodeState(ctx, weights, {
    B, S: encRun.S, maxSteps: steps, ...opts,
  });
  let uniformsBefore = null;
  let armDrained = false;
  try {
    // Match translateBatch's WebKit path while timing routing decisions. Queue
    // ordering makes a bank reusable after submit: the next writeBuffer is
    // ordered after that submit even though autotune has no readback to await.
    const paramPool = paramPoolSession.poolFor(state);
    if (paramPool) uniformsBefore = getDispatchStats(device).uniformBuffersCreated;
    const t0 = performance.now();
    for (let g = 0; g < steps; g += 8) {
      const tEnd = Math.min(g + 8, steps);
      const bank = (g / 8) % 2;
      let poolFrameActive = false;
      const scratch = [];
      try {
        if (paramPool) {
          paramPool.begin(bank);
          poolFrameActive = true;
        }
        const encoder = device.createCommandEncoder({ label: `autotune ${g}..${tEnd}` });
        const pass = encoder.beginComputePass();
        for (let t = g; t < tEnd; t++) {
          scratch.push(...encodeDecodeStep(ctx, weights, encRun, state, t, pass).scratch);
        }
        pass.end();
        if (paramPool) {
          paramPool.flush();
          poolFrameActive = false;
        }
        device.queue.submit([encoder.finish()]);
        if (paramPool) paramPool.release(bank);
      } catch (err) {
        if (poolFrameActive) paramPool.abort();
        throw err;
      } finally {
        for (const b of scratch) b.destroy();
      }
    }
    await device.queue.onSubmittedWorkDone();
    armDrained = true;
    return Number((((performance.now() - t0) * 1000) / steps).toFixed(1));
  } finally {
    try {
      if (uniformsBefore !== null) {
        const uniformsAfter = getDispatchStats(device).uniformBuffersCreated;
        paramPoolSession.recordDecodeUniformDelta(uniformsAfter - uniformsBefore);
      }
    } finally {
      try {
        // A session intentionally keeps its bank buffers across arms, but a
        // fresh decode state gives every arm a new resource generation.
        // Retire the old bind groups before destroying that state so WebKit
        // can release its backing buffers immediately.
        if (armDrained) paramPoolSession.invalidateBindings();
      } finally {
        state.destroy();
      }
    }
  }
}

export function cachedAutotuneMatches(cachedUs, currentUs, tolerance = 0.15) {
  if (!(Number.isFinite(cachedUs) && cachedUs > 0 && Number.isFinite(currentUs) && currentUs > 0)) {
    return false;
  }
  return Math.abs((currentUs / cachedUs) - 1) <= tolerance + Number.EPSILON;
}

export function cachedPolicyMatchesSafety(tuned, safety) {
  if (!tuned || !safety) return false;
  if (safety.forceSg) return tuned.sg === 'on' && tuned.sgForced === true;
  if (!safety.sgOk) return tuned.sg !== 'on';
  return true;
}

// A cache can outlive the GPU's current power/thermal regime. The RTX
// reference has two repeatable states whose B64 default step differs by
// ~25%; the b1 mega-vs-lean winner flips with that state. One warmed B64
// canary is much cheaper than a full retune and prevents a 30-day cache from
// steering today's low-clock session with a high-clock policy (or vice versa).
export async function measureAutotuneReference(
  ctx,
  weights,
  {
    B = 64, S = 96, steps = 24, rounds = 3,
    immediates = 'auto', uniformPool = false,
  } = {},
) {
  const batch = syntheticBatch(B, S);
  const encRun = await runEncoder(ctx, weights, batch, { retainEncOut: false });
  await ctx.device.queue.onSubmittedWorkDone();
  const paramPoolSession = createAutotuneParamPoolSession(ctx, uniformPool);
  try {
    await timeArm(ctx, weights, encRun, B, steps, { immediates }, paramPoolSession);
    const runs = [];
    for (let r = 0; r < rounds; r++) {
      runs.push(await timeArm(ctx, weights, encRun, B, steps, { immediates }, paramPoolSession));
    }
    return { medianUs: median(runs), runs, poolStats: paramPoolSession.snapshot() };
  } finally {
    try {
      paramPoolSession.destroy();
    } finally {
      encRun.arena.destroy();
    }
  }
}

// Measure the routing decisions on this device and return {tuned, timings}.
// tuned is a threshold table (DEFAULT_ROUTING shape); timings carries the
// per-arm medians for reporting. Weights should be the production load
// (lmHeadQ8/ffnWT/projWT) so 'auto' resolves to what the app will run.

// GPU correctness self-test for the Adreno 7xx tree-reduction miscompile
// (2026-07 Android round). The driver compiles the guarded-RMW shared-memory
// tree (`if (tid < s) { red[tid] = op(red[tid], red[tid+s]) }` + barrier)
// WRONG inside the real add_ln/attention kernels — deterministically — while
// the same idiom in a minimal shader passes; the trigger is contextual, so
// this probe replicates the smallest KNOWN-FAILING structure verbatim (the
// add_layernorm shape: strided load loop -> shared vbuf + register partials
// -> sum tree -> broadcast mean). ~5ms, no features required. Returns true
// when the device computes it wrong -> callers must force sg reductions;
// false = computes it right; null = the probe itself couldn't run.
export async function detectTreeReductionBug(device) {
  const D = 448;
  const WG = 256;
  const G = 8;
  const code = `
@group(0) @binding(0) var<storage, read_write> out: array<f32>;
const D: u32 = ${D}u;
const WG: u32 = ${WG}u;
fn vhash(i: u32, g: u32) -> f32 {
  return f32((((i + 1u) * 2654435761u) ^ ((g + 1u) * 40503u)) & 1023u);
}
var<workgroup> vbuf: array<f32, D>;
var<workgroup> scratch: array<f32, WG>;
@compute @workgroup_size(${WG})
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let g = wid.x; let tid = lid.x;
  var sum: f32 = 0.0;
  for (var i = tid; i < D; i = i + WG) {
    let v = (vhash(i, g) - 512.0) / 256.0 + (vhash(i + 7777u, g) - 512.0) / 256.0;
    vbuf[i] = v;
    sum = sum + v;
  }
  scratch[tid] = sum;
  workgroupBarrier();
  for (var s = WG / 2u; s > 0u; s = s >> 1u) {
    if (tid < s) { scratch[tid] = scratch[tid] + scratch[tid + s]; }
    workgroupBarrier();
  }
  let mu = scratch[0] / f32(D);
  workgroupBarrier();
  for (var i = tid; i < D; i = i + WG) { out[g * D + i] = vbuf[i] - mu; }
}`;
  try {
    const module = device.createShaderModule({ code });
    const pipeline = device.createComputePipeline({
      layout: 'auto', compute: { module, entryPoint: 'main' },
    });
    const outBuf = device.createBuffer({
      size: G * D * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    const staging = device.createBuffer({
      size: G * D * 4, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    const bindGroup = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [{ binding: 0, resource: { buffer: outBuf } }],
    });
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(G);
    pass.end();
    encoder.copyBufferToBuffer(outBuf, 0, staging, 0, G * D * 4);
    device.queue.submit([encoder.finish()]);
    await staging.mapAsync(GPUMapMode.READ);
    const got = new Float32Array(staging.getMappedRange().slice(0));
    staging.unmap();
    outBuf.destroy();
    staging.destroy();
    const hash = (i, g) => ((Math.imul(i + 1, 2654435761) ^ Math.imul(g + 1, 40503)) >>> 0) & 1023;
    for (let g = 0; g < G; g++) {
      const v = new Float64Array(D);
      let sum = 0;
      for (let i = 0; i < D; i++) {
        v[i] = (hash(i, g) - 512) / 256 + (hash(i + 7777, g) - 512) / 256;
        sum += v[i];
      }
      const mu = sum / D;
      for (let i = 0; i < D; i++) {
        if (!(Math.abs(got[g * D + i] - (v[i] - mu)) <= 1e-3)) return true;
      }
    }
    return false;
  } catch {
    // A probe failure must never block engine init — but a dead probe is not
    // a clean verdict either: on a tree-bug device a transient crash here
    // (memory pressure at load) used to cache a no-forced-sg table FOREVER.
    // null = inconclusive: the policy layer retries, then chooses another
    // proven-safe route or blocks initialization.
    return null;
  }
}

// GPU correctness self-test for the OTHER reduction route. The sg kernels
// (gemm_gemv IF_SG et al.) fold each TK-slice with a subgroupShuffleDown
// butterfly and assume consecutive local invocations occupy consecutive
// lanes of one subgroup. That layout is implementation-defined in WGSL —
// every shipping driver maps linearly, but nothing guarantees it (Codex
// audit 2026-07-10) — so measure it: replicate the kernel's exact butterfly
// at every power-of-two TK the dispatch gate can pick (TK ≤ subgroupMinSize)
// and check the lane-0 slice sums against CPU. Returns true = mismapped (sg
// routes are UNSAFE here), false = clean, null = probe couldn't run.
export async function detectSgShuffleMismap(device, maxTk) {
  const WG = 64;
  const G = 4;
  const tks = [4, 8, 16, 32].filter((tk) => tk <= Math.min(maxTk ?? 0, WG));
  if (!tks.length) return null;
  const hash = (i, g) => (((Math.imul(i + 1, 2654435761) ^ Math.imul(g + 1, 40503)) >>> 0) & 1023) / 64;
  try {
    for (const TK of tks) {
      const slices = WG / TK;
      const code = `
enable subgroups;
@group(0) @binding(0) var<storage, read_write> out: array<f32>;
const TK: u32 = ${TK}u;
fn vhash(i: u32, g: u32) -> f32 {
  return f32((((i + 1u) * 2654435761u) ^ ((g + 1u) * 40503u)) & 1023u) / 64.0;
}
@compute @workgroup_size(${WG})
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let g = wid.x; let tid = lid.x;
  var vr = vhash(tid, g);
  for (var s = TK / 2u; s > 0u; s = s >> 1u) {
    vr = vr + subgroupShuffleDown(vr, s);
  }
  if (tid % TK == 0u) { out[g * ${slices}u + tid / TK] = vr; }
}`;
      const module = device.createShaderModule({ code });
      const pipeline = device.createComputePipeline({
        layout: 'auto', compute: { module, entryPoint: 'main' },
      });
      const outBuf = device.createBuffer({
        size: G * slices * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
      });
      const staging = device.createBuffer({
        size: G * slices * 4, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
      });
      const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [{ binding: 0, resource: { buffer: outBuf } }],
      });
      const encoder = device.createCommandEncoder();
      const pass = encoder.beginComputePass();
      pass.setPipeline(pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.dispatchWorkgroups(G);
      pass.end();
      encoder.copyBufferToBuffer(outBuf, 0, staging, 0, G * slices * 4);
      device.queue.submit([encoder.finish()]);
      await staging.mapAsync(GPUMapMode.READ);
      const got = new Float32Array(staging.getMappedRange().slice(0));
      staging.unmap();
      outBuf.destroy();
      staging.destroy();
      for (let g = 0; g < G; g++) {
        for (let sl = 0; sl < slices; sl++) {
          let want = 0;
          for (let i = 0; i < TK; i++) want += hash(sl * TK + i, g);
          if (!(Math.abs(got[g * slices + sl] - want) <= 1e-2)) return true;
        }
      }
    }
    return false;
  } catch {
    // Same contract as detectTreeReductionBug: a dead probe is inconclusive,
    // not a verdict. The policy layer may use tree only if tree was proven
    // safe; otherwise it retries and then blocks rather than guessing.
    return null;
  }
}

export async function probeReductionSafety(
  ctx,
  { treeProbe = detectTreeReductionBug, sgProbe = detectSgShuffleMismap } = {},
) {
  const sgFeature = !!ctx.hasSubgroups && (ctx.subgroupMinSize ?? 0) >= 16;
  // Correctness gates BEFORE any perf decision — one probe per reduction
  // route. treeBug: a device with the tree-reduce miscompile must run the sg
  // variants at every reduction site (attention, add_ln, row_ln, gemv-WT,
  // decoder_mega) - sg is forced on and the b8 perf probe cannot override
  // it. sgMismap: a device whose subgroup lane layout breaks the shuffle
  // butterfly must never probe/force sg. A null verdict is retried once; after
  // that at least one route must be explicitly proven safe.
  const probe = async () => {
    const treeBug = await treeProbe(ctx.device);
    // `true` represents unavailable/unsafe when the feature itself is absent.
    const sgMismap = sgFeature
      ? await sgProbe(ctx.device, ctx.subgroupMinSize)
      : true;
    return { treeBug, sgMismap };
  };
  let verdicts = await probe();
  let safety = resolveReductionSafety({ ...verdicts, sgFeature });
  if (safety.retry) {
    await ctx.device.queue.onSubmittedWorkDone();
    verdicts = await probe();
    safety = resolveReductionSafety({ ...verdicts, sgFeature });
  }
  if (safety.status === 'unsafe' || safety.status === 'unresolved') {
    throw new ReductionSafetyError(
      `webMT: no proven-safe GPU reduction route ` +
      `(treeBug=${verdicts.treeBug}, sgMismap=${verdicts.sgMismap}, subgroups=${sgFeature})`,
      { ...verdicts, sgFeature },
    );
  }
  return { ...safety, verdicts: { ...verdicts, sgFeature } };
}

export async function autotuneRouting(ctx, weights, opts = {}) {
  const { reductionSafety = null, uniformPool = false, ...perfOpts } = opts;
  const safety = reductionSafety ?? await probeReductionSafety(ctx);
  const paramPoolSession = createAutotuneParamPoolSession(ctx, uniformPool);
  try {
    try {
      const result = await probeAndDecide(ctx, weights, perfOpts, safety, paramPoolSession);
      return { ...result, poolStats: paramPoolSession.snapshot() };
    } catch (err) {
      // The perf probe died mid-flight (device pressure, transient loss). The
      // correctness verdict above must NOT die with it: on a tree-bug device,
      // falling back to plain 'auto' routing (sg off) silently translates
      // garbage. Return defaults + the forced-sg bits; `partial` tells callers
      // this table is failure-derived and should not be cached.
      const tuned = { ...DEFAULT_ROUTING, projTiledKinds: [...DEFAULT_ROUTING.projTiledKinds] };
      if (safety.forceSg) {
        tuned.sg = 'on';
        tuned.sgForced = true;
      }
      return {
        tuned,
        timings: {},
        partial: String(err?.message ?? err),
        poolStats: paramPoolSession.snapshot(),
      };
    }
  } finally {
    paramPoolSession.destroy();
  }
}

async function probeAndDecide(
  ctx,
  weights,
  {
    S = 96, steps = 24, rounds = 3, extraRounds = 4,
    immediates = 'auto', maxB = null,
  },
  { sgOk, forceSg },
  paramPoolSession,
) {
  // One entry per batch point; arms at the same B share one encoder run.
  // 'def' is decoder.js 'auto' at that B; every alt is compared against it.
  // b1 steps are so short that per-round jitter dominates — give it extra
  // rounds (cheap: its arms run in tens of ms). At B ≤ 4 'auto' means the
  // megakernel (when eligible), so the b1 alts probe the two lean paths:
  // mega with the LN fusion question mooted, and the fully split pipeline.
  // Where sg is probeable, every b1 arm also runs WITH sg — the mega-vs-lean
  // verdict flips with it (mega's ~400 tree barriers/layer are exactly what
  // Metal pays; sg removes ~80% of them), so the b1 decision below is read
  // in whichever sg mode the b8 probe picks.
  const points = [
    {
      B: 1,
      arms: {
        def: {},
        lnOff: { decodeMega: 'off', fuseLn: 'off' },
        lnOn: { decodeMega: 'off', fuseLn: 'on' },
        ...(sgOk ? {
          defSg: { sg: 'on' },
          lnOffSg: { decodeMega: 'off', fuseLn: 'off', sg: 'on' },
          lnOnSg: { decodeMega: 'off', fuseLn: 'on', sg: 'on' },
        } : {}),
      },
      comparisons: [
        ['lnOff', 'def'], ['lnOn', 'def'],
        ...(sgOk ? [['lnOffSg', 'defSg'], ['lnOnSg', 'defSg']] : []),
      ],
    },
    { B: 16, arms: { def: {}, gemv: { lmHead: 'gemv' } }, comparisons: [['gemv', 'def']] },
    { B: 32, arms: { def: {}, all5: { tiledProj: ALL5 } }, comparisons: [['all5', 'def']] },
    {
      B: 64,
      arms: { def: {}, all5: { tiledProj: ALL5 }, fuseOff: { lmHeadFuse: 'off' } },
      comparisons: [['all5', 'def'], ['fuseOff', 'def']],
    },
  ];
  // File-mode shape: FFN_SPLITK_AUTO (sk8, an RTX constant) vs split-K off
  // at the fully-tiled route. S=64 / 12 steps keep the heaviest point's
  // encoder+probe cost bounded (~3s iGPU, ~0.7s dGPU, once per cache life).
  // The batch is resolved per device/model (see resolveSkProbeB) and the
  // point is skipped entirely where the policy can never apply.
  const skB = resolveSkProbeB(ctx, weights.dtype === 'f16' ? 2 : 4, maxB);
  if (skB) {
    points.push({
      B: skB, S: SK_PROBE_S, steps: 12,
      arms: { def: {}, sk0: { ffnSplitK: 0 } },
      comparisons: [['sk0', 'def']],
    });
  }
  // Subgroup arm at b8 — the batch point where the sg signal peaks (all
  // GEMV + add_ln sites active, mega/fused-LN out of the picture). Only
  // probed where the decode state accepts sg 'on' (feature + slice fit).
  // rounds+2 like b1: the sg arms are short and the margin on the NVIDIA
  // reference is only ~3% — a hot 3-round probe was seen flipping it on a
  // 973µs outlier run. Extra rounds keep the median honest.
  if (sgOk) {
    points.push({
      B: 8,
      arms: { def: {}, sgOn: { sg: 'on' } },
      comparisons: [['sgOn', 'def']],
      minRounds: rounds + extraRounds,
    });
  }

  // Warm the adapter with the heaviest probe first. Starting at B=1 leaves
  // an NVIDIA laptop GPU in P8/low-clock territory and makes dispatch-bound
  // mega-vs-lean routing depend on whichever clock ramp this browser launch
  // happened to get. Heavy-to-light gives every latency probe a comparable
  // post-boost state without adding another synthetic warmup workload.
  points.sort((a, b) => b.B - a.B);

  const timings = {};
  for (const point of points) {
    const batch = syntheticBatch(point.B, point.S ?? S);
    const pointSteps = point.steps ?? steps;
    const encRun = await runEncoder(ctx, weights, batch, { retainEncOut: false });
    await ctx.device.queue.onSubmittedWorkDone();
    try {
      const names = Object.keys(point.arms);
      for (const name of names) {
        await timeArm(ctx, weights, encRun, point.B, pointSteps, {
          ...point.arms[name], immediates,
        }, paramPoolSession);
      }
      const runs = Object.fromEntries(names.map((n) => [n, []]));
      const measureRound = async (r) => {
        const order = rotatedArmOrder(names, r);
        for (const name of order) {
          runs[name].push(await timeArm(ctx, weights, encRun, point.B, pointSteps, {
            ...point.arms[name], immediates,
          }, paramPoolSession));
        }
      };
      const screeningRounds = point.minRounds ?? rounds;
      for (let r = 0; r < screeningRounds; r++) await measureRound(r);
      const needsMore = point.comparisons.some(([alt, base]) =>
        pairedPerfVerdict(runs, alt, base) === 'more');
      if (needsMore && screeningRounds === rounds) {
        for (let r = screeningRounds; r < rounds + extraRounds; r++) await measureRound(r);
      }
      timings[`b${point.B}`] = Object.fromEntries(
        names.map((n) => [n, { medianUs: median(runs[n]), runs: runs[n] }]),
      );
    } finally {
      encRun.arena.destroy();
    }
  }

  const wins = (B, alt, base = 'def') => {
    const runs = Object.fromEntries(
      Object.entries(timings[`b${B}`]).map(([name, entry]) => [name, entry.runs]),
    );
    return pairedPerfVerdict(runs, alt, base) === 'win';
  };

  const tuned = { ...DEFAULT_ROUTING, projTiledKinds: [...DEFAULT_ROUTING.projTiledKinds] };
  // sg is decided FIRST (b8 probe) — the b1 mega-vs-lean question is then
  // answered under the sg mode the device will actually run.
  if (forceSg) {
    tuned.sg = 'on';
    tuned.sgForced = true; // correctness, not perf - survives any probe verdict
  } else if (timings.b8) {
    const b8Runs = Object.fromEntries(
      Object.entries(timings.b8).map(([name, entry]) => [name, entry.runs]),
    );
    if (pairedStableSgWin(b8Runs)) tuned.sg = 'on';
  }
  const sfx = tuned.sg === 'on' ? 'Sg' : '';
  // b1: 'def' is the megakernel (when eligible); a lean-arm win turns it off
  // and settles the LN-fusion question for the split pipeline at the same
  // time (lowest median of the two winning lean arms decides).
  const b1 = timings.b1;
  const b1Runs = Object.fromEntries(Object.entries(b1).map(([name, entry]) => [name, entry.runs]));
  const baseName = `def${sfx}`;
  const offName = `lnOff${sfx}`;
  const onName = `lnOn${sfx}`;
  const offRatio = median(pairedRatios(b1Runs, offName, baseName));
  const onRatio = median(pairedRatios(b1Runs, onName, baseName));
  const leanBest = offRatio <= onRatio ? offName : onName;
  if (wins(1, leanBest, `def${sfx}`)) {
    tuned.decodeMega = 'off';
    tuned.fuseLnMaxB = leanBest.startsWith('lnOff') ? 0 : DEFAULT_ROUTING.fuseLnMaxB;
  }
  if (wins(16, 'gemv')) tuned.lmHeadMinB = 32;
  if (wins(64, 'fuseOff')) tuned.lmHeadFuse = 'off';
  if (wins(32, 'all5')) {
    tuned.projTiledMinB = 32;
    tuned.projTiledKinds = [...ALL5];
  } else if (wins(64, 'all5')) {
    tuned.projTiledMinB = 64;
    tuned.projTiledKinds = [...ALL5];
  }
  if (skB && wins(skB, 'sk0')) tuned.ffnSkLarge = 'off';
  return { tuned, timings };
}

// Resolve a tuned threshold table to concrete createDecodeState/translateBatch
// options for one batch size. tuned == null falls back to plain 'auto'.
export function tunedOptions(tuned, B) {
  if (!tuned) return {};
  return {
    fuseLn: B <= tuned.fuseLnMaxB ? 'on' : 'off',
    lmHead: B >= tuned.lmHeadMinB ? 'auto' : 'gemv',
    lmHeadFuse: tuned.lmHeadFuse === 'off' ? 'off' : 'auto',
    tiledProj: B >= tuned.projTiledMinB && tuned.projTiledMinB < DEFAULT_ROUTING.projTiledMinB
      ? [...tuned.projTiledKinds]
      : 'auto',
    // Tree-bug devices: mega+sg is VERIFIED broken at hachimi dims on the
    // affected Adreno (sg golden 0/30 via mega vs 30/30 lean, 2026-07-07;
    // moxhi dims pass — likely its near-16KB shared budget). Until that
    // miscompile is bisected, correctness beats mega's small-B win there.
    decodeMega: tuned.sgForced ? 'off' : (tuned.decodeMega === 'off' ? 'off' : 'auto'),
    // Learned only where supported — 'on' would throw on a device without
    // the feature, but tuned tables are per-device (cache keyed on adapter).
    sg: tuned.sg === 'on' ? 'on' : 'off',
    // Tree-bug devices also miscompile attention_block at non-multiple-of-16
    // workgroup sizes — cap the encoder QB there (no-op where QB is already
    // a multiple of 8, e.g. every D<=64 model).
    encAttnSafe: !!tuned.sgForced,
    // Measured 'off' only applies at large batch — below FFN_SK_OFF_MIN_B
    // split-K wins on every device measured (see DEFAULT_ROUTING).
    ffnSplitK: tuned.ffnSkLarge === 'off' && B >= FFN_SK_OFF_MIN_B ? 0 : 'auto',
  };
}