@simulatte/doppler 0.1.0

package/src/gpu/kernels/kernel-configs.d.ts

@@ -0,0 +1,65 @@
+/**
+ * Kernel Configurations - Table-driven kernel metadata
+ *
+ * Contains all kernel configurations including shader files, entry points,
+ * workgroup sizes, and feature requirements.
+ *
+ * @module gpu/kernels/kernel-configs
+ */
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/** Variant-specific metadata for table-driven kernel configuration */
+export interface VariantMetadata {
+  /** Columns processed per workgroup (matmul multicol variants) */
+  colsPerWg?: number;
+  /** Tile size for M dimension (batched matmul variants) */
+  tileM?: number;
+  /** Output buffer binding index (gather F16 output variants) */
+  outputBinding?: number;
+  /** Maximum KV length for chunked attention */
+  maxKVLen?: number;
+}
+
+/** Kernel configuration */
+export interface KernelConfig {
+  shaderFile: string;
+  entryPoint: string;
+  workgroupSize: [number, number, number];
+  requires: string[];
+  validate?: (seqLen: number, numHeads: number, headDim: number) => void;
+  /** Output dtype for variants that output to specific precision */
+  outputDtype?: 'f16' | 'f32';
+  /** Tile size for attention/matmul kernels */
+  tileSize?: number;
+  /** Additional variant-specific configuration */
+  variantMetadata?: VariantMetadata;
+}
+
+// ============================================================================
+// Kernel Configurations
+// ============================================================================
+
+/** All kernel configurations by operation and variant */
+export declare const KERNEL_CONFIGS: Record<string, Record<string, KernelConfig>>;
+
+// ============================================================================
+// Config Helpers
+// ============================================================================
+
+/**
+ * Get kernel configuration
+ */
+export declare function getKernelConfig(operation: string, variant: string): KernelConfig;
+
+/**
+ * Set a validator function on a kernel config.
+ * Used to set attention validators after import to avoid circular dependencies.
+ */
+export declare function setKernelValidator(
+  operation: string,
+  variant: string,
+  validator: (seqLen: number, numHeads: number, headDim: number) => void
+): void;

package/src/gpu/kernels/kernel-configs.js

@@ -0,0 +1,50 @@
+import { loadJson } from '../../utils/load-json.js';
+
+const registry = await loadJson('../../config/kernels/registry.json', import.meta.url, 'Failed to load registry');
+import { resolveKernelConfig } from '../../config/schema/kernel-registry.schema.js';
+
+export const KERNEL_CONFIGS = Object.fromEntries(
+  Object.entries(registry.operations).map(([operation, opSchema]) => {
+    const variants = Object.fromEntries(
+      Object.entries(opSchema.variants).map(([variant, variantSchema]) => {
+        const resolved = resolveKernelConfig(operation, variant, opSchema, variantSchema);
+        return [
+          variant,
+          {
+            shaderFile: resolved.wgsl,
+            entryPoint: resolved.entryPoint,
+            workgroupSize: resolved.workgroup,
+            requires: resolved.requires,
+            bindings: resolved.bindings,
+            uniforms: resolved.uniforms,
+            wgslOverrides: resolved.wgslOverrides,
+            sharedMemory: resolved.sharedMemory,
+            outputDtype: resolved.outputDtype ?? undefined,
+            variantMetadata: resolved.variantMetadata ?? undefined,
+          },
+        ];
+      })
+    );
+
+    return [operation, variants];
+  })
+);
+
+export function getKernelConfig(operation, variant) {
+  const config = KERNEL_CONFIGS[operation]?.[variant];
+  if (!config) {
+    throw new Error(`Unknown kernel: ${operation}/${variant}`);
+  }
+  return config;
+}
+
+export function setKernelValidator(
+  operation,
+  variant,
+  validator
+) {
+  const config = KERNEL_CONFIGS[operation]?.[variant];
+  if (config) {
+    config.validate = validator;
+  }
+}

package/src/gpu/kernels/kernel-tuning.d.ts

