// WebGPU inference engine for Qwen3.5-0.8B (QwenTheBard)
// GPU handles large matrix-vector multiplies; JS handles norms, activations, attention, DeltaNet recurrence.

// ============================================================
// Model constants
// ============================================================

const VOCAB_SIZE = 248320;
const HIDDEN_SIZE = 1024;
const NUM_LAYERS = 24;
const INTERMEDIATE_SIZE = 3584;
const RMS_NORM_EPS = 1e-6;
const EOS_TOKEN = 248046;

// Full attention (layers 3,7,11,15,19,23)
const FULL_ATTN_LAYERS = new Set([3, 7, 11, 15, 19, 23]);
const NUM_HEADS = 8;
const NUM_KV_HEADS = 2;
const HEAD_DIM = 256;
const ROTARY_DIM = 64;
const ROPE_THETA = 10_000_000;
const KV_GROUPS = NUM_HEADS / NUM_KV_HEADS; // 4
const ATTN_OUTPUT_DIM = NUM_HEADS * HEAD_DIM; // 2048

// DeltaNet
const DN_NUM_HEADS = 16;
const DN_KEY_DIM = 128;
const DN_VALUE_DIM = 128;
const DN_CONV_KERNEL = 4;
const DN_CONV_DIM = DN_NUM_HEADS * (DN_KEY_DIM * 2 + DN_VALUE_DIM); // 6144
const DN_TOTAL_VALUE = DN_NUM_HEADS * DN_VALUE_DIM; // 2048
const DN_TOTAL_KEY = DN_NUM_HEADS * DN_KEY_DIM; // 2048

// Buffer sizing
const MAX_OUTPUT_SIZE = VOCAB_SIZE; // 248320 floats for logits
const MAX_INPUT_SIZE = DN_CONV_DIM; // 6144 floats (largest projection input after MLP gate*up)
const MAX_SEQ_LEN = 2048;

// ============================================================
// WGSL compute shader for Q4 matrix-vector multiply
// ============================================================

const Q4_MATVEC_SHADER = /* wgsl */ `
struct Params {
  M: u32,       // output rows
  N: u32,       // input cols
  groups_per_row: u32,
}

@group(0) @binding(0) var<storage, read> weights: array<u32>;
@group(0) @binding(1) var<storage, read> input: array<f32>;
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
@group(0) @binding(3) var<uniform> params: Params;

// Q4 group: 68 bytes = 17 u32s
// [0]: scale(f16) | zero(f16)  packed as u32
// [1..16]: 64 bytes of packed nibbles (128 elements, 2 per byte)

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let row = gid.x;
  if (row >= params.M) { return; }

  let groups_per_row = params.groups_per_row;
  let group_u32s = 17u; // 68 bytes / 4
  let row_offset = row * groups_per_row * group_u32s;

  var acc: f32 = 0.0;
  var col: u32 = 0u;

  for (var g: u32 = 0u; g < groups_per_row; g++) {
    let base = row_offset + g * group_u32s;

    // Unpack scale and zero from f16 pair
    let sz = weights[base];
    let scale_bits = sz & 0xFFFFu;
    let zero_bits = (sz >> 16u) & 0xFFFFu;
    let scale = unpack2x16float(sz).x;
    let zero = unpack2x16float(sz).y;

    // 64 bytes of nibbles = 16 u32s, each u32 has 8 nibbles = 8 values
    for (var w: u32 = 0u; w < 16u; w++) {
      let packed = weights[base + 1u + w];
      // Each u32 holds 8 nibbles (4 bytes, 2 nibbles per byte)
      for (var n: u32 = 0u; n < 8u; n++) {
        let nibble = (packed >> (n * 4u)) & 0xFu;
        let val = f32(nibble) * scale + zero;
        acc += val * input[col];
        col++;
      }
    }
  }

  output[row] = acc;
}
`;

// ============================================================
// Q4MF binary format parser
// ============================================================

interface TensorMeta {
  shape: number[];
  dtype: number; // 0=Q4, 1=F32, 2=BF16
  data: Uint8Array;
}

