@simulatte/doppler 0.1.9 → 0.2.0

package/src/gpu/kernels/gather_f16_out.wgsl

@@ -1,55 +0,0 @@
-// gather_f16_out.wgsl
-
-/**
- * Gather Kernel (F16 output)
- *
- * F32 embeddings -> F16 output for reduced memory bandwidth.
- */
-
-enable f16;
-
-// 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(4) var<storage, read_write> output_f16: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE_MAIN, 1, 1)
-fn gather_f16_out(@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;
-    }
-
-    let token_idx = tid / u.hidden_size;
-    let dim_idx = tid % u.hidden_size;
-    let token_id = indices[token_idx + u.index_offset];
-
-    if (token_id >= u.vocab_size) {
-        output_f16[tid] = f16(0.0);
-        return;
-    }
-
-    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_f16[tid] = f16(embeddings[embed_offset]);
-}

package/src/gpu/kernels/gather_f16_vec4.wgsl

@@ -1,76 +0,0 @@
-// gather_f16_vec4.wgsl
-
-/**
- * Gather Kernel (F16, vec4) - Token Embedding Lookup
- *
- * Vectorized gather for F16 embeddings -> F32 output.
- */
-
-enable f16;
-
-// Tunable workgroup size
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-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<f16>;   // F16 Embedding matrix [vocab_size, hidden_size]
-@group(0) @binding(3) var<storage, read_write> output: array<f32>; // F32 Output [num_tokens, hidden_size]
-
-// Vectorized version for better memory throughput
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn gather_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let vec4_per_row = u.hidden_size / 4u;
-    let total_vec4s = u.num_tokens * vec4_per_row;
-
-    if (tid >= total_vec4s) {
-        return;
-    }
-
-    // Compute token index and vec4 index within row
-    let token_idx = tid / vec4_per_row;
-    let vec4_idx = tid % vec4_per_row;
-
-    // Get the token ID
-    let token_id = indices[token_idx + u.index_offset];
-
-    // Bounds check
-    if (token_id >= u.vocab_size) {
-        let out_base = tid * 4u;
-        output[out_base] = 0.0;
-        output[out_base + 1u] = 0.0;
-        output[out_base + 2u] = 0.0;
-        output[out_base + 3u] = 0.0;
-        return;
-    }
-
-    // Gather 4 F16 elements and convert to F32
-    let out_base = tid * 4u;
-    let dim_base = vec4_idx * 4u;
-
-    if (u.transpose == 1u) {
-        // Transposed layout [hidden_size, vocab_size]: elements are strided by vocab_size
-        output[out_base] = f32(embeddings[(dim_base) * u.vocab_size + token_id]);
-        output[out_base + 1u] = f32(embeddings[(dim_base + 1u) * u.vocab_size + token_id]);
-        output[out_base + 2u] = f32(embeddings[(dim_base + 2u) * u.vocab_size + token_id]);
-        output[out_base + 3u] = f32(embeddings[(dim_base + 3u) * u.vocab_size + token_id]);
-    } else {
-        // Standard layout [vocab_size, hidden_size]: elements are contiguous
-        let embed_base = token_id * u.hidden_size + dim_base;
-        output[out_base] = f32(embeddings[embed_base]);
-        output[out_base + 1u] = f32(embeddings[embed_base + 1u]);
-        output[out_base + 2u] = f32(embeddings[embed_base + 2u]);
-        output[out_base + 3u] = f32(embeddings[embed_base + 3u]);
-    }
-}

package/src/gpu/kernels/gather_f16_vec4_f16_out.wgsl