@@ -0,0 +1,42 @@
+/**
+ * Kernel Tuning - Auto-tuning and kernel prewarming
+ *
+ * Provides utilities for tuning kernel workgroup sizes and
+ * prewarming kernel pipelines for optimal performance.
+ *
+ * @module gpu/kernels/kernel-tuning
+ */
+
+// ============================================================================
+// Workgroup Size Tuning
+// ============================================================================
+
+/**
+ * Get tuned workgroup size for an operation
+ */
+export declare function getTunedWorkgroupSize(
+  operation: string,
+  inputSizes?: Record<string, number>
+): Promise<[number, number, number]>;
+
+// ============================================================================
+// Auto-Tuning
+// ============================================================================
+
+/**
+ * Run auto-tuning for all kernels with given model config
+ */
+export declare function autoTuneKernels(
+  modelConfig?: Record<string, number>
+): Promise<Record<string, any>>;
+
+// ============================================================================
+// Pipeline Prewarming
+// ============================================================================
+
+/**
+ * Prewarm all supported kernel pipelines
+ */
+export declare function prewarmKernels(
+  options?: { mode?: 'parallel' | 'sequential' }
+): Promise<void>;

package/src/gpu/kernels/kernel-tuning.js

@@ -0,0 +1,149 @@
+
+
+import { getKernelCapabilities } from '../device.js';
+import { getKernelTuner, getTunerConfig } from '../kernel-tuner.js';
+import { KERNEL_CONFIGS } from './kernel-configs.js';
+import { createPipeline } from './pipeline-cache.js';
+import { hasRequiredFeatures } from './feature-check.js';
+import { log } from '../../debug/index.js';
+
+// ============================================================================
+// Workgroup Size Tuning
+// ============================================================================
+
+
+export async function getTunedWorkgroupSize(
+  operation,
+  inputSizes = {}
+) {
+  try {
+    const tuner = await getKernelTuner();
+    const result = tuner.getCachedResult(operation, inputSizes);
+
+    if (result) {
+      return result.optimalWorkgroupSize;
+    }
+
+    // Run tuning if not cached
+    const tuneResult = await tuner.tuneKernel(operation, inputSizes);
+    return tuneResult.optimalWorkgroupSize;
+  } catch (e) {
+    log.warn('KernelTuning', `Tuning failed for ${operation}, using defaults: ${e.message}`);
+    const { fallbackWorkgroupSizes } = getTunerConfig();
+    const fallback = fallbackWorkgroupSizes?.[operation] ?? fallbackWorkgroupSizes?.default;
+    if (!fallback) {
+      throw new Error(`KernelTuning: missing fallback workgroup size for "${operation}".`);
+    }
+    return fallback;
+  }
+}
+
+// ============================================================================
+// Auto-Tuning
+// ============================================================================
+
+
+export async function autoTuneKernels(
+  modelConfig = {}
+) {
+  const {
+    hiddenSize = 4096,
+    intermediateSize = 14336,
+    numHeads = 32,
+    headDim = 128,
+    maxSeqLen = 4096,
+    vocabSize = 32000,
+  } = modelConfig;
+
+  const tuner = await getKernelTuner();
+  
+  const results = {};
+
+  // Tune matmul for common sizes
+  results.matmul_hidden = await tuner.tuneKernel('matmul', {
+    M: 1, N: hiddenSize, K: hiddenSize,
+  });
+  results.matmul_ffn = await tuner.tuneKernel('matmul', {
+    M: 1, N: intermediateSize, K: hiddenSize,
+  });
+
+  // Tune attention
+  results.attention = await tuner.tuneKernel('attention', {
+    seqLen: 1, numHeads, headDim,
+  });
+
+  // Tune softmax (LM head output)
+  results.softmax = await tuner.tuneKernel('softmax', {
+    innerSize: vocabSize, outerSize: 1,
+  });
+
+  // Tune RMSNorm
+  results.rmsnorm = await tuner.tuneKernel('rmsnorm', {
+    hiddenSize, numTokens: 1,
+  });
+
+  // Tune dequant
+  results.dequant = await tuner.tuneKernel('dequant', {
+    numBlocks: 1000,
+  });
+
+  log.debug('KernelTuning', `Auto-tuning complete: ${JSON.stringify(results)}`);
+  return results;
+}
+
+// ============================================================================
+// Pipeline Prewarming
+// ============================================================================
+
+
+export async function prewarmKernels(
+  options = {}
+) {
+  const caps = getKernelCapabilities();
+  const mode = options.mode ?? 'parallel';
+  const entries = Object.entries(KERNEL_CONFIGS)
+    .sort(([a], [b]) => a.localeCompare(b))
+    .map(([operation, variants]) => [
+      operation,
+      Object.entries(variants).sort(([a], [b]) => a.localeCompare(b))
+    ]);
+
+  if (mode === 'sequential') {
+    let count = 0;
+    for (const [operation, variants] of entries) {
+      for (const [variant, cfg] of variants) {
+        if (cfg.requires && !hasRequiredFeatures(cfg.requires, caps)) {
+          continue;
+        }
+        try {
+          await createPipeline(operation, variant);
+          count += 1;
+        } catch (e) {
+          log.warn('KernelTuning', `Prewarm failed for ${operation}/${variant}: ${e.message}`);
+        }
+      }
+    }
+    log.debug('KernelTuning', `Prewarmed ${count} kernel pipelines`);
+    return;
+  }
+
+  
+  const jobs = [];
+  for (const [operation, variants] of entries) {
+    for (const [variant, cfg] of variants) {
+      if (cfg.requires && !hasRequiredFeatures(cfg.requires, caps)) {
+        continue;
+      }
+      jobs.push(
+        createPipeline(operation, variant)
+          .then(() => {}) // Ignore the pipeline result
+          .catch((e) => {
+            log.warn('KernelTuning', `Prewarm failed for ${operation}/${variant}: ${e.message}`);
+          })
+      );
+    }
+  }
+
+  await Promise.all(jobs);
+  log.debug('KernelTuning', `Prewarmed ${jobs.length} kernel pipelines`);
+}