function parseQ4MF(buffer: ArrayBuffer): Map<string, TensorMeta> {
  const view = new DataView(buffer);
  const bytes = new Uint8Array(buffer);
  let offset = 0;

  // Header: magic(u32), version(u32), num_tensors(u32), group_size(u32)
  const magic = view.getUint32(offset, true); offset += 4;
  if (magic !== 0x46344D51) throw new Error(`Invalid Q4MF magic: 0x${magic.toString(16)}`);
  const version = view.getUint32(offset, true); offset += 4;
  if (version !== 1) throw new Error(`Unsupported Q4MF version: ${version}`);
  const numTensors = view.getUint32(offset, true); offset += 4;
  /* const groupSize = */ view.getUint32(offset, true); offset += 4;

  const tensors = new Map<string, TensorMeta>();
  const decoder = new TextDecoder();

  for (let t = 0; t < numTensors; t++) {
    const nameLen = view.getUint16(offset, true); offset += 2;
    const name = decoder.decode(bytes.subarray(offset, offset + nameLen)); offset += nameLen;
    const ndim = bytes[offset]; offset += 1;
    const shape: number[] = [];
    for (let d = 0; d < ndim; d++) {
      shape.push(view.getUint32(offset, true)); offset += 4;
    }
    const dtype = bytes[offset]; offset += 1;
    const dataLen = view.getUint32(offset, true); offset += 4;
    const data = bytes.subarray(offset, offset + dataLen);
    offset += dataLen;
    tensors.set(name, { shape, dtype, data });
  }

  return tensors;
}

// ============================================================
// CPU helpers
// ============================================================

function bf16ToF32(bits: number): number {
  const buf = new ArrayBuffer(4);
  const u32 = new Uint32Array(buf);
  u32[0] = bits << 16;
  return new Float32Array(buf)[0];
}

function f16ToF32(bits: number): number {
  const sign = (bits >> 15) & 1;
  const exp = (bits >> 10) & 0x1F;
  const frac = bits & 0x3FF;
  if (exp === 0) {
    // Subnormal or zero
    return (sign ? -1 : 1) * Math.pow(2, -14) * (frac / 1024);
  }
  if (exp === 31) {
    return frac ? NaN : (sign ? -Infinity : Infinity);
  }
  return (sign ? -1 : 1) * Math.pow(2, exp - 15) * (1 + frac / 1024);
}

/** Convert a tensor's raw data to Float32Array based on dtype. */
function toF32Array(meta: TensorMeta): Float32Array {
  if (meta.dtype === 1) {
    // F32 — copy to aligned buffer since subarray offset may not be 4-byte aligned
    const out = new Float32Array(meta.data.byteLength / 4);
    out.set(new Float32Array(meta.data.slice().buffer));
    return out;
  }
  if (meta.dtype === 2) {
    // BF16
    const count = meta.data.byteLength / 2;
    const out = new Float32Array(count);
    const view = new DataView(meta.data.buffer, meta.data.byteOffset, meta.data.byteLength);
    for (let i = 0; i < count; i++) {
      out[i] = bf16ToF32(view.getUint16(i * 2, true));
    }
    return out;
  }
  throw new Error(`Cannot directly convert Q4 tensor to F32`);
}

/** Dequantize one row of a Q4 tensor on CPU. */
function dequantQ4Row(data: Uint8Array, row: number, numCols: number, groupSize: number): Float32Array {
  const groupsPerRow = Math.ceil(numCols / groupSize);
  const bytesPerGroup = 68; // 2(scale) + 2(zero) + 64(nibbles)
  const rowOffset = row * groupsPerRow * bytesPerGroup;
  const out = new Float32Array(numCols);
  const view = new DataView(data.buffer, data.byteOffset);
  let col = 0;

  for (let g = 0; g < groupsPerRow; g++) {
    const base = rowOffset + g * bytesPerGroup;
    const scale = f16ToF32(view.getUint16(base, true));
    const zero = f16ToF32(view.getUint16(base + 2, true));

    // 64 bytes of packed nibbles
    for (let b = 0; b < 64 && col < numCols; b++) {
      const byte = data[base + 4 + b];
      const lo = byte & 0xF;
      const hi = (byte >> 4) & 0xF;
      out[col++] = lo * scale + zero;
      if (col < numCols) {
        out[col++] = hi * scale + zero;
      }
    }
  }

  return out;
}

function sigmoid(x: number): number {
  return 1.0 / (1.0 + Math.exp(-x));
}

function silu(x: number): number {
  return x * sigmoid(x);
}

function softplus(x: number): number {
  return Math.log(1.0 + Math.exp(x));
}

/** RMSNorm: out[i] = x[i] * rsqrt(mean(x^2) + eps) * (1 + weight[i]) */
function rmsNorm(out: Float32Array, x: Float32Array, weight: Float32Array, eps: number, offset = 0, len?: number): void {
  const n = len ?? x.length;
  let sumSq = 0;
  for (let i = 0; i < n; i++) sumSq += x[offset + i] * x[offset + i];
  const rms = 1.0 / Math.sqrt(sumSq / n + eps);
  for (let i = 0; i < n; i++) {
    out[offset + i] = x[offset + i] * rms * (1.0 + weight[i]);
  }
}

