@simulatte/doppler 0.1.0

package/src/gpu/kernels/fused_matmul_rmsnorm_f16.wgsl

@@ -0,0 +1,303 @@
+// AUTO-GENERATED from src/gpu/kernels/fused_matmul_rmsnorm.wgsl.
+// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
+// Fused GEMV + RMSNorm Kernel (F16 variant)
+//
+// For decode (M=1), combines the down projection matmul with RMSNorm in a single kernel:
+// 1. Compute GEMV: C[1, N] = A[1, K] × B[K, N]  (down projection)
+// 2. Compute RMSNorm on C: output = C / sqrt(mean(C^2) + eps) * weight
+// 3. Optional residual: output = output + residual
+//
+// F16 variant: All buffers are f16, accumulation in f32 for stability.
+//
+// Benefits:
+// - Single GPU dispatch instead of 2
+// - No intermediate buffer for matmul output
+// - Better cache locality
+// - 2x memory bandwidth vs F32 variant
+//
+// Expected speedup: 1.3-1.8x for post-FFN normalization path
+
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+const MAX_WORKGROUP_SIZE: u32 = 256u;
+override COLS_PER_WG: u32 = 4u;
+override THREADS_PER_COL: u32 = 64u;  // WORKGROUP_SIZE / COLS_PER_WG
+override RMS_NORM_OFFSET: bool = false;
+override WEIGHT_IS_F16: bool = false;
+
+struct Uniforms {
+    N: u32,             // Output dimension (hidden_size)
+    K: u32,             // Input dimension (intermediate_size)
+    eps: f32,           // RMSNorm epsilon
+    has_residual: u32,  // 1 if residual provided, 0 otherwise
+    transpose_b: u32,   // 1 if weight is [N,K] (row-major), 0 if [K,N] (column-major)
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f16>;       // [1, K] - activation from FFN
+@group(0) @binding(2) var<storage, read> weight: array<f16>;      // [K, N] - down projection weight
+@group(0) @binding(3) var<storage, read> norm_weight: array<u32>; // [N] - RMSNorm weight (F16 or F32)
+@group(0) @binding(4) var<storage, read_write> output: array<f16>;// [1, N] - final output
+@group(0) @binding(5) var<storage, read> residual: array<f16>;    // [1, N] - optional residual
+
+// Shared memory for reduction (f32 for numerical stability)
+var<workgroup> shared_partial: array<f32, MAX_WORKGROUP_SIZE>;  // Partial dot products
+var<workgroup> shared_output: array<f32, MAX_WORKGROUP_SIZE>;   // Output values for RMSNorm
+var<workgroup> shared_sum_sq: array<f32, MAX_WORKGROUP_SIZE>;   // Sum of squares for reduction
+
+fn apply_weight(w: f32) -> f32 {
+    if (RMS_NORM_OFFSET) {
+        return 1.0 + w;
+    }
+    return w;
+}
+
+fn load_norm_weight(idx: u32) -> f32 {
+    if (WEIGHT_IS_F16) {
+        let packed = norm_weight[idx >> 1u];
+        let pair = unpack2x16float(packed);
+        return select(pair.x, pair.y, (idx & 1u) == 1u);
+    }
+    return bitcast<f32>(norm_weight[idx]);
+}
+
+// Optimized single-workgroup variant for small N (hidden_size <= WORKGROUP_SIZE)
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn gemv_rmsnorm_small(
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each thread computes one output column (for N <= WORKGROUP_SIZE)
+    var dot_sum: f32 = 0.0;
+    if (tid < N) {
+        // Unroll by 4 for better ILP
+        var k: u32 = 0u;
+        let k_aligned = (K / 4u) * 4u;
+
+        for (; k < k_aligned; k = k + 4u) {
+            let a0 = f32(input[k]);
+            let a1 = f32(input[k + 1u]);
+            let a2 = f32(input[k + 2u]);
+            let a3 = f32(input[k + 3u]);
+
+            if (u.transpose_b == 1u) {
+                // Row-major: weight[tid, k]
+                let base = tid * K + k;
+                dot_sum = dot_sum + a0 * f32(weight[base])
+                                  + a1 * f32(weight[base + 1u])
+                                  + a2 * f32(weight[base + 2u])
+                                  + a3 * f32(weight[base + 3u]);
+            } else {
+                // Column-major: weight[k, tid]
+                dot_sum = dot_sum + a0 * f32(weight[k * N + tid])
+                                  + a1 * f32(weight[(k + 1u) * N + tid])
+                                  + a2 * f32(weight[(k + 2u) * N + tid])
+                                  + a3 * f32(weight[(k + 3u) * N + tid]);
+            }
+        }
+
+        // Handle remainder
+        for (; k < K; k = k + 1u) {
+            let w_idx = select(k * N + tid, tid * K + k, u.transpose_b == 1u);
+            dot_sum = dot_sum + f32(input[k]) * f32(weight[w_idx]);
+        }
+    }
+
+    // Store output value and compute sum of squares
+    shared_output[tid] = dot_sum;
+    shared_sum_sq[tid] = select(0.0, dot_sum * dot_sum, tid < N);
+    workgroupBarrier();
+
+    // Parallel reduction for sum of squares
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Compute inverse RMS
+    let mean_sq = shared_sum_sq[0] / f32(N);
+    let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+    // Write normalized output
+    if (tid < N) {
+        var result = shared_output[tid] * inv_rms * apply_weight(load_norm_weight(tid));
+        if (u.has_residual == 1u) {
+            result = result + f32(residual[tid]);
+        }
+        output[tid] = f16(result);
+    }
+}
+
+// Medium variant for N up to ~4096 (single workgroup, multiple elements per thread)
+// Handles Gemma's hidden_size=2304 and similar models
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn gemv_rmsnorm_medium(
+    @builtin(local_invocation_id) local_id: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each thread handles multiple output columns
+    let cols_per_thread = (N + WORKGROUP_SIZE - 1u) / WORKGROUP_SIZE;
+
+    // Phase 1: Compute dot products for all columns assigned to this thread
+    var local_sum_sq: f32 = 0.0;
+
+    for (var c: u32 = 0u; c < cols_per_thread; c = c + 1u) {
+        let col = tid + c * WORKGROUP_SIZE;
+        if (col < N) {
+            // Compute dot product for this column
+            var dot_sum: f32 = 0.0;
+
+            // Unroll by 4 for better ILP
+            var k: u32 = 0u;
+            let k_aligned = (K / 4u) * 4u;
+
+            for (; k < k_aligned; k = k + 4u) {
+                let a0 = f32(input[k]);
+                let a1 = f32(input[k + 1u]);
+                let a2 = f32(input[k + 2u]);
+                let a3 = f32(input[k + 3u]);
+
+                if (u.transpose_b == 1u) {
+                    let base = col * K + k;
+                    dot_sum = dot_sum + a0 * f32(weight[base])
+                                      + a1 * f32(weight[base + 1u])
+                                      + a2 * f32(weight[base + 2u])
+                                      + a3 * f32(weight[base + 3u]);
+                } else {
+                    dot_sum = dot_sum + a0 * f32(weight[k * N + col])
+                                      + a1 * f32(weight[(k + 1u) * N + col])
+                                      + a2 * f32(weight[(k + 2u) * N + col])
+                                      + a3 * f32(weight[(k + 3u) * N + col]);
+                }
+            }
+
+            // Handle remainder
+            for (; k < K; k = k + 1u) {
+                let w_idx = select(k * N + col, col * K + k, u.transpose_b == 1u);
+                dot_sum = dot_sum + f32(input[k]) * f32(weight[w_idx]);
+            }
+
+            // Store to output buffer (will normalize in-place later)
+            output[col] = f16(dot_sum);
+            local_sum_sq = local_sum_sq + dot_sum * dot_sum;
+        }
+    }
+
+    // Phase 2: Parallel reduction for sum of squares
+    shared_sum_sq[tid] = local_sum_sq;
+    workgroupBarrier();
+
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Compute inverse RMS (all threads read the same value)
+    let mean_sq = shared_sum_sq[0] / f32(N);
+    let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+    // Phase 3: Normalize outputs in-place
+    for (var c: u32 = 0u; c < cols_per_thread; c = c + 1u) {
+        let col = tid + c * WORKGROUP_SIZE;
+        if (col < N) {
+            var result = f32(output[col]) * inv_rms * apply_weight(load_norm_weight(col));
+            if (u.has_residual == 1u) {
+                result = result + f32(residual[col]);
+            }
+            output[col] = f16(result);
+        }
+    }
+}
+
+// Main entry point - dispatches based on N size
+// Uses multiple workgroups for large N
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) wg_id: vec3<u32>,
+    @builtin(num_workgroups) num_wgs: vec3<u32>
+) {
+    let tid = local_id.x;
+    let N = u.N;
+    let K = u.K;
+
+    // Each workgroup handles COLS_PER_WG columns of output
+    let col_in_wg = tid / THREADS_PER_COL;
+    let thread_in_col = tid % THREADS_PER_COL;
+    let global_col = wg_id.x * COLS_PER_WG + col_in_wg;
+
+    // Phase 1: Compute dot product for this column
+    var dot_sum: f32 = 0.0;
+    if (global_col < N) {
+        let k_per_thread = (K + THREADS_PER_COL - 1u) / THREADS_PER_COL;
+        let k_start = thread_in_col * k_per_thread;
+        let k_end = min(k_start + k_per_thread, K);
+
+        for (var k: u32 = k_start; k < k_end; k = k + 1u) {
+            let w_idx = select(k * N + global_col, global_col * K + k, u.transpose_b == 1u);
+            dot_sum = dot_sum + f32(input[k]) * f32(weight[w_idx]);
+        }
+    }
+
+    // Store partial dot product in shared memory
+    shared_partial[tid] = dot_sum;
+    workgroupBarrier();
+
+    // Reduce within each column
+    for (var stride: u32 = THREADS_PER_COL / 2u; stride > 0u; stride = stride >> 1u) {
+        if (thread_in_col < stride) {
+            let idx = col_in_wg * THREADS_PER_COL + thread_in_col;
+            shared_partial[idx] = shared_partial[idx] + shared_partial[idx + stride];
+        }
+        workgroupBarrier();
+    }
+
+    // Thread 0 of each column has the final dot product
+    var matmul_output: f32 = 0.0;
+    if (thread_in_col == 0u && global_col < N) {
+        matmul_output = shared_partial[col_in_wg * THREADS_PER_COL];
+        shared_output[col_in_wg] = matmul_output;
+        shared_sum_sq[col_in_wg] = matmul_output * matmul_output;
+    }
+    workgroupBarrier();
+
+    // Reduce sum of squares across columns in this workgroup
+    for (var stride: u32 = COLS_PER_WG / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride && tid + stride < COLS_PER_WG) {
+            shared_sum_sq[tid] = shared_sum_sq[tid] + shared_sum_sq[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    let local_sum_sq = shared_sum_sq[0];
+
+    // Phase 3: Write normalized output
+    if (thread_in_col == 0u && global_col < N) {
+        let val = shared_output[col_in_wg];
+        let mean_sq = local_sum_sq / f32(min(COLS_PER_WG, N));
+        let inv_rms = 1.0 / sqrt(mean_sq + u.eps);
+
+        var result = val * inv_rms * apply_weight(load_norm_weight(global_col));
+
+        if (u.has_residual == 1u) {
+            result = result + f32(residual[global_col]);
+        }
+
+        output[global_col] = f16(result);
+    }
+}

package/src/gpu/kernels/fused_swiglu.wgsl

@@ -0,0 +1,63 @@
+// fused_swiglu.wgsl
+
+/**
+ * SwiGLU Activation Kernel (Row-Split Fused Gate/Up)
+ *
+ * Computes: output = SiLU(gate + gate_bias) * (up + up_bias)
+ *
+ * input layout: [num_tokens, 2 * dim] flattened, row-major
+ *   row = [gate[0..dim), up[0..dim)]
+ * bias layout: [2 * dim]
+ *   bias = [gate_bias[0..dim), up_bias[0..dim)]
+ * output layout: [num_tokens, dim]
+ */
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_tokens: u32,
+    dim: u32,
+    bias_offset: u32,  // byte offset into bias buffer (divide by 4 for F32 index)
+    clamp_max: f32,    // SwiGLU clamp (0 = disabled)
+}
+
+@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> bias: array<f32>;
+@group(0) @binding(3) var<storage, read_write> output: array<f32>;
+
+fn sigmoid(x: f32) -> f32 {
+    return 1.0 / (1.0 + exp(-x));
+}
+
+fn silu(x: f32) -> f32 {
+    return x * sigmoid(x);
+}
+
+fn clamp_swiglu(x: f32) -> f32 {
+    if (u.clamp_max <= 0.0) {
+        return x;
+    }
+    return clamp(x, -u.clamp_max, u.clamp_max);
+}
+
+@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.dim;
+    if (idx >= total) {
+        return;
+    }
+
+    let token_idx = idx / u.dim;
+    let dim_idx = idx % u.dim;
+
+    // bias_offset is in bytes, convert to F32 index
+    let bias_base = u.bias_offset / 4u;
+
+    let row_base = token_idx * u.dim * 2u;
+    let gate = input[row_base + dim_idx] + bias[bias_base + dim_idx];
+    let up = input[row_base + u.dim + dim_idx] + bias[bias_base + u.dim + dim_idx];
+
+    output[idx] = clamp_swiglu(silu(gate) * up);
+}

