用 TypeScript 写 mini-picoLM

mkdir -p minipicolm/{memory,tensor,model,runtime,tokenizer,format}
cd minipicolm

# there is nothing to install — that is the point
node --version    # 22.6+ for --experimental-strip-types

{
  "compilerOptions": {
    "strict": true,
    "noEmit": true,
    "allowImportingTsExtensions": true,
    "target": "ES2020",
    "module": "esnext",
    "moduleResolution": "bundler",
    "lib": ["dom", "esnext"]
  }
}

src/lib/minipicolm/
├── memory/          1-3    ArrayBuffer → arena → tensor views
│   ├── buffer.ts           the single allocation
│   ├── arena.ts            bump allocator, mark/release
│   └── tensor.ts           shape + a view, nothing more
├── tensor/          4,5,9,18
│   ├── matmul.ts           matmul + the FUSED quantised path
│   ├── rmsnorm.ts          x / rms(x) * w
│   ├── softmax.ts          plain, and the streaming form
│   ├── activation.ts       SwiGLU
│   └── webgpu.ts           a real compute shader, feature-detected
├── model/           6,7,10
│   ├── rope.ts             interleaved rotation + precomputed tables
│   ├── attention.ts        GQA over the cache, online softmax
│   ├── embedding.ts        one dequantised row
│   ├── mlp.ts              SwiGLU feed-forward
│   ├── block.ts            the residual pair
│   └── llama.ts            forward(token, pos)
├── runtime/         8,12,13,17
│   ├── kv-cache.ts         [layer][pos][kv_dim], written once
│   ├── workspace.ts        every scratch buffer, allocated once
│   ├── sampler.ts          temperature / top-k / top-p, seeded
│   ├── grammar.ts          JSON pushdown automaton
│   └── generation.ts       the async loop
├── tokenizer/       11
│   └── bpe.ts              SentencePiece, byte fallback
├── format/          14-16
│   ├── gguf.ts             read and write the container
│   └── quant.ts            Q4_0, Q8_0, Q4_K, Q6_K
├── main.ts                 loadFromGGUF() and createToyModel()
├── shell.ts                the commands behind the terminal
├── selftest.ts             90 assertions
├── run-selftest.mts        the headless engine gate
└── run-shell-test.mts      the headless shell gate

/**
 * mini-picoLM — step 1: the ArrayBuffer
 * =====================================
 *
 * Everything the engine touches lives in exactly one `ArrayBuffer`. Not one per
 * tensor, not one per layer — one, allocated at startup and never grown.
 *
 * The reason is the same one picoLM has in C, and it survives the move to
 * JavaScript intact: an inference engine's allocation pattern is completely
 * known before the first token. You know the model shape, you know the context
 * length, so you know every buffer you will ever need. Allocating them up front
 * turns a stream of GC-visible objects into a single flat region.
 *
 * In the browser that matters more than it does in C. A `Float32Array` created
 * per matmul is garbage the collector has to trace; a view into a pre-allocated
 * buffer is not. The difference shows up as jitter between tokens, which is the
 * one thing a streaming UI cannot hide.
 */

/** Typed-array constructors the arena knows how to hand out. */
export type TypedArrayCtor =
  | Float32Array0
  | Int32Array0
  | Uint32Array0
  | Uint8Array0
  | Int8Array0
  | Uint16Array0;

type Float32Array0 = typeof Float32Array;
type Int32Array0 = typeof Int32Array;
type Uint32Array0 = typeof Uint32Array;
type Uint8Array0 = typeof Uint8Array;
type Int8Array0 = typeof Int8Array;
type Uint16Array0 = typeof Uint16Array;

/** GPUs and SIMD both want 256-byte alignment; CPUs are happy with far less. */
export const ALIGNMENT = 256;

export function alignUp(bytes: number, alignment: number = ALIGNMENT): number {
  return Math.ceil(bytes / alignment) * alignment;
}

export function bytesHuman(n: number): string {
  if (n < 1024) return `${n} B`;
  if (n < 1024 ** 2) return `${(n / 1024).toFixed(1)} KB`;
  if (n < 1024 ** 3) return `${(n / 1024 ** 2).toFixed(1)} MB`;
  return `${(n / 1024 ** 3).toFixed(2)} GB`;
}

/**
 * Allocate the one buffer.
 *
 * Browsers cap a single ArrayBuffer well below the address space — the limit is
 * around 2 GB on 64-bit Chrome and lower elsewhere — so this fails loudly
 * rather than at some random allocation later.
 */
export function createBackingBuffer(bytes: number): ArrayBuffer {
  if (!Number.isFinite(bytes) || bytes <= 0) {
    throw new RangeError(`backing buffer size must be positive, got ${bytes}`);
  }
  try {
    return new ArrayBuffer(bytes);
  } catch (err) {
    throw new RangeError(
      `could not allocate ${bytesHuman(bytes)} in one ArrayBuffer. ` +
        `Lower the context length or use a smaller model. (${String(err)})`,
    );
  }
}

/**
 * mini-picoLM — step 2: the arena allocator
 * =========================================
 *
 * A bump allocator over the single ArrayBuffer from step 1. Allocation is one
 * addition; there is no free list, no fragmentation, and no per-object header.
 *
 * There is also no `free()`. That sounds like a limitation and is actually the
 * design: everything the engine allocates lives exactly as long as the engine
 * does. Scratch space is reused by overwriting, not by releasing. The only
 * lifetime operation that makes sense is `mark()` / `release()` — drop
 * everything allocated after a point — which is how a request-scoped scratch
 * region works.
 *
 * This is the same argument S06 makes about the KV cache from the other end.
 * There, allocation had to be dynamic because output length is unknown. Here,
 * every buffer's size is known before the first token, so the simplest possible
 * allocator is also the right one.
 */

import { ALIGNMENT, alignUp, bytesHuman, createBackingBuffer } from "./buffer.ts";

export type Allocation = {
  readonly name: string;
  readonly byteOffset: number;
  readonly byteLength: number;
};

export class Arena {
  readonly buffer: ArrayBuffer;
  readonly capacity: number;
  private offset = 0;
  private readonly log: Allocation[] = [];

  constructor(capacityBytes: number) {
    this.capacity = alignUp(capacityBytes);
    this.buffer = createBackingBuffer(this.capacity);
  }

  /** Bytes handed out so far, including alignment padding. */
  get used(): number {
    return this.offset;
  }

  get remaining(): number {
    return this.capacity - this.offset;
  }

  /** Every allocation, in order — what the S22 simulator draws. */
  get allocations(): readonly Allocation[] {
    return this.log;
  }

