@simulatte/doppler 0.1.4 → 0.1.5

package/src/gpu/kernels/relu.js

@@ -0,0 +1,45 @@
+import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../tensor.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { selectRuleValue } from './rule-registry.js';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectReluVariant(dtype) {
+  return selectRuleValue('relu', 'variant', { dtype });
+}
+
+function resolveCount(input, countOverride) {
+  if (Number.isFinite(countOverride) && countOverride > 0) {
+    return Math.floor(countOverride);
+  }
+  if (Array.isArray(input.shape) && input.shape.length > 0) {
+    return input.shape.reduce((acc, value) => acc * value, 1);
+  }
+  return Math.floor(input.buffer.size / dtypeBytes(input.dtype));
+}
+
+async function _relu(target, input, options = {}) {
+  const { count = null, outputBuffer = null } = options;
+  const size = resolveCount(input, count);
+  const variant = selectReluVariant(input.dtype);
+  const output = outputBuffer || acquireBuffer(size * dtypeBytes(input.dtype), undefined, 'relu_output');
+
+  await unifiedKernelWrapper(
+    'relu',
+    target,
+    variant,
+    [input, output],
+    { size, _pad0: 0, _pad1: 0, _pad2: 0 },
+    Math.ceil(size / WORKGROUP_SIZES.DEFAULT)
+  );
+
+  return createTensor(output, input.dtype, [...input.shape], 'relu_output');
+}
+
+export async function runReLU(input, options = {}) {
+  return _relu(null, input, options);
+}
+
+export async function recordReLU(recorder, input, options = {}) {
+  return _relu(recorder, input, options);
+}

package/src/gpu/kernels/relu.wgsl

@@ -0,0 +1,21 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    size: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= u.size) {
+        return;
+    }
+    output[idx] = max(input[idx], 0.0);
+}

package/src/gpu/kernels/relu_f16.wgsl

@@ -0,0 +1,23 @@
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    size: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f16>;
+@group(0) @binding(2) var<storage, read_write> output: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= u.size) {
+        return;
+    }
+    output[idx] = max(input[idx], 0.0h);
+}

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

@@ -0,0 +1,21 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+
+export interface RepeatChannelsOptions extends OutputBufferOptions {
+  inChannels: number;
+  height: number;
+  width: number;
+  repeats: number;
+}
+
+export declare function runRepeatChannels(
+  input: Tensor,
+  options: RepeatChannelsOptions
+): Promise<Tensor>;
+
+export declare function recordRepeatChannels(
+  recorder: CommandRecorder,
+  input: Tensor,
+  options: RepeatChannelsOptions
+): Promise<Tensor>;

package/src/gpu/kernels/repeat_channels.js

@@ -0,0 +1,60 @@
+import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../tensor.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { selectRuleValue } from './rule-registry.js';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectRepeatChannelsVariant(dtype) {
+  return selectRuleValue('repeatChannels', 'variant', { dtype });
+}
+
+async function _repeatChannels(target, input, options = {}) {
+  const {
+    inChannels,
+    height,
+    width,
+    repeats,
+    outputBuffer = null,
+  } = options;
+
+  if (
+    !Number.isFinite(inChannels) ||
+    !Number.isFinite(height) ||
+    !Number.isFinite(width) ||
+    !Number.isFinite(repeats) ||
+    repeats < 1
+  ) {
+    throw new Error('RepeatChannels requires inChannels, height, width, and repeats.');
+  }
+
+  const outChannels = inChannels * repeats;
+  const variant = selectRepeatChannelsVariant(input.dtype);
+  const bytesPerElement = dtypeBytes(input.dtype);
+  const outputSize = outChannels * height * width * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'repeat_channels_output');
+
+  await unifiedKernelWrapper(
+    'repeat_channels',
+    target,
+    variant,
+    [input, output],
+    {
+      in_channels: inChannels,
+      height,
+      width,
+      repeats,
+      _pad0: 0,
+    },
+    Math.ceil((outChannels * height * width) / WORKGROUP_SIZES.DEFAULT)
+  );
+
+  return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
+}
+
+export async function runRepeatChannels(input, options = {}) {
+  return _repeatChannels(null, input, options);
+}
+
+export async function recordRepeatChannels(recorder, input, options = {}) {
+  return _repeatChannels(recorder, input, options);
+}

