@simulatte/doppler 0.1.9 → 0.2.0

package/src/gpu/kernels/split_qg.wgsl

@@ -1,58 +0,0 @@
-// split_qg.wgsl
-
-/**
- * De-interleave Q and Gate projections from q_proj output for attentionOutputGate models.
- *
- * Models like Qwen 3.5 store q_proj weights with interleaved head layout:
- *   rows [h*headDim*2 : h*headDim*2+headDim]     = Q for head h
- *   rows [h*headDim*2+headDim : (h+1)*headDim*2] = Gate for head h
- *
- * A single full matmul over all 2*qSize rows produces interleaved output:
- *   input[token, h*headDim*2 : h*headDim*2+headDim]     = Q head h
- *   input[token, h*headDim*2+headDim : (h+1)*headDim*2] = Gate head h
- *
- * This kernel separates them into contiguous Q and G outputs:
- *   Q[token, h*headDim + dim] = input[token, h*headDim*2 + dim]
- *   G[token, h*headDim + dim] = input[token, h*headDim*2 + headDim + dim]
- *
- * Input layout  (row-major): [numTokens, numHeads * headDim * 2]
- * Output Q layout (row-major): [numTokens, numHeads * headDim]
- * Output G layout (row-major): [numTokens, numHeads * headDim]
- */
-
-struct Params {
-    num_tokens: u32,
-    num_heads: u32,
-    head_dim: u32,
-    _pad: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> params: Params;
-@group(0) @binding(1) var<storage, read> input: array<f32>;
-@group(0) @binding(2) var<storage, read_write> Q: array<f32>;
-@group(0) @binding(3) var<storage, read_write> G: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let q_size = params.num_heads * params.head_dim;
-    let total_elements = params.num_tokens * q_size;
-
-    if (idx >= total_elements) {
-        return;
-    }
-
-    let token = idx / q_size;
-    let elem = idx % q_size;
-    let head = elem / params.head_dim;
-    let dim = elem % params.head_dim;
-
-    // Input is interleaved per head: [Q_h (headDim elems), G_h (headDim elems)]
-    let src_q = token * (q_size * 2u) + head * (params.head_dim * 2u) + dim;
-    let src_g = src_q + params.head_dim;
-
-    Q[idx] = input[src_q];
-    G[idx] = input[src_g];
-}

package/src/gpu/kernels/split_qg_f16.wgsl

@@ -1,62 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/split_qg.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// split_qg_f16.wgsl
-
-/**
- * De-interleave Q and Gate projections from q_proj output for attentionOutputGate models (f16).
- *
- * Models like Qwen 3.5 store q_proj weights with interleaved head layout:
- *   rows [h*headDim*2 : h*headDim*2+headDim]     = Q for head h
- *   rows [h*headDim*2+headDim : (h+1)*headDim*2] = Gate for head h
- *
- * A single full matmul over all 2*qSize rows produces interleaved output:
- *   input[token, h*headDim*2 : h*headDim*2+headDim]     = Q head h
- *   input[token, h*headDim*2+headDim : (h+1)*headDim*2] = Gate head h
- *
- * This kernel separates them into contiguous Q and G outputs:
- *   Q[token, h*headDim + dim] = input[token, h*headDim*2 + dim]
- *   G[token, h*headDim + dim] = input[token, h*headDim*2 + headDim + dim]
- *
- * Input layout  (row-major): [numTokens, numHeads * headDim * 2]
- * Output Q layout (row-major): [numTokens, numHeads * headDim]
- * Output G layout (row-major): [numTokens, numHeads * headDim]
- */
-
-enable f16;
-
-struct Params {
-    num_tokens: u32,
-    num_heads: u32,
-    head_dim: u32,
-    _pad: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> params: Params;
-@group(0) @binding(1) var<storage, read> input: array<f16>;
-@group(0) @binding(2) var<storage, read_write> Q: array<f16>;
-@group(0) @binding(3) var<storage, read_write> G: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let q_size = params.num_heads * params.head_dim;
-    let total_elements = params.num_tokens * q_size;
-
-    if (idx >= total_elements) {
-        return;
-    }
-
-    let token = idx / q_size;
-    let elem = idx % q_size;
-    let head = elem / params.head_dim;
-    let dim = elem % params.head_dim;
-
-    // Input is interleaved per head: [Q_h (headDim elems), G_h (headDim elems)]
-    let src_q = token * (q_size * 2u) + head * (params.head_dim * 2u) + dim;
-    let src_g = src_q + params.head_dim;
-
-    Q[idx] = input[src_q];
-    G[idx] = input[src_g];
-}

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