@@ -1,68 +0,0 @@
-// gather_f16_vec4_f16_out.wgsl
-
-/**
- * Gather Kernel (F16 input -> F16 output, vec4)
- *
- * Vectorized gather for F16 activation mode with F16 embeddings.
- */
-
-enable f16;
-
-// Tunable workgroup size
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-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<f16>; // F16 Embedding matrix
-@group(0) @binding(4) var<storage, read_write> output_f16: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn gather_vec4_f16_out(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let vec4_per_row = u.hidden_size / 4u;
-    let total_vec4s = u.num_tokens * vec4_per_row;
-
-    if (tid >= total_vec4s) {
-        return;
-    }
-
-    let token_idx = tid / vec4_per_row;
-    let vec4_idx = tid % vec4_per_row;
-    let token_id = indices[token_idx + u.index_offset];
-
-    if (token_id >= u.vocab_size) {
-        let out_base = tid * 4u;
-        output_f16[out_base] = f16(0.0);
-        output_f16[out_base + 1u] = f16(0.0);
-        output_f16[out_base + 2u] = f16(0.0);
-        output_f16[out_base + 3u] = f16(0.0);
-        return;
-    }
-
-    let out_base = tid * 4u;
-    let dim_base = vec4_idx * 4u;
-
-    if (u.transpose == 1u) {
-        output_f16[out_base] = embeddings[(dim_base) * u.vocab_size + token_id];
-        output_f16[out_base + 1u] = embeddings[(dim_base + 1u) * u.vocab_size + token_id];
-        output_f16[out_base + 2u] = embeddings[(dim_base + 2u) * u.vocab_size + token_id];
-        output_f16[out_base + 3u] = embeddings[(dim_base + 3u) * u.vocab_size + token_id];
-    } else {
-        let embed_base = token_id * u.hidden_size + dim_base;
-        output_f16[out_base] = embeddings[embed_base];
-        output_f16[out_base + 1u] = embeddings[embed_base + 1u];
-        output_f16[out_base + 2u] = embeddings[embed_base + 2u];
-        output_f16[out_base + 3u] = embeddings[embed_base + 3u];
-    }
-}

package/src/gpu/kernels/gather_vec4.wgsl

@@ -1,74 +0,0 @@
-// gather_vec4.wgsl
-
-/**
- * Gather Kernel (vec4) - Token Embedding Lookup
- *
- * Vectorized gather for F32 embeddings -> F32 output.
- */
-
-// Tunable workgroup size
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-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]
-
-// Vectorized version for better memory throughput
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn gather_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let vec4_per_row = u.hidden_size / 4u;
-    let total_vec4s = u.num_tokens * vec4_per_row;
-
-    if (tid >= total_vec4s) {
-        return;
-    }
-
-    // Compute token index and vec4 index within row
-    let token_idx = tid / vec4_per_row;
-    let vec4_idx = tid % vec4_per_row;
-
-    // Get the token ID
-    let token_id = indices[token_idx + u.index_offset];
-
-    // Bounds check
-    if (token_id >= u.vocab_size) {
-        let out_base = tid * 4u;
-        output[out_base] = 0.0;
-        output[out_base + 1u] = 0.0;
-        output[out_base + 2u] = 0.0;
-        output[out_base + 3u] = 0.0;
-        return;
-    }
-
-    // Gather 4 elements
-    let out_base = tid * 4u;
-    let dim_base = vec4_idx * 4u;
-
-    if (u.transpose == 1u) {
-        // Transposed layout [hidden_size, vocab_size]: elements are strided by vocab_size
-        output[out_base] = embeddings[(dim_base) * u.vocab_size + token_id];
-        output[out_base + 1u] = embeddings[(dim_base + 1u) * u.vocab_size + token_id];
-        output[out_base + 2u] = embeddings[(dim_base + 2u) * u.vocab_size + token_id];
-        output[out_base + 3u] = embeddings[(dim_base + 3u) * u.vocab_size + token_id];
-    } else {
-        // Standard layout [vocab_size, hidden_size]: elements are contiguous
-        let embed_base = token_id * u.hidden_size + dim_base;
-        output[out_base] = embeddings[embed_base];
-        output[out_base + 1u] = embeddings[embed_base + 1u];
-        output[out_base + 2u] = embeddings[embed_base + 2u];
-        output[out_base + 3u] = embeddings[embed_base + 3u];
-    }
-}

package/src/gpu/kernels/gather_vec4_f16_out.wgsl