package/src/gpu/kernels/repeat_channels.wgsl

@@ -0,0 +1,29 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    in_channels: u32,
+    height: u32,
+    width: u32,
+    repeats: u32,
+    _pad0: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let spatial = u.height * u.width;
+    let out_channels = u.in_channels * u.repeats;
+    let total = out_channels * spatial;
+    if (idx >= total) {
+        return;
+    }
+
+    let out_channel = idx / spatial;
+    let channel = out_channel / u.repeats;
+    let spatial_idx = idx - out_channel * spatial;
+    output[idx] = input[channel * spatial + spatial_idx];
+}

package/src/gpu/kernels/repeat_channels_f16.wgsl

@@ -0,0 +1,31 @@
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    in_channels: u32,
+    height: u32,
+    width: u32,
+    repeats: u32,
+    _pad0: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f16>;
+@group(0) @binding(2) var<storage, read_write> output: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let spatial = u.height * u.width;
+    let out_channels = u.in_channels * u.repeats;
+    let total = out_channels * spatial;
+    if (idx >= total) {
+        return;
+    }
+
+    let out_channel = idx / spatial;
+    let channel = out_channel / u.repeats;
+    let spatial_idx = idx - out_channel * spatial;
+    output[idx] = input[channel * spatial + spatial_idx];
+}

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

@@ -0,0 +1,27 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+
+export interface SanaLinearAttentionOptions extends OutputBufferOptions {
+  numHeads: number;
+  headDim: number;
+  numTokens?: number;
+  hiddenSize?: number;
+  eps?: number;
+  summaryBuffer?: GPUBuffer | null;
+}
+
+export declare function runSanaLinearAttention(
+  query: Tensor,
+  key: Tensor,
+  value: Tensor,
+  options: SanaLinearAttentionOptions
+): Promise<Tensor>;
+
+export declare function recordSanaLinearAttention(
+  recorder: CommandRecorder,
+  query: Tensor,
+  key: Tensor,
+  value: Tensor,
+  options: SanaLinearAttentionOptions
+): Promise<Tensor>;

package/src/gpu/kernels/sana_linear_attention.js