package/src/gpu/kernels/kv-quantize.d.ts

@@ -0,0 +1,37 @@
+/**
+ * KV Quantization Kernels
+ *
+ * Quantize KV tensors into packed int8/int4 with per-token+head scales.
+ *
+ * @module gpu/kernels/kv-quantize
+ */
+
+export interface KVQuantizeOptions {
+  numKVHeads: number;
+  headDim: number;
+  startPos: number;
+  numTokens: number;
+  packedStride: number;
+  mode?: 'int8' | 'int4';
+}
+
+export declare function runKVQuantize(
+  keys: GPUBuffer,
+  values: GPUBuffer,
+  outputKeys: GPUBuffer,
+  outputValues: GPUBuffer,
+  scalesK: GPUBuffer,
+  scalesV: GPUBuffer,
+  options: KVQuantizeOptions
+): Promise<void>;
+
+export declare function recordKVQuantize(
+  recorder: import('../command-recorder.js').CommandRecorder,
+  keys: GPUBuffer,
+  values: GPUBuffer,
+  outputKeys: GPUBuffer,
+  outputValues: GPUBuffer,
+  scalesK: GPUBuffer,
+  scalesV: GPUBuffer,
+  options: KVQuantizeOptions
+): Promise<void>;

package/src/gpu/kernels/kv-quantize.js