package/src/gpu/kernels/fused_swiglu_f16.wgsl

@@ -0,0 +1,57 @@
+// AUTO-GENERATED from src/gpu/kernels/fused_swiglu.wgsl.
+// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
+// fused_swiglu_f16.wgsl
+
+/**
+ * SwiGLU Activation Kernel (Row-Split Fused Gate/Up, f16 IO)
+ */
+
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_tokens: u32,
+    dim: u32,
+    bias_offset: u32,
+    clamp_max: f32,
+}
+
+@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> bias: array<f16>;
+@group(0) @binding(3) var<storage, read_write> output: array<f16>;
+
+fn sigmoid(x: f32) -> f32 {
+    return 1.0 / (1.0 + exp(-x));
+}
+
+fn silu(x: f32) -> f32 {
+    return x * sigmoid(x);
+}
+
+fn clamp_swiglu(x: f32) -> f32 {
+    if (u.clamp_max <= 0.0) {
+        return x;
+    }
+    return clamp(x, -u.clamp_max, u.clamp_max);
+}
+
+@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.dim;
+    if (idx >= total) {
+        return;
+    }
+
+    let token_idx = idx / u.dim;
+    let dim_idx = idx % u.dim;
+    let bias_base = u.bias_offset / 2u;
+
+    let row_base = token_idx * u.dim * 2u;
+    let gate = f32(input[row_base + dim_idx]) + f32(bias[bias_base + dim_idx]);
+    let up = f32(input[row_base + u.dim + dim_idx]) + f32(bias[bias_base + u.dim + dim_idx]);
+
+    output[idx] = f16(clamp_swiglu(silu(gate) * up));
+}

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