@@ -0,0 +1,122 @@
+import { getDevice } from '../device.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../tensor.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { selectRuleValue } from './rule-registry.js';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectSanaLinearAttentionVariant(isF16) {
+  return selectRuleValue('sanaLinearAttention', 'variant', { isF16 });
+}
+
+async function runSummary(target, query, key, value, summaryBuffer, uniforms, variant) {
+  const summarySize = uniforms.num_heads * (uniforms.head_dim + 1) * uniforms.head_dim;
+  await unifiedKernelWrapper(
+    'sana_linear_attention_summary',
+    target,
+    variant,
+    [query, key, value, summaryBuffer],
+    {
+      num_heads: uniforms.num_heads,
+      head_dim: uniforms.head_dim,
+      num_tokens: uniforms.num_tokens,
+      hidden_size: uniforms.hidden_size,
+      _pad0: 0,
+      _pad1: 0,
+    },
+    Math.ceil(summarySize / WORKGROUP_SIZES.DEFAULT)
+  );
+}
+
+async function runApply(target, query, summaryBuffer, outputBuffer, uniforms, variant) {
+  const outputSize = uniforms.num_tokens * uniforms.hidden_size;
+  await unifiedKernelWrapper(
+    'sana_linear_attention_apply',
+    target,
+    variant,
+    [query, summaryBuffer, outputBuffer],
+    {
+      num_heads: uniforms.num_heads,
+      head_dim: uniforms.head_dim,
+      num_tokens: uniforms.num_tokens,
+      hidden_size: uniforms.hidden_size,
+      eps: uniforms.eps,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
+    Math.ceil(outputSize / WORKGROUP_SIZES.DEFAULT)
+  );
+}
+
+async function _sanaLinearAttention(target, query, key, value, options = {}) {
+  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
+  const device = target?.device || getDevice();
+  if (!device) {
+    throw new Error('SanaLinearAttention requires a WebGPU device.');
+  }
+
+  const {
+    numHeads,
+    headDim,
+    numTokens = query.shape?.[0],
+    hiddenSize = query.shape?.[1],
+    eps = 1e-15,
+    outputBuffer = null,
+    summaryBuffer = null,
+  } = options;
+
+  if (
+    !Number.isFinite(numHeads) ||
+    !Number.isFinite(headDim) ||
+    !Number.isFinite(numTokens) ||
+    !Number.isFinite(hiddenSize)
+  ) {
+    throw new Error('SanaLinearAttention requires numHeads, headDim, numTokens, and hiddenSize.');
+  }
+  if (hiddenSize !== numHeads * headDim) {
+    throw new Error(`SanaLinearAttention hiddenSize mismatch: ${hiddenSize} != ${numHeads} * ${headDim}`);
+  }
+
+  const isF16 = query.dtype === 'f16';
+  const variant = selectSanaLinearAttentionVariant(isF16);
+  const temporarySummary = summaryBuffer || acquireBuffer(
+    numHeads * (headDim + 1) * headDim * Float32Array.BYTES_PER_ELEMENT,
+    undefined,
+    'sana_linear_attention_summary'
+  );
+  const output = outputBuffer || acquireBuffer(
+    numTokens * hiddenSize * dtypeBytes(query.dtype),
+    undefined,
+    'sana_linear_attention_output'
+  );
+
+  const uniforms = {
+    num_heads: numHeads,
+    head_dim: headDim,
+    num_tokens: numTokens,
+    hidden_size: hiddenSize,
+    eps,
+  };
+
+  await runSummary(target, query, key, value, temporarySummary, uniforms, variant);
+  await runApply(target, query, temporarySummary, output, uniforms, variant);
+
+  if (!summaryBuffer) {
+    if (recorder) {
+      recorder.trackTemporaryBuffer(temporarySummary);
+    } else {
+      releaseBuffer(temporarySummary);
+    }
+  }
+
+  return createTensor(output, query.dtype, [numTokens, hiddenSize], 'sana_linear_attention_output');
+}
+
+export async function runSanaLinearAttention(query, key, value, options = {}) {
+  return _sanaLinearAttention(null, query, key, value, options);
+}
+
+export async function recordSanaLinearAttention(recorder, query, key, value, options = {}) {
+  return _sanaLinearAttention(recorder, query, key, value, options);
+}

package/src/gpu/kernels/sana_linear_attention_apply.wgsl

@@ -0,0 +1,44 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_heads: u32,
+    head_dim: u32,
+    num_tokens: u32,
+    hidden_size: u32,
+    eps: f32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> query: array<f32>;
+@group(0) @binding(2) var<storage, read> summary: array<f32>;
+@group(0) @binding(3) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let total = u.num_tokens * u.hidden_size;
+    if (idx >= total) {
+        return;
+    }
+
+    let token = idx / u.hidden_size;
+    let hidden = idx - token * u.hidden_size;
+    let head = hidden / u.head_dim;
+    let dim = hidden - head * u.head_dim;
+    let rows_per_head = u.head_dim + 1u;
+    let head_offset = head * rows_per_head * u.head_dim;
+    let hidden_base = head * u.head_dim;
+
+    var numerator: f32 = 0.0;
+    var denominator: f32 = 0.0;
+    for (var i: u32 = 0u; i < u.head_dim; i = i + 1u) {
+        let q_value = max(query[token * u.hidden_size + hidden_base + i], 0.0);
+        numerator = numerator + summary[head_offset + dim * u.head_dim + i] * q_value;
+        denominator = denominator + summary[head_offset + u.head_dim * u.head_dim + i] * q_value;
+    }
+
+    output[idx] = numerator / (denominator + u.eps);
+}

package/src/gpu/kernels/sana_linear_attention_apply_f16.wgsl

