@simulatte/doppler 0.1.8 → 0.2.0

package/src/gpu/kernels/matmul.d.ts

@@ -1,111 +0,0 @@
-/**
- * Matrix Multiplication Kernels
- *
- * Provides optimized matmul operations with support for:
- * - F16/F32 inputs and outputs
- * - Mixed precision (F16 weights, F32 activations)
- * - Tiled and naive variants
- * - Command recording for batched execution
- */
-
-import type { Tensor, TensorDtype } from '../tensor.js';
-import type { WeightBuffer } from '../weight-buffer.js';
-import type { CommandRecorder } from '../command-recorder.js';
-import type { OutputBufferOptions, OutputDtypeOptions, Vec4Options } from './types.js';
-import type { KernelPathSchema } from '../../config/schema/index.js';
-
-/** Matmul kernel options */
-export interface MatmulOptions extends OutputBufferOptions, OutputDtypeOptions, Vec4Options {
-  alpha?: number;
-  /** Optional matmul role for kernel path overrides (e.g., 'q_proj', 'ffn_gate', 'lm_head') */
-  role?: string;
-  /** Layer index for kernel path layer overrides */
-  layerIdx?: number;
-  /** Explicit kernel path context for variant selection (avoids global path state). */
-  kernelPath?: KernelPathSchema | null;
-  /** Optional explicit phase for kernel-path lookup when the runtime rewrites rows (for example prefill last-position logits). */
-  phaseOverride?: 'decode' | 'prefill' | null;
-  /**
-   * Whether B matrix is stored transposed.
-   * - true: B is [N,K] (SafeTensors/row-major), needs transpose
-   * - false: B is [K,N] (column-major/pre-transposed), direct access
-   * - 'auto': Auto-detect from buffer layout metadata (default)
-   */
-  transposeB?: boolean | 'auto';
-  aOffset?: number;
-  bOffset?: number;
-  cOffset?: number;
-  aDtype?: 'f16' | 'f32' | null;
-  bDtype?: 'f16' | 'f32' | 'q4k' | null;
-  preferF16?: boolean;
-  /** WGSL override constants for pipeline creation */
-  constants?: Record<string, number | boolean>;
-}
-
-/** Context for base matmul kernel selection rules. */
-export interface MatmulKernelSelectionContext {
-  useF16Matmul: boolean;
-  useF16wF32a: boolean;
-  useVec4: boolean;
-}
-
-/** Context for fused Q4K variant selection rules. */
-export interface MatmulQ4KFusedContext {
-  useF16A: boolean;
-  useF16Out: boolean;
-  isM1: boolean;
-}
-
-/** Context for GEMV variant selection rules. */
-export interface GemvSelectionContext {
-  hasSubgroups: boolean;
-  useF16Gemv: boolean;
-  useF32Gemv: boolean;
-  useMulticol: boolean;
-}
-
-/** Context for top-level matmul variant selection rules. */
-export interface MatmulVariantSelectionContext {
-  canFused: boolean;
-  useGemv: boolean;
-}
-
-/**
- * Check if fused Q4K kernels are disabled.
- */
-export declare function isFusedQ4KDisabled(): boolean;
-
-/**
- * Select the best matmul kernel variant
- */
-export declare function selectMatmulKernel(options?: MatmulOptions): string;
-
-/**
- * Create bind group layout for matmul operation
- */
-export declare function createMatmulBindGroupLayout(): GPUBindGroupLayout;
-
-/**
- * Run matrix multiplication
- */
-export declare function runMatmul(
-  A: Tensor,
-  B: GPUBuffer | WeightBuffer,
-  M: number,
-  N: number,
-  K: number,
-  options?: MatmulOptions
-): Promise<Tensor>;
-
-/**
- * Record matrix multiplication (batched, no submit)
- */
-export declare function recordMatmul(
-  recorder: CommandRecorder,
-  A: Tensor,
-  B: GPUBuffer | WeightBuffer,
-  M: number,
-  N: number,
-  K: number,
-  options?: MatmulOptions
-): Promise<Tensor>;