  /**
   * Carve out `count` elements of `Ctor` and return a view onto them.
   *
   * The view aliases the arena; it is not a copy. Two views can legitimately
   * overlap, which is how the workspace reuses one scratch region for several
   * differently-shaped intermediates.
   */
  alloc<T extends ArrayBufferView>(
    Ctor: {
      new (buffer: ArrayBuffer, byteOffset: number, length: number): T;
      BYTES_PER_ELEMENT: number;
    },
    count: number,
    name = "anon",
  ): T {
    const byteLength = count * Ctor.BYTES_PER_ELEMENT;
    const byteOffset = alignUp(this.offset, ALIGNMENT);
    if (byteOffset + byteLength > this.capacity) {
      throw new RangeError(
        `arena exhausted allocating ${bytesHuman(byteLength)} for "${name}": ` +
          `${bytesHuman(this.remaining)} left of ${bytesHuman(this.capacity)}`,
      );
    }
    this.offset = byteOffset + byteLength;
    this.log.push({ name, byteOffset, byteLength });
    return new Ctor(this.buffer, byteOffset, count);
  }

  f32(count: number, name?: string): Float32Array {
    return this.alloc(Float32Array, count, name);
  }

  i32(count: number, name?: string): Int32Array {
    return this.alloc(Int32Array, count, name);
  }

  u8(count: number, name?: string): Uint8Array {
    return this.alloc(Uint8Array, count, name);
  }

  /** Remember the current watermark. */
  mark(): number {
    return this.offset;
  }

  /**
   * Drop everything allocated since `mark`. The bytes are not cleared — the
   * next allocation simply overwrites them, which is the whole point.
   */
  release(mark: number): void {
    if (mark < 0 || mark > this.offset) {
      throw new RangeError(`bad arena mark ${mark} (offset is ${this.offset})`);
    }
    this.offset = mark;
    while (this.log.length && this.log[this.log.length - 1].byteOffset >= mark) {
      this.log.pop();
    }
  }

  summary(): string {
    return `${bytesHuman(this.used)} / ${bytesHuman(this.capacity)} in ${this.log.length} allocations`;
  }
}

/**
 * mini-picoLM — step 3: tensor views
 * ==================================
 *
 * A tensor here is a shape plus a `Float32Array` that already points into the
 * arena. It owns nothing and allocates nothing.
 *
 * Note what is deliberately missing: no autograd, no broadcasting, no dtype
 * promotion, no lazy graph. Inference needs none of it, and every one of those
 * features costs an allocation or an indirection per operation. The whole type
 * is twenty lines because the whole job is "interpret these bytes as a matrix".
 */

import type { Arena } from "./arena.ts";

export type Shape = readonly number[];

export class Tensor {
  readonly data: Float32Array;
  readonly shape: Shape;
  readonly name: string;

  constructor(data: Float32Array, shape: Shape, name = "anon") {
    const n = shape.reduce((a, b) => a * b, 1);
    if (data.length !== n) {
      throw new RangeError(
        `tensor "${name}": shape [${shape}] needs ${n} elements, view has ${data.length}`,
      );
    }
    this.data = data;
    this.shape = shape;
    this.name = name;
  }

  static zeros(arena: Arena, shape: Shape, name = "anon"): Tensor {
    const n = shape.reduce((a, b) => a * b, 1);
    return new Tensor(arena.f32(n, name), shape, name);
  }

  /** Wrap an existing view without copying — the common case. */
  static of(data: Float32Array, shape: Shape, name = "anon"): Tensor {
    return new Tensor(data, shape, name);
  }

  get rank(): number {
    return this.shape.length;
  }

  get length(): number {
    return this.data.length;
  }

  get rows(): number {
    return this.rank >= 2 ? this.shape[0] : 1;
  }

  get cols(): number {
    return this.shape[this.rank - 1];
  }

  /** Row `i` of a 2-D tensor, as a view. Still no copy. */
  row(i: number): Float32Array {
    const c = this.cols;
    return this.data.subarray(i * c, (i + 1) * c);
  }

  /** A sub-range of the last axis, as a view — used to slice heads apart. */
  slice(start: number, end: number): Float32Array {
    return this.data.subarray(start, end);
  }

  fill(v: number): this {
    this.data.fill(v);
    return this;
  }

  copyFrom(src: ArrayLike<number>): this {
    this.data.set(src);
    return this;
  }

  toString(): string {
    return `Tensor(${this.name}, [${this.shape}])`;
  }
}

/**
 * mini-picoLM — step 4: matmul
 * ============================
 *
 * Decode is one token at a time, so every matmul in this engine is really a
 * matrix-vector product: `out[d] = dot(W[d], x)`. There is no batch dimension,
 * exactly as in picoLM, and for the same reason — one user, one core.
 *
 * That makes this function completely memory-bound, which is S05's whole point
 * arriving in JavaScript. The arithmetic is trivial; the cost is streaming `W`
 * through the CPU once per token. Two consequences shape the code below:
 *
 *   * **Row-major, sequential.** `W` is walked in the order it is stored, so
 *     the prefetcher can keep up. Transposing the loops costs ~5× on the same
 *     arithmetic.
 *   * **No allocation.** The output view is passed in. A function that returned
 *     a fresh Float32Array would allocate `n_layers × 7` arrays per token.
 *
 * The quantised variants dequantise inline and never materialise a float copy
 * of `W` — the fused path S08 argues for and the Python port of S21 could not
 * express. TypeScript can, because a hand-written loop over a Uint8Array is
 * exactly as fast as any other loop here.
 */

import { dequantBlockInto, blockSize, type QuantType } from "../format/quant.ts";

/** out[d] = dot(W[d, :], x), W row-major [d, n], float32. */
export function matmul(
  out: Float32Array,
  x: Float32Array,
  w: Float32Array,
  n: number,
  d: number,
): void {
  for (let i = 0; i < d; i++) {
    const base = i * n;
    let acc = 0;
    // Four-way unrolled: JS engines do not vectorise this for us, but they do
    // keep four independent accumulators in registers, which hides the latency
    // of the dependent add chain.
    let j = 0;
    let a0 = 0;
    let a1 = 0;
    let a2 = 0;
    let a3 = 0;
    for (; j + 3 < n; j += 4) {
      a0 += w[base + j] * x[j];
      a1 += w[base + j + 1] * x[j + 1];
      a2 += w[base + j + 2] * x[j + 2];
      a3 += w[base + j + 3] * x[j + 3];
    }
    acc = a0 + a1 + a2 + a3;
    for (; j < n; j++) acc += w[base + j] * x[j];
    out[i] = acc;
  }
}

/**
 * The fused path: `W` is still quantised bytes, and one block is decoded into a
 * small stack scratch buffer immediately before it is consumed.
 *
 * Nothing full-precision the size of `W` is ever stored, so the bytes crossing
 * the memory bus are the quantised ones — which is the entire reason to
 * quantise. Compare `dequantizeAll` + `matmul`, which reads the int4 weights
 * AND writes and re-reads a float32 copy: strictly more traffic than never
 * quantising at all.
 */