@@ -0,0 +1,47 @@
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_heads: u32,
+    head_dim: u32,
+    num_tokens: u32,
+    hidden_size: u32,
+    eps: f32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> query: array<f16>;
+@group(0) @binding(2) var<storage, read> summary: array<f32>;
+@group(0) @binding(3) var<storage, read_write> output: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let total = u.num_tokens * u.hidden_size;
+    if (idx >= total) {
+        return;
+    }
+
+    let token = idx / u.hidden_size;
+    let hidden = idx - token * u.hidden_size;
+    let head = hidden / u.head_dim;
+    let dim = hidden - head * u.head_dim;
+    let rows_per_head = u.head_dim + 1u;
+    let head_offset = head * rows_per_head * u.head_dim;
+    let hidden_base = head * u.head_dim;
+
+    var numerator: f32 = 0.0;
+    var denominator: f32 = 0.0;
+    for (var i: u32 = 0u; i < u.head_dim; i = i + 1u) {
+        let q_value = max(f32(query[token * u.hidden_size + hidden_base + i]), 0.0);
+        numerator = numerator + summary[head_offset + dim * u.head_dim + i] * q_value;
+        denominator = denominator + summary[head_offset + u.head_dim * u.head_dim + i] * q_value;
+    }
+
+    let result = numerator / (denominator + u.eps);
+    output[idx] = f16(clamp(result, -65504.0, 65504.0));
+}

package/src/gpu/kernels/sana_linear_attention_summary.wgsl

@@ -0,0 +1,47 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_heads: u32,
+    head_dim: u32,
+    num_tokens: u32,
+    hidden_size: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> query: array<f32>;
+@group(0) @binding(2) var<storage, read> key: array<f32>;
+@group(0) @binding(3) var<storage, read> value: array<f32>;
+@group(0) @binding(4) var<storage, read_write> summary: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let rows_per_head = u.head_dim + 1u;
+    let head_span = rows_per_head * u.head_dim;
+    let total = u.num_heads * head_span;
+    if (idx >= total) {
+        return;
+    }
+
+    let head = idx / head_span;
+    let rem = idx - head * head_span;
+    let row = rem / u.head_dim;
+    let col = rem - row * u.head_dim;
+    let hidden_base = head * u.head_dim;
+
+    var acc: f32 = 0.0;
+    for (var token: u32 = 0u; token < u.num_tokens; token = token + 1u) {
+        let key_idx = token * u.hidden_size + hidden_base + col;
+        let key_value = max(key[key_idx], 0.0);
+        let value_value = select(
+            value[token * u.hidden_size + hidden_base + row],
+            1.0,
+            row == u.head_dim
+        );
+        acc = acc + value_value * key_value;
+    }
+
+    summary[idx] = acc;
+}

package/src/gpu/kernels/sana_linear_attention_summary_f16.wgsl

@@ -0,0 +1,49 @@
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_heads: u32,
+    head_dim: u32,
+    num_tokens: u32,
+    hidden_size: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> query: array<f16>;
+@group(0) @binding(2) var<storage, read> key: array<f16>;
+@group(0) @binding(3) var<storage, read> value: array<f16>;
+@group(0) @binding(4) var<storage, read_write> summary: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let rows_per_head = u.head_dim + 1u;
+    let head_span = rows_per_head * u.head_dim;
+    let total = u.num_heads * head_span;
+    if (idx >= total) {
+        return;
+    }
+
+    let head = idx / head_span;
+    let rem = idx - head * head_span;
+    let row = rem / u.head_dim;
+    let col = rem - row * u.head_dim;
+    let hidden_base = head * u.head_dim;
+
+    var acc: f32 = 0.0;
+    for (var token: u32 = 0u; token < u.num_tokens; token = token + 1u) {
+        let key_idx = token * u.hidden_size + hidden_base + col;
+        let key_value = max(f32(key[key_idx]), 0.0);
+        let value_value = select(
+            f32(value[token * u.hidden_size + hidden_base + row]),
+            1.0,
+            row == u.head_dim
+        );
+        acc = acc + value_value * key_value;
+    }
+
+    summary[idx] = acc;
+}

package/src/index-browser.d.ts

@@ -1,5 +1,5 @@
 export declare const DOPPLER_VERSION: string;
-export { doppler } from './client/doppler-api.js';
+export { doppler } from './client/doppler-api.browser.js';
 
 export {
   DopplerLoader,

package/src/index-browser.js

@@ -1,5 +1,5 @@
-export const DOPPLER_VERSION = '0.1.0';
-export { doppler } from './client/doppler-api.js';
+export { DOPPLER_VERSION } from './version.js';
+export { doppler } from './client/doppler-api.browser.js';
 
 // Core loaders
 export {

package/src/index.js

@@ -1,4 +1,4 @@
-export const DOPPLER_VERSION = '0.1.0';
+export { DOPPLER_VERSION } from './version.js';
 export { doppler } from './client/doppler-api.js';
 
 // Core loaders