web/gpu.js

// Thin WebGPU helpers shared by tests and the model runtime.

import { setSubgroups, setF16 } from "./kernels.js";

export async function initDevice() {
  const adapter = await navigator.gpu?.requestAdapter({ powerPreference: "high-performance" })
    ?? await navigator.gpu?.requestAdapter();
  if (!adapter) throw new Error("WebGPU is not available — use Chrome/Edge 113+ or Safari 26+");
  const params = typeof location !== "undefined" ? new URLSearchParams(location.search) : new URLSearchParams();
  // shader-f16 is a fast path, not a requirement: without it (Dawn gates it
  // on NVIDIA/Linux; Qualcomm mobiles lack it) f16 tensors are converted to
  // f32 at load time and the KV cache lives in f32
  const hasF16 = !params.has("nof16") && adapter.features.has("shader-f16");
  setF16(hasF16);
  const features = hasF16 ? ["shader-f16"] : [];
  // subgroups are a fast path, not a requirement: without them (iOS Safari)
  // the matvec kernels fall back to workgroup-shared reductions
  const hasSubgroups = !params.has("nosg") && adapter.features.has("subgroups");
  if (hasSubgroups) features.push("subgroups");
  setSubgroups(hasSubgroups);
  if (adapter.features.has("timestamp-query")) features.push("timestamp-query");
  if (hasSubgroups && adapter.features.has("chromium-experimental-subgroup-matrix"))
    features.push("chromium-experimental-subgroup-matrix");
  const device = await adapter.requestDevice({
    requiredFeatures: features,
    requiredLimits: {
      maxBufferSize: adapter.limits.maxBufferSize,
      maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
      maxStorageBuffersPerShaderStage: adapter.limits.maxStorageBuffersPerShaderStage,
      maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize,
    },
  });
  device.hasSubgroups = hasSubgroups;
  device.hasF16 = hasF16;
  // fused decode kernels replay a 32-lane subgroupAdd (qkvRope head norm)
  // bit-exactly — only enabled when the subgroup size is fixed at 32
  device.sg32 = hasSubgroups && adapter.info?.subgroupMinSize === 32 && adapter.info?.subgroupMaxSize === 32;
  device.lost.then((info) => console.error("GPU device lost:", info.reason, info.message));
  return device;
}

export function createBuffer(device, sizeBytes, { storage = true, uniform = false } = {}) {
  let usage = GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC;
  if (storage) usage |= GPUBufferUsage.STORAGE;
  if (uniform) usage |= GPUBufferUsage.UNIFORM;
  return device.createBuffer({ size: Math.max(sizeBytes, 16), usage });
}

export function uploadF32(device, data) {
  const buf = createBuffer(device, data.byteLength);
  device.queue.writeBuffer(buf, 0, data);
  return buf;
}

export function makePipeline(device, code, label = "") {
  const module = device.createShaderModule({ code, label });
  return device.createComputePipeline({ label, layout: "auto", compute: { module } });
}

export function bind(device, pipeline, entries) {
  return device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: Object.entries(entries).map(([binding, buffer]) => ({
      binding: Number(binding),
      resource: { buffer },
    })),
  });
}

export async function readback(device, buffer, sizeBytes, offset = 0) {
  const rb = device.createBuffer({
    size: sizeBytes,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });
  const enc = device.createCommandEncoder();
  enc.copyBufferToBuffer(buffer, offset, rb, 0, sizeBytes);
  device.queue.submit([enc.finish()]);
  await rb.mapAsync(GPUMapMode.READ);
  const data = new Float32Array(rb.getMappedRange().slice(0));
  rb.destroy();
  return data;
}

// f16 helpers (Float16Array is available in Chrome 135+)
export function toF16(f32arr) {
  return new Float16Array(f32arr);
}

export function uploadF16(device, f32arr) {
  const f16 = new Float16Array(f32arr);
  const buf = createBuffer(device, f16.byteLength);
  device.queue.writeBuffer(buf, 0, f16.buffer, f16.byteOffset, f16.byteLength);
  return buf;
}