export function matmulQuant(
  out: Float32Array,
  x: Float32Array,
  w: Uint8Array,
  n: number,
  d: number,
  qtype: QuantType,
  scratch: Float32Array,
): void {
  const bs = blockSize(qtype);
  if (n % bs !== 0) {
    throw new RangeError(`matmulQuant: ${n} is not a multiple of the ${bs}-weight block`);
  }
  const blocksPerRow = n / bs;

  for (let i = 0; i < d; i++) {
    let acc = 0;
    for (let b = 0; b < blocksPerRow; b++) {
      // decode 32 or 256 weights into registers-worth of scratch, use, discard
      dequantBlockInto(w, i * blocksPerRow + b, qtype, scratch);
      const off = b * bs;
      for (let j = 0; j < bs; j++) acc += scratch[j] * x[off + j];
    }
    out[i] = acc;
  }
}

/** out += a (in place) — the residual stream's only operation. */
export function addInto(out: Float32Array, a: Float32Array): void {
  for (let i = 0; i < out.length; i++) out[i] += a[i];
}

/** out[i] = a[i] * b[i] — SwiGLU's gate. */
export function mulInto(out: Float32Array, a: Float32Array, b: Float32Array): void {
  for (let i = 0; i < out.length; i++) out[i] = a[i] * b[i];
}

export function dot(a: Float32Array, b: Float32Array, n: number): number {
  let acc = 0;
  for (let i = 0; i < n; i++) acc += a[i] * b[i];
  return acc;
}

/**
 * mini-picoLM — step 5: RMSNorm
 * =============================
 *
 * LayerNorm with the parts that turned out not to matter removed: no mean
 * subtraction, no bias. `x / rms(x) * weight`.
 *
 * The epsilon is inside the square root and is not decorative. It is model
 * configuration, shipped in GGUF metadata: TinyLlama and Llama-2 train with
 * 1e-5, the original LLaMA with 1e-6 (S03's value). Load a model with the
 * wrong one and every activation shifts slightly, compounding over 22 layers —
 * the kind of mismatch that produces a model which works and is quietly a
 * little worse.
 */

export const RMS_EPS = 1e-5;

export function rmsnorm(
  out: Float32Array,
  x: Float32Array,
  weight: Float32Array,
  eps: number = RMS_EPS,
): void {
  const n = x.length;
  let ss = 0;
  for (let i = 0; i < n; i++) ss += x[i] * x[i];
  const scale = 1 / Math.sqrt(ss / n + eps);
  for (let i = 0; i < n; i++) out[i] = x[i] * scale * weight[i];
}

/**
 * mini-picoLM — softmax, twice
 * ============================
 *
 * Two shapes of the same function, and the difference is the subject of S12.
 *
 * `softmax` is the ordinary one: find the max, exponentiate, normalise. It
 * needs the whole row in memory at once.
 *
 * `OnlineSoftmax` never sees the whole row. It keeps a running maximum and a
 * running sum, and rescales both whenever a larger value arrives. That is what
 * lets attention stream over a KV cache of any length without ever allocating
 * an array of scores — the single most important structural idea in a decode
 * kernel, and it is nine lines.
 */

/** In-place softmax over `x[0..n)`. Subtracting the max is what stops exp() overflowing. */
export function softmax(x: Float32Array, n: number = x.length): void {
  let max = -Infinity;
  for (let i = 0; i < n; i++) if (x[i] > max) max = x[i];
  let sum = 0;
  for (let i = 0; i < n; i++) {
    const e = Math.exp(x[i] - max);
    x[i] = e;
    sum += e;
  }
  const inv = 1 / sum;
  for (let i = 0; i < n; i++) x[i] *= inv;
}

/**
 * The streaming form. Feed it one score and one value vector at a time; the
 * accumulator holds the correctly-weighted partial result at every point.
 *
 *   m  running maximum
 *   l  running sum of exp(score - m)
 *   acc  running sum of exp(score - m) * v
 *
 * When a score exceeds `m`, everything accumulated so far was scaled against
 * the old maximum, so it is corrected by exp(m_old - m_new) — one multiply per
 * element, and the invariant holds again.
 */
export class OnlineSoftmax {
  private m = -Infinity;
  private l = 0;
  private readonly acc: Float32Array;

  private readonly dim: number;

  constructor(dim: number, scratch: Float32Array) {
    this.dim = dim;
    if (scratch.length < dim) {
      throw new RangeError(`OnlineSoftmax needs ${dim} scratch elements, got ${scratch.length}`);
    }
    this.acc = scratch.subarray(0, dim);
  }

  reset(): void {
    this.m = -Infinity;
    this.l = 0;
    this.acc.fill(0);
  }

  /** Fold one (score, value) pair into the accumulator. */
  push(score: number, v: Float32Array, vOffset: number): void {
    // A -Infinity score is a masked position and its weight is exp(-inf) = 0,
    // so drop it at the door. Without this, a fully-masked stream leaves m at
    // -Infinity and the else-branch below computes exp(-inf - -inf) = NaN —
    // the classic fully-masked-tile bug. Finite scores need no guard at all:
    // exp(-inf - score) is already 0.
    if (score === -Infinity) return;
    const { acc, dim } = this;
    if (score > this.m) {
      const correction = Math.exp(this.m - score);
      this.l = this.l * correction + 1;
      for (let i = 0; i < dim; i++) acc[i] = acc[i] * correction + v[vOffset + i];
      this.m = score;
    } else {
      const wgt = Math.exp(score - this.m);
      this.l += wgt;
      for (let i = 0; i < dim; i++) acc[i] += wgt * v[vOffset + i];
    }
  }

  /** Normalise into `out`. Safe to call once per query head. */
  finish(out: Float32Array, outOffset: number): void {
    const inv = this.l === 0 ? 0 : 1 / this.l;
    for (let i = 0; i < this.dim; i++) out[outOffset + i] = this.acc[i] * inv;
  }
}

/**
 * mini-picoLM — step 9: SwiGLU
 * ============================
 *
 * `SwiGLU(x) = (SiLU(x·W1) * (x·W3)) · W2`
 *
 * Three matrices where a classic FFN has two. The extra one is the *gate*: its
 * output multiplies the activation elementwise, so the network can suppress a
 * channel rather than merely fail to excite it.
 *
 * The cost is the reason `d_ff` looks odd in Llama configs. A two-matrix FFN at
 * `4 * d_model` and a three-matrix one at `(8/3) * d_model` hold the same
 * parameter count, which is why you see 5632 next to 2048 rather than 8192.
 */

/** SiLU, a.k.a. swish: x * sigmoid(x). */
export function silu(x: Float32Array, n: number = x.length): void {
  for (let i = 0; i < n; i++) x[i] = x[i] / (1 + Math.exp(-x[i]));
}

export function siluScalar(v: number): number {
  return v / (1 + Math.exp(-v));
}

