// Pipeline plumbing shared by all kernels: WGSL template substitution,
// per-device pipeline caching, and dispatch helpers.

import gemmSource from './kernels/gemm.wgsl.js';
import gemvSource from './kernels/gemm_gemv.wgsl.js';
import gemmTiledSource from './kernels/gemm_tiled.wgsl.js';
import gemmTiled2Source from './kernels/gemm_tiled2.wgsl.js';
import attentionSource from './kernels/attention.wgsl.js';
import attentionBlockSource from './kernels/attention_block.wgsl.js';
import addLnSource from './kernels/add_layernorm.wgsl.js';
import gemmRowLnSource from './kernels/gemm_row_ln.wgsl.js';
import gemmReduceSource from './kernels/gemm_reduce.wgsl.js';
import embedSource from './kernels/embed.wgsl.js';
import scatterRowsSource from './kernels/scatter_rows.wgsl.js';
import compactGatherSource from './kernels/compact_gather.wgsl.js';
import decoderMegaSource from './kernels/decoder_mega.wgsl.js';
import kvAppendSource from './kernels/kv_append.wgsl.js';
import argmaxSource from './kernels/argmax_penalty.wgsl.js';
import argmaxReduceSource from './kernels/argmax_reduce.wgsl.js';
import {
  D_MODEL, HEADS, HEAD_DIM, FFN, SCORES_CAP, ATTN_SCALE, LN_EPS, EMBED_SCALE, DECODER_START,
  VOCAB, EOS, PAD, REP_PENALTY, BITMASK_WORDS, DECODE_CAP,
} from './constants.js';

// Substitute template placeholders in WGSL source.
//
// flags:
//   t       'f16'|'f32'  storage type of inputs ({{T}}), default 'f32'
//   outT    'f16'|'f32'  storage type of output ({{OUT_T}}), default = t
//   wg      number       workgroup size ({{WG}}), default 64
//   bias    bool         {{IF_BIAS}}...{{/IF_BIAS}} block
//   silu    bool         {{IF_SILU}}...{{/IF_SILU}} block
//   wt      bool         {{IF_WT}}...{{/IF_WT}} block (gemm: W stored [N,K])
//   defines {}           extra placeholders for later kernels: boolean values
//                        drive {{IF_NAME}} blocks, everything else substitutes
//                        {{NAME}} scalars (keys are uppercased).
//
// {{ENABLE_F16}} becomes 'enable f16;' iff t or outT is f16. Conditional
// blocks do not nest. Unknown placeholders throw (typo guard).

// Single source of truth for flag defaults, shared by buildShader and the
// pipeline cache key — equivalent flag spellings ({}, {t:'f32'}, different
// key order) resolve to one normalized shape and thus one compiled pipeline.
function normalizeFlags(flags = {}) {
  const defines = flags.defines ?? {};
  return {
    t: flags.t ?? 'f32',
    outT: flags.outT ?? flags.t ?? 'f32',
    wg: flags.wg ?? 64,
    bias: !!flags.bias,
    silu: !!flags.silu,
    wt: !!flags.wt,
    sg: !!flags.sg,
    immediate: !!flags.immediate,
    // Sorted keys so {a, b} and {b, a} serialize to the same cache key.
    defines: Object.fromEntries(Object.keys(defines).sort().map((k) => [k, defines[k]])),
  };
}

export function buildShader(source, flags = {}) {
  const { t, outT, wg, bias, silu, wt, sg, immediate, defines } = normalizeFlags(flags);
  const values = {
    T: t,
    OUT_T: outT,
    WG: String(wg),
    ENABLE_F16: t === 'f16' || outT === 'f16' ? 'enable f16;' : '',
    // Subgroup reductions — kernels opt in with {{ENABLE_SG}} + IF_SG/IF_NOSG.
    // Callers gate flags.sg on ctx.hasSubgroups AND slice width ≤
    // ctx.subgroupMinSize (see initDevice): SG variants assume a reduction
    // slice never straddles a subgroup.
    ENABLE_SG: sg ? 'enable subgroups;' : '',
    ENABLE_IMMEDIATE: immediate ? 'requires immediate_address_space;' : '',
    PARAM_BINDING: immediate ? '' : '@group(0) @binding(0) ',
    PARAM_ADDRESS: immediate ? 'immediate' : 'uniform',
  };
  const conds = { BIAS: !!bias, SILU: !!silu, WT: !!wt, SG: !!sg, NOSG: !sg };
  for (const [name, value] of Object.entries(defines)) {
    if (typeof value === 'boolean') conds[name.toUpperCase()] = value;
    else values[name.toUpperCase()] = String(value);
  }
  // Conditionals may nest (e.g. IF_BIAS inside gemm_gemv's IF_WT): replaced
  // bodies are not re-scanned by String.replace, so iterate to a fixed point.
  // Same-name nesting is still unsupported (the non-greedy match would pair
  // the outer open with the inner close).
  let code = source;
  for (let prev = null; prev !== code;) {
    prev = code;
    code = code.replace(/\{\{IF_([A-Z0-9_]+)\}\}([\s\S]*?)\{\{\/IF_\1\}\}/g, (_m, name, body) => {
      if (!(name in conds)) throw new Error(`buildShader: unknown conditional {{IF_${name}}}`);
      return conds[name] ? body : '';
    });
  }
  code = code.replace(/\{\{([A-Z0-9_/]+)\}\}/g, (_m, name) => {
    if (!(name in values)) throw new Error(`buildShader: unresolved placeholder {{${name}}}`);
    return values[name];
  });
  return code;
}

// Per-device pipeline cache:
// WeakMap<GPUDevice, Map<cacheKey, {pipeline, source}>>. The source template
// is stored per entry so a later kernel accidentally reusing a key name fails
// loudly instead of silently returning the wrong pipeline.
const pipelineCache = new WeakMap();

const DEFAULT_BIND_GROUP_CACHE_LIMIT = 256;
const DEFAULT_UNIFORM_POOL_BANK_BYTES = 256 * 1024;
export const MAX_UNIFORM_POOL_BATCH = 64;
const dispatchStates = new WeakMap();
const objectIds = new WeakMap();
const pooledUniformBuffers = new WeakSet();
let nextObjectId = 1;

function objectId(object) {
  let id = objectIds.get(object);
  if (!id) {
    id = nextObjectId++;
    objectIds.set(object, id);
  }
  return id;
}

function dispatchState(device) {
  let state = dispatchStates.get(device);
  if (!state) {
    state = {
      bindGroups: new Map(),
      bindGroupLimit: DEFAULT_BIND_GROUP_CACHE_LIMIT,
      activeUniformPools: 0,
      uniformPoolOriginalBindGroupLimit: null,
      dummyStorage: null,
      uniformFrame: null,
      stats: {
        uniformBuffersCreated: 0,
        uniformPoolBuffersCreated: 0,
        uniformPoolBuffersDestroyed: 0,
        uniformPoolFramesBegun: 0,
        uniformPoolFramesFlushed: 0,
        uniformPoolBlocks: 0,
        uniformPoolBytes: 0,
        uniformPoolBindGroupCacheHits: 0,
        uniformPoolWarmBindGroupLookups: 0,
        uniformPoolWarmBindGroupCacheHits: 0,
        uniformPoolWarmBindGroupResets: 0,
        uniformPoolGenerationInvalidations: 0,
        uniformPoolCachePurges: 0,
        dummyBuffersCreated: 0,
        bindGroupsCreated: 0,
        bindGroupCacheHits: 0,
        bindGroupEvictions: 0,
        bindGroupTargetedPurgeCalls: 0,
        bindGroupTargetedPurges: 0,
        immediateSets: 0,
      },
    };
    dispatchStates.set(device, state);
  }
  return state;
}

export function getDispatchStats(device) {
  const state = dispatchState(device);
  return {
    ...state.stats,
    bindGroupCacheSize: state.bindGroups.size,
    bindGroupCacheLimit: state.bindGroupLimit,
  };
}