@@ -1,51 +0,0 @@
-/**
- * Split QKV Kernel
- *
- * Splits fused QKV projection output into separate Q, K, V buffers.
- * Used for 3→1 matmul optimization in attention.
- */
-
-import type { Tensor } from '../tensor.js';
-import type { CommandRecorder } from '../command-recorder.js';
-
-/** Split QKV options */
-export interface SplitQKVOptions {
-  numTokens: number;
-  qSize: number;  // numHeads * headDim
-  kSize: number;  // numKVHeads * headDim
-  vSize: number;  // numKVHeads * headDim
-  /** Pre-allocated Q output tensor */
-  qTensor?: Tensor | null;
-  /** Pre-allocated K output tensor */
-  kTensor?: Tensor | null;
-  /** Pre-allocated V output tensor */
-  vTensor?: Tensor | null;
-}
-
-/** Split QKV result */
-export interface SplitQKVResult {
-  Q: Tensor;
-  K: Tensor;
-  V: Tensor;
-}
-
-/**
- * Split fused QKV output into separate Q, K, V tensors.
- *
- * @param qkvTensor - Fused QKV output [numTokens, qSize + kSize + vSize]
- * @param options - Split configuration
- * @returns Separate Q, K, V tensors
- */
-export declare function runSplitQKV(
-  qkvTensor: Tensor,
-  options: SplitQKVOptions
-): Promise<SplitQKVResult>;
-
-/**
- * Record split QKV (batched, no submit).
- */
-export declare function recordSplitQKV(
-  recorder: CommandRecorder,
-  qkvTensor: Tensor,
-  options: SplitQKVOptions
-): Promise<SplitQKVResult>;

package/src/gpu/kernels/split_qkv.js

@@ -1,51 +0,0 @@
-
-import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
-import { createTensor, dtypeBytes } from '../tensor.js';
-import { WORKGROUP_SIZES } from './constants.js';
-import { unifiedKernelWrapper } from './utils.js';
-import { selectRuleValue } from './rule-registry.js';
-
-async function _splitQKV(target, qkvTensor, options) {
-  const { numTokens, qSize, kSize, vSize, qTensor = null, kTensor = null, vTensor = null } = options;
-  const ownsQ = qTensor == null;
-  const ownsK = kTensor == null;
-  const ownsV = vTensor == null;
-
-  const outputDtype = qkvTensor.dtype;
-  const pipelineVariant = selectRuleValue('splitQkv', 'variant', { outputDtype });
-  const bytesPerElement = dtypeBytes(outputDtype);
-
-  const qBuffer = qTensor?.buffer || acquireBuffer(numTokens * qSize * bytesPerElement, undefined, 'Q');
-  const kBuffer = kTensor?.buffer || acquireBuffer(numTokens * kSize * bytesPerElement, undefined, 'K');
-  const vBuffer = vTensor?.buffer || acquireBuffer(numTokens * vSize * bytesPerElement, undefined, 'V');
-
-  const totalElements = numTokens * (qSize + kSize + vSize);
-
-  try {
-    await unifiedKernelWrapper(
-      'split_qkv', target, pipelineVariant,
-      [qkvTensor, qBuffer, kBuffer, vBuffer],
-      { num_tokens: numTokens, q_size: qSize, k_size: kSize, v_size: vSize },
-      Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT)
-    );
-
-    const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
-    const K = kTensor || createTensor(kBuffer, outputDtype, [numTokens, kSize], 'K');
-    const V = vTensor || createTensor(vBuffer, outputDtype, [numTokens, vSize], 'V');
-
-    return { Q, K, V };
-  } catch (error) {
-    if (ownsQ) releaseBuffer(qBuffer);
-    if (ownsK) releaseBuffer(kBuffer);
-    if (ownsV) releaseBuffer(vBuffer);
-    throw error;
-  }
-}
-
-export async function runSplitQKV(qkvTensor, options) {
-  return _splitQKV(null, qkvTensor, options);
-}
-
-export async function recordSplitQKV(recorder, qkvTensor, options) {
-  return _splitQKV(recorder, qkvTensor, options);
-}