@@ -0,0 +1,138 @@
+
+import { getDevice, getKernelCapabilities } from '../device.js';
+import { dispatch, recordDispatch } from './dispatch.js';
+import { createPipeline, createUniformBufferWithView, getKernelConfig, hasRequiredFeatures } from './utils.js';
+import { selectRuleValue as selectKernelRuleValue } from './rule-registry.js';
+
+
+function resolveQuantizeVariant(mode) {
+  return selectKernelRuleValue('kv_quantize', 'variant', { mode });
+}
+
+
+function createQuantizeUniformBuffer(device, recorder, params) {
+  return createUniformBufferWithView(
+    'kv_quantize_uniforms',
+    32,
+    (view) => {
+      view.setUint32(0, params.numKVHeads, true);
+      view.setUint32(4, params.headDim, true);
+      view.setUint32(8, params.startPos, true);
+      view.setUint32(12, params.numTokens, true);
+      view.setUint32(16, params.packedStride, true);
+      view.setUint32(20, 0, true);
+    },
+    recorder,
+    device
+  );
+}
+
+
+export async function runKVQuantize(
+  keys,
+  values,
+  outputKeys,
+  outputValues,
+  scalesK,
+  scalesV,
+  options = {}
+) {
+  const device = getDevice();
+  const {
+    numKVHeads,
+    headDim,
+    startPos,
+    numTokens,
+    packedStride,
+    mode = 'int8',
+  } = options;
+
+  const variant = resolveQuantizeVariant(mode);
+  const config = getKernelConfig('kv_quantize', variant);
+  const caps = getKernelCapabilities();
+  if (!hasRequiredFeatures(config.requires, caps)) {
+    throw new Error(`KV quantize kernel "${variant}" requires unsupported GPU features.`);
+  }
+
+  const pipeline = await createPipeline('kv_quantize', variant);
+  const uniformBuffer = createQuantizeUniformBuffer(device, null, {
+    numKVHeads,
+    headDim,
+    startPos,
+    numTokens,
+    packedStride,
+  });
+
+  const bindGroup = device.createBindGroup({
+    label: 'kv_quantize_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: keys } },
+      { binding: 2, resource: { buffer: values } },
+      { binding: 3, resource: { buffer: outputKeys } },
+      { binding: 4, resource: { buffer: outputValues } },
+      { binding: 5, resource: { buffer: scalesK } },
+      { binding: 6, resource: { buffer: scalesV } },
+    ],
+  });
+
+  const workgroups = [numKVHeads, numTokens, 1];
+  dispatch(device, pipeline, bindGroup, workgroups, 'kv_quantize');
+  uniformBuffer.destroy();
+}
+
+
+export async function recordKVQuantize(
+  recorder,
+  keys,
+  values,
+  outputKeys,
+  outputValues,
+  scalesK,
+  scalesV,
+  options = {}
+) {
+  const device = recorder.device;
+  const {
+    numKVHeads,
+    headDim,
+    startPos,
+    numTokens,
+    packedStride,
+    mode = 'int8',
+  } = options;
+
+  const variant = resolveQuantizeVariant(mode);
+  const config = getKernelConfig('kv_quantize', variant);
+  const caps = getKernelCapabilities();
+  if (!hasRequiredFeatures(config.requires, caps)) {
+    throw new Error(`KV quantize kernel "${variant}" requires unsupported GPU features.`);
+  }
+
+  const pipeline = await createPipeline('kv_quantize', variant);
+  const uniformBuffer = createQuantizeUniformBuffer(device, recorder, {
+    numKVHeads,
+    headDim,
+    startPos,
+    numTokens,
+    packedStride,
+  });
+
+  const bindGroup = device.createBindGroup({
+    label: 'kv_quantize_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: keys } },
+      { binding: 2, resource: { buffer: values } },
+      { binding: 3, resource: { buffer: outputKeys } },
+      { binding: 4, resource: { buffer: outputValues } },
+      { binding: 5, resource: { buffer: scalesK } },
+      { binding: 6, resource: { buffer: scalesV } },
+    ],
+  });
+
+  const workgroups = [numKVHeads, numTokens, 1];
+  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'kv_quantize');
+}

package/src/gpu/kernels/kv_quantize_int4.wgsl