export function resetDispatchStats(device, { clearCache = false } = {}) {
  const state = dispatchState(device);
  for (const key of Object.keys(state.stats)) state.stats[key] = 0;
  if (clearCache) state.bindGroups.clear();
}

export function setBindGroupCacheLimit(device, limit) {
  if (!Number.isInteger(limit) || limit < 0) {
    throw new Error(`bind-group cache limit must be a non-negative integer, got ${limit}`);
  }
  const state = dispatchState(device);
  state.bindGroupLimit = limit;
  while (state.bindGroups.size > limit) {
    state.bindGroups.delete(state.bindGroups.keys().next().value);
    state.stats.bindGroupEvictions++;
  }
}

// WebKit can defer releasing GPUBindGroup-owned backing allocations even
// after the cache entry is removed and every referenced GPUBuffer is
// destroyed. Repeated large file batches therefore use the established
// transient-uniform path; the pool stays enabled only through the largest
// batch size proven stable on the affected iPhone.
export function shouldUseUniformParamPool(enabled, { B, immediate } = {}) {
  return !!enabled
    && immediate === false
    && Number.isInteger(B)
    && B >= 1
    && B <= MAX_UNIFORM_POOL_BATCH;
}

// Run-scoped uniform-parameter arenas for browsers without WebGPU immediates.
// One bank belongs to one in-flight decode group until its readback completes.
// Parameter bindings keep their ordinary auto-layout interface; stable buffer
// identities + aligned offsets merely make the existing bind groups reusable.
export function createUniformParamPool(device, {
  banks = 2,
  bankBytes = DEFAULT_UNIFORM_POOL_BANK_BYTES,
  alignment = device?.limits?.minUniformBufferOffsetAlignment ?? 256,
} = {}) {
  if (!Number.isInteger(banks) || banks < 1) {
    throw new Error(`uniform pool banks must be a positive integer, got ${banks}`);
  }
  if (!Number.isInteger(alignment) || alignment < 16 || alignment % 16 !== 0) {
    throw new Error(`uniform pool alignment must be a positive multiple of 16, got ${alignment}`);
  }
  if (!Number.isInteger(bankBytes) || bankBytes < alignment) {
    throw new Error(`uniform pool bankBytes must be an integer >= alignment, got ${bankBytes}`);
  }
  bankBytes = Math.ceil(bankBytes / alignment) * alignment;

  const state = dispatchState(device);
  const statsAtCreate = {
    uniformBuffersCreated: state.stats.uniformBuffersCreated,
    bindGroupsCreated: state.stats.bindGroupsCreated,
    bindGroupCacheHits: state.stats.bindGroupCacheHits,
    pooledBindGroupCacheHits: state.stats.uniformPoolBindGroupCacheHits,
    warmBindGroupLookups: state.stats.uniformPoolWarmBindGroupLookups,
    warmBindGroupCacheHits: state.stats.uniformPoolWarmBindGroupCacheHits,
    warmBindGroupResets: state.stats.uniformPoolWarmBindGroupResets,
    generationInvalidations: state.stats.uniformPoolGenerationInvalidations,
    cachePurges: state.stats.uniformPoolCachePurges,
  };
  const cacheLimitStats = { highWater: state.bindGroupLimit };
  const poolBanks = Array.from({ length: banks }, (_, i) => {
    const buffer = device.createBuffer({
      label: `uniform params bank ${i}`,
      size: bankBytes,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    pooledUniformBuffers.add(buffer);
    return {
      buffer,
      cpu: new Uint32Array(bankBytes / 4),
      cursor: 0,
      blocks: 0,
      busy: false,
      flushed: false,
      highWaterBytes: 0,
      highWaterBlocks: 0,
    };
  });
  if (state.activeUniformPools === 0) {
    state.uniformPoolOriginalBindGroupLimit = state.bindGroupLimit;
  }
  state.activeUniformPools++;
  const cacheEnabled = state.uniformPoolOriginalBindGroupLimit > 0;
  state.stats.uniformPoolBuffersCreated += poolBanks.length;
  const bankIds = poolBanks.map((bank) => objectId(bank.buffer));

  // A cached bind group owns strong references to every bound GPUBuffer, not
  // just the small uniform bank. WebKit keeps those backing allocations alive
  // after GPUBuffer.destroy() while the bind group remains reachable. Purge
  // entries for this pool whenever their resource generation is retired.
  const purgeBindings = () => {
    let removed = 0;
    for (const key of [...state.bindGroups.keys()]) {
      if (bankIds.some((id) => key.includes(`|${id}@`))) {
        state.bindGroups.delete(key);
        state.stats.uniformPoolCachePurges++;
        removed++;
      }
    }
    return removed;
  };

  let destroyed = false;
  const needLive = () => {
    if (destroyed) throw new Error('uniform pool is destroyed');
  };
  const needBank = (index) => {
    if (!Number.isInteger(index) || index < 0 || index >= poolBanks.length) {
      throw new Error(`uniform pool bank ${index} out of range 0..${poolBanks.length - 1}`);
    }
    return poolBanks[index];
  };

  const pool = {
    begin(index) {
      needLive();
      if (state.uniformFrame) throw new Error('uniform pool frame already active');
      const bank = needBank(index);
      if (bank.busy) throw new Error(`uniform pool bank ${index} reused while busy`);
      bank.cursor = 0;
      bank.blocks = 0;
      bank.busy = true;
      bank.flushed = false;
      state.uniformFrame = {
        pool,
        bank,
        index,
        alignment,
        bankBytes,
        cacheBanks: poolBanks.length,
        cacheEnabled,
        cacheLimitStats,
        warm: bank.highWaterBlocks > 0,
      };
      state.stats.uniformPoolFramesBegun++;
      return index;
    },

    flush() {
      needLive();
      const frame = state.uniformFrame;
      if (!frame || frame.pool !== pool) throw new Error('uniform pool has no active frame to flush');
      const { bank, index } = frame;
      const usedBytes = Math.ceil(bank.cursor / 4) * 4;
      if (usedBytes > 0) {
        device.queue.writeBuffer(bank.buffer, 0, bank.cpu.buffer, bank.cpu.byteOffset, usedBytes);
      }
      bank.flushed = true;
      bank.highWaterBytes = Math.max(bank.highWaterBytes, usedBytes);
      bank.highWaterBlocks = Math.max(bank.highWaterBlocks, bank.blocks);
      state.uniformFrame = null;
      state.stats.uniformPoolFramesFlushed++;
      state.stats.uniformPoolBlocks += bank.blocks;
      state.stats.uniformPoolBytes += usedBytes;
      return { bank: index, blocks: bank.blocks, usedBytes };
    },

    abort() {
      needLive();
      const frame = state.uniformFrame;
      if (!frame || frame.pool !== pool) throw new Error('uniform pool has no active frame to abort');
      frame.bank.cursor = 0;
      frame.bank.blocks = 0;
      frame.bank.busy = false;
      frame.bank.flushed = false;
      state.uniformFrame = null;
    },

    release(index) {
      needLive();
      const bank = needBank(index);
      if (state.uniformFrame?.bank === bank) {
        throw new Error(`uniform pool bank ${index} released while its frame is active`);
      }
      if (!bank.busy || !bank.flushed) {
        throw new Error(`uniform pool bank ${index} released before a flushed submission`);
      }
      bank.busy = false;
      bank.flushed = false;
    },

    // Drop bind groups for the resource generation that has just drained,
    // while retaining the two stable uniform banks for the next generation.
    // Every bank must be idle: callers invalidate immediately before they
    // destroy/replace state buffers referenced by those bind groups.
    invalidateBindings() {
      needLive();
      if (state.uniformFrame?.pool === pool) {
        throw new Error('uniform pool generation invalidated while its frame is active');
      }
      const busy = poolBanks.findIndex((bank) => bank.busy);
      if (busy >= 0) {
        throw new Error(`uniform pool bank ${busy} is busy during generation invalidation`);
      }
      state.stats.uniformPoolGenerationInvalidations++;
      return purgeBindings();
    },

    snapshot() {
      return {
        banks: poolBanks.length,
        bankBytes,
        alignment,
        bindGroupCacheLimit: state.bindGroupLimit,
        bindGroupCacheLimitHighWater: cacheLimitStats.highWater,
        busyBanks: poolBanks.filter((bank) => bank.busy).length,
        highWaterBytes: Math.max(0, ...poolBanks.map((bank) => bank.highWaterBytes)),
        highWaterBlocks: Math.max(0, ...poolBanks.map((bank) => bank.highWaterBlocks)),
        transientUniformBuffersCreated:
          state.stats.uniformBuffersCreated - statsAtCreate.uniformBuffersCreated,
        bindGroupsCreated: state.stats.bindGroupsCreated - statsAtCreate.bindGroupsCreated,
        bindGroupCacheHits: state.stats.bindGroupCacheHits - statsAtCreate.bindGroupCacheHits,
        pooledBindGroupCacheHits:
          state.stats.uniformPoolBindGroupCacheHits - statsAtCreate.pooledBindGroupCacheHits,
        warmBindGroupLookups:
          state.stats.uniformPoolWarmBindGroupLookups - statsAtCreate.warmBindGroupLookups,
        warmBindGroupCacheHits:
          state.stats.uniformPoolWarmBindGroupCacheHits - statsAtCreate.warmBindGroupCacheHits,
        warmBindGroupResets:
          state.stats.uniformPoolWarmBindGroupResets - statsAtCreate.warmBindGroupResets,
        generationInvalidations:
          state.stats.uniformPoolGenerationInvalidations - statsAtCreate.generationInvalidations,
        bindGroupsPurged:
          state.stats.uniformPoolCachePurges - statsAtCreate.cachePurges,
      };
    },

    destroy() {
      if (destroyed) return;
      if (state.uniformFrame?.pool === pool) {
        const bank = state.uniformFrame.bank;
        bank.busy = false;
        bank.flushed = false;
        state.uniformFrame = null;
      }
      purgeBindings();
      for (const bank of poolBanks) {
        pooledUniformBuffers.delete(bank.buffer);
        bank.busy = false;
        bank.flushed = false;
        bank.buffer.destroy();
      }
      state.activeUniformPools--;
      if (state.activeUniformPools === 0) {
        state.bindGroupLimit = state.uniformPoolOriginalBindGroupLimit;
        state.uniformPoolOriginalBindGroupLimit = null;
        while (state.bindGroups.size > state.bindGroupLimit) {
          state.bindGroups.delete(state.bindGroups.keys().next().value);
          state.stats.bindGroupEvictions++;
        }
      }
      state.stats.uniformPoolBuffersDestroyed += poolBanks.length;
      destroyed = true;
    },
  };
  return pool;
}

// `key` must uniquely identify the source template (the source text itself is
// not part of the cache key, but collisions are detected on hit).
export function getPipeline(device, key, source, flags = {}) {
  let map = pipelineCache.get(device);
  if (!map) {
    map = new Map();
    pipelineCache.set(device, map);
  }
  // normalizeFlags builds the object literally, so JSON key order is stable.
  const cacheKey = `${key}:${JSON.stringify(normalizeFlags(flags))}`;
  let entry = map.get(cacheKey);
  if (entry) {
    if (entry.source !== source) throw new Error(`pipeline cache key collision: ${cacheKey}`);
    return entry.pipeline;
  }
  const module = device.createShaderModule({ label: cacheKey, code: buildShader(source, flags) });
  const pipeline = device.createComputePipeline({
    label: cacheKey,
    layout: 'auto',
    compute: { module, entryPoint: 'main' },
  });
  map.set(cacheKey, { pipeline, source });
  return pipeline;
}

// Buffer arguments to the dispatch helpers below may be either a plain
// GPUBuffer or a binding descriptor {buffer, offset, size} (the shape
// weights.bindingFor returns). Normalize to a bind-group resource.
function asResource(buf) {
  return buf.buffer ? buf : { buffer: buf };
}

// Small per-call uniform buffer written at creation. Returned buffers belong
// in the caller's scratch list (safe to destroy after submit).
function makeUniform(device, label, vals) {
  dispatchState(device).stats.uniformBuffersCreated++;
  const size = Math.max(16, Math.ceil((vals.length * 4) / 16) * 16);
  const buf = device.createBuffer({ label, size, usage: GPUBufferUsage.UNIFORM, mappedAtCreation: true });
  new Uint32Array(buf.getMappedRange()).set(vals);
  buf.unmap();
  return buf;
}

function dummyStorage(device) {
  const state = dispatchState(device);
  if (!state.dummyStorage) {
    state.dummyStorage = device.createBuffer({
      label: 'shared dummy storage', size: 4, usage: GPUBufferUsage.STORAGE,
    });
    state.stats.dummyBuffersCreated++;
  }
  return state.dummyStorage;
}

// Small parameter blocks use WebGPU immediates when the shader variant asks
// for them. The compatibility route remains the original mapped uniform.
function makeParams(device, label, vals, immediate) {
  if (immediate) {
    return { resource: null, values: Uint32Array.from(vals), scratch: [] };
  }
  const state = dispatchState(device);
  const frame = state.uniformFrame;
  if (frame) {
    const size = Math.max(16, Math.ceil((vals.length * 4) / 16) * 16);
    const offset = Math.ceil(frame.bank.cursor / frame.alignment) * frame.alignment;
    const end = offset + size;
    if (end > frame.bankBytes) {
      throw new Error(
        `uniform pool bank ${frame.index} overflow: need ${end} bytes, cap ${frame.bankBytes}`,
      );
    }
    frame.bank.cpu.fill(0, offset / 4, end / 4);
    frame.bank.cpu.set(vals, offset / 4);
    frame.bank.cursor = end;
    frame.bank.blocks++;
    // A pooled key includes its bank buffer + parameter offset. Grow the
    // bounded LRU before record() inserts this block so the first large frame
    // cannot evict itself. Reserve the same number of slots for every bank.
    if (frame.cacheEnabled) {
      state.bindGroupLimit = Math.max(
        state.bindGroupLimit,
        frame.bank.blocks * frame.cacheBanks,
      );
      frame.cacheLimitStats.highWater = Math.max(
        frame.cacheLimitStats.highWater,
        state.bindGroupLimit,
      );
    }
    return {
      resource: { buffer: frame.bank.buffer, offset, size },
      values: null,
      scratch: [],
    };
  }
  const resource = makeUniform(device, label, vals);
  return { resource, values: null, scratch: [resource] };
}

function paramResources(params, resources) {
  return params.resource ? [params.resource, ...resources] : resources;
}

function bindGroupKey(pipeline, resources, immediate) {
  const resourceKeys = resources.map((resource) => {
    const normalized = asResource(resource);
    return `${objectId(normalized.buffer)}@${normalized.offset ?? 0}:${normalized.size ?? '*'}`;
  });
  return `${immediate ? 'i' : 'u'}|p${objectId(pipeline)}|${resourceKeys.join('|')}`;
}

function record(pass, pipeline, device, resources, wgX, wgY = 1, wgZ = 1, immediateValues = null) {
  const firstBinding = immediateValues ? 1 : 0;
  const state = dispatchState(device);
  const pooledUniform = !immediateValues && resources.length > 0
    && pooledUniformBuffers.has(asResource(resources[0]).buffer);
  let bindGroup = null;
  let cacheKey = null;
  if ((immediateValues || pooledUniform) && state.bindGroupLimit > 0) {
    cacheKey = bindGroupKey(pipeline, resources, !!immediateValues);
    bindGroup = state.bindGroups.get(cacheKey) ?? null;
    const warmPooledLookup = pooledUniform && state.uniformFrame?.warm;
    if (warmPooledLookup) {
      state.stats.uniformPoolWarmBindGroupLookups++;
    }
    if (bindGroup) {
      state.bindGroups.delete(cacheKey);
      state.bindGroups.set(cacheKey, bindGroup);
      state.stats.bindGroupCacheHits++;
      if (pooledUniform) {
        state.stats.uniformPoolBindGroupCacheHits++;
        if (state.uniformFrame?.warm) state.stats.uniformPoolWarmBindGroupCacheHits++;
      }
    } else if (warmPooledLookup) {
      // Compaction swaps the decode resource set, so both banks need one cold
      // frame for the new generation. Stop classifying the rest of this frame
      // as warm after its first miss; the next reuse of this bank is warm.
      state.uniformFrame.warm = false;
      state.stats.uniformPoolWarmBindGroupResets++;
    }
  }
  if (!bindGroup) {
    bindGroup = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: resources.map((r, binding) => ({
        binding: binding + firstBinding,
        resource: asResource(r),
      })),
    });
    state.stats.bindGroupsCreated++;
    if (cacheKey) {
      state.bindGroups.set(cacheKey, bindGroup);
      if (state.bindGroups.size > state.bindGroupLimit) {
        state.bindGroups.delete(state.bindGroups.keys().next().value);
        state.stats.bindGroupEvictions++;
      }
    }
  }
  pass.setPipeline(pipeline);
  if (immediateValues) {
    if (typeof pass.setImmediates !== 'function') {
      throw new Error('WebGPU immediate shader selected but pass.setImmediates is unavailable');
    }
    pass.setImmediates(0, immediateValues);
    state.stats.immediateSets++;
  }
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(wgX, wgY, wgZ);
}

