@simulatte/doppler 0.1.3 → 0.1.5

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,4 +1,5 @@
 export declare const DOPPLER_VERSION: string;
+export { doppler } from './client/doppler-api.browser.js';
 
 export {
   DopplerLoader,

package/src/index-browser.js

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

package/src/index.d.ts

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

package/src/index.js

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