package/src/gpu/kernels/split_qkv.wgsl

@@ -1,71 +0,0 @@
-// split_qkv.wgsl
-
-/**
- * Split fused QKV output into separate Q, K, V buffers.
- *
- * Input layout (row-major):
- *   [numTokens, qSize + kSize + vSize]
- *   Each row: [q_values, k_values, v_values]
- *
- * Output layout:
- *   Q: [numTokens, qSize]
- *   K: [numTokens, kSize]
- *   V: [numTokens, vSize]
- */
-
-struct Params {
-    numTokens: u32,
-    qSize: u32,
-    kSize: u32,
-    vSize: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> params: Params;
-@group(0) @binding(1) var<storage, read> input: array<f32>;
-@group(0) @binding(2) var<storage, read_write> Q: array<f32>;
-@group(0) @binding(3) var<storage, read_write> K: array<f32>;
-@group(0) @binding(4) var<storage, read_write> V: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let numTokens = params.numTokens;
-    let qSize = params.qSize;
-    let kSize = params.kSize;
-    let vSize = params.vSize;
-    let qkvSize = qSize + kSize + vSize;
-
-    // Each thread handles one element from Q, K, or V
-    let totalQ = numTokens * qSize;
-    let totalK = numTokens * kSize;
-    let totalV = numTokens * vSize;
-    let totalElements = totalQ + totalK + totalV;
-
-    if (idx >= totalElements) {
-        return;
-    }
-
-    if (idx < totalQ) {
-        // Q output
-        let token = idx / qSize;
-        let elem = idx % qSize;
-        let srcIdx = token * qkvSize + elem;
-        Q[idx] = input[srcIdx];
-    } else if (idx < totalQ + totalK) {
-        // K output
-        let kIdx = idx - totalQ;
-        let token = kIdx / kSize;
-        let elem = kIdx % kSize;
-        let srcIdx = token * qkvSize + qSize + elem;
-        K[kIdx] = input[srcIdx];
-    } else {
-        // V output
-        let vIdx = idx - totalQ - totalK;
-        let token = vIdx / vSize;
-        let elem = vIdx % vSize;
-        let srcIdx = token * qkvSize + qSize + kSize + elem;
-        V[vIdx] = input[srcIdx];
-    }
-}

package/src/gpu/kernels/split_qkv_f16.wgsl

@@ -1,75 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/split_qkv.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// split_qkv_f16.wgsl
-
-/**
- * Split fused QKV output into separate Q, K, V buffers (f16).
- *
- * Input layout (row-major):
- *   [numTokens, qSize + kSize + vSize]
- *   Each row: [q_values, k_values, v_values]
- *
- * Output layout:
- *   Q: [numTokens, qSize]
- *   K: [numTokens, kSize]
- *   V: [numTokens, vSize]
- */
-
-enable f16;
-
-struct Params {
-    numTokens: u32,
-    qSize: u32,
-    kSize: u32,
-    vSize: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> params: Params;
-@group(0) @binding(1) var<storage, read> input: array<f16>;
-@group(0) @binding(2) var<storage, read_write> Q: array<f16>;
-@group(0) @binding(3) var<storage, read_write> K: array<f16>;
-@group(0) @binding(4) var<storage, read_write> V: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let numTokens = params.numTokens;
-    let qSize = params.qSize;
-    let kSize = params.kSize;
-    let vSize = params.vSize;
-    let qkvSize = qSize + kSize + vSize;
-
-    // Each thread handles one element from Q, K, or V
-    let totalQ = numTokens * qSize;
-    let totalK = numTokens * kSize;
-    let totalV = numTokens * vSize;
-    let totalElements = totalQ + totalK + totalV;
-
-    if (idx >= totalElements) {
-        return;
-    }
-
-    if (idx < totalQ) {
-        // Q output
-        let token = idx / qSize;
-        let elem = idx % qSize;
-        let srcIdx = token * qkvSize + elem;
-        Q[idx] = input[srcIdx];
-    } else if (idx < totalQ + totalK) {
-        // K output
-        let kIdx = idx - totalQ;
-        let token = kIdx / kSize;
-        let elem = kIdx % kSize;
-        let srcIdx = token * qkvSize + qSize + elem;
-        K[kIdx] = input[srcIdx];
-    } else {
-        // V output
-        let vIdx = idx - totalQ - totalK;
-        let token = vIdx / vSize;
-        let elem = vIdx % vSize;
-        let srcIdx = token * qkvSize + qSize + kSize + elem;
-        V[vIdx] = input[srcIdx];
-    }
-}