// Record a GEMM dispatch into an existing compute pass. Y[m,n] = X·W (+B),
// X [M,K], W [K,N], Y [M,N], all row-major (W layout: K_in × N_out).
// flags.wt flips the W layout to TRANSPOSED [N,K] row-major (LM head reads
// shared.weight [24000,448] directly).
//
// x/w/b/y are GPUBuffers or {buffer, offset, size} binding descriptors (b may
// be null when flags.bias is falsy — a 4-byte dummy is bound). flags:
// {t, outT, wg, silu, wt} as in buildShader; bias is derived from the
// presence of b.
//
// Creates a tiny per-call Dims uniform (and possibly a dummy B) — fine for
// tests; engine decode paths later manage their own uniforms. Returns
// {pipeline, scratch} where scratch lists buffers safe to destroy after the
// encoder is submitted.
// flags.gemv additionally routes to the GEMV-style kernel (gemm_gemv.wgsl —
// small-M decode projections; see the kernel header for the layout rules:
// wt requires K%4 == 0, non-wt requires N%4 == 0). flags.tk/flags.tn override
// the tile shape (defaults TK=16 k-lanes; TN=8 outputs wt / 4 quads non-wt).
// storeKV = {kCache, vCache, t, Lmax} (gemv non-wt only) additionally
// scatters the k|v slices of a fused QKV output into the decode caches from
// the epilogue (replaces a kv_append dispatch).
// flags.tiled routes to the shared-memory tiled kernel (gemm_tiled.wgsl —
// large-M sites: the encoder GEMMs). flags.bm/bn/bkk/tm/tn override the tile
// geometry (defaults 64×64×16 block, 4×4 register subtile → 256 threads).
// fusedArgmax = {partials, lbias, seen} (tiled v2 only) replaces the Y store
// with the fused greedy-argmax epilogue: per-row (val, idx) partials land in
// `partials` [M, ceil(N/BN)] vec2<u32> — finish with dispatchArgmaxReduce.
// lbias is final_logits_bias (f32 [N]), seen the repetition bitmask (read).
// y is ignored (pass null).
// splitK = {parts, sk} (tiled v2 only) partitions K over grid.z for starved
// small-N sites: RAW f32 partials land in `parts` [nz, M, N] and bias/SiLU
// are deferred — finish with dispatchGemmReduce (y is ignored; b/flags.silu
// belong to the reduce call). nz = splitKParts(K, sk, BK) ≤ sk.
export function dispatchGemm(device, pass, { x, w, b = null, y, M, K, N, storeKV = null, scales = null, fusedArgmax = null, splitK = null, flags = {} }) {
  if (flags.tiled && flags.gemv) throw new Error('flags.tiled and flags.gemv are exclusive');
  if (flags.tiled) {
    if (storeKV) throw new Error('storeKV requires flags.gemv');
    return dispatchGemmTiled(device, pass, { x, w, b, y, M, K, N, scales, fusedArgmax, splitK, flags });
  }
  if (fusedArgmax) throw new Error('fusedArgmax requires flags.tiled');
  if (splitK) throw new Error('splitK requires flags.tiled');
  if (flags.gemv) return dispatchGemv(device, pass, { x, w, b, y, M, K, N, storeKV, scales, flags });
  if (scales) throw new Error('scales (wq8) requires flags.tiled or flags.gemv');
  if (storeKV) throw new Error('storeKV requires flags.gemv');
  const wg = flags.wg ?? 64;
  const pipeline = getPipeline(device, 'gemm', gemmSource, { ...flags, bias: !!b });

  const dims = makeParams(device, 'gemm dims', [M, K, N, 0], !!flags.immediate);
  const scratch = [...dims.scratch];
  let bias = b;
  if (!bias) {
    bias = dummyStorage(device);
  }

  // dispatchWorkgroups per-dimension limit is 65535. x: ceil(24000/64)=375,
  // fine. y: one workgroup per row — fine for this model's M (decode rows /
  // sentence-length prefill), would need chunking for M > 65535.
  record(pass, pipeline, device, paramResources(dims, [x, w, bias, y]),
    Math.ceil(N / wg), M, 1, dims.values);
  return { pipeline, scratch };
}