/** Gated RMSNorm: out = rmsNorm(x, weight_direct) * silu(gate). weight used directly (NOT 1+w). */
function rmsNormGated(out: Float32Array, outOff: number, x: Float32Array, xOff: number,
  gate: Float32Array, gateOff: number, weight: Float32Array, eps: number, len: number): void {
  let sumSq = 0;
  for (let i = 0; i < len; i++) sumSq += x[xOff + i] * x[xOff + i];
  const rms = 1.0 / Math.sqrt(sumSq / len + eps);
  for (let i = 0; i < len; i++) {
    out[outOff + i] = x[xOff + i] * rms * weight[i] * silu(gate[gateOff + i]);
  }
}

/** In-place L2 normalization. */
function l2Norm(x: Float32Array, offset: number, len: number): void {
  let sumSq = 0;
  for (let i = 0; i < len; i++) sumSq += x[offset + i] * x[offset + i];
  const inv = 1.0 / (Math.sqrt(sumSq) + 1e-12);
  for (let i = 0; i < len; i++) x[offset + i] *= inv;
}

/** Apply rotary position embedding (rotate_half style) to first rotaryDim elements. */
function applyRope(x: Float32Array, xOff: number, cos: Float32Array, sin: Float32Array,
  cosOff: number, rotaryDim: number): void {
  const half = rotaryDim / 2;
  for (let i = 0; i < half; i++) {
    const x0 = x[xOff + i];
    const x1 = x[xOff + i + half];
    x[xOff + i] = x0 * cos[cosOff + i] - x1 * sin[cosOff + i];
    x[xOff + i + half] = x1 * cos[cosOff + i + half] + x0 * sin[cosOff + i + half];
  }
}

// ============================================================
// GPUEngine
// ============================================================

const GROUP_SIZE = 128;

export class GPUEngine {
  // WebGPU objects
  private device!: GPUDevice;
  private q4Pipeline!: GPUComputePipeline;
  private bindGroupLayout!: GPUBindGroupLayout;

  // Weight buffers on GPU (Q4 data)
  private weightBuffers = new Map<string, GPUBuffer>();

  // Parsed weight metadata
  private weightMeta = new Map<string, { shape: number[]; dtype: number }>();

  // Small unquantized weights on CPU
  private cpuWeights = new Map<string, Float32Array>();

  // Embedding Q4 data kept on CPU for row-level dequant
  private embedData!: Uint8Array;
  private embedCols!: number;

  // Working GPU buffers
  private inputBuf!: GPUBuffer;
  private outputBufs: GPUBuffer[] = [];
  private paramsBufs: GPUBuffer[] = [];
  private stagingBuf!: GPUBuffer;

  // CPU-side state vectors
  private hidden = new Float32Array(HIDDEN_SIZE);
  private residual = new Float32Array(HIDDEN_SIZE);
  private normed = new Float32Array(HIDDEN_SIZE);

  // KV caches for full attention layers (6 layers)
  // Each stores [maxSeq * kvHeads * headDim]
  private kvKeys: Float32Array[] = [];
  private kvValues: Float32Array[] = [];
  private kvLen: number[] = [];

  // DeltaNet states (18 layers)
  // Recurrent state: [heads * keyDim * valueDim] per layer
  private deltaStates: Float32Array[] = [];
  // Conv state: [convDim * (kernel-1)] per layer
  private convStates: Float32Array[] = [];

  // RoPE tables
  private ropeCos!: Float32Array;
  private ropeSin!: Float32Array;

  private constructor() {}

  /** Initialize WebGPU and create pipeline. Returns null if WebGPU unavailable. */
  static async create(): Promise<GPUEngine | null> {
    if (!navigator.gpu) {
      console.warn("[gpu] navigator.gpu not available (need secure context: localhost or HTTPS)");
      return null;
    }
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) {
      console.warn("[gpu] No GPU adapter found — check chrome://gpu and chrome://flags/#enable-unsafe-webgpu");
      return null;
    }
    const device = await adapter.requestDevice({
      requiredLimits: {
        maxStorageBufferBindingSize: 256 * 1024 * 1024, // 256MB for large Q4 weight tensors
        maxBufferSize: 256 * 1024 * 1024, // 256MB
      },
    });

    const engine = new GPUEngine();
    engine.device = device;