@@ -0,0 +1,64 @@
+/**
+ * Gather (Embedding Lookup) Kernels
+ *
+ * Provides token embedding lookups from embedding tables.
+ */
+
+import type { CommandRecorder } from '../command-recorder.js';
+import type { Tensor, TensorDtype } from '../tensor.js';
+import type { OutputBufferOptions } from './types.js';
+
+/** Gather kernel options */
+export interface GatherOptions extends OutputBufferOptions {
+  useVec4?: boolean;
+  embeddingDtype?: 'f16' | 'f32';
+  /**
+   * Output dtype. When 'f16', converts F32 embeddings to F16 output.
+   * Default: 'f32'
+   */
+  outputDtype?: 'f16' | 'f32';
+  /**
+   * True if embeddings are stored as [hidden_size, vocab_size] (GGUF layout).
+   * False if embeddings are stored as [vocab_size, hidden_size] (PyTorch layout).
+   * Default: false (RDRR format uses PyTorch layout from SafeTensors).
+   */
+  transpose?: boolean;
+  /** Optional index offset into the token indices buffer. */
+  indexOffset?: number;
+  /** Optional indirect dispatch buffer for GPU-driven workgroup counts. */
+  indirectBuffer?: GPUBuffer | null;
+  /** Byte offset into indirect dispatch buffer (default: 0). */
+  indirectOffset?: number;
+}
+
+/** Context for gather variant selection rules. */
+export interface GatherSelectionContext {
+  useF16Input: boolean;
+  useF16Output: boolean;
+  useVec4: boolean;
+}
+
+/**
+ * Run gather/embedding lookup
+ */
+export declare function runGather(
+  indices: GPUBuffer,
+  embeddings: GPUBuffer,
+  numTokens: number,
+  hiddenSize: number,
+  vocabSize: number,
+  options?: GatherOptions
+): Promise<Tensor>;
+
+/**
+ * Record gather (batched, no submit)
+ */
+export declare function recordGather(
+  recorder: CommandRecorder,
+  indices: GPUBuffer,
+  embeddings: GPUBuffer,
+  numTokens: number,
+  hiddenSize: number,
+  vocabSize: number,
+  options?: GatherOptions
+): Promise<Tensor>;