// GEMV-style GEMM (see gemm_gemv.wgsl). Workgroup = TK k-lanes × TN outputs
// (wt: scalars, non-wt: quads of 4). Bind-group shape matches dispatchGemm.
function dispatchGemv(device, pass, { x, w, b, y, M, K, N, storeKV = null, scales = null, flags }) {
  const wt = !!flags.wt;
  // wq8 (W8A16 int8 weights): WT layout only, scales at binding 5 — which
  // storeKV also claims, so the two are mutually exclusive (never needed
  // together: q8 sites are lm_head/FFN, storeKV is self_qkv).
  const wq8 = !!flags.wq8;
  if (wq8 && (!wt || !scales || storeKV)) {
    throw new Error('gemv wq8: needs wt layout and scales, excludes storeKV');
  }
  if (wt && K % 4 !== 0) throw new Error(`gemv wt requires K%4==0, got K=${K}`);
  if (!wt && N % 4 !== 0) throw new Error(`gemv requires N%4==0, got N=${N}`);
  if (storeKV && N % 3 !== 0) throw new Error('storeKV requires fused QKV (N=3·H·D)');
  const TK = flags.tk ?? 16;
  const TN = flags.tn ?? (wt ? 8 : 4);
  const MT = wt ? (flags.mt ?? 8) : 1; // wt: rows served per workgroup (W-tile reuse)
  const pipeline = getPipeline(device, 'gemm_gemv', gemvSource, {
    t: flags.t, outT: flags.outT, wg: TK * TN, bias: !!b, silu: flags.silu, wt,
    immediate: !!flags.immediate,
    // SG (subgroup reduction) exists on the WT path only; the caller gates
    // flags.sg on ctx.hasSubgroups and TK ≤ ctx.subgroupMinSize.
    sg: !!flags.sg && wt,
    defines: { TK, TN, NWT: !wt, STORE_KV: !!storeKV, WQ8: wq8, WQF: !wq8, ...(wt ? { MT } : {}) },
  });

  const dims = makeParams(device, 'gemv dims', storeKV
    ? [M, K, N, 0, storeKV.t, storeKV.Lmax, 0, 0]
    : [M, K, N, 0], !!flags.immediate);
  const scratch = [...dims.scratch];
  let bias = b;
  if (!bias) {
    bias = dummyStorage(device);
  }
  const wgX = wt ? Math.ceil(N / TN) : Math.ceil(N / (4 * TN));
  const wgY = wt ? Math.ceil(M / MT) : M;
  const resources = [x, w, bias, y];
  if (storeKV) resources.push(storeKV.kCache, storeKV.vCache);
  if (wq8) resources.push(scales);
  record(pass, pipeline, device, paramResources(dims, resources), wgX, wgY, 1, dims.values);
  return { pipeline, scratch };
}