/**
 * Fuse the activation and the gate into one pass.
 *
 * `hb` holds the gate projection, `hb2` the up projection. Doing this in one
 * loop rather than two keeps `hb` in cache between the two reads — the same
 * argument as kernel fusion in S13, at a scale where you can see the whole
 * thing.
 */
export function siluMulInto(hb: Float32Array, hb2: Float32Array, n: number = hb.length): void {
  for (let i = 0; i < n; i++) {
    const g = hb[i];
    hb[i] = (g / (1 + Math.exp(-g))) * hb2[i];
  }
}

/**
 * mini-picoLM — steps 15 & 16: Q8 and Q4
 * ======================================
 *
 * The GGUF block formats, in TypeScript. These are the same algorithms the S21
 * Python port implements, and that port was compiled against picoLM's own
 * `quant.c` and compared byte for byte on random blocks — so the numbers here
 * have a verified reference. `selftest.ts` re-checks them against blocks
 * written by hand from the spec.
 *
 * Two families:
 *
 *   **Q8_0 / Q4_0** — 32 weights, one fp16 scale. Simple, symmetric, and what
 *   "4-bit" meant before 2023.
 *
 *   **Q4_K / Q6_K** — 256-weight super-blocks with *nested* scales. Q4_K packs
 *   a 6-bit scale and min per 32-weight sub-block under one fp16 pair; Q6_K
 *   carries sixteen signed 8-bit scales, one per 16-weight sub-block, under a
 *   single fp16 scale. That nesting is what tracks outliers at four bits, and
 *   it is why a real 4-bit model is 4.5 bits per weight rather than 4.0.
 *
 * The awkward bit-packing in `getScaleMin` is the format, not the port: for the
 * first four sub-blocks the 6-bit fields sit in their own bytes, and for the
 * last four they straddle byte boundaries.
 */

export type QuantType = "F32" | "F16" | "Q4_0" | "Q8_0" | "Q4_K" | "Q6_K";

/** GGUF's numeric tensor-type tags. */
export const GGUF_TYPE: Record<number, QuantType> = {
  0: "F32",
  1: "F16",
  2: "Q4_0",
  8: "Q8_0",
  12: "Q4_K",
  14: "Q6_K",
};

const BLOCK: Record<QuantType, number> = {
  F32: 1,
  F16: 1,
  Q4_0: 32,
  Q8_0: 32,
  Q4_K: 256,
  Q6_K: 256,
};

const QSIZE: Record<QuantType, number> = {
  F32: 4,
  F16: 2,
  Q4_0: 18,
  Q8_0: 34,
  Q4_K: 144,
  Q6_K: 210,
};

export function blockSize(t: QuantType): number {
  return BLOCK[t];
}

export function quantSize(t: QuantType): number {
  return QSIZE[t];
}

export function rowSize(t: QuantType, n: number): number {
  return (n / BLOCK[t]) * QSIZE[t];
}

export function bitsPerWeight(t: QuantType): number {
  return (8 * QSIZE[t]) / BLOCK[t];
}

/* ------------------------------------------------------------------ fp16 */

/**
 * JavaScript has no Float16Array in every engine we care about yet, so decode
 * IEEE-754 binary16 by hand — the same bit surgery picoLM writes in C because
 * its target CPUs lack hardware fp16.
 */
export function fp16ToFp32(h: number): number {
  const sign = (h >> 15) & 1;
  const exp = (h >> 10) & 0x1f;
  const mant = h & 0x3ff;
  let value: number;
  if (exp === 0) {
    value = mant === 0 ? 0 : mant * 2 ** -24; // subnormal
  } else if (exp === 31) {
    value = mant === 0 ? Infinity : NaN;
  } else {
    value = (mant + 1024) * 2 ** (exp - 25);
  }
  return sign ? -value : value;
}

const _f32 = new Float32Array(1);
const _u32 = new Uint32Array(_f32.buffer);

export function fp32ToFp16(f: number): number {
  _f32[0] = f;
  const bits = _u32[0];
  const sign = (bits >>> 16) & 0x8000;
  const exp = ((bits >>> 23) & 0xff) - 127 + 15;
  const mant = bits & 0x7fffff;
  if (((bits >>> 23) & 0xff) === 0xff) return sign | 0x7c00 | (mant ? 0x200 : 0);
  if (exp >= 31) return sign | 0x7c00;
  if (exp <= 0) {
    if (exp < -10) return sign;
    const m = (mant | 0x800000) >>> (14 - exp);
    return sign | m;
  }
  return sign | (exp << 10) | (mant >>> 13);
}

/* ------------------------------------------------- one block at a time */

/** Q4_K's six-bit scale/min pair for sub-block `j`, unpacked from 12 bytes. */
function getScaleMin(q: Uint8Array, base: number, j: number): [number, number] {
  if (j < 4) {
    return [q[base + j] & 63, q[base + j + 4] & 63];
  }
  const sc = (q[base + j + 4] & 0xf) | ((q[base + j - 4] >> 6) << 4);
  const mn = (q[base + j + 4] >> 4) | ((q[base + j] >> 6) << 4);
  return [sc, mn];
}

/**
 * Decode block `index` of a quantised tensor into `out[0..blockSize)`.
 *
 * This is the primitive the fused matmul calls. It writes into a caller-owned
 * scratch buffer so that nothing the size of the weight matrix is ever
 * allocated — the whole point of fusing.
 */