package/src/gpu/kernels/topk.wgsl

@@ -1,243 +0,0 @@
-// topk.wgsl
-
-/**
- * Top-K Selection Kernel for MoE Routing
- *
- * Selects top-k experts for each token based on router logits.
- * Optimized for small k (typically 2) and small n (typically 8 experts).
- *
- * Input: softmax probabilities [numTokens, numExperts]
- * Output:
- *   - indices [numTokens, topK] (u32)
- *   - weights [numTokens, topK] (f32, renormalized)
- */
-
-override WORKGROUP_SIZE: u32 = 256u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    num_tokens: u32,     // Number of tokens
-    num_experts: u32,    // Number of experts (typically 8)
-    top_k: u32,          // Number of experts to select (typically 2)
-    normalize: u32,      // Whether to renormalize weights (1 = yes)
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> probs: array<f32>;           // [numTokens, numExperts]
-@group(0) @binding(2) var<storage, read_write> out_indices: array<u32>; // [numTokens, topK]
-@group(0) @binding(3) var<storage, read_write> out_weights: array<f32>; // [numTokens, topK]
-
-// Workgroup shared memory for sorting
-// Supports up to WORKGROUP_SIZE experts (default 256; covers DeepSeek-V2's 160, Snowflake Arctic's 128, etc.)
-var<workgroup> shared_probs: array<f32, MAX_WORKGROUP_SIZE>;
-var<workgroup> shared_indices: array<u32, MAX_WORKGROUP_SIZE>;
-
-// Main kernel: one workgroup per token
-// Workgroup size WORKGROUP_SIZE to support loading up to WORKGROUP_SIZE experts in parallel
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let token_idx = wg_id.x;
-    let thread_idx = local_id.x;
-    let num_experts = u.num_experts;
-    let top_k = u.top_k;
-
-    if (token_idx >= u.num_tokens) {
-        return;
-    }
-
-    let base_offset = token_idx * num_experts;
-
-    // Load probabilities into shared memory (first numExperts threads)
-    if (thread_idx < num_experts) {
-        shared_probs[thread_idx] = probs[base_offset + thread_idx];
-        shared_indices[thread_idx] = thread_idx;
-    }
-    workgroupBarrier();
-
-    // Simple selection sort for top-k (efficient for small k and n)
-    // Only thread 0 does the sorting to avoid race conditions
-    if (thread_idx == 0u) {
-        // Find top-k by partial selection sort
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            var max_idx = k;
-            var max_val = shared_probs[k];
-
-            // Find maximum in remaining elements
-            for (var i: u32 = k + 1u; i < num_experts; i = i + 1u) {
-                if (shared_probs[i] > max_val) {
-                    max_val = shared_probs[i];
-                    max_idx = i;
-                }
-            }
-
-            // Swap if needed
-            if (max_idx != k) {
-                let tmp_prob = shared_probs[k];
-                let tmp_idx = shared_indices[k];
-                shared_probs[k] = shared_probs[max_idx];
-                shared_indices[k] = shared_indices[max_idx];
-                shared_probs[max_idx] = tmp_prob;
-                shared_indices[max_idx] = tmp_idx;
-            }
-        }
-
-        // Compute weight sum for normalization
-        var weight_sum: f32 = 0.0;
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            weight_sum = weight_sum + shared_probs[k];
-        }
-
-        // Write output indices and weights
-        let out_base = token_idx * top_k;
-        let inv_sum = select(1.0, 1.0 / weight_sum, u.normalize == 1u && weight_sum > 0.0);
-
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            out_indices[out_base + k] = shared_indices[k];
-            out_weights[out_base + k] = shared_probs[k] * inv_sum;
-        }
-    }
-}
-
-// Optimized version for topK=2, numExperts<=8
-// Each thread handles one token
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn topk_2_small(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let token_idx = gid.x;
-
-    if (token_idx >= u.num_tokens) {
-        return;
-    }
-
-    let num_experts = u.num_experts;
-    let base_offset = token_idx * num_experts;
-
-    // Find top 2 in a single pass
-    var top1_idx: u32 = 0u;
-    var top1_val: f32 = probs[base_offset];
-    var top2_idx: u32 = 1u;
-    var top2_val: f32 = probs[base_offset + 1u];
-
-    // Ensure top1 >= top2
-    if (top2_val > top1_val) {
-        let tmp_idx = top1_idx;
-        let tmp_val = top1_val;
-        top1_idx = top2_idx;
-        top1_val = top2_val;
-        top2_idx = tmp_idx;
-        top2_val = tmp_val;
-    }
-
-    // Scan remaining experts
-    for (var i: u32 = 2u; i < num_experts; i = i + 1u) {
-        let val = probs[base_offset + i];
-        if (val > top1_val) {
-            top2_idx = top1_idx;
-            top2_val = top1_val;
-            top1_idx = i;
-            top1_val = val;
-        } else if (val > top2_val) {
-            top2_idx = i;
-            top2_val = val;
-        }
-    }
-
-    // Renormalize weights
-    let weight_sum = top1_val + top2_val;
-    let inv_sum = select(1.0, 1.0 / weight_sum, u.normalize == 1u && weight_sum > 0.0);
-
-    // Write output
-    let out_base = token_idx * 2u;
-    out_indices[out_base] = top1_idx;
-    out_indices[out_base + 1u] = top2_idx;
-    out_weights[out_base] = top1_val * inv_sum;
-    out_weights[out_base + 1u] = top2_val * inv_sum;
-}
-
-// Fused softmax + top-k for efficiency
-// Avoids separate softmax kernel call
-// Workgroup size WORKGROUP_SIZE to support up to WORKGROUP_SIZE experts
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn softmax_topk(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let token_idx = wg_id.x;
-    let thread_idx = local_id.x;
-    let num_experts = u.num_experts;
-    let top_k = u.top_k;
-
-    if (token_idx >= u.num_tokens) {
-        return;
-    }
-
-    let base_offset = token_idx * num_experts;
-
-    // Load logits and find max (for numerical stability)
-    if (thread_idx < num_experts) {
-        shared_probs[thread_idx] = probs[base_offset + thread_idx];
-        shared_indices[thread_idx] = thread_idx;
-    }
-    workgroupBarrier();
-
-    // Thread 0 does softmax + top-k
-    if (thread_idx == 0u) {
-        // Find max
-        var max_val: f32 = shared_probs[0];
-        for (var i: u32 = 1u; i < num_experts; i = i + 1u) {
-            max_val = max(max_val, shared_probs[i]);
-        }
-
-        // Compute exp and sum
-        var exp_sum: f32 = 0.0;
-        for (var i: u32 = 0u; i < num_experts; i = i + 1u) {
-            let exp_val = exp(shared_probs[i] - max_val);
-            shared_probs[i] = exp_val;
-            exp_sum = exp_sum + exp_val;
-        }
-
-        // Normalize to get probabilities
-        let inv_exp_sum = 1.0 / exp_sum;
-        for (var i: u32 = 0u; i < num_experts; i = i + 1u) {
-            shared_probs[i] = shared_probs[i] * inv_exp_sum;
-        }
-
-        // Partial selection sort for top-k
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            var max_idx = k;
-            var max_prob = shared_probs[k];
-
-            for (var i: u32 = k + 1u; i < num_experts; i = i + 1u) {
-                if (shared_probs[i] > max_prob) {
-                    max_prob = shared_probs[i];
-                    max_idx = i;
-                }
-            }
-
-            if (max_idx != k) {
-                let tmp_prob = shared_probs[k];
-                let tmp_idx = shared_indices[k];
-                shared_probs[k] = shared_probs[max_idx];
-                shared_indices[k] = shared_indices[max_idx];
-                shared_probs[max_idx] = tmp_prob;
-                shared_indices[max_idx] = tmp_idx;
-            }
-        }
-
-        // Renormalize top-k weights
-        var weight_sum: f32 = 0.0;
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            weight_sum = weight_sum + shared_probs[k];
-        }
-
-        let out_base = token_idx * top_k;
-        let inv_sum = select(1.0, 1.0 / weight_sum, u.normalize == 1u && weight_sum > 0.0);
-
-        for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-            out_indices[out_base + k] = shared_indices[k];
-            out_weights[out_base + k] = shared_probs[k] * inv_sum;
-        }
-    }
-}