// Shared-memory tiled GEMM. Workgroup = BM×BN output tile, K walked in BK
// slices through workgroup memory. Two kernel versions:
//   v2 (gemm_tiled2.wgsl, default when eligible): vec4 global loads + vec4
//      shared arrays + optional 8×4 subtile (flags.tm8). Needs K%4==0, and
//      N%4==0 when !wt (vec4 reads must not straddle row boundaries).
//      flags.sh16 stores f16 in the shared tiles (bit-exact for f16 data);
//      flags.dbuf double-buffers the tiles — one barrier per K-slice.
//   v1 (gemm_tiled.wgsl): scalar loads, 4×4 subtile — fallback for shapes v2
//      can't take, and the sweep control (force with flags.tiledV: 1).
// Bind-group shape matches dispatchGemm. Geometry constraints checked here so
// a bad override fails at dispatch, not as a cryptic WGSL compile error.
// Split-K partition arithmetic: KSL = the BK-aligned K range per grid.z
// slice, nz = how many slices actually cover K (≤ sk when K is small
// relative to sk·BK — the reduce must fold exactly nz, never sk).
export function splitKParts(K, sk, BK = 16) {
  const KSL = Math.ceil(K / sk / BK) * BK;
  return { KSL, nz: Math.ceil(K / KSL) };
}

function dispatchGemmTiled(device, pass, { x, w, b, y, M, K, N, scales = null, fusedArgmax = null, splitK = null, flags }) {
  const BM = flags.bm ?? 64;
  const BN = flags.bn ?? 64;
  const BK = flags.bkk ?? 16;
  if (BM % 4 !== 0 || BN % 4 !== 0) {
    throw new Error(`gemm_tiled: BM/BN must be multiples of 4 (${BM}, ${BN})`);
  }
  // v3 staging flags (v2 only, checked below): sh16 stores the native f16 in
  // the shared tiles (bit-exact for f16 data — f32→f16 round-trip of
  // f16-origin values; int8 q values ≤127 are also exact); dbuf double-
  // buffers the tiles for one barrier per K-slice. Fused argmax excludes
  // both: its pVal partials alias Xs as raw f32 lanes.
  const sh16 = !!flags.sh16;
  const dbuf = !!flags.dbuf;
  if (fusedArgmax && (sh16 || dbuf)) {
    throw new Error('gemm_tiled2 fused argmax: sh16/dbuf unsupported (pVal aliases f32 Xs)');
  }
  // Split-K: raw partials only — the bias/SiLU epilogue moves to
  // dispatchGemmReduce, so accepting them here would silently drop them.
  if (splitK) {
    if (fusedArgmax) throw new Error('gemm_tiled2 splitK: exclusive with fusedArgmax');
    if (!(splitK.sk >= 2)) throw new Error(`gemm_tiled2 splitK: sk must be >= 2, got ${splitK.sk}`);
    if (b || flags.silu) throw new Error('gemm_tiled2 splitK: pass bias/silu to dispatchGemmReduce, not the GEMM');
  }
  // Fused argmax adds the pIdx array (BM·BN/4 u32); pVal aliases Xs, which
  // requires the Xs lane count BK·BM to cover the BM·BN/4 partial slots.
  const fusedShared = fusedArgmax ? BM * (BN / 4) * 4 : 0;
  const laneBytes = sh16 && flags.t === 'f16' ? 2 : 4;
  const sharedBytes = (BM + BN) * BK * laneBytes * (dbuf ? 2 : 1) + fusedShared;
  if (sharedBytes > 16384) {
    throw new Error(`gemm_tiled: shared memory ${sharedBytes} bytes > 16384 limit`);
  }
  if (fusedArgmax && BK < BN / 4) {
    throw new Error(`gemm_tiled2 fused argmax: BK=${BK} < BN/4=${BN / 4} — pVal cannot alias Xs`);
  }
  // wq8 (W8A16 int8 weights — lm_head, decode FFN): v2-only, [N,K]-packed
  // like wt, per-N scales in their own binding (bias/silu still available).
  const wq8 = !!flags.wq8;
  if (wq8 && (!scales || !flags.wt)) {
    throw new Error('gemm_tiled2 wq8: needs scales and wt layout');
  }
  const v2Eligible = K % 4 === 0 && BK % 4 === 0 && (flags.wt || N % 4 === 0);
  const useV2 = (flags.tiledV ?? (v2Eligible ? 2 : 1)) === 2;
  if (useV2 && !v2Eligible) {
    throw new Error(`gemm_tiled2: shape M=${M} K=${K} N=${N} wt=${!!flags.wt} BK=${BK} not vec4-eligible`);
  }
  if (wq8 && !useV2) throw new Error('gemm_tiled2 wq8: v1 fallback has no int8 path');
  if (fusedArgmax && !useV2) throw new Error('gemm_tiled2 fused argmax: v2 only');
  if ((sh16 || dbuf) && !useV2) throw new Error('gemm_tiled2 sh16/dbuf: v2 only');
  if (splitK && (!useV2 || wq8)) throw new Error('gemm_tiled2 splitK: v2 only, no wq8');
  const kp = splitK ? splitKParts(K, splitK.sk, BK) : null;
  const TM = useV2 && flags.tm8 ? 8 : 4;
  if (BM % TM !== 0) throw new Error(`gemm_tiled: BM=${BM} not a multiple of TM=${TM}`);
  const threads = (BM / TM) * (BN / 4);
  if (threads > 256) throw new Error(`gemm_tiled: ${threads} threads > 256 workgroup limit`);
  const pipeline = useV2
    ? getPipeline(device, 'gemm_tiled2', gemmTiled2Source, {
      t: flags.t, outT: flags.outT, wg: threads, bias: !!b, silu: flags.silu,
      immediate: !!flags.immediate,
      wt: flags.wt && !wq8, // wq8 has its own [N,K] staging block
      defines: {
        BM, BN, BK, TM8: TM === 8, WNT: !flags.wt && !wq8, WQ8: wq8, WQF: !wq8,
        STORE_Y: !fusedArgmax && !splitK, ARGMAX: !!fusedArgmax,
        SPLITK: !!splitK, NOSPLITK: !splitK,
        SH16: sh16, SH32: !sh16, DBUF: dbuf, SBUF: !dbuf,
        ...(splitK ? { KSL: kp.KSL } : {}),
        ...(fusedArgmax ? { PENALTY: REP_PENALTY, MASK_WORDS: fusedArgmax.maskWords ?? BITMASK_WORDS } : {}),
      },
    })
    : getPipeline(device, 'gemm_tiled', gemmTiledSource, {
      t: flags.t, outT: flags.outT, wg: threads, bias: !!b, silu: flags.silu, wt: flags.wt,
      immediate: !!flags.immediate,
      defines: { BM, BN, BK },
    });

  const dims = makeParams(device, 'gemm_tiled dims', [M, K, N, 0], !!flags.immediate);
  const scratch = [...dims.scratch];
  let bias = b;
  if (!bias) {
    bias = dummyStorage(device);
  }
  // Slot 4 is Y (plain), the argmax partials (fused), or the split-K raw
  // partials; lbias/seen trail the optional wq8 scales so binding numbers
  // stay consecutive in every mode.
  const resources = [x, w, bias, fusedArgmax?.partials ?? splitK?.parts ?? y];
  if (wq8) resources.push(scales);
  if (fusedArgmax) resources.push(fusedArgmax.lbias, fusedArgmax.seen);
  record(pass, pipeline, device, paramResources(dims, resources),
    Math.ceil(N / BN), Math.ceil(M / BM), kp?.nz ?? 1, dims.values);
  return { pipeline, scratch };
}