export function dequantBlockInto(
  raw: Uint8Array,
  index: number,
  qtype: QuantType,
  out: Float32Array,
): void {
  const off = index * QSIZE[qtype];
  const dv = new DataView(raw.buffer, raw.byteOffset, raw.byteLength);

  switch (qtype) {
    case "F32": {
      out[0] = dv.getFloat32(off, true);
      return;
    }
    case "F16": {
      out[0] = fp16ToFp32(dv.getUint16(off, true));
      return;
    }
    case "Q8_0": {
      const d = fp16ToFp32(dv.getUint16(off, true));
      for (let i = 0; i < 32; i++) out[i] = d * dv.getInt8(off + 2 + i);
      return;
    }
    case "Q4_0": {
      const d = fp16ToFp32(dv.getUint16(off, true));
      for (let i = 0; i < 16; i++) {
        const b = raw[off + 2 + i];
        out[i] = d * ((b & 0xf) - 8);
        out[i + 16] = d * ((b >> 4) - 8);
      }
      return;
    }
    case "Q4_K": {
      const d = fp16ToFp32(dv.getUint16(off, true));
      const dmin = fp16ToFp32(dv.getUint16(off + 2, true));
      const sBase = off + 4;
      const qBase = off + 16;
      let o = 0;
      let qb = 0;
      for (let j = 0; j < 4; j++) {
        const [sc1, mn1] = getScaleMin(raw, sBase, j * 2);
        const [sc2, mn2] = getScaleMin(raw, sBase, j * 2 + 1);
        const d1 = d * sc1;
        const m1 = dmin * mn1;
        const d2 = d * sc2;
        const m2 = dmin * mn2;
        for (let l = 0; l < 32; l++) {
          const b = raw[qBase + qb + l];
          out[o + l] = d1 * (b & 0xf) - m1;
          out[o + l + 32] = d2 * (b >> 4) - m2;
        }
        o += 64;
        qb += 32;
      }
      return;
    }
    case "Q6_K": {
      const d = fp16ToFp32(dv.getUint16(off + 208, true));
      const qlBase = off;
      const qhBase = off + 128;
      const scBase = off + 192;
      for (let c = 0; c < 2; c++) {
        const ql = qlBase + c * 64;
        const qh = qhBase + c * 32;
        const isBase = c * 8;
        const y = c * 128;
        for (let l = 0; l < 32; l++) {
          const isL = isBase + (l >> 4);
          const lo = raw[ql + l];
          const hi = raw[ql + l + 32];
          const h = raw[qh + l];
          out[y + l] = d * dv.getInt8(scBase + isL) * (((lo & 0xf) | (((h >> 0) & 3) << 4)) - 32);
          out[y + l + 32] =
            d * dv.getInt8(scBase + isL + 2) * (((hi & 0xf) | (((h >> 2) & 3) << 4)) - 32);
          out[y + l + 64] =
            d * dv.getInt8(scBase + isL + 4) * (((lo >> 4) | (((h >> 4) & 3) << 4)) - 32);
          out[y + l + 96] =
            d * dv.getInt8(scBase + isL + 6) * (((hi >> 4) | (((h >> 6) & 3) << 4)) - 32);
        }
      }
      return;
    }
    default: {
      const never: never = qtype;
      throw new Error(`unsupported quant type ${String(never)}`);
    }
  }
}

/**
 * Decode `n` consecutive weights.
 *
 * Convenient, and exactly what a hot loop must not do — it materialises the
 * float copy that fusing exists to avoid. Used for the norm weights (tiny) and
 * by the tests.
 */
export function dequantizeRow(raw: Uint8Array, n: number, qtype: QuantType): Float32Array {
  const bs = BLOCK[qtype];
  if (n % bs !== 0) {
    throw new RangeError(`${n} weights is not a whole number of ${qtype} blocks`);
  }
  const out = new Float32Array(n);
  const scratch = new Float32Array(bs);
  const nb = n / bs;
  for (let b = 0; b < nb; b++) {
    dequantBlockInto(raw, b, qtype, scratch);
    out.set(scratch, b * bs);
  }
  return out;
}

/**
 * mini-picoLM — step 6: RoPE
 * ==========================
 *
 * Rotary position embedding. Instead of adding a position vector, rotate each
 * pair of dimensions by an angle proportional to the position. The dot product
 * of two rotated vectors then depends only on the *gap* between their
 * positions — which is exactly why a KV cache is sound: a cached key's rotation
 * is fixed forever, because its absolute position never changes.
 *
 * There are two incompatible conventions, both called RoPE:
 *
 *   **interleaved** — pair `x[2i]` with `x[2i+1]`. llama.cpp, picoLM, and every
 *   Llama GGUF file.
 *
 *   **NeoX / split-half** — pair `x[i]` with `x[i + D/2]`. GPT-NeoX, and the
 *   one S03 teaches.
 *
 * Pick the wrong one and the model runs, sounds fluent, and is subtly wrong.
 * This is the interleaved one, because this engine loads GGUF.
 */

/** Precomputed cos/sin for every position — keeps trig out of the token loop. */
export class RopeTables {
  readonly cos: Float32Array;
  readonly sin: Float32Array;
  readonly half: number;

  constructor(maxSeq: number, headDim: number, freqBase = 10000) {
    this.half = headDim >> 1;
    this.cos = new Float32Array(maxSeq * this.half);
    this.sin = new Float32Array(maxSeq * this.half);
    for (let pos = 0; pos < maxSeq; pos++) {
      for (let i = 0; i < this.half; i++) {
        const theta = pos / Math.pow(freqBase, (2 * i) / headDim);
        this.cos[pos * this.half + i] = Math.cos(theta);
        this.sin[pos * this.half + i] = Math.sin(theta);
      }
    }
  }
}

/**
 * Rotate `x` in place: `nHeads` heads of `headDim` laid out contiguously.
 *
 * Applied to Q and K, never to V — V carries content, not position.
 */
export function ropeInPlace(
  x: Float32Array,
  nHeads: number,
  headDim: number,
  tables: RopeTables,
  pos: number,
): void {
  const half = tables.half;
  const base = pos * half;
  for (let h = 0; h < nHeads; h++) {
    const o = h * headDim;
    for (let i = 0; i < half; i++) {
      const c = tables.cos[base + i];
      const s = tables.sin[base + i];
      const x0 = x[o + i * 2];
      const x1 = x[o + i * 2 + 1];
      x[o + i * 2] = x0 * c - x1 * s;
      x[o + i * 2 + 1] = x0 * s + x1 * c;
    }
  }
}

/**
 * mini-picoLM — step 8: the KV cache
 * ==================================
 *
 * One flat region of the arena, indexed `[layer][position][kv_dim]`. Writing
 * position `p` is a `set()` into a subarray; reading positions `0..p` is a
 * subarray. No allocation on either path.
 *
 * Two properties do all the work:
 *
 *   **Keys are stored after RoPE.** A cached token's position never changes, so
 *   its rotation never changes. Storing pre-rotation would force re-rotating
 *   the whole cache every step — replacing one compute cost with a smaller one
 *   and defeating most of the point.
 *
 *   **Entries are immutable once written.** Writing position N must not disturb
 *   anything before it. `selftest.ts` asserts this, because a cache that
 *   silently corrupts history produces text that degrades over hundreds of
 *   tokens with no error anywhere.
 *
 * Unlike picoLM this stores float32, not fp16 — `Float16Array` is not yet
 * available across the browsers this has to run in. That doubles the cache, and
 * since the cache is the resident cost, it is the main reason the context
 * length here is smaller than picoLM's on the same device. The arithmetic is in
 * the S22 simulator.
 */

import type { Arena } from "../memory/arena.ts";

export class KVCache {
  readonly key: Float32Array;
  readonly value: Float32Array;
  readonly nLayers: number;
  readonly maxSeq: number;
  readonly kvDim: number;
  private _length = 0;

  constructor(arena: Arena, nLayers: number, maxSeq: number, kvDim: number) {
    this.nLayers = nLayers;
    this.maxSeq = maxSeq;
    this.kvDim = kvDim;
    const n = nLayers * maxSeq * kvDim;
    this.key = arena.f32(n, "kv.key");
    this.value = arena.f32(n, "kv.value");
  }

  /** Positions written so far. */
  get length(): number {
    return this._length;
  }

  get bytes(): number {
    return this.key.byteLength + this.value.byteLength;
  }