@@ -0,0 +1,119 @@
+// KV Quantization Kernel (int4)
+//
+// Packs 8 int4 values into u32, with per-token+head scale (f16).
+
+enable f16;
+
+const MAX_WORKGROUP_SIZE: u32 = 256u;
+const PACK_FACTOR: u32 = 8u;
+const QMAX: f32 = 7.0;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_kv_heads: u32,
+    head_dim: u32,
+    start_pos: u32,
+    num_tokens: u32,
+    packed_stride: u32,
+    _pad: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input_k: array<f16>;
+@group(0) @binding(2) var<storage, read> input_v: array<f16>;
+@group(0) @binding(3) var<storage, read_write> output_k: array<u32>;
+@group(0) @binding(4) var<storage, read_write> output_v: array<u32>;
+@group(0) @binding(5) var<storage, read_write> scales_k: array<f16>;
+@group(0) @binding(6) var<storage, read_write> scales_v: array<f16>;
+
+var<workgroup> shared_k: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_v: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_abs_k: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_abs_v: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_qk: array<i32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_qv: array<i32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_scale_k: f32;
+var<workgroup> shared_scale_v: f32;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) workgroup_id: vec3<u32>,
+) {
+    let head_idx = workgroup_id.x;
+    let token_idx = workgroup_id.y;
+    let tid = local_id.x;
+    let head_dim = u.head_dim;
+
+    if (head_idx >= u.num_kv_heads || token_idx >= u.num_tokens) {
+        return;
+    }
+    if (head_dim > WORKGROUP_SIZE || WORKGROUP_SIZE > MAX_WORKGROUP_SIZE) {
+        return;
+    }
+
+    let valid = tid < head_dim;
+    let token_base = token_idx * u.num_kv_heads + head_idx;
+    let input_base = token_base * head_dim;
+
+    if (valid) {
+        let idx = input_base + tid;
+        let k_val = f32(input_k[idx]);
+        let v_val = f32(input_v[idx]);
+        shared_k[tid] = k_val;
+        shared_v[tid] = v_val;
+        shared_abs_k[tid] = abs(k_val);
+        shared_abs_v[tid] = abs(v_val);
+    }
+    workgroupBarrier();
+
+    if (tid == 0u) {
+        var max_k: f32 = 0.0;
+        var max_v: f32 = 0.0;
+        for (var d: u32 = 0u; d < head_dim; d++) {
+            max_k = max(max_k, shared_abs_k[d]);
+            max_v = max(max_v, shared_abs_v[d]);
+        }
+        var scale_k = max_k / QMAX;
+        var scale_v = max_v / QMAX;
+        if (scale_k == 0.0) { scale_k = 1.0; }
+        if (scale_v == 0.0) { scale_v = 1.0; }
+        shared_scale_k = scale_k;
+        shared_scale_v = scale_v;
+        let scale_idx = (u.start_pos + token_idx) * u.num_kv_heads + head_idx;
+        scales_k[scale_idx] = f16(scale_k);
+        scales_v[scale_idx] = f16(scale_v);
+    }
+    workgroupBarrier();
+
+    if (valid) {
+        let scale_k = shared_scale_k;
+        let scale_v = shared_scale_v;
+        let qk = i32(round(clamp(shared_k[tid] / scale_k, -QMAX, QMAX)));
+        let qv = i32(round(clamp(shared_v[tid] / scale_v, -QMAX, QMAX)));
+        shared_qk[tid] = qk;
+        shared_qv[tid] = qv;
+    }
+    workgroupBarrier();
+
+    if ((tid % PACK_FACTOR) == 0u) {
+        let pack_idx = tid / PACK_FACTOR;
+        if (pack_idx < u.packed_stride) {
+            var packed_k: u32 = 0u;
+            var packed_v: u32 = 0u;
+            for (var i: u32 = 0u; i < PACK_FACTOR; i++) {
+                let lane = tid + i;
+                let qk = select(0, shared_qk[lane], lane < head_dim);
+                let qv = select(0, shared_qv[lane], lane < head_dim);
+                let uk = u32(qk) & 0xfu;
+                let uv = u32(qv) & 0xfu;
+                packed_k = packed_k | (uk << (4u * i));
+                packed_v = packed_v | (uv << (4u * i));
+            }
+            let packed_offset = ((u.start_pos + token_idx) * u.num_kv_heads + head_idx) * u.packed_stride + pack_idx;
+            output_k[packed_offset] = packed_k;
+            output_v[packed_offset] = packed_v;
+        }
+    }
+}