    // Create bind group layout
    engine.bindGroupLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
        { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
      ],
    });

    // Create compute pipeline
    const shaderModule = device.createShaderModule({ code: Q4_MATVEC_SHADER });
    engine.q4Pipeline = device.createComputePipeline({
      layout: device.createPipelineLayout({ bindGroupLayouts: [engine.bindGroupLayout] }),
      compute: { module: shaderModule, entryPoint: 'main' },
    });

    // Pre-allocate working buffers
    engine.inputBuf = device.createBuffer({
      size: MAX_INPUT_SIZE * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });
    // Pool of output + param buffers for batched matvec (up to 4 concurrent ops)
    for (let i = 0; i < 4; i++) {
      engine.outputBufs.push(device.createBuffer({
        size: MAX_OUTPUT_SIZE * 4,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
      }));
      engine.paramsBufs.push(device.createBuffer({
        size: 16,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
      }));
    }
    engine.stagingBuf = device.createBuffer({
      size: MAX_OUTPUT_SIZE * 4,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });

    // Precompute RoPE tables
    engine.ropeCos = new Float32Array(MAX_SEQ_LEN * ROTARY_DIM);
    engine.ropeSin = new Float32Array(MAX_SEQ_LEN * ROTARY_DIM);
    const halfDim = ROTARY_DIM / 2;
    for (let pos = 0; pos < MAX_SEQ_LEN; pos++) {
      for (let i = 0; i < halfDim; i++) {
        const freq = 1.0 / Math.pow(ROPE_THETA, (2 * i) / ROTARY_DIM);
        const angle = pos * freq;
        const c = Math.cos(angle);
        const s = Math.sin(angle);
        engine.ropeCos[pos * ROTARY_DIM + i] = c;
        engine.ropeCos[pos * ROTARY_DIM + halfDim + i] = c;
        engine.ropeSin[pos * ROTARY_DIM + i] = s;
        engine.ropeSin[pos * ROTARY_DIM + halfDim + i] = s;
      }
    }

    // Initialize KV caches and DeltaNet states
    for (let layer = 0; layer < NUM_LAYERS; layer++) {
      if (FULL_ATTN_LAYERS.has(layer)) {
        engine.kvKeys.push(new Float32Array(MAX_SEQ_LEN * NUM_KV_HEADS * HEAD_DIM));
        engine.kvValues.push(new Float32Array(MAX_SEQ_LEN * NUM_KV_HEADS * HEAD_DIM));
        engine.kvLen.push(0);
      } else {
        engine.deltaStates.push(new Float32Array(DN_NUM_HEADS * DN_KEY_DIM * DN_VALUE_DIM));
        engine.convStates.push(new Float32Array(DN_CONV_DIM * (DN_CONV_KERNEL - 1)));
      }
    }

    return engine;
  }

  /** Parse Q4MF binary and upload Q4 weights to GPU, keep small weights on CPU. */
  async loadWeights(q4mfData: ArrayBuffer): Promise<void> {
    console.log(`[gpu] Parsing Q4MF (${(q4mfData.byteLength / 1024 / 1024).toFixed(1)} MB)...`);
    const tensors = parseQ4MF(q4mfData);
    console.log(`[gpu] ${tensors.size} tensors parsed, uploading to GPU...`);
    const prefix = 'model.language_model.';
    let gpuBytes = 0;
    let cpuCount = 0;

    // Small Q4 tensors that are accessed element-wise on CPU (not via GPU matvec)
    const cpuDequantSuffixes = ['conv1d.weight', 'A_log', 'dt_bias'];

    for (const [name, meta] of tensors) {
      this.weightMeta.set(name, { shape: meta.shape, dtype: meta.dtype });

      if (meta.dtype === 0) {
        // Check if this Q4 tensor should be dequantized for CPU access
        const needsCpu = cpuDequantSuffixes.some(s => name.endsWith(s))
          || name === prefix + 'embed_tokens.weight';

        if (needsCpu && name !== prefix + 'embed_tokens.weight') {
          // Small Q4 tensor used element-wise — dequant all rows to CPU
          const rows = meta.shape[0];
          const cols = meta.shape.length > 1 ? meta.shape.reduce((a, b) => a * b) / rows : 1;
          const totalElements = meta.shape.reduce((a, b) => a * b);
          const flat = new Float32Array(totalElements);
          for (let r = 0; r < rows; r++) {
            const row = dequantQ4Row(meta.data, r, cols, GROUP_SIZE);
            flat.set(row, r * cols);
          }
          this.cpuWeights.set(name, flat);
          cpuCount++;
        }

        // Upload to GPU for matvec
        const bufSize = Math.ceil(meta.data.byteLength / 4) * 4;
        const gpuBuf = this.device.createBuffer({
          size: bufSize,
          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
          mappedAtCreation: true,
        });
        new Uint8Array(gpuBuf.getMappedRange()).set(meta.data);
        gpuBuf.unmap();
        this.weightBuffers.set(name, gpuBuf);
        gpuBytes += bufSize;

        // Keep embed_tokens data on CPU for row dequant
        if (name === prefix + 'embed_tokens.weight') {
          this.embedData = meta.data;
          this.embedCols = meta.shape[1];
        }
      } else {
        // F32 or BF16 — keep on CPU
        this.cpuWeights.set(name, toF32Array(meta));
        cpuCount++;
      }
    }
    console.log(`[gpu] Weights loaded: ${this.weightBuffers.size} on GPU (${(gpuBytes / 1024 / 1024).toFixed(1)} MB), ${cpuCount} on CPU`);
  }

  /** Dequantize one embedding row on CPU. */
  private getEmbedding(tokenId: number): Float32Array {
    return dequantQ4Row(this.embedData, tokenId, this.embedCols, GROUP_SIZE);
  }

  /** Batched GPU Q4 matrix-vector multiply: runs multiple matvecs in a single GPU submission + readback. */
  async matvecBatch(
    input: Float32Array,
    ops: Array<{ weightName: string; outputSize: number }>,
  ): Promise<Float32Array[]> {
    // Upload input once
    this.device.queue.writeBuffer(this.inputBuf, 0, input.buffer, input.byteOffset, input.byteLength);

    // Upload per-op params
    for (let i = 0; i < ops.length; i++) {
      const meta = this.weightMeta.get(ops[i].weightName)!;
      const M = ops[i].outputSize;
      const N = meta.shape[meta.shape.length - 1];
      this.device.queue.writeBuffer(this.paramsBufs[i], 0, new Uint32Array([M, N, Math.ceil(N / GROUP_SIZE)]));
    }

    const commandEncoder = this.device.createCommandEncoder();
    let totalFloats = 0;
    const offsets: number[] = [];

    for (let i = 0; i < ops.length; i++) {
      const wBuf = this.weightBuffers.get(ops[i].weightName);
      if (!wBuf) throw new Error(`Weight not found on GPU: ${ops[i].weightName}`);

      const bindGroup = this.device.createBindGroup({
        layout: this.bindGroupLayout,
        entries: [
          { binding: 0, resource: { buffer: wBuf } },
          { binding: 1, resource: { buffer: this.inputBuf } },
          { binding: 2, resource: { buffer: this.outputBufs[i] } },
          { binding: 3, resource: { buffer: this.paramsBufs[i] } },
        ],
      });

      const pass = commandEncoder.beginComputePass();
      pass.setPipeline(this.q4Pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.dispatchWorkgroups(Math.ceil(ops[i].outputSize / 256));
      pass.end();

      offsets.push(totalFloats);
      totalFloats += ops[i].outputSize;
    }

    // Copy all outputs to staging buffer at different offsets
    for (let i = 0; i < ops.length; i++) {
      commandEncoder.copyBufferToBuffer(
        this.outputBufs[i], 0,
        this.stagingBuf, offsets[i] * 4,
        ops[i].outputSize * 4,
      );
    }

    this.device.queue.submit([commandEncoder.finish()]);

    // Single readback for all results
    const totalBytes = totalFloats * 4;
    await this.stagingBuf.mapAsync(GPUMapMode.READ, 0, totalBytes);
    const mapped = this.stagingBuf.getMappedRange(0, totalBytes);
    const allData = new Float32Array(totalFloats);
    allData.set(new Float32Array(mapped));
    this.stagingBuf.unmap();

    const results: Float32Array[] = [];
    for (let i = 0; i < ops.length; i++) {
      results.push(allData.slice(offsets[i], offsets[i] + ops[i].outputSize));
    }
    return results;
  }

  /** GPU Q4 matrix-vector multiply: output = weights[M, N] @ input[N]. */
  async matvec(weightName: string, input: Float32Array, outputSize: number): Promise<Float32Array> {
    const [result] = await this.matvecBatch(input, [{ weightName, outputSize }]);
    return result;
  }

  /** Helper to get a CPU weight by name, with prefix. */
  private cpu(name: string): Float32Array {
    const w = this.cpuWeights.get(name);
    if (!w) throw new Error(`CPU weight not found: ${name}`);
    return w;
  }

  /** Helper to get the full weight name for a layer tensor. */
  private layerName(layer: number, suffix: string): string {
    return `model.language_model.layers.${layer}.${suffix}`;
  }

  /** Run forward pass for a single token. Returns logits [VOCAB_SIZE]. */
  async forward(tokenId: number, pos: number): Promise<Float32Array> {
    const prefix = 'model.language_model.';

    // 1. Embedding lookup (CPU dequant of one row)
    const embed = this.getEmbedding(tokenId);
    this.hidden.set(embed);

    // Track which full-attn and delta-net layer index we're on
    let fullAttnIdx = 0;
    let deltaIdx = 0;

    // 2. Layer loop
    for (let layer = 0; layer < NUM_LAYERS; layer++) {
      const isFullAttn = FULL_ATTN_LAYERS.has(layer);

      // Save residual
      this.residual.set(this.hidden);

      // Pre-attention RMSNorm
      const lnW = this.cpu(this.layerName(layer, 'input_layernorm.weight'));
      rmsNorm(this.normed, this.hidden, lnW, RMS_NORM_EPS);

      if (isFullAttn) {
        await this.fullAttentionForward(layer, fullAttnIdx, pos);
        fullAttnIdx++;
      } else {
        await this.deltaNetForward(layer, deltaIdx);
        deltaIdx++;
      }

      // Residual add
      for (let i = 0; i < HIDDEN_SIZE; i++) this.hidden[i] += this.residual[i];

      // Save residual for MLP
      this.residual.set(this.hidden);

      // Post-attention RMSNorm
      const postLnW = this.cpu(this.layerName(layer, 'post_attention_layernorm.weight'));
      rmsNorm(this.normed, this.hidden, postLnW, RMS_NORM_EPS);

      // MLP (SwiGLU)
      await this.mlpForward(layer);

      // Residual add
      for (let i = 0; i < HIDDEN_SIZE; i++) this.hidden[i] += this.residual[i];
    }

    // 3. Final RMSNorm
    const finalNormW = this.cpu(prefix + 'norm.weight');
    rmsNorm(this.normed, this.hidden, finalNormW, RMS_NORM_EPS);

    // 4. LM head (tied weights — reuse embed_tokens)
    const logits = await this.matvec(prefix + 'embed_tokens.weight', this.normed, VOCAB_SIZE);
    return logits;
  }

  /** Full attention forward for one layer. Writes result into this.hidden. */
  private async fullAttentionForward(layer: number, kvIdx: number, pos: number): Promise<void> {
    const ln = (s: string) => this.layerName(layer, s);

    // Batch Q/K/V projections (all from same normed input → single GPU submission)
    const [qBuf, kBuf, vBuf] = await this.matvecBatch(this.normed, [
      { weightName: ln('self_attn.q_proj.weight'), outputSize: NUM_HEADS * HEAD_DIM * 2 },
      { weightName: ln('self_attn.k_proj.weight'), outputSize: NUM_KV_HEADS * HEAD_DIM },
      { weightName: ln('self_attn.v_proj.weight'), outputSize: NUM_KV_HEADS * HEAD_DIM },
    ]);

    // Split Q into query and gate
    const query = new Float32Array(ATTN_OUTPUT_DIM);
    const gate = new Float32Array(ATTN_OUTPUT_DIM);
    for (let h = 0; h < NUM_HEADS; h++) {
      const srcOff = h * HEAD_DIM * 2;
      const dstOff = h * HEAD_DIM;
      for (let d = 0; d < HEAD_DIM; d++) {
        query[dstOff + d] = qBuf[srcOff + d];
        gate[dstOff + d] = qBuf[srcOff + HEAD_DIM + d];
      }
    }

    // QK-norm per head
    const qNormW = this.cpu(ln('self_attn.q_norm.weight'));
    const kNormW = this.cpu(ln('self_attn.k_norm.weight'));
    for (let h = 0; h < NUM_HEADS; h++) {
      rmsNorm(query, query, qNormW, RMS_NORM_EPS, h * HEAD_DIM, HEAD_DIM);
    }
    for (let h = 0; h < NUM_KV_HEADS; h++) {
      rmsNorm(kBuf, kBuf, kNormW, RMS_NORM_EPS, h * HEAD_DIM, HEAD_DIM);
    }

    // Apply partial RoPE
    const cosOff = pos * ROTARY_DIM;
    for (let h = 0; h < NUM_HEADS; h++) {
      applyRope(query, h * HEAD_DIM, this.ropeCos, this.ropeSin, cosOff, ROTARY_DIM);
    }
    for (let h = 0; h < NUM_KV_HEADS; h++) {
      applyRope(kBuf, h * HEAD_DIM, this.ropeCos, this.ropeSin, cosOff, ROTARY_DIM);
    }

    // Append K, V to cache
    const kvKeyDim = NUM_KV_HEADS * HEAD_DIM;
    const seqPos = this.kvLen[kvIdx];
    this.kvKeys[kvIdx].set(kBuf, seqPos * kvKeyDim);
    this.kvValues[kvIdx].set(vBuf, seqPos * kvKeyDim);
    this.kvLen[kvIdx] = seqPos + 1;
    const seqLen = this.kvLen[kvIdx];

    // Compute attention per query head (GQA)
    const attnOut = new Float32Array(ATTN_OUTPUT_DIM);
    const scale = 1.0 / Math.sqrt(HEAD_DIM);
    const scores = new Float32Array(seqLen);

    for (let h = 0; h < NUM_HEADS; h++) {
      const kvH = Math.floor(h / KV_GROUPS);
      const qOff = h * HEAD_DIM;

      // Attention scores
      let maxScore = -Infinity;
      for (let t = 0; t < seqLen; t++) {
        const kOff = t * kvKeyDim + kvH * HEAD_DIM;
        let dot = 0;
        for (let d = 0; d < HEAD_DIM; d++) {
          dot += query[qOff + d] * this.kvKeys[kvIdx][kOff + d];
        }
        dot *= scale;
        scores[t] = dot;
        if (dot > maxScore) maxScore = dot;
      }

      // Softmax
      let sumExp = 0;
      for (let t = 0; t < seqLen; t++) {
        scores[t] = Math.exp(scores[t] - maxScore);
        sumExp += scores[t];
      }
      const invSum = sumExp > 0 ? 1.0 / sumExp : 0;
      for (let t = 0; t < seqLen; t++) scores[t] *= invSum;

      // Weighted sum of values
      for (let d = 0; d < HEAD_DIM; d++) attnOut[qOff + d] = 0;
      for (let t = 0; t < seqLen; t++) {
        const vOff = t * kvKeyDim + kvH * HEAD_DIM;
        const w = scores[t];
        for (let d = 0; d < HEAD_DIM; d++) {
          attnOut[qOff + d] += w * this.kvValues[kvIdx][vOff + d];
        }
      }
    }

    // Apply output gate: attnOut *= sigmoid(gate)
    for (let i = 0; i < ATTN_OUTPUT_DIM; i++) {
      attnOut[i] *= sigmoid(gate[i]);
    }

    // Output projection
    this.hidden.set(await this.matvec(ln('self_attn.o_proj.weight'), attnOut, HIDDEN_SIZE));
  }

  /** DeltaNet forward for one layer. Writes result into this.hidden. */
  private async deltaNetForward(layer: number, dnIdx: number): Promise<void> {
    const ln = (s: string) => this.layerName(layer, s);

    // Batch all DeltaNet projections (all from same normed input → single GPU submission)
    const [qkv, z, aRaw, bRaw] = await this.matvecBatch(this.normed, [
      { weightName: ln('linear_attn.in_proj_qkv.weight'), outputSize: DN_CONV_DIM },
      { weightName: ln('linear_attn.in_proj_z.weight'), outputSize: DN_TOTAL_VALUE },
      { weightName: ln('linear_attn.in_proj_a.weight'), outputSize: DN_NUM_HEADS },
      { weightName: ln('linear_attn.in_proj_b.weight'), outputSize: DN_NUM_HEADS },
    ]);

    // Causal conv1d step
    const convW = this.cpu(ln('linear_attn.conv1d.weight'));
    const convState = this.convStates[dnIdx];
    const stateCols = DN_CONV_KERNEL - 1;

    for (let ch = 0; ch < DN_CONV_DIM; ch++) {
      const csBase = ch * stateCols;

      // Compute conv output before updating state
      let convOut = 0;
      for (let k = 0; k < stateCols; k++) {
        convOut += convState[csBase + k] * convW[ch * DN_CONV_KERNEL + k];
      }
      convOut += qkv[ch] * convW[ch * DN_CONV_KERNEL + stateCols];

      // Shift state left, append new input
      for (let j = 0; j + 1 < stateCols; j++) {
        convState[csBase + j] = convState[csBase + j + 1];
      }
      convState[csBase + stateCols - 1] = qkv[ch];

      // SiLU activation
      qkv[ch] = silu(convOut);
    }

    // Split QKV
    const qStart = 0;
    const kStart = DN_TOTAL_KEY;
    const vStart = DN_TOTAL_KEY * 2;

    // Compute decay and beta
    const aLog = this.cpu(ln('linear_attn.A_log'));
    const dtBias = this.cpu(ln('linear_attn.dt_bias'));
    const gVals = new Float32Array(DN_NUM_HEADS);
    const betaVals = new Float32Array(DN_NUM_HEADS);
    for (let h = 0; h < DN_NUM_HEADS; h++) {
      gVals[h] = -Math.exp(aLog[h]) * softplus(aRaw[h] + dtBias[h]);
      betaVals[h] = sigmoid(bRaw[h]);
    }

    // DeltaNet recurrence per head
    const stateMatrix = this.deltaStates[dnIdx];
    const dnOut = new Float32Array(DN_TOTAL_VALUE);
    const scaleQ = 1.0 / Math.sqrt(DN_KEY_DIM);
    const qHead = new Float32Array(DN_KEY_DIM);
    const kHead = new Float32Array(DN_KEY_DIM);

    for (let h = 0; h < DN_NUM_HEADS; h++) {
      // Extract Q and K, L2-normalize
      for (let d = 0; d < DN_KEY_DIM; d++) {
        qHead[d] = qkv[qStart + h * DN_KEY_DIM + d];
        kHead[d] = qkv[kStart + h * DN_KEY_DIM + d];
      }
      l2Norm(qHead, 0, DN_KEY_DIM);
      l2Norm(kHead, 0, DN_KEY_DIM);

      // Scale query
      for (let d = 0; d < DN_KEY_DIM; d++) qHead[d] *= scaleQ;

      const smBase = h * DN_KEY_DIM * DN_VALUE_DIM;

      // Decay state
      const decay = Math.exp(gVals[h]);
      for (let i = 0; i < DN_KEY_DIM * DN_VALUE_DIM; i++) {
        stateMatrix[smBase + i] *= decay;
      }

      // kv_mem[v] = sum_d(S[d,v] * k[d])
      const kvMem = new Float32Array(DN_VALUE_DIM);
      for (let v = 0; v < DN_VALUE_DIM; v++) {
        let sum = 0;
        for (let d = 0; d < DN_KEY_DIM; d++) {
          sum += stateMatrix[smBase + d * DN_VALUE_DIM + v] * kHead[d];
        }
        kvMem[v] = sum;
      }

      // delta = (v - kv_mem) * beta
      const beta = betaVals[h];
      for (let v = 0; v < DN_VALUE_DIM; v++) {
        const delta = (qkv[vStart + h * DN_VALUE_DIM + v] - kvMem[v]) * beta;
        // Update state: S[d,v] += k[d] * delta
        for (let d = 0; d < DN_KEY_DIM; d++) {
          stateMatrix[smBase + d * DN_VALUE_DIM + v] += kHead[d] * delta;
        }
      }

      // Output: out[v] = sum_d(S[d,v] * q[d])
      const outOff = h * DN_VALUE_DIM;
      for (let v = 0; v < DN_VALUE_DIM; v++) {
        let sum = 0;
        for (let d = 0; d < DN_KEY_DIM; d++) {
          sum += stateMatrix[smBase + d * DN_VALUE_DIM + v] * qHead[d];
        }
        dnOut[outOff + v] = sum;
      }
    }

    // Gated RMS norm per head
    const normW = this.cpu(ln('linear_attn.norm.weight'));
    for (let h = 0; h < DN_NUM_HEADS; h++) {
      const off = h * DN_VALUE_DIM;
      rmsNormGated(dnOut, off, dnOut, off, z, off, normW, RMS_NORM_EPS, DN_VALUE_DIM);
    }

    // Output projection
    this.hidden.set(await this.matvec(ln('linear_attn.out_proj.weight'), dnOut, HIDDEN_SIZE));
  }

  /** SwiGLU MLP forward. Writes result into this.hidden. */
  private async mlpForward(layer: number): Promise<void> {
    const ln = (s: string) => this.layerName(layer, s);

    // Batch gate + up projections (both from same normed input → single GPU submission)
    const [gateOut, upOut] = await this.matvecBatch(this.normed, [
      { weightName: ln('mlp.gate_proj.weight'), outputSize: INTERMEDIATE_SIZE },
      { weightName: ln('mlp.up_proj.weight'), outputSize: INTERMEDIATE_SIZE },
    ]);

    // SwiGLU: silu(gate) * up
    for (let i = 0; i < INTERMEDIATE_SIZE; i++) {
      gateOut[i] = silu(gateOut[i]) * upOut[i];
    }

    this.hidden.set(await this.matvec(ln('mlp.down_proj.weight'), gateOut, HIDDEN_SIZE));
  }

  /** Reset all caches for a new generation. */
  reset(): void {
    for (let i = 0; i < this.kvKeys.length; i++) {
      this.kvKeys[i].fill(0);
      this.kvValues[i].fill(0);
      this.kvLen[i] = 0;
    }
    for (let i = 0; i < this.deltaStates.length; i++) {
      this.deltaStates[i].fill(0);
      this.convStates[i].fill(0);
    }
  }
}