  private offset(layer: number, pos: number): number {
    return (layer * this.maxSeq + pos) * this.kvDim;
  }

  /** Append this layer's K and V for position `pos`. */
  write(layer: number, pos: number, k: Float32Array, v: Float32Array): void {
    if (pos >= this.maxSeq) {
      throw new RangeError(`position ${pos} past the ${this.maxSeq}-token cache`);
    }
    const o = this.offset(layer, pos);
    this.key.set(k, o);
    this.value.set(v, o);
    if (layer === this.nLayers - 1) this._length = Math.max(this._length, pos + 1);
  }

  /** Byte offset of one head's K vector — attention reads through this. */
  headOffset(layer: number, pos: number, kvHead: number, headDim: number): number {
    return this.offset(layer, pos) + kvHead * headDim;
  }

  reset(): void {
    this._length = 0;
    this.key.fill(0);
    this.value.fill(0);
  }
}

/**
 * mini-picoLM — step 7: attention
 * ===============================
 *
 * Grouped-query attention over the KV cache, computed with an online softmax.
 *
 * The shape of this function is the whole argument of S12. There is no `scores`
 * array. At no context length does this allocate anything proportional to the
 * sequence — it streams over the cache once per head, folding each (score,
 * value) pair into a running accumulator. Peak intermediate memory is
 * `head_dim` floats, whatever the context.
 *
 * GQA is the `kvHead = h / kvMul` line. Several query heads read the same K and
 * V rows, so the cache shrinks by that factor while the query side keeps its
 * full expressiveness. It is a memory optimisation, not a quality one.
 */

import { OnlineSoftmax } from "../tensor/softmax.ts";
import type { KVCache } from "../runtime/kv-cache.ts";

export type AttentionShape = {
  nHeads: number;
  nKvHeads: number;
  headDim: number;
};

/**
 * out[0..nHeads*headDim) = attention(q, cache[layer][0..=pos])
 *
 * `scratch` must hold at least `headDim` floats and is reused across heads.
 */
export function attention(
  out: Float32Array,
  q: Float32Array,
  cache: KVCache,
  layer: number,
  pos: number,
  shape: AttentionShape,
  scratch: Float32Array,
): void {
  const { nHeads, nKvHeads, headDim } = shape;
  const kvMul = nHeads / nKvHeads;
  const inv = 1 / Math.sqrt(headDim);
  const online = new OnlineSoftmax(headDim, scratch);

  for (let h = 0; h < nHeads; h++) {
    const kvHead = Math.floor(h / kvMul);
    const qo = h * headDim;
    online.reset();

    for (let t = 0; t <= pos; t++) {
      const ko = cache.headOffset(layer, t, kvHead, headDim);
      let score = 0;
      for (let i = 0; i < headDim; i++) score += q[qo + i] * cache.key[ko + i];
      online.push(score * inv, cache.value, ko);
    }

    online.finish(out, qo);
  }
}

/**
 * mini-picoLM — the embedding lookup
 * ==================================
 *
 * A row of the embedding table, dequantised on demand.
 *
 * Worth noticing how little happens here. `token_embd.weight` is usually the
 * largest single tensor in the file — 32000 × 2048 for TinyLlama — and a decode
 * step touches exactly one of its rows. Dequantising the whole table at load
 * time would cost hundreds of megabytes to save a few microseconds per token.
 * So it stays quantised, mapped, and untouched except for the one row.
 */

import { dequantBlockInto, blockSize, rowSize, type QuantType } from "../format/quant.ts";

/**
 * Write row `token` of a [vocab, dim] quantised table into `out`.
 *
 * `out` is a caller-owned view — the residual stream itself, in practice, so
 * the embedding is written straight into the buffer the first layer reads.
 */
export function embedInto(
  out: Float32Array,
  table: Uint8Array,
  token: number,
  dim: number,
  qtype: QuantType,
  scratch: Float32Array,
): void {
  const bs = blockSize(qtype);
  const perRow = rowSize(qtype, dim);
  const rowStart = token * perRow;
  const blocksPerRow = dim / bs;
  // `dequantBlockInto` indexes blocks from the start of the view it is given,
  // so hand it a view that starts at this row.
  const row = table.subarray(rowStart, rowStart + perRow);

  for (let b = 0; b < blocksPerRow; b++) {
    dequantBlockInto(row, b, qtype, scratch);
    out.set(scratch.subarray(0, bs), b * bs);
  }
}

/**
 * mini-picoLM — the feed-forward network
 * ======================================
 *
 * SwiGLU: gate and up in parallel, activate-and-multiply, then down.
 *
 * This is where most of the model's parameters live — roughly 70% of a Llama
 * block — which makes it the dominant term in decode bandwidth. Everything the
 * quantisation chapters buy, they buy mostly here.
 */

import { siluMulInto } from "../tensor/activation.ts";
import type { QuantWeight } from "./llama.ts";
import type { Workspace } from "../runtime/workspace.ts";

export function feedForward(
  out: Float32Array,
  x: Float32Array,
  gate: QuantWeight,
  up: QuantWeight,
  down: QuantWeight,
  dim: number,
  hidden: number,
  ws: Workspace,
): void {
  ws.matmul(ws.hb, x, gate, dim, hidden);
  ws.matmul(ws.hb2, x, up, dim, hidden);
  siluMulInto(ws.hb, ws.hb2, hidden);
  ws.matmul(out, ws.hb, down, hidden, dim);
}

/**
 * mini-picoLM — the workspace
 * ===========================
 *
 * Every scratch buffer the forward pass needs, allocated once from the arena.
 *
 * The list is short and fixed because the shapes are known before the first
 * token: one residual stream, two same-size temporaries, a query vector, two
 * FFN-width buffers, a logits row, and one block of dequantisation scratch.
 * That is the entire working set of a decode step.
 *
 * `matmul` lives here rather than in `tensor/` because it needs the scratch
 * buffer to fuse against. Routing every weight through this one method also
 * makes the WebGPU backend a single substitution rather than a rewrite.
 */

import type { Arena } from "../memory/arena.ts";
import { matmul, matmulQuant } from "../tensor/matmul.ts";
import { blockSize } from "../format/quant.ts";
import type { QuantWeight } from "../model/llama.ts";

export class Workspace {
  readonly x: Float32Array; // residual stream
  readonly xb: Float32Array; // post-norm / attention out
  readonly xb2: Float32Array; // second temporary
  readonly q: Float32Array;
  readonly k: Float32Array;
  readonly v: Float32Array;
  readonly hb: Float32Array; // ffn hidden (gate)
  readonly hb2: Float32Array; // ffn hidden (up)
  readonly logits: Float32Array;
  readonly attScratch: Float32Array;
  readonly qScratch: Float32Array; // one dequantised block

  /** Counted so the simulator can show where the time goes. */
  matmulCalls = 0;
  matmulMacs = 0;

