webgpu

package
v1.4.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 22, 2026 License: MIT Imports: 7 Imported by: 0

Documentation

Rendered for js/wasm

Overview

Package webgpu is a thin, hand-written Go wrapper over the browser's native WebGPU API (accessed via syscall/js), mirroring how internal/metal wraps Apple's Metal API for the darwin/cgo build. No third-party JS/WASM code is used anywhere in this package — only the browser's own navigator.gpu surface.

Index

Constants

View Source 
const Q4KMatvecWGSL = wgslCommon + `
@group(0) @binding(0) var<storage, read> weights: array<u32>;
@group(0) @binding(1) var<storage, read> x: array<f32>;
@group(0) @binding(2) var<storage, read_write> out: array<f32>;
@group(0) @binding(3) var<storage, read> params: array<u32>;

var<workgroup> partial: array<f32, 64>;

const Q4K_BLOCK_BYTES: u32 = 144u;

@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let row = wg.x;
  let localId = lid.x;
  let cols = params[0];
  let rowOffset = params[1];
  let numBlocks = cols / 256u;
  let rowByteOffset = row * numBlocks * Q4K_BLOCK_BYTES;

  var acc: f32 = 0.0;
  var blockIdx = localId;
  loop {
    if (blockIdx >= numBlocks) { break; }
    let base = rowByteOffset + blockIdx * Q4K_BLOCK_BYTES;
    let d = f16_to_f32(read_u16(base));
    let dmin = f16_to_f32(read_u16(base + 2u));
    let scalesBase = base + 4u;
    let qBase = base + 16u;
    let yoff = blockIdx * 256u;

    var jj: u32 = 0u;
    loop {
      if (jj >= 4u) { break; }
      let is0 = jj * 2u;
      let sm1 = get_scale_min_k4(is0, scalesBase);
      let sm2 = get_scale_min_k4(is0 + 1u, scalesBase);
      let d1 = d * f32(sm1.x);
      let d2 = d * f32(sm2.x);
      let min1 = dmin * f32(sm1.y);
      let min2 = dmin * f32(sm2.y);
      let qOff = qBase + jj * 32u;
      let xOff = yoff + jj * 64u;

      var l: u32 = 0u;
      loop {
        if (l >= 32u) { break; }
        let qByte = read_u8(qOff + l);
        let lo = qByte & 0x0fu;
        let hi = qByte >> 4u;
        let vLo = d1 * f32(lo) - min1;
        let vHi = d2 * f32(hi) - min2;
        acc = acc + vLo * x[xOff + l];
        acc = acc + vHi * x[xOff + 32u + l];
        l = l + 1u;
      }
      jj = jj + 1u;
    }
    blockIdx = blockIdx + 64u;
  }

  partial[localId] = acc;
  workgroupBarrier();
  var stride: u32 = 32u;
  loop {
    if (stride == 0u) { break; }
    if (localId < stride) {
      partial[localId] = partial[localId] + partial[localId + stride];
    }
    workgroupBarrier();
    stride = stride / 2u;
  }
  if (localId == 0u) {
    out[rowOffset + row] = partial[0];
  }
}
`