@@ -1,68 +0,0 @@
-// gather_vec4_f16_out.wgsl
-
-/**
- * Gather Kernel (vec4, F16 output)
- *
- * Vectorized gather for F32 embeddings -> F16 output.
- */
-
-enable f16;
-
-// Tunable workgroup size
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-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(4) var<storage, read_write> output_f16: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn gather_vec4_f16_out(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let tid = gid.x;
-    let vec4_per_row = u.hidden_size / 4u;
-    let total_vec4s = u.num_tokens * vec4_per_row;
-
-    if (tid >= total_vec4s) {
-        return;
-    }
-
-    let token_idx = tid / vec4_per_row;
-    let vec4_idx = tid % vec4_per_row;
-    let token_id = indices[token_idx + u.index_offset];
-
-    if (token_id >= u.vocab_size) {
-        let out_base = tid * 4u;
-        output_f16[out_base] = f16(0.0);
-        output_f16[out_base + 1u] = f16(0.0);
-        output_f16[out_base + 2u] = f16(0.0);
-        output_f16[out_base + 3u] = f16(0.0);
-        return;
-    }
-
-    let out_base = tid * 4u;
-    let dim_base = vec4_idx * 4u;
-
-    if (u.transpose == 1u) {
-        output_f16[out_base] = f16(embeddings[(dim_base) * u.vocab_size + token_id]);
-        output_f16[out_base + 1u] = f16(embeddings[(dim_base + 1u) * u.vocab_size + token_id]);
-        output_f16[out_base + 2u] = f16(embeddings[(dim_base + 2u) * u.vocab_size + token_id]);
-        output_f16[out_base + 3u] = f16(embeddings[(dim_base + 3u) * u.vocab_size + token_id]);
-    } else {
-        let embed_base = token_id * u.hidden_size + dim_base;
-        output_f16[out_base] = f16(embeddings[embed_base]);
-        output_f16[out_base + 1u] = f16(embeddings[embed_base + 1u]);
-        output_f16[out_base + 2u] = f16(embeddings[embed_base + 2u]);
-        output_f16[out_base + 3u] = f16(embeddings[embed_base + 3u]);
-    }
-}

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

@@ -1,33 +0,0 @@
-/**
- * GeLU Activation Kernels
- *
- * Provides GeLU activation: x * Phi(x) where Phi is the CDF of standard normal distribution.
- */
-
-import type { Tensor } from '../tensor.js';
-import type { CommandRecorder } from '../command-recorder.js';
-import type { OutputBufferOptions } from './types.js';
-
-/** GeLU kernel options */
-export interface GeLUOptions extends OutputBufferOptions {
-  size?: number | null;
-  gate?: Tensor | null;
-}
-
-/**
- * Run GeLU activation
- */
-export declare function runGeLU(
-  input: Tensor,
-  options?: GeLUOptions
-): Promise<Tensor>;
-
-/**
- * Record GeLU (batched, no submit)
- * Supports gated variant (GeGLU) when options.gate is provided.
- */
-export declare function recordGeLU(
-  recorder: CommandRecorder,
-  input: Tensor,
-  options?: GeLUOptions
-): Promise<Tensor>;

package/src/gpu/kernels/gelu.js

@@ -1,55 +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';
-
-function selectGeluVariant(isF16) {
-  return selectRuleValue('gelu', 'variant', { isF16 });
-}
-
-function resolveOverrides(context) {
-  const overrides = selectRuleValue('gelu', 'overrides', context);
-  return overrides && Object.keys(overrides).length > 0 ? overrides : null;
-}
-
-async function _gelu(target, input, options = {}) {
-  const { size, gate = null, outputBuffer = null } = options;
-
-  const isF16 = input.dtype === 'f16';
-  const bytesPerElement = dtypeBytes(input.dtype);
-  const variant = selectGeluVariant(isF16);
-  const overrides = resolveOverrides({ hasGate: Boolean(gate), useRowsplit: false });
-
-  const inferredSize = size || (input.buffer.size / bytesPerElement);
-  const outputSize = inferredSize * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gelu_output');
-  const gateBuffer = gate ?? input;
-  const ownedOutput = outputBuffer ? null : output;
-
-  try {
-    await unifiedKernelWrapper(
-      'gelu', target, variant,
-      [input, output, gateBuffer],
-      { size: inferredSize, rowsplit_dim: 0 },
-      Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT),
-      overrides
-    );
-
-    return createTensor(output, input.dtype, [inferredSize], 'gelu_output');
-  } catch (error) {
-    if (ownedOutput) {
-      releaseBuffer(ownedOutput);
-    }
-    throw error;
-  }
-}
-
-export async function runGeLU(input, options = {}) {
-  return _gelu(null, input, options);
-}
-
-export async function recordGeLU(recorder, input, options = {}) {
-  return _gelu(recorder, input, options);
-}

package/src/gpu/kernels/gelu.wgsl