  constructor(arena: Arena, dim: number, hidden: number, kvDim: number, vocab: number) {
    this.x = arena.f32(dim, "ws.x");
    this.xb = arena.f32(dim, "ws.xb");
    this.xb2 = arena.f32(dim, "ws.xb2");
    this.q = arena.f32(dim, "ws.q");
    this.k = arena.f32(kvDim, "ws.k");
    this.v = arena.f32(kvDim, "ws.v");
    this.hb = arena.f32(hidden, "ws.hb");
    this.hb2 = arena.f32(hidden, "ws.hb2");
    this.logits = arena.f32(vocab, "ws.logits");
    this.attScratch = arena.f32(Math.max(dim, 256), "ws.att");
    this.qScratch = arena.f32(256, "ws.dequant"); // one Q4_K/Q6_K super-block
  }

  /**
   * The single matmul entry point. Dispatches on how the weight is stored.
   *
   * Synchronous by design — see `tensor/webgpu.ts` for why the GPU path cannot
   * live behind this signature, and why that is the honest answer rather than a
   * limitation to route around.
   */
  matmul(out: Float32Array, x: Float32Array, w: QuantWeight, n: number, d: number): void {
    this.matmulCalls++;
    this.matmulMacs += n * d;

    if (w.kind === "f32") {
      matmul(out, x, w.data, n, d);
      return;
    }
    const bs = blockSize(w.qtype);
    matmulQuant(out, x, w.raw, n, d, w.qtype, this.qScratch.subarray(0, bs));
  }

  resetCounters(): void {
    this.matmulCalls = 0;
    this.matmulMacs = 0;
  }
}

/**
 * mini-picoLM — step 10: the transformer block
 * ============================================
 *
 * Seven operations, twice through the residual stream:
 *
 *     x += attn(rmsnorm(x))
 *     x += ffn(rmsnorm(x))
 *
 * Everything before this file was a part; this is where they become a layer.
 * Note that `x` is never copied — both sublayers read it, write into scratch,
 * and add back. The residual stream is one buffer for the whole forward pass.
 */

import { rmsnorm } from "../tensor/rmsnorm.ts";
import { addInto } from "../tensor/matmul.ts";
import { ropeInPlace, type RopeTables } from "./rope.ts";
import { attention } from "./attention.ts";
import { feedForward } from "./mlp.ts";
import type { LayerWeights, ModelConfig } from "./llama.ts";
import type { Workspace } from "../runtime/workspace.ts";
import type { KVCache } from "../runtime/kv-cache.ts";

export function forwardBlock(
  layer: number,
  pos: number,
  cfg: ModelConfig,
  w: LayerWeights,
  ws: Workspace,
  cache: KVCache,
  rope: RopeTables,
): void {
  const { dim, nHeads, nKvHeads, headDim, kvDim, hidden } = cfg;

  /* ---- attention ---- */
  rmsnorm(ws.xb, ws.x, w.attnNorm);

  ws.matmul(ws.q, ws.xb, w.attnQ, dim, dim);
  ws.matmul(ws.k, ws.xb, w.attnK, dim, kvDim);
  ws.matmul(ws.v, ws.xb, w.attnV, dim, kvDim);

  // Rotate before caching: a cached key's position is fixed forever.
  ropeInPlace(ws.q, nHeads, headDim, rope, pos);
  ropeInPlace(ws.k, nKvHeads, headDim, rope, pos);

  cache.write(layer, pos, ws.k, ws.v);

  attention(ws.xb, ws.q, cache, layer, pos, { nHeads, nKvHeads, headDim }, ws.attScratch);

  ws.matmul(ws.xb2, ws.xb, w.attnOut, dim, dim);
  addInto(ws.x, ws.xb2);

  /* ---- feed-forward ---- */
  rmsnorm(ws.xb, ws.x, w.ffnNorm);
  feedForward(ws.xb2, ws.xb, w.ffnGate, w.ffnUp, w.ffnDown, dim, hidden, ws);
  addInto(ws.x, ws.xb2);
}

/**
 * mini-picoLM — the model
 * =======================
 *
 * Config, weights, and `forward(token, pos)`.
 *
 * That signature is the whole course in one line. One token per call is the S01
 * loop; `pos` indexing a persistent cache is S05; the weights arriving as
 * quantised bytes is S08; the online softmax inside attention is S12;
 * `nKvHeads < nHeads` is S03's GQA.
 *
 * There is no batch dimension anywhere, which is why there is no scheduler, no
 * block manager and no continuous batching. One user, one tab.
 */

import { Arena } from "../memory/arena.ts";
import { rmsnorm } from "../tensor/rmsnorm.ts";
import { embedInto } from "./embedding.ts";
import { forwardBlock } from "./block.ts";
import { RopeTables } from "./rope.ts";
import { KVCache } from "../runtime/kv-cache.ts";
import { Workspace } from "../runtime/workspace.ts";
import { blockSize, rowSize, type QuantType } from "../format/quant.ts";

export type ModelConfig = {
  dim: number;
  hidden: number;
  nHeads: number;
  nKvHeads: number;
  nLayers: number;
  vocabSize: number;
  maxSeq: number;
  ropeFreqBase: number;
  headDim: number;
  kvDim: number;
};

export function makeConfig(
  p: Omit<ModelConfig, "headDim" | "kvDim"> & Partial<Pick<ModelConfig, "headDim" | "kvDim">>,
): ModelConfig {
  const headDim = p.headDim ?? p.dim / p.nHeads;
  return { ...p, headDim, kvDim: p.kvDim ?? p.nKvHeads * headDim };
}

/** A weight is either a plain float view or quantised bytes plus its type. */
export type QuantWeight =
  | { kind: "f32"; data: Float32Array }
  | { kind: "quant"; raw: Uint8Array; qtype: QuantType };

export type LayerWeights = {
  attnNorm: Float32Array;
  ffnNorm: Float32Array;
  attnQ: QuantWeight;
  attnK: QuantWeight;
  attnV: QuantWeight;
  attnOut: QuantWeight;
  ffnGate: QuantWeight;
  ffnUp: QuantWeight;
  ffnDown: QuantWeight;
};

export type ModelWeights = {
  tokenEmbd: { raw: Uint8Array; qtype: QuantType };
  outputNorm: Float32Array;
  output: QuantWeight;
  layers: LayerWeights[];
};

export class LlamaModel {
  readonly cfg: ModelConfig;
  readonly weights: ModelWeights;
  readonly arena: Arena;
  readonly cache: KVCache;
  readonly ws: Workspace;
  readonly rope: RopeTables;