// Split-K fold (gemm_reduce.wgsl): Y[m,n] = Σ_z parts[z,m,n] (+B[n], SiLU) —
// the deferred epilogue of a splitK dispatchGemm. nz MUST be splitKParts'
// nz for the same (K, sk, BK), not sk — trailing slices may not exist.
// storeKV = {kCache, vCache, t, Lmax} (split-K self_qkv): the fused row's
// k|v slices additionally scatter into the decode caches, bit-identical to
// Y's slices (same {{OUT_T}} value — the kv_append contract).
export function dispatchGemmReduce(device, pass, { parts, b = null, y, M, N, nz, storeKV = null, flags = {} }) {
  if (storeKV && N % 3 !== 0) throw new Error('gemm_reduce storeKV requires fused QKV (N=3·H·D)');
  const wg = flags.wg ?? 128;
  const pipeline = getPipeline(device, 'gemm_reduce', gemmReduceSource, {
    t: flags.t, outT: flags.outT, wg, bias: !!b, silu: !!flags.silu,
    immediate: !!flags.immediate,
    defines: { STORE_KV: !!storeKV },
  });
  const dims = makeParams(device, 'gemm_reduce dims',
    [M, N, nz, storeKV?.t ?? 0, storeKV?.Lmax ?? 0, 0, 0, 0], !!flags.immediate);
  const scratch = [...dims.scratch];
  let bias = b;
  if (!bias) {
    bias = dummyStorage(device);
  }
  const resources = [parts, bias, y];
  if (storeKV) resources.push(storeKV.kCache, storeKV.vCache);
  record(pass, pipeline, device, paramResources(dims, resources),
    Math.ceil((M * N) / wg), 1, 1, dims.values);
  return { pipeline, scratch };
}

// Blocked-attention tile chooser: solves the two constraints the kernel is
// compiled against — QB·D4 within the 256-thread workgroup, and the shared
// take within the base WebGPU 16384B budget.
// Default QB: the largest ≤ 16 that fits 256 threads at this head dim — 16
// for D ≤ 64 (Moxhi's 56 keeps its measured tile), 14 for Hachimi-60's D=72.
// qbAlign8 (Adreno tree-bug devices, 2026-07): that driver also miscompiles
// this kernel when the workgroup size (QB·D4) is not a multiple of 16 —
// measured surface: 112/144/224 threads correct, 196/216/252 wrong (~1e-2
// errors). Find the largest QB whose ACTUAL workgroup size QB·D4 is a
// multiple of 16. This preserves already-aligned shapes such as D4=20/QB=12
// and, critically, never rounds a small QB upward past the 256-thread cap.
// K/V tiles are staged in the weights' native dtype, so the f32 fallback
// (adapters without shader-f16 — first seen: Colab T4 via Vulkan, 2026-07)
// doubles kvBytes and the f16-measured default JB=32 no longer fits: the
// default JB halves (floor 8) until the budget holds. Explicit qb/jb
// override the solver (and may throw at the dispatch guard).
export function attnBlockTile({ D4, t, qb = null, jb = null, qbAlign8 = false }) {
  let QB = qb ?? Math.max(1, Math.min(16, Math.floor(256 / D4)));
  if (qbAlign8 && qb == null && (QB * D4) % 16 !== 0) {
    while (QB > 1 && (QB * D4) % 16 !== 0) QB--;
    if ((QB * D4) % 16 !== 0) {
      throw new Error(`attention block: no QB <= 256 threads aligns D4=${D4} to 16 threads`);
    }
  }
  const kvBytes = t === 'f16' ? 8 : 16;
  // Shared budget: Qs f32 quads + Ks/Vs native-T quads + p tile scores +
  // 3 per-query f32 arrays.
  const sharedFor = (j) => QB * D4 * 16 + 2 * j * D4 * kvBytes + QB * j * 4 + 3 * QB * 4;
  let JB = jb ?? 32;
  if (jb == null) while (JB > 8 && sharedFor(JB) > 16384) JB >>= 1;
  return { QB, JB, shared: sharedFor(JB) };
}

// Record a unified-attention dispatch (grid B·M × H). q/k/v/y GPUBuffers or
// binding descriptors; q, k and v may all alias one fused buffer — the
// strides/offsets (elements, not bytes) select the slices (see
// attention.wgsl). lens is required when lenMode is 1; a dummy is bound for
// lenMode 0. flags.t picks the storage type. flags.block routes to the
// blocked encoder kernel (attention_block.wgsl, lenMode 1 only) with tile
// shape flags.qb × flags.jb (defaults from attnBlockTile). flags.packed
// (block only) switches to the row-packed layout: Q/K/V/Y hold T = Σ lens
// rows and `starts` (u32 [B], required) carries each sequence's first packed
// row. Returns {pipeline, scratch}.
export function dispatchAttention(device, pass, {
  q, k, v, lens = null, y, B, M, L, lenMode, step = 0, starts = null,
  qStride = HEADS * HEAD_DIM, qOff = 0, kvStride = HEADS * HEAD_DIM, kOff = 0, vOff = 0,
  flags = {},
}) {
  // Q/K/V are bound as vec4 arrays (see attention.wgsl): every stride/offset
  // must be vec4-aligned. HEAD_DIM%4 == 0 is enforced by applyModelConfig.
  for (const [name, val] of [['qStride', qStride], ['qOff', qOff], ['kvStride', kvStride],
    ['kOff', kOff], ['vOff', vOff], ['HEAD_DIM', HEAD_DIM]]) {
    if (val % 4 !== 0) throw new Error(`attention: ${name}=${val} not vec4-aligned`);
  }
  if (flags.block) {
    // Blocked encoder path (attention_block.wgsl): QB query rows per
    // workgroup, K/V tiles staged in shared. lenMode-1 only — decode's M=1
    // gains nothing from query blocking and keeps the unblocked kernel.
    if (lenMode !== 1) throw new Error('attention block: lenMode must be 1');
    if (!lens) throw new Error('attention block: lens buffer required');
    if (flags.packed && !starts) throw new Error('attention block: packed needs a starts buffer');
    const D4 = HEAD_DIM / 4;
    const { QB, JB, shared } = attnBlockTile({
      D4, t: flags.t, qb: flags.qb ?? null, jb: flags.jb ?? null,
      qbAlign8: !!flags.qbAlign8,
    });
    const threads = QB * D4;
    if (threads > 256) throw new Error(`attention block: QB=${QB} needs ${threads} > 256 threads`);
    if (shared > 16384) throw new Error(`attention block: QB=${QB} JB=${JB} needs ${shared}B shared > 16384`);
    const pipeline = getPipeline(device, 'attention_block', attentionBlockSource, {
      t: flags.t, immediate: !!flags.immediate,
      defines: {
        H: HEADS, D: HEAD_DIM, QB, JB, ATTN_SCALE,
        Q_STRIDE: qStride, Q_OFF: qOff, KV_STRIDE: kvStride, K_OFF: kOff, V_OFF: vOff,
        PACKED: !!flags.packed, NOPACKED: !flags.packed,
      },
    });
    const params = makeParams(device, 'attn params',
      [B, M, L, lenMode, step, 0, 0, 0], !!flags.immediate);
    const resources = [q, k, v, lens, y];
    if (flags.packed) resources.push(starts);
    record(pass, pipeline, device, paramResources(params, resources),
      Math.ceil(M / QB), HEADS, B, params.values);
    return { pipeline, scratch: params.scratch };
  }
  const pipeline = getPipeline(device, 'attention', attentionSource, {
    t: flags.t, wg: flags.wg ?? 128, sg: !!flags.sg,
    immediate: !!flags.immediate,
    defines: {
      H: HEADS, D: HEAD_DIM, SCORES_CAP, ATTN_SCALE,
      Q_STRIDE: qStride, Q_OFF: qOff, KV_STRIDE: kvStride, K_OFF: kOff, V_OFF: vOff,
    },
  });
  const params = makeParams(device, 'attn params',
    [B, M, L, lenMode, step, 0, 0, 0], !!flags.immediate);
  const scratch = [...params.scratch];
  let lensBuf = lens;
  if (!lensBuf) {
    lensBuf = dummyStorage(device);
  }
  record(pass, pipeline, device, paramResources(params, [q, k, v, lensBuf, y]),
    B * M, HEADS, 1, params.values);
  return { pipeline, scratch };
}