package/src/gpu/kernels/matmul.js

@@ -1,315 +0,0 @@
-import { getDevice, getKernelCapabilities } from '../device.js';
-import { createTensor } from '../tensor.js';
-import { getBuffer, getLayout, getWeightDtype } from '../weight-buffer.js';
-import { log, trace, isTraceEnabled } from '../../debug/index.js';
-import { releaseBuffer } from '../../memory/buffer-pool.js';
-import { releaseUniformBuffer } from '../uniform-cache.js';
-import { castF16ToF32, recordCastF16ToF32 } from './cast.js';
-import {
-  resolveMatmulPhase,
-  resolveMatmulConstants,
-  getMatmulConfig,
-  isFusedQ4KDisabled,
-  toMatmulDtype,
-  resolveTransposeB,
-  validateMatmulDimensions,
-  validateMatmulOffsets,
-  getMatmulBindingSizes,
-  requiresF32Input,
-  selectMatmulVariantAndFlags,
-  resolveMatmulOutput,
-  selectMatmulKernel,
-} from './matmul-selection.js';
-import {
-  MatmulKernel,
-  calculateMatmulDispatch,
-  createMatmulUniformBuffer,
-  createMatmulBindGroupLayout,
-  getMatmulPipeline,
-} from './matmul-dispatch.js';
-
-export { isFusedQ4KDisabled, selectMatmulKernel };
-export { createMatmulBindGroupLayout };
-
-let _runMatmulDebugCount = 0;
-let _recordMatmulDebugCount = 0;
-
-function getDebugCounter(isRecord) {
-  return isRecord ? _recordMatmulDebugCount : _runMatmulDebugCount;
-}
-
-function incrementDebugCounter(isRecord) {
-  if (isRecord) {
-    _recordMatmulDebugCount += 1;
-    return;
-  }
-  _runMatmulDebugCount += 1;
-}
-
-function buildProfileLabel(options = {}) {
-  const layerLabel = Number.isFinite(options.layerIdx) ? `:L${options.layerIdx}` : '';
-  const roleLabel = options.role ? `:${options.role}` : '';
-  return `matmul${roleLabel}${layerLabel}`;
-}
-
-function assertBindGroupBuffer(kernelName, variant, bindingIndex, bindingLabel, buffer, details = []) {
-  const isGpuBuffer = buffer && (
-    typeof GPUBuffer === 'undefined'
-      ? true
-      : buffer instanceof GPUBuffer
-  );
-  if (isGpuBuffer) {
-    return;
-  }
-  const detailText = details.filter(Boolean).join(', ');
-  throw new Error(
-    `[${kernelName}] variant="${variant}" binding ${bindingIndex} "${bindingLabel}" requires a GPUBuffer` +
-    (detailText ? ` (${detailText})` : '') +
-    '.'
-  );
-}
-
-function createMatmulBindGroupEntries(variant, uniformBuffer, matmulInput, bBuffer, outputBuffer, offsets, bindingSizes) {
-  const isQ4KF16 = variant === 'q4_fused_multicol_f16'
-    || variant === 'q4_fused_f16a'
-    || variant === 'q4_fused_batched_f16'
-    || variant === 'q4_fused_multicol_f16a'
-    || variant === 'q4_fused_batched_f16a';
-
-  assertBindGroupBuffer('matmul', variant, 0, 'uniforms', uniformBuffer);
-  assertBindGroupBuffer('matmul', variant, 1, 'input', matmulInput?.buffer, [
-    `inputLabel=${matmulInput?.label ?? 'unknown'}`,
-    `inputDtype=${matmulInput?.dtype ?? 'unknown'}`,
-  ]);
-  assertBindGroupBuffer('matmul', variant, 2, 'weights', bBuffer);
-  assertBindGroupBuffer('matmul', variant, isQ4KF16 ? 4 : 3, 'output', outputBuffer);
-
-  const entries = [
-    { binding: 0, resource: { buffer: uniformBuffer } },
-    { binding: 1, resource: { buffer: matmulInput.buffer, offset: offsets.aOffset, size: bindingSizes.aBindingSize } },
-    { binding: 2, resource: { buffer: bBuffer, offset: offsets.bOffset, size: bindingSizes.bBindingSize } },
-  ];
-
-  if (isQ4KF16) {
-    entries.push({
-      binding: 4,
-      resource: { buffer: outputBuffer, offset: offsets.cOffset, size: bindingSizes.cBindingSize },
-    });
-  } else {
-    entries.push({
-      binding: 3,
-      resource: { buffer: outputBuffer, offset: offsets.cOffset, size: bindingSizes.cBindingSize },
-    });
-  }
-
-  return entries;
-}
-
-async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
-  const isRecord = Boolean(recorder);
-  const mode = isRecord ? 'record' : 'run';
-  const opLabel = isRecord ? 'recordMatmul' : 'runMatmul';
-  const device = recorder?.device || getDevice();
-  const capabilities = getKernelCapabilities();
-
-  const {
-    alpha = 1.0,
-    outputBuffer = null,
-    transposeB: transposeBOption = true,
-    aOffset = 0,
-    bOffset = 0,
-    cOffset = 0,
-  } = options;
-
-  const bBuffer = getBuffer(B);
-  const weightDtype = getWeightDtype(B);
-  const weightLabel = (B && typeof B === 'object' ? B.label : null) ?? bBuffer?.label ?? null;
-  const weightLayout = getLayout(B);
-  const weightShape = B?.shape ? `[${B.shape.join(', ')}]` : null;
-
-  if (isTraceEnabled('kernels') && getDebugCounter(isRecord) < 20) {
-    incrementDebugCounter(isRecord);
-    const modeLabel = isRecord ? 'recordMatmul' : 'runMatmul';
-    trace.kernels(`${modeLabel}: M=${M}, N=${N}, K=${K}, transposeBOption=${transposeBOption}, weightLayout=${weightLayout}, weightDtype=${weightDtype}`);
-  }
-
-  const transposeB = resolveTransposeB(B, transposeBOption);
-  validateMatmulDimensions(opLabel, M, N, K);
-
-  const aDtype = toMatmulDtype(A.dtype);
-  const bDtype = toMatmulDtype(weightDtype ?? options.bDtype);
-  const requestedOutputDtype = options.outputDtype || A.dtype;
-
-  if (bDtype === 'f16' && capabilities?.hasF16 !== true) {
-    throw new Error(`[${opLabel}] f16 weights require shader-f16 support.`);
-  }
-  if (requestedOutputDtype === 'f16' && capabilities?.hasF16 !== true) {
-    throw new Error(`[${opLabel}] f16 output requires shader-f16 support.`);
-  }
-
-  if (!isRecord && isTraceEnabled('kernels') && !weightDtype && !options.bDtype && M <= 2) {
-    log.warn('Matmul', `runMatmul: B buffer dtype unknown! size=${bBuffer.size}, M=${M}, N=${N}, K=${K}. Assuming f32.`);
-  }
-
-  validateMatmulOffsets(opLabel, aOffset, bOffset, cOffset);
-
-  const { variant, useQ4KFused, useGemv } = selectMatmulVariantAndFlags(
-    mode,
-    M,
-    N,
-    K,
-    aDtype,
-    bDtype,
-    transposeB,
-    requestedOutputDtype,
-    options
-  );
-
-  const phase = resolveMatmulPhase(M, options.phaseOverride ?? null);
-  const constants = resolveMatmulConstants(options, phase);
-
-  let matmulInput = A;
-  let matmulADtype = aDtype;
-  let castedInput = null;
-  if (matmulADtype === 'f16' && requiresF32Input(variant)) {
-    if (isTraceEnabled('kernels')) {
-      trace.kernels(`Matmul: casting f16 activations to f32 for variant=${variant}`);
-    }
-    if (isRecord) {
-      castedInput = await recordCastF16ToF32(recorder, A);
-      recorder.trackTemporaryBuffer(castedInput.buffer);
-    } else {
-      castedInput = await castF16ToF32(A);
-    }
-    matmulInput = castedInput;
-    matmulADtype = 'f32';
-  }
-
-  let bindingSizes;
-  try {
-    bindingSizes = getMatmulBindingSizes(
-      opLabel,
-      matmulInput.buffer,
-      bBuffer,
-      M,
-      N,
-      K,
-      matmulADtype,
-      bDtype,
-      transposeB,
-      aOffset,
-      bOffset
-    );
-  } catch (err) {
-    if (!isRecord && err instanceof Error && err.message.includes('B buffer too small')) {
-      const detailParts = [];
-      if (weightLabel) detailParts.push(`label=${weightLabel}`);
-      if (weightDtype) detailParts.push(`weightDtype=${weightDtype}`);
-      if (weightLayout) detailParts.push(`layout=${weightLayout}`);
-      if (weightShape) detailParts.push(`shape=${weightShape}`);
-      if (Number.isFinite(bBuffer?.size)) detailParts.push(`bSize=${bBuffer.size}`);
-      const detail = detailParts.length ? ` (${detailParts.join(', ')})` : '';
-      throw new Error(`${err.message}${detail}`);
-    }
-    throw err;
-  }
-
-  if (!isRecord && isTraceEnabled('kernels') && bDtype === 'q4k') {
-    if (useQ4KFused) {
-      trace.kernels(`Q4K FUSED: M=${M}, N=${N}, K=${K}, variant=${variant} (WARNING: 2.3x slower than dequant)`);
-    } else {
-      trace.kernels(`Q4K DEQUANT: M=${M}, N=${N}, K=${K}, will dequant first then matmul with variant=${variant}`);
-    }
-  }
-
-  if (!isRecord && isTraceEnabled('kernels') && N > 100000) {
-    trace.kernels(`MATMUL_LARGE: N=${N}, variant=${variant}, aDtype=${aDtype}, bDtype=${bDtype}, transposeB=${transposeB}`);
-  }
-
-  const config = getMatmulConfig(variant, constants);
-  const kernel = new MatmulKernel(device);
-  const pipeline = await getMatmulPipeline(variant, constants);
-
-  const { output: C, outputSize, cBindingSize, actualOutputDtype } = resolveMatmulOutput(
-    variant,
-    M,
-    N,
-    outputBuffer
-  );
-  const ownsOutput = outputBuffer == null;
-
-  if (!Number.isFinite(outputSize) || outputSize <= 0) {
-    throw new Error(`[${opLabel}] Invalid output size: ${outputSize} (M=${M}, N=${N})`);
-  }
-
-  const cRequired = cOffset + cBindingSize;
-  if (C.size < cRequired) {
-    throw new Error(`[${opLabel}] Output buffer too small: ${C.size} < ${cRequired} (M=${M}, N=${N})`);
-  }
-
-  const dispatchPlan = calculateMatmulDispatch(variant, useQ4KFused, useGemv, M, N, config);
-  let uniformBuffer = null;
-  let completed = false;
-  try {
-    uniformBuffer = createMatmulUniformBuffer(
-      'matmul_uniforms',
-      M,
-      N,
-      K,
-      alpha,
-      useQ4KFused,
-      transposeB,
-      dispatchPlan.uniformWorkgroupsX,
-      recorder || null,
-      device
-    );
-
-    const entries = createMatmulBindGroupEntries(
-      variant,
-      uniformBuffer,
-      matmulInput,
-      bBuffer,
-      C,
-      { aOffset, bOffset, cOffset },
-      {
-        aBindingSize: bindingSizes.aBindingSize,
-        bBindingSize: bindingSizes.bBindingSize,
-        cBindingSize,
-      }
-    );
-
-    const bindGroup = device.createBindGroup({
-      label: 'matmul_bind_group',
-      layout: pipeline.getBindGroupLayout(0),
-      entries,
-    });
-
-    if (isRecord) {
-      kernel.record(recorder, pipeline, bindGroup, dispatchPlan.workgroups, buildProfileLabel(options));
-    } else {
-      kernel.dispatch(pipeline, bindGroup, dispatchPlan.workgroups);
-    }
-    completed = true;
-    return createTensor(C, actualOutputDtype, [M, N], 'matmul_output');
-  } finally {
-    if (!isRecord && uniformBuffer) {
-      releaseUniformBuffer(uniformBuffer);
-    }
-    if (!isRecord && castedInput) {
-      releaseBuffer(castedInput.buffer);
-    }
-    if (!completed && ownsOutput) {
-      releaseBuffer(C);
-    }
-  }
-}
-
-
-export async function runMatmul(A, B, M, N, K, options = {}) {
-  return executeMatmul(null, A, B, M, N, K, options);
-}
-
-
-export async function recordMatmul(recorder, A, B, M, N, K, options = {}) {
-  return executeMatmul(recorder, A, B, M, N, K, options);
-}

