@simulatte/doppler 0.1.9 → 0.2.0

package/src/gpu/kernels/sample_f16.wgsl

@@ -1,361 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/sample.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// sample_f16.wgsl
-
-/**
- * GPU-Side Sampling Kernel (f16 logits)
- *
- * Same as sample.wgsl but logits are f16.
- * Computation is performed in f32 for stability.
- */
-
-enable f16;
-
-// Configuration
-override WORKGROUP_SIZE: u32 = 256u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-const MAX_TOP_K: u32 = 128u;
-
-struct Uniforms {
-    vocab_size: u32,
-    top_k: u32,
-    temperature: f32,
-    random_value: f32,
-    pad_token_id: u32,
-    logit_softcap: f32,
-    output_index: u32,
-    pad0: u32,
-}
-
-fn apply_softcap(x: f32, softcap: f32) -> f32 {
-    if (softcap <= 0.0) {
-        return x;
-    }
-    return softcap * tanh(x / softcap);
-}
-
-fn candidate_beats(candidate_value: f32, candidate_index: u32, best_value: f32, best_index: u32) -> bool {
-    if (candidate_value > best_value) {
-        return true;
-    }
-    if (candidate_value < best_value) {
-        return false;
-    }
-    return candidate_index < best_index;
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> logits: array<f16>;
-@group(0) @binding(2) var<storage, read_write> output: array<u32>;
-@group(0) @binding(3) var<storage, read_write> topk_indices: array<u32>;
-@group(0) @binding(4) var<storage, read_write> topk_logits: array<f32>;
-
-var<workgroup> shared_values: array<f32, MAX_WORKGROUP_SIZE>;
-var<workgroup> shared_indices: array<u32, MAX_WORKGROUP_SIZE>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn find_topk_phase1(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>,
-    @builtin(workgroup_id) wgid: vec3<u32>,
-    @builtin(num_workgroups) num_wg: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let global_idx = gid.x;
-    let vocab_size = u.vocab_size;
-    let temperature = u.temperature;
-    let pad_id = u.pad_token_id;
-    let softcap = u.logit_softcap;
-
-    if (num_wg.x == 1u) {
-        var val: f32 = -3.402823e+38;
-        if (global_idx < vocab_size && global_idx != pad_id) {
-            val = apply_softcap(f32(logits[global_idx]), softcap) / temperature;
-        }
-        topk_logits[thread_idx] = val;
-        topk_indices[thread_idx] = global_idx;
-        return;
-    }
-
-    var local_max: f32 = -3.402823e+38;
-    var local_max_idx: u32 = 0u;
-
-    var idx = global_idx;
-    while (idx < vocab_size) {
-        if (idx != pad_id) {
-            let val = apply_softcap(f32(logits[idx]), softcap) / temperature;
-            if (candidate_beats(val, idx, local_max, local_max_idx)) {
-                local_max = val;
-                local_max_idx = idx;
-            }
-        }
-        idx = idx + WORKGROUP_SIZE * num_wg.x;
-    }
-
-    shared_values[thread_idx] = local_max;
-    shared_indices[thread_idx] = local_max_idx;
-    workgroupBarrier();
-
-    var stride = WORKGROUP_SIZE / 2u;
-    while (stride > 0u) {
-        if (thread_idx < stride) {
-            if (candidate_beats(
-                shared_values[thread_idx + stride],
-                shared_indices[thread_idx + stride],
-                shared_values[thread_idx],
-                shared_indices[thread_idx]
-            )) {
-                shared_values[thread_idx] = shared_values[thread_idx + stride];
-                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
-            }
-        }
-        workgroupBarrier();
-        stride = stride / 2u;
-    }
-
-    if (thread_idx == 0u) {
-        topk_logits[wgid.x] = shared_values[0];
-        topk_indices[wgid.x] = shared_indices[0];
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn find_topk_phase2(
-    @builtin(local_invocation_id) lid: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let top_k = u.top_k;
-    let num_groups = min(WORKGROUP_SIZE, (u.vocab_size + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE);
-    let num_candidates = select(num_groups, min(u.vocab_size, WORKGROUP_SIZE), num_groups == 1u);
-
-    if (thread_idx < WORKGROUP_SIZE) {
-        if (thread_idx < num_candidates) {
-            shared_values[thread_idx] = topk_logits[thread_idx];
-            shared_indices[thread_idx] = topk_indices[thread_idx];
-        } else {
-            shared_values[thread_idx] = -3.402823e+38;
-            shared_indices[thread_idx] = 0u;
-        }
-    }
-    workgroupBarrier();
-
-    if (thread_idx == 0u) {
-        for (var k: u32 = 0u; k < top_k && k < num_candidates; k = k + 1u) {
-            var max_idx = k;
-            var max_val = shared_values[k];
-
-            for (var i: u32 = k + 1u; i < num_candidates; i = i + 1u) {
-                if (candidate_beats(shared_values[i], shared_indices[i], max_val, shared_indices[max_idx])) {
-                    max_val = shared_values[i];
-                    max_idx = i;
-                }
-            }
-
-            if (max_idx != k) {
-                let tmp_val = shared_values[k];
-                let tmp_idx = shared_indices[k];
-                shared_values[k] = shared_values[max_idx];
-                shared_indices[k] = shared_indices[max_idx];
-                shared_values[max_idx] = tmp_val;
-                shared_indices[max_idx] = tmp_idx;
-            }
-        }
-
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            topk_logits[k] = shared_values[k];
-            topk_indices[k] = shared_indices[k];
-        }
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn softmax_and_sample(
-    @builtin(local_invocation_id) lid: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let top_k = u.top_k;
-    let random_val = u.random_value;
-
-    if (thread_idx < top_k) {
-        shared_values[thread_idx] = topk_logits[thread_idx];
-        shared_indices[thread_idx] = topk_indices[thread_idx];
-    }
-    workgroupBarrier();
-
-    if (thread_idx == 0u) {
-        var max_val: f32 = shared_values[0];
-        for (var i: u32 = 1u; i < top_k; i = i + 1u) {
-            max_val = max(max_val, shared_values[i]);
-        }
-
-        var exp_sum: f32 = 0.0;
-        for (var i: u32 = 0u; i < top_k; i = i + 1u) {
-            let exp_val = exp(shared_values[i] - max_val);
-            shared_values[i] = exp_val;
-            exp_sum = exp_sum + exp_val;
-        }
-
-        let inv_sum = 1.0 / exp_sum;
-        var cum_prob: f32 = 0.0;
-        var selected_token: u32 = shared_indices[top_k - 1u];
-
-        for (var i: u32 = 0u; i < top_k; i = i + 1u) {
-            let prob = shared_values[i] * inv_sum;
-            cum_prob = cum_prob + prob;
-            if (cum_prob >= random_val) {
-                selected_token = shared_indices[i];
-                break;
-            }
-        }
-
-        output[u.output_index] = selected_token;
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn sample_single_pass(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>,
-    @builtin(num_workgroups) num_wg: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let vocab_size = u.vocab_size;
-    let top_k = min(u.top_k, MAX_TOP_K);
-    let temperature = u.temperature;
-    let random_val = u.random_value;
-    let pad_id = u.pad_token_id;
-    let softcap = u.logit_softcap;
-
-    var local_max: f32 = -3.402823e+38;
-    var local_max_idx: u32 = 0u;
-
-    var idx = gid.x;
-    while (idx < vocab_size) {
-        if (idx != pad_id) {
-            let val = apply_softcap(f32(logits[idx]), softcap) / temperature;
-            if (candidate_beats(val, idx, local_max, local_max_idx)) {
-                local_max = val;
-                local_max_idx = idx;
-            }
-        }
-        idx = idx + num_wg.x * WORKGROUP_SIZE;
-    }
-
-    shared_values[thread_idx] = local_max;
-    shared_indices[thread_idx] = local_max_idx;
-    workgroupBarrier();
-
-    var stride = WORKGROUP_SIZE / 2u;
-    while (stride > 0u) {
-        if (thread_idx < stride) {
-            if (candidate_beats(
-                shared_values[thread_idx + stride],
-                shared_indices[thread_idx + stride],
-                shared_values[thread_idx],
-                shared_indices[thread_idx]
-            )) {
-                shared_values[thread_idx] = shared_values[thread_idx + stride];
-                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
-            }
-        }
-        workgroupBarrier();
-        stride = stride / 2u;
-    }
-
-    if (thread_idx == 0u && num_wg.x == 1u) {
-        output[u.output_index] = shared_indices[0];
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn argmax(
-    @builtin(global_invocation_id) gid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>,
-    @builtin(workgroup_id) wgid: vec3<u32>,
-    @builtin(num_workgroups) num_wg: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let global_idx = gid.x;
-    let vocab_size = u.vocab_size;
-    let pad_id = u.pad_token_id;
-    let softcap = u.logit_softcap;
-
-    var local_max: f32 = -3.402823e+38;
-    var local_max_idx: u32 = 0u;
-
-    var idx = global_idx;
-    while (idx < vocab_size) {
-        if (idx != pad_id) {
-            let val = apply_softcap(f32(logits[idx]), softcap);
-            if (candidate_beats(val, idx, local_max, local_max_idx)) {
-                local_max = val;
-                local_max_idx = idx;
-            }
-        }
-        idx = idx + num_wg.x * WORKGROUP_SIZE;
-    }
-
-    shared_values[thread_idx] = local_max;
-    shared_indices[thread_idx] = local_max_idx;
-    workgroupBarrier();
-
-    var stride = WORKGROUP_SIZE / 2u;
-    while (stride > 0u) {
-        if (thread_idx < stride) {
-            if (candidate_beats(
-                shared_values[thread_idx + stride],
-                shared_indices[thread_idx + stride],
-                shared_values[thread_idx],
-                shared_indices[thread_idx]
-            )) {
-                shared_values[thread_idx] = shared_values[thread_idx + stride];
-                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
-            }
-        }
-        workgroupBarrier();
-        stride = stride / 2u;
-    }
-
-    if (thread_idx == 0u) {
-        topk_logits[wgid.x] = shared_values[0];
-        topk_indices[wgid.x] = shared_indices[0];
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn argmax_reduce(
-    @builtin(local_invocation_id) lid: vec3<u32>
-) {
-    let thread_idx = lid.x;
-    let num_groups = min(WORKGROUP_SIZE, (u.vocab_size + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE);
-
-    if (thread_idx < num_groups) {
-        shared_values[thread_idx] = topk_logits[thread_idx];
-        shared_indices[thread_idx] = topk_indices[thread_idx];
-    } else {
-        shared_values[thread_idx] = -3.402823e+38;
-        shared_indices[thread_idx] = 0u;
-    }
-    workgroupBarrier();
-
-    var stride = WORKGROUP_SIZE / 2u;
-    while (stride > 0u) {
-        if (thread_idx < stride) {
-            if (candidate_beats(
-                shared_values[thread_idx + stride],
-                shared_indices[thread_idx + stride],
-                shared_values[thread_idx],
-                shared_indices[thread_idx]
-            )) {
-                shared_values[thread_idx] = shared_values[thread_idx + stride];
-                shared_indices[thread_idx] = shared_indices[thread_idx + stride];
-            }
-        }
-        workgroupBarrier();
-        stride = stride / 2u;
-    }
-
-    if (thread_idx == 0u) {
-        output[u.output_index] = shared_indices[0];
-    }
-}

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

@@ -1,27 +0,0 @@
-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

@@ -1,129 +0,0 @@
-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) {
-  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(uniforms.hidden_size / WORKGROUP_SIZES.DEFAULT), uniforms.num_tokens, 1]
-  );
-}
-
-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;
-  const ownsSummary = summaryBuffer == null;
-  const ownsOutput = outputBuffer == null;
-
-  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,
-  };
-
-  try {
-    await runSummary(target, query, key, value, temporarySummary, uniforms, variant);
-    await runApply(target, query, temporarySummary, output, uniforms, variant);
-    return createTensor(output, query.dtype, [numTokens, hiddenSize], 'sana_linear_attention_output');
-  } catch (error) {
-    if (ownsOutput) {
-      releaseBuffer(output);
-    }
-    throw error;
-  } finally {
-    if (ownsSummary) {
-      if (recorder) {
-        recorder.trackTemporaryBuffer(temporarySummary);
-      } else {
-        releaseBuffer(temporarySummary);
-      }
-    }
-  }
-}
-
-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

@@ -1,43 +0,0 @@
-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 hidden = gid.x;
-    let token = gid.y;
-    if (token >= u.num_tokens || hidden >= u.hidden_size) {
-        return;
-    }
-
-    let idx = token * u.hidden_size + hidden;
-    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

@@ -1,46 +0,0 @@
-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 hidden = gid.x;
-    let token = gid.y;
-    if (token >= u.num_tokens || hidden >= u.hidden_size) {
-        return;
-    }
-
-    let idx = token * u.hidden_size + hidden;
-    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

@@ -1,51 +0,0 @@
-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 query_value = query[token * u.hidden_size + hidden_base + col];
-        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
-        );
-        if (u.hidden_size == 0u) {
-            acc = acc + query_value;
-        }
-        acc = acc + value_value * key_value;
-    }
-
-    summary[idx] = acc;
-}