@@ -1,64 +0,0 @@
-// GeLU Activation Kernel
-//
-// GeLU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
-// Variants:
-// - geglu: GeLU(gate) * up (separate gate buffer)
-// - geglu_rowsplit: input = [numTokens, 2*dim] row-split
-
-override WORKGROUP_SIZE: u32 = 256u;
-override HAS_GATE: bool = false;
-override USE_ROWSPLIT: bool = false;
-
-struct Uniforms {
-    size: u32,          // Total output elements
-    rowsplit_dim: u32,  // Dim for rowsplit variants (0 when unused)
-    _pad0: u32,
-    _pad1: u32,
-}
-
-@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_write> output: array<f32>;
-@group(0) @binding(3) var<storage, read> gate: array<f32>;
-
-fn gelu(x: f32) -> f32 {
-    let sqrt_2_over_pi: f32 = 0.7978845608;
-    let c: f32 = 0.044715;
-    let inner = sqrt_2_over_pi * (x + c * x * x * x);
-    let inner_clamped = clamp(inner, -15.0, 15.0);
-    return 0.5 * x * (1.0 + tanh(inner_clamped));
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(global_invocation_id) global_id: vec3<u32>
-) {
-    let idx = global_id.x;
-    if (idx >= u.size) {
-        return;
-    }
-
-    if (USE_ROWSPLIT) {
-        if (u.rowsplit_dim == 0u) {
-            return;
-        }
-        let dim = u.rowsplit_dim;
-        let token_idx = idx / dim;
-        let dim_idx = idx % dim;
-        let row_base = token_idx * dim * 2u;
-        let g = input[row_base + dim_idx];
-        let up = input[row_base + dim + dim_idx];
-        output[idx] = gelu(g) * up;
-        return;
-    }
-
-    if (HAS_GATE) {
-        let up = input[idx];
-        let g = gate[idx];
-        output[idx] = gelu(g) * up;
-        return;
-    }
-
-    let x = input[idx];
-    output[idx] = gelu(x);
-}

package/src/gpu/kernels/gelu_f16.wgsl

@@ -1,66 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/gelu.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// GeLU Activation Kernel with F16 Input/Output
-//
-// F16 variant for reduced memory bandwidth when using F16 activations.
-// Intermediate computations use F32 for precision.
-
-enable f16;
-
-override WORKGROUP_SIZE: u32 = 256u;
-override HAS_GATE: bool = false;
-override USE_ROWSPLIT: bool = false;
-
-struct Uniforms {
-    size: u32,          // Total output elements
-    rowsplit_dim: u32,  // Dim for rowsplit variants (0 when unused)
-    _pad0: u32,
-    _pad1: u32,
-}
-
-@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_write> output: array<f16>;
-@group(0) @binding(3) var<storage, read> gate: array<f16>;
-
-fn gelu(x: f32) -> f32 {
-    let sqrt_2_over_pi: f32 = 0.7978845608;
-    let c: f32 = 0.044715;
-    let inner = sqrt_2_over_pi * (x + c * x * x * x);
-    let inner_clamped = clamp(inner, -15.0, 15.0);
-    return 0.5 * x * (1.0 + tanh(inner_clamped));
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(global_invocation_id) global_id: vec3<u32>
-) {
-    let idx = global_id.x;
-    if (idx >= u.size) {
-        return;
-    }
-
-    if (USE_ROWSPLIT) {
-        if (u.rowsplit_dim == 0u) {
-            return;
-        }
-        let dim = u.rowsplit_dim;
-        let token_idx = idx / dim;
-        let dim_idx = idx % dim;
-        let row_base = token_idx * dim * 2u;
-        let g = f32(input[row_base + dim_idx]);
-        let up = f32(input[row_base + dim + dim_idx]);
-        output[idx] = f16(gelu(g) * up);
-        return;
-    }
-
-    if (HAS_GATE) {
-        let up = f32(input[idx]);
-        let g = f32(gate[idx]);
-        output[idx] = f16(gelu(g) * up);
-        return;
-    }
-
-    let x = f32(input[idx]);
-    output[idx] = f16(gelu(x));
-}

package/src/gpu/kernels/gptoss_mxfp4_expert_fused.wgsl