package/src/gpu/kernels/gather.js

@@ -0,0 +1,119 @@
+import { getKernelCapabilities } from '../device.js';
+import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { WORKGROUP_SIZES, VEC4_ELEMENTS_PER_WG } from './constants.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { trace } from '../../debug/index.js';
+import { createTensor } from '../tensor.js';
+import { DTYPE_SIZES, padToQ4KBlock } from '../../config/schema/index.js';
+import { selectRuleValue as selectKernelRuleValue } from './rule-registry.js';
+import { selectRuleValue as selectSharedRuleValue } from '../../rules/rule-registry.js';
+
+function selectGatherVariant(useF16Input, useF16Output, useVec4) {
+  return selectKernelRuleValue(
+    'gather',
+    'variant',
+    { useF16Input, useF16Output, useVec4 }
+  );
+}
+
+async function _gather(
+  target,
+  indices,
+  embeddings,
+  numTokens,
+  hiddenSize,
+  vocabSize,
+  options = {}
+) {
+  const {
+    useVec4 = true,
+    outputBuffer = null,
+    embeddingDtype,
+    outputDtype,
+    transpose = false,
+    indexOffset = 0,
+    indirectBuffer = null,
+    indirectOffset = 0,
+  } = options;
+
+  const caps = getKernelCapabilities();
+  if (embeddingDtype == null) {
+    throw new Error('[Gather] embeddingDtype is required.');
+  }
+  if (outputDtype == null) {
+    throw new Error('[Gather] outputDtype is required.');
+  }
+
+  const useF16Input = embeddingDtype === 'f16' && caps.hasF16;
+  const useF16Output = outputDtype === 'f16' && caps.hasF16;
+
+  trace.embed(
+    `Gather: numTokens=${numTokens}, hiddenSize=${hiddenSize}, vocabSize=${vocabSize}, ` +
+    `transpose=${transpose}, indexOffset=${indexOffset}, ` +
+    `embeddingDtype=${embeddingDtype}, outputDtype=${outputDtype}, ` +
+    `useF16Input=${useF16Input}, useF16Output=${useF16Output}`
+  );
+
+  const variant = selectGatherVariant(useF16Input, useF16Output, useVec4);
+  trace.embed(`Gather variant: ${variant}`);
+
+  // Pad hiddenSize to Q4K alignment for downstream fused Q4K matmul kernels
+  // that read 256-element blocks. Extra padding elements stay zero.
+  const actualDtype = selectSharedRuleValue('shared', 'dtype', 'f16OrF32', { useF16: useF16Output });
+  const bytesPerElement = DTYPE_SIZES[actualDtype];
+  const paddedHiddenSize = padToQ4KBlock(hiddenSize);
+  const outputSize = numTokens * paddedHiddenSize * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gather_output');
+
+  const uniforms = {
+    num_tokens: numTokens,
+    hidden_size: hiddenSize,
+    vocab_size: vocabSize,
+    transpose: transpose ? 1 : 0,
+    index_offset: indexOffset,
+    _pad0: 0,
+    _pad1: 0,
+    _pad2: 0,
+  };
+
+  const workgroups = indirectBuffer
+    ? { indirectBuffer, indirectOffset }
+    : (useVec4
+      ? Math.ceil((numTokens * hiddenSize) / VEC4_ELEMENTS_PER_WG)
+      : Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT));
+
+  await unifiedKernelWrapper(
+    'gather',
+    target,
+    variant,
+    [indices, embeddings, output],
+    uniforms,
+    workgroups
+  );
+
+  return createTensor(output, actualDtype, [numTokens, hiddenSize], 'gather_output');
+}
+
+export async function runGather(
+  indices,
+  embeddings,
+  numTokens,
+  hiddenSize,
+  vocabSize,
+  options = {}
+) {
+  return _gather(null, indices, embeddings, numTokens, hiddenSize, vocabSize, options);
+}
+
+export async function recordGather(
+  recorder,
+  indices,
+  embeddings,
+  numTokens,
+  hiddenSize,
+  vocabSize,
+  options = {}
+) {
+  return _gather(recorder, indices, embeddings, numTokens, hiddenSize, vocabSize, options);
+}
+