// Record an add+LayerNorm dispatch: y = LN(x + r), one workgroup per row.
// x/r/gamma/beta/y GPUBuffers or binding descriptors; y must not alias x or r
// (read/read_write usage conflict). Returns {pipeline, scratch}.
export function dispatchAddLn(device, pass, { x, r, gamma, beta, y, rows, flags = {} }) {
  const pipeline = getPipeline(device, 'add_ln', addLnSource, {
    t: flags.t, wg: flags.wg ?? 256, sg: !!flags.sg, immediate: !!flags.immediate,
    defines: { D: D_MODEL, EPS: LN_EPS },
  });
  const params = makeParams(device, 'add_ln params', [rows, 0, 0, 0], !!flags.immediate);
  record(pass, pipeline, device, paramResources(params, [x, r, gamma, beta, y]),
    rows, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Record a FUSED projection + residual add + LayerNorm dispatch (one
// workgroup per row — see gemm_row_ln.wgsl): y = LN(x·W + b + r)·gamma+beta.
// W is the [K,N] row-major NWT tensor; K and N must be multiples of 4 and
// K is baked into the pipeline (shared-memory X row). Small-B decode only —
// M workgroups can't feed the GPU at large batch.
export function dispatchGemmRowLn(device, pass, { x, w, b, r, gamma, beta, y, M, K, N, flags = {} }) {
  if (K % 4 !== 0 || N % 4 !== 0) throw new Error(`gemm_row_ln: K=${K}/N=${N} must be vec4-aligned`);
  if ((K + N) * 4 + (flags.wg ?? 128) * 4 > 16384) {
    throw new Error(`gemm_row_ln: shared memory over budget at K=${K}, N=${N}`);
  }
  const pipeline = getPipeline(device, 'gemm_row_ln', gemmRowLnSource, {
    t: flags.t, wg: flags.wg ?? 128, sg: !!flags.sg, immediate: !!flags.immediate,
    defines: { KDIM: K, D: N, EPS: LN_EPS },
  });
  const params = makeParams(device, 'gemm_row_ln params', [M, 0, 0, 0], !!flags.immediate);
  record(pass, pipeline, device, paramResources(params, [x, w, b, r, gamma, beta, y]),
    M, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Record an embedding dispatch (one workgroup per row). mode 'src' (encoder,
// ids [B·S], pos = row % s) or 'decode' (ids = token ring, pos = step).
// packed ('src' only): ids holds T = Σ lens row-packed words
// (pos << 16 | id) — requires s < 2^16 (id always fits: VOCAB 24000).
// Returns {pipeline, scratch}.
export function dispatchEmbed(device, pass, { ids, table, posEmbed, y, mode, nRows, step = 0, batch, s = 0, packed = false, flags = {} }) {
  if (packed && (mode !== 'src' || s > 0xffff)) {
    throw new Error(`embed: packed needs mode 'src' and s < 65536 (got ${mode}, s=${s})`);
  }
  const pipeline = getPipeline(device, 'embed', embedSource, {
    t: flags.t, wg: flags.wg ?? 224, immediate: !!flags.immediate,
    defines: {
      D: D_MODEL, EMBED_SCALE, SRC_IDS: mode === 'src', DECODE: mode === 'decode', DECODER_START,
      PACKED: !!packed, NOPACKED: !packed,
    },
  });
  const params = makeParams(device, 'embed params',
    [nRows, step, batch, s], !!flags.immediate);
  record(pass, pipeline, device, paramResources(params, [ids, table, posEmbed, y]),
    nRows, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Decode-megakernel workgroup-shared budget at the ACTIVE model dims — the
// kernel's xs4/tmp4/out4/scores/red arrays (see decoder_mega.wgsl). Exported
// so createDecodeState can keep 'auto' off models that cannot compile it.
export function decodeMegaSharedBytes(wg = 256) {
  const hd4 = (HEADS * HEAD_DIM) / 4;
  const tmp4 = Math.max(FFN / 4, hd4 + wg);
  // wg·4 is the NOSG red array; SG shrinks it to wg·2 — this stays the
  // conservative bound so eligibility never depends on the sg flag.
  return (2 * hd4 + tmp4) * 16 + SCORES_CAP * 4 + wg * 4;
}

// Record a decode-step MEGAKERNEL dispatch (decoder_mega.wgsl): one
// workgroup per batch row computes the row's whole decoder layer (embed
// folded in when `embed` — layer 0). Reads the ORIGINAL [K,N] .weight
// tensors (gemm_row_ln access pattern — no transposed copies needed);
// every tensor is addressed inside the ONE weights buffer via compile-time
// vec4 offsets (byteOffset/8) — one pipeline per layer. Grid: (B). x is the
// global hidden buffer the layer reads (embed: ignored) and writes back.
export function dispatchDecoderMega(device, pass, {
  weights, layer, embed = false, ring, kCache, vCache, crossKV, lens, x,
  B, t, S, kvCapacity = DECODE_CAP, flags = {},
}) {
  if (weights.dtype !== 'f16') throw new Error('decoder mega: needs f16 weights');
  const shared = decodeMegaSharedBytes(flags.wg ?? 256);
  if (shared > 16384) {
    throw new Error(`decoder mega: shared memory ${shared} bytes > 16384 limit at these dims`);
  }
  const off4 = (name) => {
    const ten = weights.tensors.get(name);
    if (!ten) throw new Error(`decoder mega: missing tensor ${name}`);
    if (ten.byteOffset % 8 !== 0) throw new Error(`decoder mega: ${name} offset not vec4-aligned`);
    return ten.byteOffset / 8;
  };
  const p = (n) => `dec.${layer}.${n}`;
  // flags.sg swaps the tree wgMax/wgSum for subgroup reductions (~5× fewer
  // barriers — the Metal lever). Callers gate it like every sg site; the
  // kernel itself only needs subgroup size ≥ 4.
  const pipeline = getPipeline(device, 'decoder_mega', decoderMegaSource, {
    t: 'f16', wg: flags.wg ?? 256, sg: !!flags.sg, immediate: !!flags.immediate,
    defines: {
      EMBED: !!embed, NOEMBED: !embed,
      ...(embed ? {
        TABLE4: off4('shared.weight'), POS4: off4('pos_embed'),
        EMBED_SCALE, DECODER_START,
      } : {}),
      H: HEADS, D: HEAD_DIM, FFN4: FFN / 4,
      LMAX: kvCapacity, SCORES_CAP, ATTN_SCALE, EPS: LN_EPS,
      QKVW4: off4(p('self_qkv.weight')), QKVB4: off4(p('self_qkv.bias')),
      OUTW4: off4(p('self_out.weight')), OUTB4: off4(p('self_out.bias')),
      LN1G4: off4(p('ln1.weight')), LN1B4: off4(p('ln1.bias')),
      CQW4: off4(p('cross_q.weight')), CQB4: off4(p('cross_q.bias')),
      COW4: off4(p('cross_out.weight')), COB4: off4(p('cross_out.bias')),
      LN2G4: off4(p('ln2.weight')), LN2B4: off4(p('ln2.bias')),
      FC1W4: off4(p('fc1.weight')), FC1B4: off4(p('fc1.bias')),
      FC2W4: off4(p('fc2.weight')), FC2B4: off4(p('fc2.bias')),
      LN3G4: off4(p('ln3.weight')), LN3B4: off4(p('ln3.bias')),
    },
  });
  const params = makeParams(device, 'mega params', [B, t, S, 0], !!flags.immediate);
  record(pass, pipeline, device,
    paramResources(params, [weights.buffer, ring, kCache, vCache, crossKV, lens, x]),
    B, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Record a row-scatter dispatch (encoder row-packing): packed activations
// [T, N] → padded [B·S, N] via starts/lens (see scatter_rows.wgsl). Padding
// rows are left untouched (zero-initialized arena buffers read as zeros).
// N must be vec4-aligned. Returns {pipeline, scratch}.
export function dispatchScatterRows(device, pass, { x, y, starts, lens, B, S, N, flags = {} }) {
  if (N % 4 !== 0) throw new Error(`scatter_rows: N=${N} not vec4-aligned`);
  const pipeline = getPipeline(device, 'scatter_rows', scatterRowsSource, {
    t: flags.t, wg: flags.wg ?? 128, immediate: !!flags.immediate,
  });
  const params = makeParams(device, 'scatter_rows params',
    [B, S, N / 4, 0], !!flags.immediate);
  record(pass, pipeline, device, paramResources(params, [starts, lens, x, y]),
    B * S, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Drop cached bind groups that retain any of the supplied GPUBuffer objects.
// Cache keys delimit every resource id as `|<id>@`, so matching that complete
// token cannot confuse (for example) buffer 12 with buffer 112. Submitted
// command buffers retain their own internal references; deleting the JS-side
// cache entry is safe even while a uniform-pool bank is still in flight and is
// required before a replaced resource generation is destroyed on WebKit.
export function purgeBindGroupsForBuffers(device, buffers) {
  if (!Array.isArray(buffers)) {
    throw new Error('bind-group targeted purge needs an array of buffers');
  }
  const state = dispatchState(device);
  state.stats.bindGroupTargetedPurgeCalls++;
  const ids = new Set();
  for (const item of buffers) {
    if (!item) continue;
    const buffer = item.buffer ?? item;
    const id = objectIds.get(buffer);
    if (id) ids.add(id);
  }
  if (ids.size === 0) return 0;
  const needles = [...ids].map((id) => `|${id}@`);
  let removed = 0;
  for (const key of [...state.bindGroups.keys()]) {
    if (!needles.some((needle) => key.includes(needle))) continue;
    state.bindGroups.delete(key);
    removed++;
  }
  state.stats.bindGroupTargetedPurges += removed;
  return removed;
}

// Record one same-buffer live-row gather. `params` is a caller-owned aligned
// uniform binding containing [rows, rowStride, copyLen, 0]; the decode state
// keeps it persistent so compaction allocates no transient buffer. This
// dispatch intentionally bypasses the bind-group cache: its first resource is
// neither an immediate block nor a pooled-uniform bank.
export function dispatchCompactGather(device, pass, {
  data, map, params, rowStrideU32, copyLenU32, flags = {},
}) {
  if (!Number.isInteger(rowStrideU32) || rowStrideU32 < 1
      || !Number.isInteger(copyLenU32) || copyLenU32 < 1
      || copyLenU32 > rowStrideU32) {
    throw new Error(
      `compact_gather: bad row shape stride=${rowStrideU32} copy=${copyLenU32}`,
    );
  }
  const pipeline = getPipeline(device, 'compact_gather', compactGatherSource, {
    wg: flags.wg ?? 256,
  });
  record(pass, pipeline, device, [params, map, data], 1, 1, 1, null);
  return { pipeline, scratch: [] };
}

// Record a kv_append APPEND dispatch: scatter the k|v slices of a fused QKV
// projection output [B, 3·H·D] into the [B, Lmax, H, D] K/V caches at decode
// position t (see kv_append.wgsl). Returns {pipeline, scratch}.
export function dispatchKvAppend(device, pass, { fused, kCache, vCache, B, t, Lmax, flags = {} }) {
  const wg = flags.wg ?? 128;
  const pipeline = getPipeline(device, 'kv_append', kvAppendSource, {
    t: flags.t, wg, immediate: !!flags.immediate,
    defines: { H: HEADS, D: HEAD_DIM, APPEND: true, SPLIT: false },
  });
  const params = makeParams(device, 'kv_append params', [B, t, Lmax, 0], !!flags.immediate);
  record(pass, pipeline, device, paramResources(params, [fused, kCache, vCache]),
    Math.ceil((B * HEADS * HEAD_DIM) / wg), 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Record a repetition-penalty + greedy-argmax + token-writeback dispatch (one
// workgroup per batch row). logits f32 [B·V]; bias is final_logits_bias (f32
// [V] — added to the raw logits BEFORE the penalty, matching HF); tokens is
// the ring [T_max·B] written at slot t·B+b. Returns {pipeline, scratch}.
// short = {n, maskWords, idmap, gmask}: shortlisted lm_head — logits/bias/
// bitmask are local-space [n]; the epilogue maps the winner through idmap and
// mirrors the seen bit into the vocab-space gmask (see argmax_penalty.wgsl).
export function dispatchArgmaxPenalty(device, pass, { logits, bias, bitmask, done, tokens, B, t, short = null, flags = {} }) {
  const pipeline = getPipeline(device, 'argmax_penalty', argmaxSource, {
    wg: flags.wg ?? 256, immediate: !!flags.immediate,
    defines: {
      V: short?.n ?? VOCAB, EOS, PAD, PENALTY: REP_PENALTY,
      MASK_WORDS: short?.maskWords ?? BITMASK_WORDS,
      SHORT: !!short, NOSHORT: !short,
      ...(short ? { GMASK_WORDS: BITMASK_WORDS } : {}),
    },
  });
  const params = makeParams(device, 'argmax params', [B, t, 0, 0], !!flags.immediate);
  const resources = [logits, bias, bitmask, done, tokens];
  if (short) resources.push(short.idmap, short.gmask);
  record(pass, pipeline, device, paramResources(params, resources), B, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Record the fused-argmax finish dispatch (one workgroup per batch row): fold
// the NT = ceil(V/BN) per-tile (val, idx) partials a fused gemm_tiled2 wrote
// and run argmax_penalty's token/done/bitmask epilogue (see
// argmax_reduce.wgsl). Returns {pipeline, scratch}.
// short as in dispatchArgmaxPenalty: bitmask is the LOCAL mask (the fused
// epilogue's read side), idmap/gmask translate the winner to vocab space.
export function dispatchArgmaxReduce(device, pass, { partials, bitmask, done, tokens, B, t, NT, short = null, flags = {} }) {
  const pipeline = getPipeline(device, 'argmax_reduce', argmaxReduceSource, {
    wg: flags.wg ?? 256, immediate: !!flags.immediate,
    defines: {
      EOS, PAD, MASK_WORDS: short?.maskWords ?? BITMASK_WORDS,
      SHORT: !!short, NOSHORT: !short,
      ...(short ? { GMASK_WORDS: BITMASK_WORDS } : {}),
    },
  });
  const params = makeParams(device, 'argmax_reduce params', [B, t, NT, 0], !!flags.immediate);
  const resources = [partials, bitmask, done, tokens];
  if (short) resources.push(short.idmap, short.gmask);
  record(pass, pipeline, device, paramResources(params, resources), B, 1, 1, params.values);
  return { pipeline, scratch: params.scratch };
}

// Test convenience: run a single GEMM in its own encoder/pass and submit.
export function runGemmOnce(device, opts) {
  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  const { scratch } = dispatchGemm(device, pass, opts);
  pass.end();
  device.queue.submit([encoder.finish()]);
  for (const buf of scratch) buf.destroy(); // safe post-submit
}