import {
  sampleVocabulary,
  candidateDistribution,
  drawDistribution,
  verifySampledDraft,
  randomStream,
} from "./sampling.js";
import { TokenStream } from "./token-stream.js";
import { Tokenizer } from "@huggingface/tokenizers";
import { renderMiniCPMPrompt } from "./prompt.js";
import { inspectGGUF } from "./gguf.js";
import {
  buildMiniCPMKernels,
  shardVocabularyKernel,
  H,
  F,
  VOCAB,
  LAYERS,
  KV,
  MAX_CHUNKS,
  PREFILL_BATCH,
} from "./minicpm-kernels.js";
const STRIDE = 256,
  DECODE_BATCH = 4,
  EOS = new Set([1, 130073]);
export const layoutTypes = {
  saveCandidates: ["r", "w"],
  embed: ["r", "r", "w"],
  norm: ["w", "r", "r", "w"],
  qkv: ["r", "r", "w"],
  output: ["r", "r", "w"],
  down: ["r", "r", "w"],
  head: ["r", "r", "w"],
  upGate: ["r", "r", "w"],
  qkvBatch: ["r", "r", "w"],
  outputBatch: ["r", "r", "w"],
  downBatch: ["r", "r", "w"],
  upGateBatch: ["r", "r", "w"],
  silu: ["r", "w"],
  rope: ["r", "w", "w", "w"],
  store: ["r", "r", "w", "w"],
  attention: ["r", "r", "r", "w", "w"],
  combine: ["r", "r", "w"],
  sortFirst64: ["r", "w", "r", "u"],
  sortFirst256: ["r", "w", "r", "u"],
  sortMerge64: ["r", "w", "u"],
  sortMerge256: ["r", "w", "u"],
  sample: ["r", "w", "u"],
};
export function cacheBytes(context, format) {
  return (
    context *
    KV *
    ({ f32: 4, f16: 2, q8_0: 36 / 32, q4_0: 20 / 32 }[format] || 2)
  );
}
export class MiniCPMRuntime {
  constructor(onStatus) {
    this.onStatus = onStatus;
    this.buffers = new Set();
    this.resident = [];
    this.residentRun = 0;
    this.pipelines = {};
    this.layouts = {};
    this.layers = [];
  }
  buffer(size, label = "", usage) {
    const b = this.device.createBuffer({
      size: Math.ceil(size / 4) * 4,
      label,
      usage:
        usage ??
        GPUBufferUsage.STORAGE |
          GPUBufferUsage.COPY_SRC |
          GPUBufferUsage.COPY_DST,
    });
    this.buffers.add(b);
    return b;
  }
  uniform(data, label) {
    const b = this.buffer(
      Math.max(16, data.byteLength),
      label,
      GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    );
    this.device.queue.writeBuffer(b, 0, data);
    return b;
  }
  group(name, buffers) {
    return this.device.createBindGroup({
      layout: this.layouts[name],
      entries: buffers.map((b, binding) => ({
        binding,
        resource: b.buffer ? b : { buffer: b },
      })),
    });
  }
  async load(blobs, settings, signal) {
    this.settings = settings;
    this.maxContext = settings.context;
    this.adapter = await navigator.gpu.requestAdapter({
      powerPreference: "high-performance",
    });
    if (!this.adapter) throw new Error("No WebGPU adapter available.");
    this.subgroups =
      this.adapter.features.has("subgroups") &&
      (this.adapter.info?.subgroupMinSize || 0) >= 32 &&
      !(
        import.meta.env.DEV &&
        new URLSearchParams(location.search).has("qa-portable")
      );
    const limits = this.adapter.limits;
    this.sharded =
      limits.maxStorageBufferBindingSize < 284098560 ||
      limits.maxBufferSize < 284098560 ||
      (import.meta.env.DEV &&
        new URLSearchParams(location.search).has("qa-sharded"));
    this.smallWorkgroups =
      limits.maxComputeWorkgroupStorageSize < 25600 ||
      (import.meta.env.DEV &&
        new URLSearchParams(location.search).has("qa-small-workgroups"));
    this.device = await this.adapter.requestDevice({
      requiredFeatures: this.subgroups ? ["subgroups"] : [],
      requiredLimits: {
        maxStorageBufferBindingSize: Math.min(
          limits.maxStorageBufferBindingSize,
          536870912,
        ),
        maxBufferSize: Math.min(limits.maxBufferSize, 536870912),
        maxComputeWorkgroupStorageSize: Math.min(
          limits.maxComputeWorkgroupStorageSize,
          32768,
        ),
      },
    });
    this.device.lost.then((info) => {
      this.failure = new Error("GPU device lost: " + info.message);
    });
    this.device.addEventListener("uncapturederror", (event) => {
      this.failure = new Error(event.error.message);
      console.error(this.failure);
    });
    this.stepLayout = this.device.createBindGroupLayout({
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "uniform",
            hasDynamicOffset: true,
            minBindingSize: 16,
          },
        },
      ],
    });
    const kernels = buildMiniCPMKernels({
      subgroups: this.subgroups,
      kv: settings.cacheK,
      smallWorkgroups: this.smallWorkgroups,
    });
    if (this.sharded)
      for (const name of ["embed", "head", "headBlock"])
        kernels[name] = shardVocabularyKernel(kernels[name], name === "embed");
    const tokenizerPromise = this.loadTokenizer();
    await Promise.all(
      Object.entries(kernels).map(async ([name, code]) => {
        const module = this.device.createShaderModule({ code, label: name });
        const info = await module.getCompilationInfo();
        const errors = info.messages.filter((m) => m.type === "error");
        if (errors.length)
          throw new Error(
            name +
              ": " +
              errors.map((e) => e.message + " at " + e.lineNum).join("; "),
          );
        const layout = this.device.createBindGroupLayout({
          entries: this.layoutFor(name).map((type, binding) => ({
            binding,
            visibility: GPUShaderStage.COMPUTE,
            buffer: {
              type:
                type === "u"
                  ? "uniform"
                  : type === "w"
                    ? "storage"
                    : "read-only-storage",
            },
          })),
        });
        this.layouts[name] = layout;
        this.pipelines[name] = await this.device.createComputePipelineAsync({
          label: name,
          layout: this.device.createPipelineLayout({
            bindGroupLayouts: [this.stepLayout, layout],
          }),
          compute: { module, entryPoint: "main" },
        });
      }),
    );
    const tensors = new Map();
    for (const blob of blobs) {
      signal?.throwIfAborted();
      const header = await inspectGGUF(blob);
      for (const tensor of header.tensors) {
        if (tensors.has(tensor.name))
          throw new Error("Duplicate weight " + tensor.name);
        tensors.set(tensor.name, { ...tensor, blob });
      }
    }
    if (tensors.size !== 381)
      throw new Error("MiniCPM5 requires all 381 original weight tensors.");
    const upload = async (names, label) => {
      const list = names.map((name) => {
        const t = tensors.get(name);
        if (!t) throw new Error("Missing tensor: " + name);
        return t;
      });
      const b = this.buffer(
        list.reduce((n, t) => n + t.bytes, 0),
        label,
      );
      let offset = 0;
      for (const t of list) {
        for (let start = 0; start < t.bytes; start += 16 * 1024 * 1024) {
          signal?.throwIfAborted();
          const bytes = await t.blob
            .slice(
              t.offset + start,
              t.offset + Math.min(t.bytes, start + 16 * 1024 * 1024),
            )
            .arrayBuffer();
          this.device.queue.writeBuffer(b, offset + start, bytes);
        }
        offset += t.bytes;
      }
      return b;
    };
    const w = (this.work = {});
    for (const [name, size] of Object.entries({
      x: PREFILL_BATCH * H,
      norm: PREFILL_BATCH * H,
      temp: PREFILL_BATCH * H,
      zeros: PREFILL_BATCH * H,
      qkv: PREFILL_BATCH * (H + KV * 2),
      q: PREFILL_BATCH * H,
      k: PREFILL_BATCH * KV,
      v: PREFILL_BATCH * KV,
      attention: PREFILL_BATCH * H,
      upGate: PREFILL_BATCH * F * 2,
      hidden: PREFILL_BATCH * F,
      partial: PREFILL_BATCH * 16 * MAX_CHUNKS * 128,
      md: PREFILL_BATCH * 16 * MAX_CHUNKS * 2,
      logits: VOCAB,
      tokens: this.maxContext + DECODE_BATCH + 1,
      sortA: 128 * 256 * 2,
      sortB: 128 * 256 * 2,
    }))
      w[name] = this.buffer(size * 4, name);
    w.candidates = this.buffer(15 * 256 * 8, "speculative candidates");
    w.candidatesRead = this.buffer(
      15 * 256 * 8,
      "candidate readback",
      GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    );
    w.verifyLogits = this.buffer(8 * VOCAB * 4, "verification logits");
    w.verifyRead = this.buffer(
      64,
      "verification readback",
      GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    );
    w.sampleSteps = this.buffer(
      (this.maxContext + 9) * STRIDE,
      "sampler steps",
      GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    );
    w.steps = this.buffer(
      (this.maxContext + DECODE_BATCH + 2) * STRIDE,
      "steps",
      GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    );
    w.read = this.buffer(
      4 * DECODE_BATCH,
      "sampled tokens",
      GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    );
    w.logitsRead = this.buffer(
      VOCAB * 4,
      "logits readback",
      GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    );
    w.penalty = this.uniform(new Float32Array([1, 0, 0, 0]), "repeat penalty");
    w.sampler = this.uniform(new Float32Array(4), "sampler");
    this.sampleStepGroup = this.device.createBindGroup({
      layout: this.stepLayout,
      entries: [{ binding: 0, resource: { buffer: w.sampleSteps, size: 16 } }],
    });
    this.stepGroup = this.device.createBindGroup({
      layout: this.stepLayout,
      entries: [{ binding: 0, resource: { buffer: w.steps, size: 16 } }],
    });
    const uploadVocab = async (name, label) => {
      if (!this.sharded) return [await upload([name], label)];
      const t = tensors.get(name);
      if (t.bytes !== 284098560)
        throw new Error("Unexpected vocabulary weight size.");
      const parts = [];
      for (let part = 0; part < 4; part++) {
        const size = 71024640,
          b = this.buffer(size, label + " " + part);
        for (let at = 0; at < size; at += 16 * 1024 * 1024) {
          signal?.throwIfAborted();
          const start = t.offset + part * size + at;
          this.device.queue.writeBuffer(
            b,
            at,
            await t.blob
              .slice(start, start + Math.min(size - at, 16 * 1024 * 1024))
              .arrayBuffer(),
          );
        }
        parts.push(b);
      }
      return parts;
    };
    this.embedding = await uploadVocab("token_embd.weight", "embeddings");
    this.head = await uploadVocab("output.weight", "output weights");
    this.finalNorm = await upload(["output_norm.weight"], "output norm");
    for (let i = 0; i < LAYERS; i++) {
      const prefix = "blk." + i + ".";
      const get = (names, label) =>
        upload(
          names.map((n) => prefix + n + ".weight"),
          label + " " + i,
        );
      const [norm, post, qkv, o, upGate, down] = await Promise.all([
        get(["attn_norm"], "norm"),
        get(["ffn_norm"], "post norm"),
        get(["attn_q", "attn_k", "attn_v"], "qkv"),
        get(["attn_output"], "attention output"),
        get(["ffn_gate", "ffn_up"], "gate/up"),
        get(["ffn_down"], "down"),
      ]);
      const k = this.buffer(
          cacheBytes(settings.context, settings.cacheK),
          "key cache " + i,
        ),
        v = this.buffer(
          cacheBytes(settings.context, settings.cacheK),
          "value cache " + i,
        );
      const groups = {
        input: this.group("norm", [w.x, i ? w.temp : w.zeros, norm, w.norm]),
        post: this.group("norm", [w.x, w.temp, post, w.norm]),
        qkv: this.group("qkv", [w.norm, qkv, w.qkv]),
        qkvBatch: this.group("qkvBatch", [w.norm, qkv, w.qkv]),
        rope: this.group("rope", [w.qkv, w.q, w.k, w.v]),
        store: this.group("store", [w.k, w.v, k, v]),
        attention: this.group("attention", [w.q, k, v, w.partial, w.md]),
        combine: this.group("combine", [w.partial, w.md, w.attention]),
        output: this.group("output", [w.attention, o, w.temp]),
        outputBatch: this.group("outputBatch", [w.attention, o, w.temp]),
        upGate: this.group("upGate", [w.norm, upGate, w.hidden]),
        upGateBatch: this.group("upGateBatch", [w.norm, upGate, w.upGate]),
        silu: this.group("silu", [w.upGate, w.hidden]),
        down: this.group("down", [w.hidden, down, w.temp]),
        downBatch: this.group("downBatch", [w.hidden, down, w.temp]),
      };
      this.layers.push({ index: i, groups, k, v });
      await this.device.queue.onSubmittedWorkDone();
      this.check();
      this.onStatus?.({
        diagnostic: `Dedicated WebGPU: layer ${i + 1} / ${LAYERS} uploaded`,
        level: "info",
      });
    }
    this.groups = {
      embed: this.group("embed", [...this.embedding, w.tokens, w.x]),
      finalNorm: this.group("norm", [w.x, w.temp, this.finalNorm, w.norm]),
    };
    this.headGroups = Array.from({ length: PREFILL_BATCH }, (_, i) =>
      this.group("head", [
        { buffer: w.norm, offset: i * H * 4, size: H * 4 },
        ...this.head,
        w.logits,
      ]),
    );
    this.verifyHead = this.group("headBlock", [
      w.norm,
      ...this.head,
      w.verifyLogits,
    ]);
    this.sortGroups = {};
    for (const keep of [64, 256]) {
      const rounds = [];
      let count = 128 * keep,
        source = w.sortA,
        target = w.sortB;
      while (count > keep) {
        const n = Math.ceil(count / 1024);
        rounds.push({
          n,
          group: this.group("sortMerge" + keep, [
            source,
            target,
            this.uniform(new Uint32Array([count, 0, 0, 0]), "sort size"),
          ]),
        });
        count = n * keep;
        [source, target] = [target, source];
      }
      this.sortGroups[keep] = {
        source,
        save: Array.from({ length: 15 }, (_, i) =>
          this.group("saveCandidates", [
            source,
            { buffer: w.candidates, offset: i * 256 * 8, size: 256 * 8 },
          ]),
        ),
        verify: Array.from({ length: 8 }, (_, i) =>
          this.group("sortFirst" + keep, [
            { buffer: w.verifyLogits, offset: i * VOCAB * 4, size: VOCAB * 4 },
            w.sortA,
            w.tokens,
            w.penalty,
          ]),
        ),
        first: this.group("sortFirst" + keep, [
          w.logits,
          w.sortA,
          w.tokens,
          w.penalty,
        ]),
        rounds,
        sample: this.group("sample", [source, w.tokens, w.sampler]),
      };
    }
    await tokenizerPromise;
    await this.device.queue.onSubmittedWorkDone();
    this.check();
    await this.warmup();
    this.ready = true;
  }
  layoutFor(name) {
    return this.sharded && ["embed", "head", "headBlock"].includes(name)
      ? ["r", "r", "r", "r", "r", "w"]
      : layoutTypes[name] || layoutTypes[name.replace(/Block$/, "")];
  }
  async warmup() {
    this.sampleK = 40;
    this.device.queue.writeBuffer(this.work.tokens, 0, new Uint32Array([0]));
    const values = new ArrayBuffer(16),
      view = new DataView(values);
    view.setFloat32(4, 0.95, true);
    view.setUint32(8, 40, true);
    this.device.queue.writeBuffer(this.work.sampler, 0, values);
    for (let round = 0; round < 2; round++) {
      const start = performance.now(),
        encoder = this.device.createCommandEncoder();
      for (let i = 0; i < 4; i++) this.encode(encoder, i, 1, true);
      await this.readTokens(encoder, 1, 4);
      this.baselineMs = (performance.now() - start) / 4;
    }
  }
  async loadTokenizer() {
    const [json, config] = await Promise.all(
      ["tokenizer.json", "tokenizer_config.json"].map(async (name) => {
        const r = await fetch(new URL("../tokenizer/" + name, import.meta.url));
        if (!r.ok) throw new Error("Could not load tokenizer.");
        return r.json();
      }),
    );
    this.tokenizer = new Tokenizer(json, config);
    this.tokenizerConfig = config;
  }
  check() {
    if (this.failure) throw this.failure;
  }
  prepareStep(position, count = 1) {
    const tokens = position + count;
    let chunks = Math.min(MAX_CHUNKS, Math.max(1, Math.ceil(tokens / 64)));
    const chunk = Math.ceil(Math.ceil(tokens / chunks) / 4) * 4;
    chunks = Math.ceil(tokens / chunk);
    this.device.queue.writeBuffer(
      this.work.steps,
      position * STRIDE,
      new Uint32Array([position, count, chunk, chunks]),
    );
    return { chunk, chunks };
  }
  dispatch(pass, name, group, x = 1, y = 1, z = 1) {
    pass.setPipeline(this.pipelines[name]);
    pass.setBindGroup(1, group);
    pass.dispatchWorkgroups(x, y, z);
  }
  encode(encoder, position, count = 1, sample = true, updateContext = true) {
    const plan = this.prepareStep(position, count),
      pass = encoder.beginComputePass();
    pass.setBindGroup(0, this.stepGroup, [position * STRIDE]);
    const run = (name, group, x = 1, y = 1, z = 1) =>
      this.dispatch(pass, name, group, x, y, z);
    const batch = count > 1,
      block = batch && count <= 8,
      suffix = block ? "Block" : batch ? "Batch" : "",
      divisor = batch && !block ? 32 : 8;
    run("embed", this.groups.embed, H / 256, count);
    for (const { index, groups: g } of this.layers) {
      run("norm", g.input, count);
      run("qkv" + suffix, batch ? g.qkvBatch : g.qkv, (H + KV * 2) / divisor);
      run("rope", g.rope, (H + KV) / 256, count);
      run("store", g.store, KV / 32, count, 2);
      run("attention", g.attention, plan.chunks, 16, count);
      run("combine", g.combine, 16, count);
      run("output" + suffix, batch ? g.outputBatch : g.output, H / divisor);
      run("norm", g.post, count);
      run(
        "upGate" + suffix,
        batch ? g.upGateBatch : g.upGate,
        batch ? (F * 2) / divisor : F / 8,
      );
      if (batch) run("silu", g.silu, F / 256, count);
      run("down" + suffix, batch ? g.downBatch : g.down, H / divisor);
      if (this.draftActive) this.draft.capture(pass, index, count);
    }
    if (sample) {
      run("norm", this.groups.finalNorm, count);
      run("head", this.headGroups[count - 1], VOCAB / 8);
      if (!this.fullSampler) {
        this.prepareSample(pass, position + count - 1);
        this.encodeSample(pass);
      }
    }
    pass.end();
    if (this.draftActive && updateContext)
      this.draft.updateContext(encoder, position, count);
  }
  prepareSample(pass, position) {
    this.device.queue.writeBuffer(
      this.work.sampleSteps,
      position * STRIDE,
      new Uint32Array([position, 1, 0, 1]),
    );
    pass.setBindGroup(0, this.sampleStepGroup, [position * STRIDE]);
  }
  encodeSample(
    pass,
    first,
    g = this.sortGroups[this.sampleK <= 64 ? 64 : 256],
    sample = true,
  ) {
    const keep = this.sampleK <= 64 ? 64 : 256;
    this.dispatch(pass, "sortFirst" + keep, first || g.first, 128);
    for (const round of g.rounds)
      this.dispatch(pass, "sortMerge" + keep, round.group, round.n);
    if (sample) this.dispatch(pass, "sample", g.sample);
  }
  async speculate(position) {
    const started = performance.now();
    const encoder = this.device.createCommandEncoder();
    this.draft.propose(encoder, position);
    this.encode(encoder, position, 8, false, false);
    const pass = encoder.beginComputePass();
    pass.setBindGroup(0, this.stepGroup, [position * STRIDE]);
    this.dispatch(pass, "norm", this.groups.finalNorm, 8);
    this.dispatch(pass, "headBlock", this.verifyHead, VOCAB / 8);
    const groups = this.sortGroups[this.sampleK <= 64 ? 64 : 256];
    for (let i = 0; i < 8; i++) {
      this.prepareSample(pass, position + i);
      this.encodeSample(
        pass,
        groups.verify[i],
        groups,
        this.samplingSettings.temperature <= 0,
      );
      if (this.samplingSettings.temperature > 0)
        this.dispatch(
          pass,
          "saveCandidates",
          groups.save[i],
          this.sampleK <= 64 ? 1 : 4,
        );
    }
    pass.end();
    encoder.copyBufferToBuffer(
      this.work.tokens,
      (position + 1) * 4,
      this.work.verifyRead,
      0,
      32,
    );
    encoder.copyBufferToBuffer(
      this.draft.work.tokens,
      (position + 1) * 4,
      this.work.verifyRead,
      32,
      28,
    );
    if (this.samplingSettings.temperature > 0)
      encoder.copyBufferToBuffer(
        this.work.candidates,
        0,
        this.work.candidatesRead,
        0,
        15 * 256 * 8,
      );
    this.device.queue.submit([encoder.finish()]);
    await this.work.verifyRead.mapAsync(GPUMapMode.READ);
    const data = new Uint32Array(this.work.verifyRead.getMappedRange()).slice();
    this.work.verifyRead.unmap();
    this.check();
    let accepted = 0,
      result;
    if (this.samplingSettings.temperature > 0) {
      await this.work.candidatesRead.mapAsync(GPUMapMode.READ);
      const words = new Uint32Array(this.work.candidatesRead.getMappedRange());
      result = [];
      for (let i = 0; i < 8; i++) {
        const p = candidateDistribution(
            words.subarray(i * 512, (i + 1) * 512),
            this.samplingSettings,
          ),
          random = (n) =>
            randomStream(this.samplingSettings.seed, position + i, n);
        if (i === 7) {
          result.push(drawDistribution(p, random(2)));
          break;
        }
        const q = candidateDistribution(
            words.subarray((8 + i) * 512, (9 + i) * 512),
            this.samplingSettings,
          ),
          check = verifySampledDraft(p, q, data[8 + i], random(1), random(2));
        result.push(check.token);
        if (!check.accepted) break;
        accepted++;
      }

      this.work.candidatesRead.unmap();
      this.device.queue.writeBuffer(
        this.work.tokens,
        (position + 1) * 4,
        new Uint32Array(result),
      );
    } else {
      while (accepted < 7 && data[accepted] === data[8 + accepted]) accepted++;
      result = Array.from(data.slice(0, accepted + 1));
    }
    const count = accepted + 1;
    const context = this.device.createCommandEncoder();
    this.draft.updateContext(context, position, count);
    this.device.queue.submit([context.finish()]);
    this.specStats.steps++;
    this.specStats.accepted += accepted;
    this.specStats.milliseconds += performance.now() - started;
    if (
      this.settings.dspark !== "on" &&
      ((this.specStats.steps >= 3 &&
        this.specStats.milliseconds /
          (this.specStats.steps + this.specStats.accepted) >
          this.baselineMs * 1.5) ||
        (this.specStats.steps >= 6 &&
          this.specStats.milliseconds /
            (this.specStats.steps + this.specStats.accepted) >
            this.baselineMs * 1.1))
    ) {
      this.draftActive = false;
      this.specStats.fallback = true;
    }
    return result;
  }
  async readTokens(encoder, from, count) {
    if (this.fullSampler) {
      encoder.copyBufferToBuffer(
        this.work.logits,
        0,
        this.work.logitsRead,
        0,
        VOCAB * 4,
      );
      this.device.queue.submit([encoder.finish()]);
      await this.work.logitsRead.mapAsync(GPUMapMode.READ);
      const values = new Float32Array(this.work.logitsRead.getMappedRange());
      const token = sampleVocabulary(
        values,
        this.samplingSettings,
        this.resident,
        from - 1,
      );
      this.work.logitsRead.unmap();
      this.check();
      this.device.queue.writeBuffer(
        this.work.tokens,
        from * 4,
        new Uint32Array([token]),
      );
      return [token];
    }
    encoder.copyBufferToBuffer(
      this.work.tokens,
      from * 4,
      this.work.read,
      0,
      count * 4,
    );
    this.device.queue.submit([encoder.finish()]);
    await this.work.read.mapAsync(GPUMapMode.READ, 0, count * 4);
    const ids = Array.from(
      new Uint32Array(this.work.read.getMappedRange(0, count * 4)),
    );
    this.work.read.unmap();
    this.check();
    return ids;
  }
  renderPrompt(messages, settings, tools) {
    return renderMiniCPMPrompt(messages, {
      thinking: settings.thinking,
      tools,
    });
  }

  // Keep streaming callback-based: Safari's reported FTL crash passed through
  // nested async-generator resumption. Await callbacks to preserve backpressure.
  async generate(prompt, settings, signal, onToken, onProgress) {
    this.check();
    const ids = this.tokenizer.encode(prompt, {
      add_special_tokens: false,
    }).ids;
    if (ids.length + 1 >= this.maxContext)
      throw new Error(
        `The prompt needs ${ids.length} tokens and exceeds the ${this.maxContext}-token context window.`,
      );
    this.fullSampler = settings.topK === 0 && settings.temperature > 0;
    this.sampleK = settings.topK || 1;
    this.samplingSettings = {
      ...settings,
      seed:
        settings.seed < 0
          ? crypto.getRandomValues(new Uint32Array(1))[0]
          : settings.seed,
    };
    const config = new ArrayBuffer(16),
      v = new DataView(config);
    v.setFloat32(0, settings.temperature, true);
    v.setFloat32(4, settings.topP, true);
    v.setUint32(8, this.sampleK, true);
    v.setUint32(12, this.samplingSettings.seed, true);
    this.device.queue.writeBuffer(this.work.sampler, 0, config);
    this.device.queue.writeBuffer(
      this.work.penalty,
      0,
      new Float32Array([settings.repeatPenalty, 0, 0, 0]),
    );
    if (this.draft) {
      v.setUint32(12, (this.samplingSettings.seed ^ 0x9e3779b9) >>> 0, true);
      this.device.queue.writeBuffer(this.draft.work.sampler, 0, config);
    }
    this.specStats = { steps: 0, accepted: 0, milliseconds: 0 };
    this.draftActive =
      !this.fullSampler && !!this.draft && settings.dspark !== "off";
    this.settings.dspark = settings.dspark || "auto";
    let reuse = 0;
    while (
      reuse < Math.min(ids.length - 1, this.residentRun) &&
      ids[reuse] === this.resident[reuse]
    )
      reuse++;
    if (this.draftActive) reuse = Math.min(reuse, this.draft.contextRun || 0);
    this.device.queue.writeBuffer(this.work.tokens, 0, new Uint32Array(ids));
    this.resident = [...ids];
    const start = performance.now();
    let ready = [];
    for (let p = reuse; p < ids.length; ) {
      signal?.throwIfAborted();
      const count = Math.min(
          settings.performance === "balanced" ? 8 : PREFILL_BATCH,
          ids.length - p,
        ),
        last = p + count === ids.length,
        encoder = this.device.createCommandEncoder();
      this.encode(encoder, p, count, last);
      if (last) ready = await this.readTokens(encoder, ids.length, 1);
      else {
        this.device.queue.submit([encoder.finish()]);
        await this.device.queue.onSubmittedWorkDone();
        this.check();
      }
      p += count;
      onProgress?.({ total: ids.length, processed: p, cache: reuse });
      await onToken({
        progress: { total: ids.length, processed: p, cache: reuse },
      });
    }
    this.residentRun = ids.length;
    const prefillMs = performance.now() - start;
    const decoder = new TokenStream(this.tokenizer);
    let text = "",
      position = ids.length;
    const max = Math.min(settings.maxTokens, this.maxContext - ids.length);
    for (let emitted = 0; emitted < max; ) {
      signal?.throwIfAborted();
      if (
        !ready.length &&
        this.draftActive &&
        position + 8 < this.maxContext &&
        max - emitted >= 8
      ) {
        ready = await this.speculate(position);
        this.residentRun = position + ready.length;
        position += ready.length;
      }
      if (!ready.length) {
        const n = Math.min(this.fullSampler ? 1 : DECODE_BATCH, max - emitted),
          encoder = this.device.createCommandEncoder();
        for (let i = 0; i < n; i++) this.encode(encoder, position + i, 1, true);
        ready = await this.readTokens(encoder, position + 1, n);
        this.residentRun = position + n;
        position += n;
      }
      const token = ready.shift();
      this.resident.push(token);
      if (EOS.has(token)) break;
      emitted++;
      const delta = decoder.push(token);
      text += delta;
      const elapsed = performance.now() - start;
      await onToken({
        token,
        delta,
        text,
        index: emitted,
        prefillMs,
        elapsed,
        promptTokens: ids.length,
        cacheTokens: reuse,
        finish: emitted >= max ? "length" : null,
      });
    }
  }
  async exit() {
    this.ready = false;
    if (this.draft) {
      for (const b of this.draft.buffers) b.destroy();
      this.draft.buffers.clear();
    }
    for (const b of this.buffers) b.destroy();
    this.buffers.clear();
    this.device?.destroy();
  }
}