Q4KMatvecWGSL is a from-scratch WGSL reimplementation of this project's own scalar Q4_K matvec kernel (simd.go's DequantRowQ4KInto/ getScaleMinK4), authored directly from the GGUF Q4_K block layout (a public file-format fact, not copyrightable expression): 144 bytes/256 elements = f16 d (2B) + f16 dmin (2B) + 12B of packed 6-bit scale/min pairs for 8 sub-blocks of 32 + 128B of nibble-packed 4-bit quants (low nibble = even sub-block, high nibble = odd sub-block within each 64-wide byte-pair span). Dequant per element: value = d*scale - dmin*min, using that element's sub-block's 6-bit scale/min (getScaleMinK4's packing trick fits 8x6-bit scales + 8x6-bit mins into 12 bytes).

One workgroup computes one output row: each of 64 threads accumulates a partial dot product over the row's blocks (blockIdx = localId, localId+64, ...), then a standard workgroup-shared-memory tree reduction sums the 64 partials into the final value, written by thread 0. Bindings: 0=weights (this dispatch's row-chunk's raw Q4_K bytes, array<u32>), 1=x (activation vector, length cols), 2=out (output vector, write out[rowOffset+row]), 3=params (array<u32>: [cols, rowOffset]).

View Source 
const Q6KMatvecWGSL = wgslCommon + `
@group(0) @binding(0) var<storage, read> weights: array<u32>;
@group(0) @binding(1) var<storage, read> x: array<f32>;
@group(0) @binding(2) var<storage, read_write> out: array<f32>;
@group(0) @binding(3) var<storage, read> params: array<u32>;

var<workgroup> partial6: array<f32, 64>;

const Q6K_BLOCK_BYTES: u32 = 210u;

fn read_i8(byteOffset: u32) -> i32 {
  let v = read_u8(byteOffset);
  return select(i32(v), i32(v) - 256, v >= 128u);
}

@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let row = wg.x;
  let localId = lid.x;
  let cols = params[0];
  let rowOffset = params[1];
  let numBlocks = cols / 256u;
  let rowByteOffset = row * numBlocks * Q6K_BLOCK_BYTES;

  var acc: f32 = 0.0;
  var blockIdx = localId;
  loop {
    if (blockIdx >= numBlocks) { break; }
    let base = rowByteOffset + blockIdx * Q6K_BLOCK_BYTES;
    let d = f16_to_f32(read_u16(base + 208u));
    let yoff = blockIdx * 256u;

    var n: u32 = 0u;
    loop {
      if (n >= 2u) { break; }
      let qlBase = base + n * 64u;
      let qhBase = base + 128u + n * 32u;
      let scBase = base + 192u + n * 8u;
      let xOff = yoff + n * 128u;

      var l: u32 = 0u;
      loop {
        if (l >= 32u) { break; }
        let is0 = l / 16u;
        let qlL = read_u8(qlBase + l);
        let qlL32 = read_u8(qlBase + l + 32u);
        let qhL = read_u8(qhBase + l);

        let q1 = i32((qlL & 0x0fu) | ((qhL & 0x03u) << 4u)) - 32;
        let q2 = i32((qlL32 & 0x0fu) | (((qhL >> 2u) & 0x03u) << 4u)) - 32;
        let q3 = i32((qlL >> 4u) | (((qhL >> 4u) & 0x03u) << 4u)) - 32;
        let q4 = i32((qlL32 >> 4u) | (((qhL >> 6u) & 0x03u) << 4u)) - 32;

        let sc0 = f32(read_i8(scBase + is0));
        let sc2 = f32(read_i8(scBase + is0 + 2u));
        let sc4 = f32(read_i8(scBase + is0 + 4u));
        let sc6 = f32(read_i8(scBase + is0 + 6u));

        acc = acc + (d * sc0 * f32(q1)) * x[xOff + l];
        acc = acc + (d * sc2 * f32(q2)) * x[xOff + 32u + l];
        acc = acc + (d * sc4 * f32(q3)) * x[xOff + 64u + l];
        acc = acc + (d * sc6 * f32(q4)) * x[xOff + 96u + l];
        l = l + 1u;
      }
      n = n + 1u;
    }
    blockIdx = blockIdx + 64u;
  }

  partial6[localId] = acc;
  workgroupBarrier();
  var stride: u32 = 32u;
  loop {
    if (stride == 0u) { break; }
    if (localId < stride) {
      partial6[localId] = partial6[localId] + partial6[localId + stride];
    }
    workgroupBarrier();
    stride = stride / 2u;
  }
  if (localId == 0u) {
    out[rowOffset + row] = partial6[0];
  }
}
`

Q6KMatvecWGSL is a from-scratch WGSL reimplementation of this project's own scalar Q6_K matvec kernel (simd.go's DequantRowQ6KInto), authored directly from the GGUF Q6_K block layout: 210 bytes/256 elements = 128B low nibbles (ql) + 64B high 2-bit planes (qh) + 16B signed int8 sub-block scales (sc) + f16 d (2B, at the very end of the block). 256 values split into two 128-wide halves; within each half, for l in [0,32) four 6-bit quant values are reassembled per l from ql/qh (each qh byte supplies the top 2 bits for 4 different 6-bit values spread 32 apart; each ql byte supplies two 4-bit low halves). Each reassembled 6-bit value (range [0,63]) has 32 subtracted (signed range [-32,31]), multiplied by a signed int8 scale for its sub-block, then by the shared f16 d: value = d * sc[k] * (q6bits - 32).