  constructor(cfg: ModelConfig, weights: ModelWeights, arena?: Arena) {
    this.cfg = cfg;
    this.weights = weights;
    this.arena = arena ?? new Arena(LlamaModel.arenaBytes(cfg));
    this.cache = new KVCache(this.arena, cfg.nLayers, cfg.maxSeq, cfg.kvDim);
    this.ws = new Workspace(this.arena, cfg.dim, cfg.hidden, cfg.kvDim, cfg.vocabSize);
    this.rope = new RopeTables(cfg.maxSeq, cfg.headDim, cfg.ropeFreqBase);
  }

  /** Exactly how much the arena needs — computed, not guessed. */
  static arenaBytes(cfg: ModelConfig): number {
    const kv = 2 * cfg.nLayers * cfg.maxSeq * cfg.kvDim * 4;
    const ws = (4 * cfg.dim + 2 * cfg.kvDim + 2 * cfg.hidden + cfg.vocabSize + 512) * 4;
    const slack = 64 * 1024; // alignment padding, one page of headroom
    return kv + ws + slack;
  }

  get kvBytes(): number {
    return this.cache.bytes;
  }

  reset(): void {
    this.cache.reset();
    this.ws.x.fill(0);
  }

  /** One token in, one row of logits out. */
  forward(token: number, pos: number): Float32Array {
    const { cfg, ws, weights } = this;

    embedInto(
      ws.x,
      weights.tokenEmbd.raw,
      token,
      cfg.dim,
      weights.tokenEmbd.qtype,
      ws.qScratch.subarray(0, blockSize(weights.tokenEmbd.qtype)),
    );

    for (let l = 0; l < cfg.nLayers; l++) {
      forwardBlock(l, pos, cfg, weights.layers[l], ws, this.cache, this.rope);
    }

    rmsnorm(ws.xb, ws.x, weights.outputNorm);
    ws.matmul(ws.logits, ws.xb, weights.output, cfg.dim, cfg.vocabSize);
    return ws.logits;
  }
}

/** Bytes a quantised tensor of `n` weights occupies — used by the loader. */
export function weightBytes(qtype: QuantType, n: number): number {
  return rowSize(qtype, n);
}

/**
 * mini-picoLM — step 11: the BPE tokenizer
 * ========================================
 *
 * SentencePiece-style BPE, the family Llama GGUF files carry. Every token has a
 * **score**, and the merge with the highest score wins — not the lowest rank,
 * which is the GPT/tiktoken convention.
 *
 * Three details that are easy to get wrong and hard to notice:
 *
 *   * spaces become `▁` (U+2581) and one is *prepended*, so "Once" and " Once"
 *     tokenize identically;
 *   * anything not in the vocabulary falls back to a `<0xHH>` byte token, which
 *     is what makes the vocabulary total;
 *   * detokenization has to be **incremental and byte-aware**. A multi-byte
 *     character can span several tokens, so emitting each token's bytes as they
 *     arrive produces replacement characters in the UI. `StreamDecoder` holds
 *     partial sequences back.
 */

export const SPACE = "▁";
const BYTE_TOKEN = /^<0x([0-9A-Fa-f]{2})>$/;

export class Tokenizer {
  readonly tokens: string[];
  readonly scores: Float32Array;
  readonly bosId: number;
  readonly eosId: number;
  private readonly lookup = new Map<string, number>();
  private readonly byteOf = new Map<number, number>();
  private readonly encoder = new TextEncoder();

  constructor(tokens: string[], scores: Float32Array, bosId = 1, eosId = 2) {
    this.tokens = tokens;
    this.scores = scores;
    this.bosId = bosId;
    this.eosId = eosId;
    for (let i = 0; i < tokens.length; i++) {
      if (!this.lookup.has(tokens[i])) this.lookup.set(tokens[i], i);
      const m = BYTE_TOKEN.exec(tokens[i]);
      if (m) this.byteOf.set(i, parseInt(m[1], 16));
    }
  }

  get vocabSize(): number {
    return this.tokens.length;
  }

  /** One token per character, with a byte fallback for anything unknown. */
  private seed(text: string): number[] {
    const out: number[] = [];
    for (const ch of text) {
      const direct = this.lookup.get(ch);
      if (direct !== undefined) {
        out.push(direct);
        continue;
      }
      for (const b of this.encoder.encode(ch)) {
        const bt = this.lookup.get(`<0x${b.toString(16).toUpperCase().padStart(2, "0")}>`);
        if (bt !== undefined) out.push(bt);
      }
    }
    return out;
  }

  encode(text: string, addBos = true): number[] {
    const norm = SPACE + text.replaceAll(" ", SPACE);
    const parts = this.seed(norm);

    // Greedy merge. Quadratic in the number of pieces, which is fine: it runs
    // once per request on a short string.
    for (;;) {
      let bestScore = -Infinity;
      let bestIdx = -1;
      let bestId = -1;
      for (let i = 0; i < parts.length - 1; i++) {
        const merged = this.tokens[parts[i]] + this.tokens[parts[i + 1]];
        const id = this.lookup.get(merged);
        if (id !== undefined && this.scores[id] > bestScore) {
          bestScore = this.scores[id];
          bestIdx = i;
          bestId = id;
        }
      }
      if (bestIdx < 0) break;
      parts.splice(bestIdx, 2, bestId);
    }

    return addBos ? [this.bosId, ...parts] : parts;
  }

  /** The raw bytes a token contributes. */
  pieceBytes(id: number): Uint8Array {
    const b = this.byteOf.get(id);
    if (b !== undefined) return new Uint8Array([b]);
    return this.encoder.encode(this.tokens[id].replaceAll(SPACE, " "));
  }

  /** Human-readable text for a token — what the grammar masks against. */
  piece(id: number): string {
    const b = this.byteOf.get(id);
    if (b !== undefined) return String.fromCharCode(b);
    return this.tokens[id].replaceAll(SPACE, " ");
  }

  decode(ids: number[], stripLeadingSpace = true): string {
    let total = 0;
    const parts = ids.map((i) => {
      const p = this.pieceBytes(i);
      total += p.length;
      return p;
    });
    const buf = new Uint8Array(total);
    let o = 0;
    for (const p of parts) {
      buf.set(p, o);
      o += p.length;
    }
    const text = new TextDecoder().decode(buf);
    return stripLeadingSpace && text.startsWith(" ") ? text.slice(1) : text;
  }
}

/**
 * Emits text only once the pending bytes form a complete character.
 *
 * `TextDecoder` with `{ stream: true }` does exactly this natively — it buffers
 * an incomplete sequence internally and emits nothing until the continuation
 * bytes arrive. Using it is both shorter and more correct than hand-rolling.
 */
export class StreamDecoder {
  private readonly decoder = new TextDecoder("utf-8");
  private readonly tok: Tokenizer;

  constructor(tok: Tokenizer) {
    this.tok = tok;
  }

  push(id: number): string {
    return this.decoder.decode(this.tok.pieceBytes(id), { stream: true });
  }

  /** Flush anything still buffered at end of stream. */
  finish(): string {
    return this.decoder.decode(new Uint8Array(0));
  }
}