package/src/gpu/kernels/gather.wgsl

@@ -0,0 +1,61 @@
+// gather.wgsl
+
+/**
+ * Gather Kernel - Token Embedding Lookup
+ *
+ * Gathers rows from an embedding matrix based on token indices.
+ * Used for efficient embedding lookup on GPU without CPU readback.
+ */
+
+// Tunable workgroup size
+override WORKGROUP_SIZE_MAIN: u32 = 256u;
+
+struct Uniforms {
+    num_tokens: u32,      // Number of tokens to gather
+    hidden_size: u32,     // Embedding dimension
+    vocab_size: u32,      // Vocabulary size (for bounds checking)
+    transpose: u32,       // 1 if embeddings are [hidden_size, vocab_size] (GGUF layout), 0 otherwise
+    index_offset: u32,    // Starting index into indices buffer
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> indices: array<u32>;      // Token IDs [num_tokens]
+@group(0) @binding(2) var<storage, read> embeddings: array<f32>;   // Embedding matrix [vocab_size, hidden_size]
+@group(0) @binding(3) var<storage, read_write> output: array<f32>; // Output [num_tokens, hidden_size]
+
+@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let tid = gid.x;
+    let total_elements = u.num_tokens * u.hidden_size;
+
+    if (tid >= total_elements) {
+        return;
+    }
+
+    // Compute token index and dimension index
+    let token_idx = tid / u.hidden_size;
+    let dim_idx = tid % u.hidden_size;
+
+    // Get the token ID (with bounds check)
+    let token_id = indices[token_idx + u.index_offset];
+
+    // Bounds check on vocab
+    if (token_id >= u.vocab_size) {
+        output[tid] = 0.0;
+        return;
+    }
+
+    // Gather from embedding matrix
+    // For GGUF layout [hidden_size, vocab_size]: offset = dim_idx * vocab_size + token_id
+    // For standard layout [vocab_size, hidden_size]: offset = token_id * hidden_size + dim_idx
+    var embed_offset: u32;
+    if (u.transpose == 1u) {
+        embed_offset = dim_idx * u.vocab_size + token_id;
+    } else {
+        embed_offset = token_id * u.hidden_size + dim_idx;
+    }
+    output[tid] = embeddings[embed_offset];
+}