package/src/gpu/kernels/matmul_f16.wgsl

@@ -1,170 +0,0 @@
-// Matrix Multiplication Kernel - FP16 (Half Precision)
-//
-// Tiled matrix multiplication using FP16 for improved throughput.
-// Requires 'shader-f16' feature enabled.
-// C[M,N] = A[M,K] * B[K,N]  (or B^T when transpose_b=1)
-//
-// Uses FP16 for storage and computation, with optional FP32 accumulation
-// for better numerical stability.
-
-enable f16;
-
-// Tile dimensions - can use larger tiles with f16 due to smaller footprint
-const MAX_TILE_SIZE: u32 = 16u;
-const TILE_AREA: u32 = MAX_TILE_SIZE * MAX_TILE_SIZE;
-
-override TILE_SIZE: u32 = 16u;  // Must be <= MAX_TILE_SIZE
-
-// Uniforms for matrix dimensions
-struct Uniforms {
-    M: u32,           // Rows of A and C
-    N: u32,           // Cols of B and C (or rows of B when transposed)
-    K: u32,           // Cols of A, Rows of B (or cols of B when transposed)
-    alpha: f32,       // Scaling factor
-    transpose_b: u32, // 0 = normal, 1 = B is stored transposed [N,K] -> treat as [K,N]
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> A: array<f16>;
-@group(0) @binding(2) var<storage, read> B: array<f16>;
-@group(0) @binding(3) var<storage, read_write> C: array<f16>;
-
-// Shared memory tiles - f16 allows 2x data in same space
-var<workgroup> tileA: array<f16, TILE_AREA>;
-var<workgroup> tileB: array<f16, TILE_AREA>;
-var<workgroup> tileB_vec4: array<vec4<f16>, TILE_AREA>;
-
-@compute @workgroup_size(TILE_SIZE, TILE_SIZE, 1)
-fn main(
-    @builtin(global_invocation_id) global_id: vec3<u32>,
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) workgroup_id: vec3<u32>
-) {
-    if (TILE_SIZE > MAX_TILE_SIZE) {
-        return;
-    }
-
-    let row = global_id.x;
-    let col = global_id.y;
-    let local_row = local_id.x;
-    let local_col = local_id.y;
-
-    // Use f32 accumulator for numerical stability during summation
-    var sum: f32 = 0.0;
-
-    let num_tiles = (u.K + TILE_SIZE - 1u) / TILE_SIZE;
-    let tile_idx = local_row * TILE_SIZE + local_col;
-
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        let a_col = t * TILE_SIZE + local_col;
-        let b_row = t * TILE_SIZE + local_row;
-
-        // Load tile from A
-        if (row < u.M && a_col < u.K) {
-            tileA[tile_idx] = A[row * u.K + a_col];
-        } else {
-            tileA[tile_idx] = f16(0.0);
-        }
-
-        // Load tile from B (handle transpose)
-        if (b_row < u.K && col < u.N) {
-            if (u.transpose_b == 0u) {
-                tileB[tile_idx] = B[b_row * u.N + col];
-            } else {
-                // B is [N, K], access element [col, b_row]
-                tileB[tile_idx] = B[col * u.K + b_row];
-            }
-        } else {
-            tileB[tile_idx] = f16(0.0);
-        }
-
-        workgroupBarrier();
-
-        // Compute partial dot product with f32 accumulation
-        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
-            let a_val = f32(tileA[local_row * TILE_SIZE + k]);
-            let b_val = f32(tileB[k * TILE_SIZE + local_col]);
-            sum = sum + a_val * b_val;
-        }
-
-        workgroupBarrier();
-    }
-
-    // Write result back as f16
-    if (row < u.M && col < u.N) {
-        C[row * u.N + col] = f16(sum * u.alpha);
-    }
-}
-
-// Alternative entry point for vec4 column groups (best when N is a multiple of 4)
-@compute @workgroup_size(TILE_SIZE, TILE_SIZE, 1)
-fn main_vec4(
-    @builtin(global_invocation_id) global_id: vec3<u32>,
-    @builtin(local_invocation_id) local_id: vec3<u32>
-) {
-    if (TILE_SIZE > MAX_TILE_SIZE) {
-        return;
-    }
-
-    let row = global_id.x;
-    let col_base = global_id.y * 4u;  // Each thread handles 4 columns
-    let local_row = local_id.x;
-    let local_col = local_id.y;
-
-    var sum: vec4<f32> = vec4<f32>(0.0);
-
-    let num_tiles = (u.K + TILE_SIZE - 1u) / TILE_SIZE;
-    let tile_idx = local_row * TILE_SIZE + local_col;
-
-    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
-        let a_col = t * TILE_SIZE + local_col;
-        let b_row = t * TILE_SIZE + local_row;
-
-        // Load A tile (single element per thread)
-        if (row < u.M && a_col < u.K) {
-            tileA[tile_idx] = A[row * u.K + a_col];
-        } else {
-            tileA[tile_idx] = f16(0.0);
-        }
-
-        // Load B tile (handle transpose)
-        var b_vec: vec4<f16> = vec4<f16>(f16(0.0));
-        if (b_row < u.K) {
-            if (u.transpose_b == 0u) {
-                if (col_base + 0u < u.N) { b_vec.x = B[b_row * u.N + col_base]; }
-                if (col_base + 1u < u.N) { b_vec.y = B[b_row * u.N + col_base + 1u]; }
-                if (col_base + 2u < u.N) { b_vec.z = B[b_row * u.N + col_base + 2u]; }
-                if (col_base + 3u < u.N) { b_vec.w = B[b_row * u.N + col_base + 3u]; }
-            } else {
-                if (col_base + 0u < u.N) { b_vec.x = B[(col_base + 0u) * u.K + b_row]; }
-                if (col_base + 1u < u.N) { b_vec.y = B[(col_base + 1u) * u.K + b_row]; }
-                if (col_base + 2u < u.N) { b_vec.z = B[(col_base + 2u) * u.K + b_row]; }
-                if (col_base + 3u < u.N) { b_vec.w = B[(col_base + 3u) * u.K + b_row]; }
-            }
-        }
-        tileB_vec4[tile_idx] = b_vec;
-
-        workgroupBarrier();
-
-        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
-            let a_val = f32(tileA[local_row * TILE_SIZE + k]);
-            let b_val = tileB_vec4[k * TILE_SIZE + local_col];
-            sum = sum + a_val * vec4<f32>(b_val);
-        }
-
-        workgroupBarrier();
-    }
-
-    // Write results
-    if (row < u.M && col_base < u.N) {
-        let base = row * u.N + col_base;
-        let scaled = sum * u.alpha;
-        if (col_base + 0u < u.N) { C[base] = f16(scaled.x); }
-        if (col_base + 1u < u.N) { C[base + 1u] = f16(scaled.y); }
-        if (col_base + 2u < u.N) { C[base + 2u] = f16(scaled.z); }
-        if (col_base + 3u < u.N) { C[base + 3u] = f16(scaled.w); }
-    }
-}