Same one-workgroup-per-row, 64-thread partial-sum + workgroup tree reduction structure as Q4KMatvecWGSL. Bindings identical: 0=weights, 1=x, 2=out, 3=params ([cols, rowOffset]).

Variables

This section is empty.

Functions

func Available 

func Available() bool

Available reports whether navigator.gpu exists at all -- a cheap synchronous check, safe to call before attempting the async RequestAdapterAndDevice.

Types

type Buffer 

type Buffer struct {
	// contains filtered or unexported fields
}

Buffer wraps one GPUBuffer.

func (*Buffer) Destroy 

func (b *Buffer) Destroy()

Destroy releases the underlying GPUBuffer's memory. No method may use b afterward.

func (*Buffer) Size 

func (b *Buffer) Size() int

Size returns the buffer's byte size.

type ComputePipeline 

type ComputePipeline struct {
	// contains filtered or unexported fields
}

ComputePipeline wraps one GPUComputePipeline.

type Device 

type Device struct {
	// contains filtered or unexported fields
}

Device wraps a GPUDevice and its default GPUQueue.

func RequestAdapterAndDevice 

func RequestAdapterAndDevice(ctx context.Context) (*Device, error)

RequestAdapterAndDevice acquires a GPU adapter and device. ctx cancels the wait (WebGPU's own APIs have no native cancellation, so this races the result against ctx.Done(); the underlying JS Promise keeps running to completion regardless, same as any other unawaited Promise -- acceptable since adapter/device acquisition is normally a few milliseconds).

func (*Device) CreateBindGroup 

func (d *Device) CreateBindGroup(pipeline *ComputePipeline, buffers ...*Buffer) js.Value

CreateBindGroup builds a group-0 bind group for pipeline's auto-inferred layout from buffers, bound at consecutive binding indices starting at 0 in the given order -- callers must match this order to their WGSL shader's @binding declarations.

func (*Device) CreateComputePipeline 

func (d *Device) CreateComputePipeline(wgslSource, entryPoint string) (*ComputePipeline, error)

CreateComputePipeline compiles wgslSource (hand-written WGSL -- see the root package's *.wgsl.go files for the actual kernels) and creates a compute pipeline with the given entry point, using WebGPU's "auto" pipeline layout (bind group layouts are inferred from the shader's own @group/@binding declarations). Shader compilation errors are collected via getCompilationInfo and returned as a Go error rather than left to silently produce a pipeline that does nothing -- with no external WGSL reference to check against, this is the primary debugging signal while authoring new kernels.

func (*Device) CreateReadbackBuffer added in v1.1.0 

func (d *Device) CreateReadbackBuffer(size int) (*Buffer, error)

CreateReadbackBuffer allocates a MAP_READ staging buffer. Compute outputs cannot be mapped directly, so callers copy into one of these buffers after ending their compute pass and then call ReadMappedBufferInto. Keeping a staging buffer alive across decode steps is materially cheaper than allocating/destroying a GPUBuffer for every matvec result.

func (*Device) CreateStorageBuffer 

func (d *Device) CreateStorageBuffer(size int, data []byte) (*Buffer, error)

CreateStorageBuffer allocates a GPUBuffer readable by compute shaders (STORAGE usage) plus COPY_DST/COPY_SRC so it can be uploaded to and, if ever needed, copied from directly. If data is non-nil its bytes are uploaded immediately via WriteBuffer. The underlying GPU allocation is rounded up to a multiple of 4 bytes (see roundUpTo4); Buffer.Size still reports the caller's logical (unrounded) size, and the shaders in this package never read past it, so the extra padding is inert.

func (*Device) CreateUniformBuffer 

func (d *Device) CreateUniformBuffer(size int) (*Buffer, error)

CreateUniformBuffer allocates a small GPUBuffer with UNIFORM usage, for per-dispatch scalar parameters (row offsets, dimensions) a WGSL shader reads via a uniform binding. See CreateStorageBuffer's doc comment for why the allocation is rounded up to a multiple of 4.

func (*Device) Destroy 

func (d *Device) Destroy()

Destroy releases the underlying GPUDevice. No method may be called on d or anything created from it afterward.

func (*Device) Limits 

func (d *Device) Limits() DeviceLimits

Limits returns the device's real reported limits.

func (*Device) NewEncoder 

func (d *Device) NewEncoder() *Encoder

NewEncoder starts recording a new command buffer.

func (*Device) PrepareWeight 

func (d *Device) PrepareWeight(data []byte, rows, cols int) (*PreparedWeight, error)

PrepareWeight uploads data (one tensor's raw on-disk quantized bytes, in the on-disk Q4_K/Q6_K block layout the shaders expect) into a new GPU storage buffer once. Returns an error if data exceeds this device's maxStorageBufferBindingSize -- splitting an oversized tensor across multiple buffers (needed for very large tensors like a tied output/ embedding projection) is not implemented yet; callers should catch this and fall back to the CPU path for that specific tensor rather than failing the whole model load.

func (*Device) ReadBuffer 

func (d *Device) ReadBuffer(src *Buffer, byteOffset, size int) ([]byte, error)

ReadBuffer copies size bytes starting at byteOffset out of src (which must have been created with COPY_SRC usage) via a temporary staging buffer with MAP_READ usage -- a compute-shader-writable STORAGE buffer cannot itself be mapped for reading under WebGPU, so results must first be copied into a MAP_READ-capable buffer. Every size WebGPU touches here (buffer/copy/map/range size) must be a multiple of 4 (see CreateStorageBuffer's doc comment); the caller's requested size need not be, so the extra padding is read into the staging buffer and trimmed back off before returning.

func (*Device) ReadMappedBufferInto added in v1.1.0 

func (d *Device) ReadMappedBufferInto(src *Buffer, byteOffset int, dst []byte) error

ReadMappedBufferInto maps an already-submitted MAP_READ buffer and copies its bytes into dst. The caller is responsible for recording and submitting the copy into src first. dst must be 4-byte aligned in length, matching the WebGPU map range requirement. This deliberately accepts a caller-owned destination so hot paths can reuse Go memory instead of allocating a fresh byte slice per GPU dispatch.

func (*Device) WriteBuffer 

func (d *Device) WriteBuffer(buf *Buffer, byteOffset int, data []byte) error

WriteBuffer uploads data into buf at byteOffset via the device queue. queue.writeBuffer is synchronous from the caller's perspective (the browser copies immediately) -- no map/unmap dance needed for uploads, only for reading results back (see ReadBuffer). data is zero-padded up to a multiple of 4 bytes before the call if needed (see roundUpTo4); buf's own allocation is already padded to match by CreateStorageBuffer/ CreateUniformBuffer, so the padded write always stays in bounds.

func (*Device) WriteBufferView added in v1.1.0 

func (d *Device) WriteBufferView(buf *Buffer, byteOffset int, view js.Value, size int)

WriteBufferView writes size bytes from an existing Uint8Array view to buf without allocating another JS typed array. It is intended for hot paths that keep a capacity-sized upload view around (activation vectors and scalar parameters in MatvecKernel). Because the view is Uint8Array, WebGPU interprets its dataOffset and size arguments in bytes. Supplying size is important after a buffer has grown for a large FFN activation: a later smaller projection must not re-upload the unused tail capacity.

type DeviceLimits 

type DeviceLimits struct {
	MaxStorageBufferBindingSize       int
	MaxStorageBuffersPerShaderStage   int
	MinStorageBufferOffsetAlignment   int
	MaxComputeWorkgroupsPerDimension  int
	MaxComputeInvocationsPerWorkgroup int
}

DeviceLimits mirrors the subset of GPUSupportedLimits this package's callers need, read from the real device at acquisition time -- never hardcoded, since these vary by hardware/driver (see RequestAdapterAndDevice).

type Encoder 

type Encoder struct {
	// contains filtered or unexported fields
}

Encoder records one or more compute dispatches into a single GPUCommandEncoder, so a whole transformer layer's matmuls can be submitted to the GPU queue once instead of once per matrix -- the CPU<->GPU round-trip latency this avoids, not raw compute time, is what dominates a naive per-matrix-submit design (see the project plan's CPU/GPU split rationale).

func (*Encoder) BeginCompute 

func (e *Encoder) BeginCompute()

BeginCompute opens a compute pass. Must be paired with EndCompute before Submit (or before another BeginCompute) -- WebGPU allows only one open pass per encoder at a time.

func (*Encoder) CopyBuffer added in v1.1.0 

func (e *Encoder) CopyBuffer(src *Buffer, srcOffset int, dst *Buffer, dstOffset, size int)

CopyBuffer records a buffer-to-buffer copy in this encoder. Call it after EndCompute when copying a compute output into a MAP_READ staging buffer; keeping both operations in one command buffer avoids an extra queue submit and an avoidable CPU/GPU synchronization point per matvec.

func (*Encoder) Dispatch 

func (e *Encoder) Dispatch(pipeline *ComputePipeline, bindGroup js.Value, x, y, z int)

Dispatch binds pipeline and bindGroup, then dispatches an x*y*z workgroup grid. Must be called between BeginCompute and EndCompute.

func (*Encoder) EndCompute 

func (e *Encoder) EndCompute()

EndCompute closes the currently open compute pass.

func (*Encoder) Submit 

func (e *Encoder) Submit()

Submit finishes recording and submits the command buffer to the device queue. The Encoder must not be reused afterward.

type MatvecKernel 

type MatvecKernel struct {
	// contains filtered or unexported fields
}

MatvecKernel is a compiled dequantizing matvec compute pipeline (Q4_K or Q6_K -- see shader_q4k.go/shader_q6k.go), reusable across many Run/ RunPrepared calls against the same Device so the (relatively expensive) shader compilation happens once.

func NewQ4KMatvecKernel 

func NewQ4KMatvecKernel(dev *Device) (*MatvecKernel, error)

NewQ4KMatvecKernel compiles the Q4_K dequantizing matvec kernel.

func NewQ6KMatvecKernel 

func NewQ6KMatvecKernel(dev *Device) (*MatvecKernel, error)

NewQ6KMatvecKernel compiles the Q6_K dequantizing matvec kernel.

func (*MatvecKernel) Destroy 

func (k *MatvecKernel) Destroy()

Destroy releases reusable per-dispatch buffers. PreparedWeight values own their own buffers and must be released separately.

func (*MatvecKernel) Run 

func (k *MatvecKernel) Run(weightBytes []byte, x []float32, rows, cols int) ([]float32, error)

Run is a one-shot convenience wrapper for testing/benchmarking: prepares weightBytes as a fresh buffer, runs once, and discards it. Real decode should call PrepareWeight once at load time and reuse the result across many RunPrepared calls instead (see PreparedWeight's doc comment).

func (*MatvecKernel) RunPrepared 

func (k *MatvecKernel) RunPrepared(w *PreparedWeight, x []float32) ([]float32, error)

RunPrepared computes out[i] = dot(row_i, x) against an already-uploaded weight (see PrepareWeight), uploading only the small activation vector x fresh on every call -- the cheap per-token cost real decode should pay. It is the allocating convenience form of RunPreparedInto.

func (*MatvecKernel) RunPreparedInto added in v1.1.0 

func (k *MatvecKernel) RunPreparedInto(w *PreparedWeight, x, out []float32) error

RunPreparedInto is RunPrepared without a result allocation. out must have room for w.rows values. Real model decode should use this method because its destination is already a reusable activation scratch slice; avoiding the intermediate []float32 and copy removes a full output-vector allocation from every WebGPU matvec.

type PreparedWeight 

type PreparedWeight struct {
	// contains filtered or unexported fields
}

PreparedWeight is one tensor's raw quantized bytes already resident in a GPU buffer, created once (see Device.PrepareWeight) and reused for every subsequent RunPrepared call. This distinction matters a lot in practice: decode calls a matvec fresh for every generated token, so re-uploading a multi-megabyte weight matrix's bytes on every single call (as the one-shot Run convenience method below does) would make the GPU path slower than the CPU path it's meant to replace, not faster.

func (*PreparedWeight) Destroy 

func (w *PreparedWeight) Destroy()

Destroy releases the GPU buffer backing this prepared weight. No method may use w afterward.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.