@@ -1,127 +0,0 @@
-// gptoss_mxfp4_expert_fused.wgsl
-
-/**
- * GPT-OSS MXFP4 expert dequant kernel with vec4-oriented write path.
- *
- * This keeps the same bind-group contract as dequant_mxfp4_expert_f16.wgsl,
- * but emits f16 and processes contiguous groups to reduce scalar indexing
- * overhead on subgroup-capable backends.
- */
-
-enable f16;
-
-override WORKGROUP_SIZE_EXPERT: u32 = 256u;
-
-struct ExpertUniforms {
-    expert_idx: u32,
-    num_experts: u32,
-    out_dim: u32,
-    num_groups: u32,
-    total_output: u32,
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-@group(0) @binding(0) var<uniform> eu: ExpertUniforms;
-@group(0) @binding(1) var<storage, read> expert_blocks: array<u32>;
-@group(0) @binding(2) var<storage, read> expert_scales: array<u32>;
-@group(0) @binding(3) var<storage, read_write> expert_output: array<f16>;
-
-fn unpack_nibble(byte_data: u32, nibble_idx: u32) -> f32 {
-    let byte_idx = nibble_idx / 2u;
-    let is_high = nibble_idx % 2u;
-    let byte_val = (byte_data >> (byte_idx * 8u)) & 0xFFu;
-
-    var nibble: u32;
-    if (is_high == 1u) {
-        nibble = (byte_val >> 4u) & 0xFu;
-    } else {
-        nibble = byte_val & 0xFu;
-    }
-
-    let sign_bit = (nibble >> 3u) & 1u;
-    let exp = (nibble >> 1u) & 3u;
-    let mantissa = nibble & 1u;
-
-    var value: f32;
-    if (exp == 0u) {
-        value = f32(mantissa) * 0.5;
-    } else {
-        let m = 1.0 + f32(mantissa) * 0.5;
-        value = m * pow(2.0, f32(exp) - 1.0);
-    }
-
-    if (sign_bit == 1u) {
-        value = -value;
-    }
-    return value;
-}
-
-fn read_scale(scale_data: u32, idx: u32) -> f32 {
-    let byte_idx = idx % 4u;
-    let scale_byte = (scale_data >> (byte_idx * 8u)) & 0xFFu;
-    if (scale_byte == 0u || scale_byte == 255u) {
-        return 0.0;
-    }
-    return exp2(f32(i32(scale_byte) - 127));
-}
-
-fn compute_output_value(out_elem: u32) -> f16 {
-    let row_idx = out_elem / (eu.num_groups * 32u);
-    let col_idx = out_elem % (eu.num_groups * 32u);
-    let group_in_row = col_idx / 32u;
-    let elem_in_group = col_idx % 32u;
-
-    let expert_offset = eu.expert_idx;
-    let blocks_per_expert = eu.out_dim * eu.num_groups * 4u;
-    let blocks_per_row = eu.num_groups * 4u;
-
-    let block_word_idx = expert_offset * blocks_per_expert
-                       + row_idx * blocks_per_row
-                       + group_in_row * 4u
-                       + (elem_in_group / 8u);
-
-    let block_word = expert_blocks[block_word_idx];
-    let nibble_in_word = elem_in_group % 8u;
-    let nibble_val = unpack_nibble(block_word, nibble_in_word);
-
-    let scales_bytes_per_expert = eu.out_dim * eu.num_groups;
-    let scales_bytes_per_row = eu.num_groups;
-    let scale_byte_offset = expert_offset * scales_bytes_per_expert
-                          + row_idx * scales_bytes_per_row
-                          + group_in_row;
-    let scale_word_idx = scale_byte_offset / 4u;
-    let scale_byte_in_word = scale_byte_offset % 4u;
-    let scale = read_scale(expert_scales[scale_word_idx], scale_byte_in_word);
-
-    return f16(nibble_val * scale);
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE_EXPERT, 1, 1)
-fn main_expert(
-    @builtin(global_invocation_id) global_id: vec3<u32>
-) {
-    let base = global_id.x * 4u;
-
-    if (base >= eu.total_output) {
-        return;
-    }
-
-    expert_output[base] = compute_output_value(base);
-
-    let i1 = base + 1u;
-    if (i1 < eu.total_output) {
-        expert_output[i1] = compute_output_value(i1);
-    }
-
-    let i2 = base + 2u;
-    if (i2 < eu.total_output) {
-        expert_output[i2] = compute_output_value(i2);
-    }
-
-    let i3 = base + 3u;
-    if (i3 < eu.total_output) {
-        expert_output[i3] = compute_output_value(i3);
-    }
-}