@simulatte/doppler 0.1.8 → 0.2.0

package/src/gpu/kernels/moe_gather.wgsl

@@ -1,170 +0,0 @@
-// moe_gather.wgsl
-
-/**
- * MoE Gather Kernel - Gather tokens by expert for batched execution
- *
- * Groups tokens by their selected experts so that each expert can
- * process its assigned tokens in a single batched operation.
- *
- * NOTE: This kernel requires EXPLICIT bind group layout (not 'auto')
- * because count_and_map and gather_tokens use different subsets of bindings.
- *
- * Input:
- *   - hidden_states [num_tokens, hidden_size]
- *   - indices [num_tokens, top_k] - selected expert indices per token
- *
- * Output:
- *   - gathered [num_experts, max_tokens_per_expert, hidden_size]
- *   - token_counts [num_experts] - actual token count per expert
- *   - token_map [num_experts, max_tokens_per_expert] - original token index mapping
- */
-
-// Tunable workgroup sizes
-override WORKGROUP_SIZE_MAIN: u32 = 256u;
-
-struct Uniforms {
-    num_tokens: u32,            // Number of input tokens
-    hidden_size: u32,           // Hidden dimension
-    num_experts: u32,           // Number of experts
-    top_k: u32,                 // Number of experts per token
-    max_tokens_per_expert: u32, // Max tokens any expert can receive
-    threads_per_row: u32,       // For 2D dispatch: dispatchSizeX * WORKGROUP_SIZE_MAIN
-    _pad2: u32,
-    _pad3: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> hidden_states: array<f32>;      // [num_tokens, hidden_size]
-@group(0) @binding(2) var<storage, read> expert_indices: array<u32>;     // [num_tokens, top_k]
-@group(0) @binding(3) var<storage, read_write> gathered: array<f32>;     // [num_experts, max_tokens_per_expert, hidden_size]
-@group(0) @binding(4) var<storage, read_write> token_counts: array<atomic<u32>>; // [num_experts]
-@group(0) @binding(5) var<storage, read_write> token_map: array<u32>;    // [num_experts, max_tokens_per_expert, 2] (token_idx, k_idx)
-
-// Phase 1: Count tokens per expert and build token map
-// Run with num_tokens * top_k threads
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn count_and_map(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let total_slots = u.num_tokens * u.top_k;
-
-    if (tid >= total_slots) {
-        return;
-    }
-
-    let token_idx = tid / u.top_k;
-    let k_idx = tid % u.top_k;
-    let expert_idx = expert_indices[tid];
-
-    // Atomically increment token count for this expert and get slot
-    let slot = atomicAdd(&token_counts[expert_idx], 1u);
-
-    // Store mapping: which original token goes to this slot
-    if (slot < u.max_tokens_per_expert) {
-        let map_base = expert_idx * u.max_tokens_per_expert * 2u + slot * 2u;
-        token_map[map_base] = token_idx;
-        token_map[map_base + 1u] = k_idx;
-    }
-}
-
-// Phase 2: Gather hidden states based on token map
-// Run with num_experts * max_tokens_per_expert * hidden_size threads
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn gather_tokens(@builtin(global_invocation_id) gid: vec3<u32>) {
-    // Support 2D dispatch for large token counts (>65535 workgroups)
-    let tid = gid.x + gid.y * u.threads_per_row;
-    let hidden_size = u.hidden_size;
-    let max_tokens_per_expert = u.max_tokens_per_expert;
-    let num_experts = u.num_experts;
-
-    let elements_per_expert = max_tokens_per_expert * hidden_size;
-    let total_elements = num_experts * elements_per_expert;
-
-    if (tid >= total_elements) {
-        return;
-    }
-
-    // Decode position
-    let expert_idx = tid / elements_per_expert;
-    let within_expert = tid % elements_per_expert;
-    let slot_idx = within_expert / hidden_size;
-    let dim_idx = within_expert % hidden_size;
-
-    // Check if this slot is valid (within actual token count)
-    let actual_count = atomicLoad(&token_counts[expert_idx]);
-    if (slot_idx >= actual_count) {
-        // Zero out unused slots
-        gathered[tid] = 0.0;
-        return;
-    }
-
-    // Look up original token index from map
-    let map_base = expert_idx * max_tokens_per_expert * 2u + slot_idx * 2u;
-    let token_idx = token_map[map_base];
-
-    // Gather from original hidden states
-    let src_idx = token_idx * hidden_size + dim_idx;
-    gathered[tid] = hidden_states[src_idx];
-}
-
-// Combined single-pass version for small models
-// Each workgroup handles one expert
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn gather_single_pass(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let expert_idx = wg_id.x;
-    let thread_idx = local_id.x;
-    let hidden_size = u.hidden_size;
-    let num_tokens = u.num_tokens;
-    let top_k = u.top_k;
-    let max_tokens_per_expert = u.max_tokens_per_expert;
-
-    if (expert_idx >= u.num_experts) {
-        return;
-    }
-
-    // Phase 1: Count tokens for this expert (thread 0 only)
-    var token_count: u32 = 0u;
-    if (thread_idx == 0u) {
-        for (var t: u32 = 0u; t < num_tokens; t = t + 1u) {
-            for (var k: u32 = 0u; k < top_k; k = k + 1u) {
-                if (expert_indices[t * top_k + k] == expert_idx) {
-                    if (token_count < max_tokens_per_expert) {
-                        let map_base = expert_idx * max_tokens_per_expert * 2u + token_count * 2u;
-                        token_map[map_base] = t;
-                        token_map[map_base + 1u] = k;
-                        token_count = token_count + 1u;
-                    }
-                }
-            }
-        }
-        atomicStore(&token_counts[expert_idx], token_count);
-    }
-
-    workgroupBarrier();
-
-    // Phase 2: Gather (all threads participate)
-    let actual_count = atomicLoad(&token_counts[expert_idx]);
-    let elements_per_slot = hidden_size;
-    let total_work = actual_count * elements_per_slot;
-    let work_per_thread = (total_work + WORKGROUP_SIZE_MAIN - 1u) / WORKGROUP_SIZE_MAIN;
-
-    for (var i: u32 = 0u; i < work_per_thread; i = i + 1u) {
-        let work_idx = thread_idx * work_per_thread + i;
-        if (work_idx >= total_work) {
-            break;
-        }
-
-        let slot_idx = work_idx / hidden_size;
-        let dim_idx = work_idx % hidden_size;
-
-        let map_base = expert_idx * max_tokens_per_expert * 2u + slot_idx * 2u;
-        let token_idx = token_map[map_base];
-
-        let src_idx = token_idx * hidden_size + dim_idx;
-        let dst_idx = expert_idx * max_tokens_per_expert * hidden_size + slot_idx * hidden_size + dim_idx;
-
-        gathered[dst_idx] = hidden_states[src_idx];
-    }
-}

package/src/gpu/kernels/moe_gather_f16.wgsl

@@ -1,82 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/moe_gather.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// moe_gather_f16.wgsl
-
-/**
- * MoE Gather Kernel - f16 inputs/outputs
- */
-
-enable f16;
-
-override WORKGROUP_SIZE_MAIN: u32 = 256u;
-
-struct Uniforms {
-    num_tokens: u32,
-    hidden_size: u32,
-    num_experts: u32,
-    top_k: u32,
-    max_tokens_per_expert: u32,
-    threads_per_row: u32,  // For 2D dispatch: dispatchSizeX * WORKGROUP_SIZE_MAIN
-    _pad2: u32,
-    _pad3: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> hidden_states: array<f16>;
-@group(0) @binding(2) var<storage, read> expert_indices: array<u32>;
-@group(0) @binding(3) var<storage, read_write> gathered: array<f16>;
-@group(0) @binding(4) var<storage, read_write> token_counts: array<atomic<u32>>;
-@group(0) @binding(5) var<storage, read_write> token_map: array<u32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn count_and_map(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let total_slots = u.num_tokens * u.top_k;
-
-    if (tid >= total_slots) {
-        return;
-    }
-
-    let expert_idx = expert_indices[tid];
-    let slot = atomicAdd(&token_counts[expert_idx], 1u);
-
-    if (slot < u.max_tokens_per_expert) {
-        let map_base = expert_idx * u.max_tokens_per_expert * 2u + slot * 2u;
-        let token_idx = tid / u.top_k;
-        let k_idx = tid % u.top_k;
-        token_map[map_base] = token_idx;
-        token_map[map_base + 1u] = k_idx;
-    }
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn gather_tokens(@builtin(global_invocation_id) gid: vec3<u32>) {
-    // Support 2D dispatch for large token counts (>65535 workgroups)
-    let tid = gid.x + gid.y * u.threads_per_row;
-    let hidden_size = u.hidden_size;
-    let max_tokens_per_expert = u.max_tokens_per_expert;
-    let num_experts = u.num_experts;
-
-    let elements_per_expert = max_tokens_per_expert * hidden_size;
-    let total_elements = num_experts * elements_per_expert;
-
-    if (tid >= total_elements) {
-        return;
-    }
-
-    let expert_idx = tid / elements_per_expert;
-    let within_expert = tid % elements_per_expert;
-    let slot_idx = within_expert / hidden_size;
-    let dim_idx = within_expert % hidden_size;
-
-    let actual_count = atomicLoad(&token_counts[expert_idx]);
-    if (slot_idx >= actual_count) {
-        gathered[tid] = f16(0.0);
-        return;
-    }
-
-    let map_base = expert_idx * max_tokens_per_expert * 2u + slot_idx * 2u;
-    let token_idx = token_map[map_base];
-    let src_idx = token_idx * hidden_size + dim_idx;
-    gathered[tid] = hidden_states[src_idx];
-}

package/src/gpu/kernels/moe_gather_vec4.wgsl

@@ -1,74 +0,0 @@
-// moe_gather_vec4.wgsl
-
-/**
- * MoE Gather Kernel (vec4) - Gather tokens by expert for batched execution
- *
- * Optimized gather with vec4 loads.
- */
-
-// Tunable workgroup size
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-struct Uniforms {
-    num_tokens: u32,            // Number of input tokens
-    hidden_size: u32,           // Hidden dimension
-    num_experts: u32,           // Number of experts
-    top_k: u32,                 // Number of experts per token
-    max_tokens_per_expert: u32, // Max tokens any expert can receive
-    _pad1: u32,
-    _pad2: u32,
-    _pad3: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> hidden_states: array<f32>;      // [num_tokens, hidden_size]
-@group(0) @binding(2) var<storage, read> expert_indices: array<u32>;     // [num_tokens, top_k]
-@group(0) @binding(3) var<storage, read_write> gathered: array<f32>;     // [num_experts, max_tokens_per_expert, hidden_size]
-@group(0) @binding(4) var<storage, read_write> token_counts: array<atomic<u32>>; // [num_experts]
-@group(0) @binding(5) var<storage, read_write> token_map: array<u32>;    // [num_experts, max_tokens_per_expert, 2]
-
-// Optimized version: Gather with vec4 loads
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn gather_tokens_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let hidden_size = u.hidden_size;
-    let max_tokens_per_expert = u.max_tokens_per_expert;
-    let num_experts = u.num_experts;
-    let vec4_per_token = hidden_size / 4u;
-
-    let vec4_per_expert = max_tokens_per_expert * vec4_per_token;
-    let total_vec4s = num_experts * vec4_per_expert;
-
-    if (tid >= total_vec4s) {
-        return;
-    }
-
-    // Decode position
-    let expert_idx = tid / vec4_per_expert;
-    let within_expert = tid % vec4_per_expert;
-    let slot_idx = within_expert / vec4_per_token;
-    let vec4_idx = within_expert % vec4_per_token;
-
-    // Check if slot is valid
-    let actual_count = atomicLoad(&token_counts[expert_idx]);
-    let dst_base = tid * 4u;
-
-    if (slot_idx >= actual_count) {
-        gathered[dst_base] = 0.0;
-        gathered[dst_base + 1u] = 0.0;
-        gathered[dst_base + 2u] = 0.0;
-        gathered[dst_base + 3u] = 0.0;
-        return;
-    }
-
-    // Look up original token
-    let map_base = expert_idx * max_tokens_per_expert * 2u + slot_idx * 2u;
-    let token_idx = token_map[map_base];
-
-    // Gather 4 elements
-    let src_base = token_idx * hidden_size + vec4_idx * 4u;
-    gathered[dst_base] = hidden_states[src_base];
-    gathered[dst_base + 1u] = hidden_states[src_base + 1u];
-    gathered[dst_base + 2u] = hidden_states[src_base + 2u];
-    gathered[dst_base + 3u] = hidden_states[src_base + 3u];
-}

package/src/gpu/kernels/moe_offsets.wgsl

@@ -1,48 +0,0 @@
-// moe_offsets.wgsl
-//
-// Build per-routing-slot token offsets on GPU:
-// token_offsets[token_idx * top_k + k_idx] = expert_idx * max_tokens_per_expert + slot_idx
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    num_tokens: u32,
-    num_experts: u32,
-    top_k: u32,
-    max_tokens_per_expert: u32,
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-    _pad3: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> token_counts: array<u32>;
-@group(0) @binding(2) var<storage, read> token_map: array<u32>;
-@group(0) @binding(3) var<storage, read_write> token_offsets: array<u32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn build_offsets(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let slot_linear = gid.x;
-    let total_slots = u.num_experts * u.max_tokens_per_expert;
-    if (slot_linear >= total_slots) {
-        return;
-    }
-
-    let expert_idx = slot_linear / u.max_tokens_per_expert;
-    let slot_idx = slot_linear % u.max_tokens_per_expert;
-    let count = token_counts[expert_idx];
-    if (slot_idx >= count) {
-        return;
-    }
-
-    let map_base = (expert_idx * u.max_tokens_per_expert + slot_idx) * 2u;
-    let token_idx = token_map[map_base];
-    let k_idx = token_map[map_base + 1u];
-    if (token_idx >= u.num_tokens || k_idx >= u.top_k) {
-        return;
-    }
-
-    let routing_idx = token_idx * u.top_k + k_idx;
-    token_offsets[routing_idx] = expert_idx * u.max_tokens_per_expert + slot_idx;
-}

package/src/gpu/kernels/pipeline-cache.d.ts

@@ -1,88 +0,0 @@
-/**
- * Pipeline Cache - GPU pipeline creation and caching
- *
- * Handles creation and caching of compute pipelines, bind group layouts,
- * and pipeline layouts for kernel execution.
- *
- * @module gpu/kernels/pipeline-cache
- */
-
-// ============================================================================
-// Bind Group Layout
-// ============================================================================
-
-/**
- * Get or create a cached bind group layout.
- */
-export declare function getOrCreateBindGroupLayout(
-  label: string,
-  entries: GPUBindGroupLayoutEntry[],
-  deviceOverride?: GPUDevice | null
-): GPUBindGroupLayout;
-
-// ============================================================================
-// Pipeline Layout
-// ============================================================================
-
-/**
- * Get or create a cached pipeline layout.
- */
-export declare function getOrCreatePipelineLayout(
-  label: string,
-  bindGroupLayouts: GPUBindGroupLayout[],
-  deviceOverride?: GPUDevice | null
-): GPUPipelineLayout;
-
-// ============================================================================
-// Pipeline Creation
-// ============================================================================
-
-/**
- * Synchronously get a cached pipeline, or null if not cached.
- * Use this for fast path when you know the pipeline should be warm.
- */
-export declare function getCachedPipeline(
-  operation: string,
-  variant: string,
-  constants?: Record<string, number | boolean> | null
-): GPUComputePipeline | null;
-
-/**
- * Get a pipeline, using synchronous cache lookup when available.
- * Falls back to async compilation if not cached.
- * This is the preferred way to get pipelines in hot paths.
- */
-export declare function getPipelineFast(
-  operation: string,
-  variant: string,
-  bindGroupLayout?: GPUBindGroupLayout | null,
-  constants?: Record<string, number | boolean> | null
-): Promise<GPUComputePipeline>;
-
-/**
- * Create a compute pipeline for a kernel
- */
-export declare function createPipeline(
-  operation: string,
-  variant: string,
-  bindGroupLayout?: GPUBindGroupLayout | null,
-  constants?: Record<string, number | boolean> | null
-): Promise<GPUComputePipeline>;
-
-// ============================================================================
-// Cache Management
-// ============================================================================
-
-/**
- * Clear the pipeline caches
- */
-export declare function clearPipelineCaches(): void;
-
-/**
- * Get pipeline cache statistics
- */
-export declare function getPipelineCacheStats(): {
-  pipelines: number;
-  bindGroupLayouts: number;
-  pipelineLayouts: number;
-};