@simulatte/doppler 0.1.0

package/src/gpu/kernels/backward/conv2d_backward_weight.wgsl

@@ -0,0 +1,70 @@
+// conv2d_backward_weight.wgsl
+//
+// Computes gradient wrt weights (dW).
+// dW [out_channels, in_channels, kernel_h, kernel_w]
+// dW[oc, ic, ky, kx] = sum_over_spatial(dY[oc, oy, ox] * X[ic, oy*stride + ky - pad, ox*stride + kx - pad])
+
+struct Uniforms {
+    in_channels: u32,
+    out_channels: u32,
+    height: u32,
+    width: u32,
+    out_height: u32,
+    out_width: u32,
+    kernel_h: u32,
+    kernel_w: u32,
+    stride: u32,
+    pad: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> grad_output: array<f32>; // dY [out_channels, out_height, out_width]
+@group(0) @binding(2) var<storage, read> input: array<f32>;       // X [in_channels, height, width]
+@group(0) @binding(3) var<storage, read_write> grad_weight: array<f32>; // dW [out_channels, in_channels, kernel_h, kernel_w]
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let weight_idx = gid.x;
+    let total_weight_elements = u.out_channels * u.in_channels * u.kernel_h * u.kernel_w;
+    
+    if (weight_idx >= total_weight_elements) {
+        return;
+    }
+
+    let weight_ic_stride = u.kernel_h * u.kernel_w;
+    let weight_oc_stride = u.in_channels * weight_ic_stride;
+
+    let oc = weight_idx / weight_oc_stride;
+    let rem_oc = weight_idx % weight_oc_stride;
+    let ic = rem_oc / weight_ic_stride;
+    let rem_ic = rem_oc % weight_ic_stride;
+    let ky = rem_ic / u.kernel_w;
+    let kx = rem_ic % u.kernel_w;
+
+    var sum: f32 = 0.0;
+    
+    let out_c_offset = oc * u.out_height * u.out_width;
+    let in_c_offset = ic * u.height * u.width;
+
+    for (var oy: u32 = 0u; oy < u.out_height; oy = oy + 1u) {
+        let iy = oy * u.stride + ky;
+        if (iy < u.pad || iy >= u.height + u.pad) { continue; }
+        let real_iy = iy - u.pad;
+
+        for (var ox: u32 = 0u; ox < u.out_width; ox = ox + 1u) {
+            let ix = ox * u.stride + kx;
+            if (ix < u.pad || ix >= u.width + u.pad) { continue; }
+            let real_ix = ix - u.pad;
+
+            let dy = grad_output[out_c_offset + oy * u.out_width + ox];
+            let x = input[in_c_offset + real_iy * u.width + real_ix];
+            sum = sum + dy * x;
+        }
+    }
+
+    grad_weight[weight_idx] = sum;
+}

package/src/gpu/kernels/backward/cross_entropy_backward.d.ts

@@ -0,0 +1,23 @@
+import type { CommandRecorder } from '../../command-recorder.js';
+import type { Tensor } from '../../tensor.js';
+import type { BackwardKernelOptions } from './utils.js';
+
+export interface CrossEntropyBackwardOptions extends BackwardKernelOptions {
+  numTokens: number;
+  vocabSize: number;
+}
+
+export declare function runCrossEntropyBackward(
+  softmax: Tensor,
+  targets: Tensor,
+  gradOutput: Tensor,
+  options: CrossEntropyBackwardOptions
+): Promise<Tensor>;
+
+export declare function recordCrossEntropyBackward(
+  recorder: CommandRecorder,
+  softmax: Tensor,
+  targets: Tensor,
+  gradOutput: Tensor,
+  options: CrossEntropyBackwardOptions
+): Promise<Tensor>;

package/src/gpu/kernels/backward/cross_entropy_backward.js

@@ -0,0 +1,29 @@
+import { dtypeBytes } from '../../tensor.js';
+import { WORKGROUP_SIZES } from '../constants.js';
+import { createBackwardKernel } from './utils.js';
+
+const { run, record } = createBackwardKernel('cross_entropy_backward', {
+  uniformSize: 16,
+  writeUniforms: (view, opts) => {
+    view.setUint32(0, opts.numTokens, true);
+    view.setUint32(4, opts.vocabSize, true);
+  },
+  calcWorkgroups: (opts) => Math.ceil((opts.numTokens * opts.vocabSize) / WORKGROUP_SIZES.DEFAULT),
+  outputBytes: (opts) => opts.numTokens * opts.vocabSize * opts._bytesPerElement,
+  outputShape: (opts) => [opts.numTokens, opts.vocabSize],
+  dtype: (opts, inputs) => inputs[0].dtype,
+  getDevice: true,
+  validate: (opts) => {
+    if (!opts.numTokens || !opts.vocabSize) throw new Error('cross entropy backward requires numTokens and vocabSize');
+  },
+});
+
+export async function runCrossEntropyBackward(softmax, targets, gradOutput, options = {}) {
+  const bytesPerElement = dtypeBytes(softmax.dtype);
+  return run(softmax, targets, gradOutput, { ...options, _bytesPerElement: bytesPerElement });
+}
+
+export async function recordCrossEntropyBackward(recorder, softmax, targets, gradOutput, options = {}) {
+  const bytesPerElement = dtypeBytes(softmax.dtype);
+  return record(recorder, softmax, targets, gradOutput, { ...options, _bytesPerElement: bytesPerElement });
+}

package/src/gpu/kernels/backward/cross_entropy_backward.wgsl

@@ -0,0 +1,39 @@
+// cross_entropy_backward.wgsl
+
+/**
+ * Cross-entropy backward kernel (expects softmax input).
+ */
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_tokens: u32,
+    vocab_size: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> softmax: array<f32>;
+@group(0) @binding(2) var<storage, read> targets: array<u32>;
+@group(0) @binding(3) var<storage, read> grad_output: array<f32>;
+@group(0) @binding(4) var<storage, read_write> output: array<f32>;
+
+@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.vocab_size;
+    if (idx >= total) {
+        return;
+    }
+
+    let token_idx = idx / u.vocab_size;
+    let class_idx = idx % u.vocab_size;
+    let target_idx = targets[token_idx];
+    var grad = softmax[idx];
+    if (class_idx == target_idx) {
+        grad = grad - 1.0;
+    }
+    let scale = grad_output[token_idx];
+    output[idx] = grad * scale;
+}

package/src/gpu/kernels/backward/embed_backward.d.ts

@@ -0,0 +1,29 @@
+import type { CommandRecorder } from '../../command-recorder.js';
+import type { Tensor } from '../../tensor.js';
+
+export declare function runEmbedBackward(
+  indices: Tensor,
+  gradOutput: Tensor,
+  options: {
+    numTokens?: number;
+    hiddenSize?: number;
+    vocabSize?: number;
+    transpose?: boolean;
+    indexOffset?: number;
+    outputBuffer?: GPUBuffer | null;
+  }
+): Promise<Tensor>;
+
+export declare function recordEmbedBackward(
+  recorder: CommandRecorder,
+  indices: Tensor,
+  gradOutput: Tensor,
+  options: {
+    numTokens?: number;
+    hiddenSize?: number;
+    vocabSize?: number;
+    transpose?: boolean;
+    indexOffset?: number;
+    outputBuffer?: GPUBuffer | null;
+  }
+): Promise<Tensor>;

package/src/gpu/kernels/backward/embed_backward.js

@@ -0,0 +1,118 @@
+import { getDevice } from '../../device.js';
+import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { createTensor } from '../../tensor.js';
+import { WORKGROUP_SIZES } from '../constants.js';
+import { createPipeline, createUniformBufferWithView } from '../utils.js';
+
+function resolveOptions(gradOutput, options) {
+  const numTokens = options.numTokens ?? gradOutput.shape?.[0] ?? null;
+  const hiddenSize = options.hiddenSize ?? gradOutput.shape?.[1] ?? null;
+  const vocabSize = options.vocabSize ?? null;
+  if (!numTokens || !hiddenSize || !vocabSize) {
+    throw new Error('embed backward requires numTokens, hiddenSize, and vocabSize');
+  }
+  const transpose = options.transpose === true;
+  const indexOffset = options.indexOffset ?? 0;
+  if (!Number.isFinite(indexOffset) || indexOffset < 0) {
+    throw new Error('embed backward requires a non-negative indexOffset');
+  }
+  return { numTokens, hiddenSize, vocabSize, transpose, indexOffset };
+}
+
+export async function runEmbedBackward(indices, gradOutput, options = {}) {
+  const device = getDevice();
+  if (gradOutput.dtype !== 'f32') {
+    throw new Error('embed backward requires f32 gradOutput');
+  }
+
+  const { numTokens, hiddenSize, vocabSize, transpose, indexOffset } = resolveOptions(gradOutput, options);
+  const { outputBuffer = null } = options;
+  const outputSize = vocabSize * hiddenSize * 4;
+  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'embed_backward_output');
+
+  const pipeline = await createPipeline('embed_backward', 'default');
+  const uniformBuffer = createUniformBufferWithView(
+    'embed_backward_uniforms',
+    32,
+    (view) => {
+      view.setUint32(0, numTokens, true);
+      view.setUint32(4, hiddenSize, true);
+      view.setUint32(8, vocabSize, true);
+      view.setUint32(12, transpose ? 1 : 0, true);
+      view.setUint32(16, indexOffset, true);
+    },
+    null,
+    device
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'embed_backward_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: indices.buffer } },
+      { binding: 2, resource: { buffer: gradOutput.buffer } },
+      { binding: 3, resource: { buffer: outputBuf } },
+    ],
+  });
+
+  const encoder = device.createCommandEncoder({ label: 'embed_backward_encoder' });
+  encoder.clearBuffer(outputBuf);
+  const pass = encoder.beginComputePass({ label: 'embed_backward_pass' });
+  pass.setPipeline(pipeline);
+  pass.setBindGroup(0, bindGroup);
+  const workgroups = Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT);
+  pass.dispatchWorkgroups(workgroups);
+  pass.end();
+  device.queue.submit([encoder.finish()]);
+  uniformBuffer.destroy();
+
+  const shape = transpose ? [hiddenSize, vocabSize] : [vocabSize, hiddenSize];
+  return createTensor(outputBuf, 'f32', shape, 'embed_backward_output');
+}
+
+export async function recordEmbedBackward(recorder, indices, gradOutput, options = {}) {
+  if (gradOutput.dtype !== 'f32') {
+    throw new Error('embed backward requires f32 gradOutput');
+  }
+
+  const { numTokens, hiddenSize, vocabSize, transpose, indexOffset } = resolveOptions(gradOutput, options);
+  const { outputBuffer = null } = options;
+  const outputSize = vocabSize * hiddenSize * 4;
+  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'embed_backward_output');
+
+  const pipeline = await createPipeline('embed_backward', 'default');
+  const uniformBuffer = createUniformBufferWithView(
+    'embed_backward_uniforms',
+    32,
+    (view) => {
+      view.setUint32(0, numTokens, true);
+      view.setUint32(4, hiddenSize, true);
+      view.setUint32(8, vocabSize, true);
+      view.setUint32(12, transpose ? 1 : 0, true);
+      view.setUint32(16, indexOffset, true);
+    },
+    recorder
+  );
+
+  const bindGroup = recorder.device.createBindGroup({
+    label: 'embed_backward_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: indices.buffer } },
+      { binding: 2, resource: { buffer: gradOutput.buffer } },
+      { binding: 3, resource: { buffer: outputBuf } },
+    ],
+  });
+
+  recorder.getEncoder().clearBuffer(outputBuf);
+  const pass = recorder.beginComputePass('embed_backward');
+  pass.setPipeline(pipeline);
+  pass.setBindGroup(0, bindGroup);
+  pass.dispatchWorkgroups(Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT));
+  pass.end();
+
+  const shape = transpose ? [hiddenSize, vocabSize] : [vocabSize, hiddenSize];
+  return createTensor(outputBuf, 'f32', shape, 'embed_backward_output');
+}

package/src/gpu/kernels/backward/embed_backward.wgsl

@@ -0,0 +1,73 @@
+// embed_backward.wgsl
+
+/**
+ * Embedding backward kernel.
+ *
+ * Accumulates dL/d(embed_table) from dL/d(embed_out) and token indices.
+ *
+ * Forward gather (non-transposed weights):
+ *   out[token, dim] = embed_table[token_id, dim]
+ *
+ * Backward:
+ *   grad_embed[token_id, dim] += grad_out[token, dim]
+ *
+ * Note: Multiple tokens can reference the same token_id, so accumulation
+ * uses atomic add via atomic<u32> compare-exchange on f32 bit patterns.
+ */
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_tokens: u32,
+    hidden_size: u32,
+    vocab_size: u32,
+    transpose: u32,
+    index_offset: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> indices: array<u32>;
+@group(0) @binding(2) var<storage, read> grad_output: array<f32>;
+@group(0) @binding(3) var<storage, read_write> grad_embeddings: array<atomic<u32>>;
+
+fn atomic_add_f32(ptr: ptr<storage, atomic<u32>, read_write>, delta: f32) {
+    var old_bits = atomicLoad(ptr);
+    loop {
+        let old_val = bitcast<f32>(old_bits);
+        let new_bits = bitcast<u32>(old_val + delta);
+        let res = atomicCompareExchangeWeak(ptr, old_bits, new_bits);
+        if (res.exchanged) {
+            break;
+        }
+        old_bits = res.old_value;
+    }
+}
+
+@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.hidden_size;
+    if (idx >= total) {
+        return;
+    }
+
+    let token_idx = idx / u.hidden_size;
+    let dim_idx = idx % u.hidden_size;
+
+    let token_id = indices[token_idx + u.index_offset];
+    if (token_id >= u.vocab_size) {
+        return;
+    }
+
+    var out_idx: u32;
+    if (u.transpose == 1u) {
+        out_idx = dim_idx * u.vocab_size + token_id;
+    } else {
+        out_idx = token_id * u.hidden_size + dim_idx;
+    }
+
+    atomic_add_f32(&grad_embeddings[out_idx], grad_output[idx]);
+}

package/src/gpu/kernels/backward/gelu_backward.d.ts

@@ -0,0 +1,16 @@
+import type { CommandRecorder } from '../../command-recorder.js';
+import type { Tensor } from '../../tensor.js';
+import type { BackwardKernelOptions } from './utils.js';
+
+export declare function runGeluBackward(
+  input: Tensor,
+  gradOutput: Tensor,
+  options?: BackwardKernelOptions
+): Promise<Tensor>;
+
+export declare function recordGeluBackward(
+  recorder: CommandRecorder,
+  input: Tensor,
+  gradOutput: Tensor,
+  options?: BackwardKernelOptions
+): Promise<Tensor>;

package/src/gpu/kernels/backward/gelu_backward.js

@@ -0,0 +1,39 @@
+import { dtypeBytes } from '../../tensor.js';
+import { WORKGROUP_SIZES } from '../constants.js';
+import { createBackwardKernel } from './utils.js';
+
+const { run, record } = createBackwardKernel('gelu_backward', {
+  uniformSize: 16,
+  writeUniforms: (view, opts) => {
+    view.setUint32(0, opts._count, true);
+  },
+  calcWorkgroups: (opts) => Math.ceil(opts._count / WORKGROUP_SIZES.DEFAULT),
+  outputBytes: (opts) => opts._count * opts._bytesPerElement,
+  outputShape: (opts) => opts._shape,
+  dtype: (opts) => opts._dtype,
+  getDevice: true,
+});
+
+export function runGeluBackward(input, gradOutput, options = {}) {
+  const bytesPerElement = dtypeBytes(gradOutput.dtype);
+  const count = options.count ?? Math.floor(gradOutput.buffer.size / bytesPerElement);
+  return run(input, gradOutput, {
+    ...options,
+    _count: count,
+    _bytesPerElement: bytesPerElement,
+    _shape: [...gradOutput.shape],
+    _dtype: gradOutput.dtype,
+  });
+}
+
+export function recordGeluBackward(recorder, input, gradOutput, options = {}) {
+  const bytesPerElement = dtypeBytes(gradOutput.dtype);
+  const count = options.count ?? Math.floor(gradOutput.buffer.size / bytesPerElement);
+  return record(recorder, input, gradOutput, {
+    ...options,
+    _count: count,
+    _bytesPerElement: bytesPerElement,
+    _shape: [...gradOutput.shape],
+    _dtype: gradOutput.dtype,
+  });
+}

package/src/gpu/kernels/backward/gelu_backward.wgsl

@@ -0,0 +1,38 @@
+// gelu_backward.wgsl
+
+/**
+ * GELU backward kernel.
+ */
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    size: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: 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> grad_output: 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 idx = gid.x;
+    if (idx >= u.size) {
+        return;
+    }
+    let x = input[idx];
+    let sqrt_2_over_pi: f32 = 0.7978845608;
+    let c: f32 = 0.044715;
+    let x3 = x * x * x;
+    let inner = sqrt_2_over_pi * (x + c * x3);
+    let inner_clamped = clamp(inner, -15.0, 15.0);
+    let tanh_inner = tanh(inner_clamped);
+    let sech2 = 1.0 - tanh_inner * tanh_inner;
+    let inner_deriv = sqrt_2_over_pi * (1.0 + 3.0 * c * x * x);
+    let deriv = 0.5 * (1.0 + tanh_inner) + 0.5 * x * sech2 * inner_deriv;
+    output[idx] = grad_output[idx] * deriv;
+}

package/src/gpu/kernels/backward/groupnorm_backward.d.ts

@@ -0,0 +1,24 @@
+import type { BackwardKernelOptions } from './utils.js';
+
+export interface GroupNormBackwardOptions extends BackwardKernelOptions {
+  channels: number;
+  height: number;
+  width: number;
+  numGroups: number;
+  eps?: number;
+}
+
+export declare function runGroupNormBackward(
+  input: any,
+  weight: any,
+  gradOutput: any,
+  options: GroupNormBackwardOptions
+): Promise<any>;
+
+export declare function recordGroupNormBackward(
+  recorder: any,
+  input: any,
+  weight: any,
+  gradOutput: any,
+  options: GroupNormBackwardOptions
+): Promise<any>;

package/src/gpu/kernels/backward/groupnorm_backward.js

@@ -0,0 +1,29 @@
+import { createBackwardKernel } from './utils.js';
+
+const { run, record } = createBackwardKernel('groupnorm_backward', {
+  uniformSize: 32,
+  writeUniforms: (view, opts) => {
+    view.setUint32(0, opts.channels, true);
+    view.setUint32(4, opts.height, true);
+    view.setUint32(8, opts.width, true);
+    view.setUint32(12, opts.numGroups, true);
+    view.setFloat32(16, opts.eps ?? 1e-5, true);
+  },
+  calcWorkgroups: (opts) => opts.numGroups,
+  outputBytes: (opts) => opts.channels * opts.height * opts.width * 4,
+  outputShape: (opts) => [opts.channels, opts.height, opts.width],
+  getDevice: true,
+  validate: (opts) => {
+    if (!opts.channels || !opts.height || !opts.width || !opts.numGroups) {
+      throw new Error('groupnorm backward requires all dimensions');
+    }
+  },
+});
+
+export async function runGroupNormBackward(input, weight, gradOutput, options = {}) {
+  return run(input, weight, gradOutput, options);
+}
+
+export async function recordGroupNormBackward(recorder, input, weight, gradOutput, options = {}) {
+  return record(recorder, input, weight, gradOutput, options);
+}

package/src/gpu/kernels/backward/groupnorm_backward.wgsl

@@ -0,0 +1,143 @@
+// groupnorm_backward.wgsl
+//
+// GroupNorm Backward Kernel (GPU)
+// Computes dInput per element. 
+// dWeight/dBias require cross-spatial reduction (left for future or handled via separate reduce).
+
+struct Uniforms {
+    channels: u32,
+    height: u32,
+    width: u32,
+    num_groups: u32,
+    eps: f32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f32>;        // x
+@group(0) @binding(2) var<storage, read> weight: array<f32>;       // gamma (per channel)
+@group(0) @binding(3) var<storage, read> grad_output: array<f32>;  // dY
+@group(0) @binding(4) var<storage, read_write> grad_input: array<f32>; // dX
+
+override WORKGROUP_SIZE: u32 = 256u;
+const MAX_WORKGROUP_SIZE: u32 = 256u;
+
+var<workgroup> shared_sum: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_sum_sq: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_dot_w: array<f32, MAX_WORKGROUP_SIZE>;
+var<workgroup> shared_dot_w_norm: array<f32, MAX_WORKGROUP_SIZE>;
+
+@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 group_idx = wg_id.x;
+    let thread_idx = local_id.x;
+    
+    if (group_idx >= u.num_groups) {
+        return;
+    }
+
+    let channels_per_group = u.channels / u.num_groups;
+    let spatial = u.height * u.width;
+    let elements_per_group = channels_per_group * spatial;
+    
+    let group_c_start = group_idx * channels_per_group;
+    
+    // -------------------------------------------------------------------------
+    // Phase 1: Recompute Mean and Variance for the group
+    // -------------------------------------------------------------------------
+    var local_sum: f32 = 0.0;
+    var local_sum_sq: f32 = 0.0;
+    
+    for (var i = thread_idx; i < elements_per_group; i += WORKGROUP_SIZE) {
+        let c_in_group = i / spatial;
+        let s_idx = i % spatial;
+        let c = group_c_start + c_in_group;
+        let idx = (c * u.height + (s_idx / u.width)) * u.width + (s_idx % u.width);
+        // Simplified mapping if NCHW is contiguous spatial
+        let real_idx = c * spatial + s_idx; 
+        
+        let x = input[real_idx];
+        local_sum += x;
+        local_sum_sq += x * x;
+    }
+
+    shared_sum[thread_idx] = local_sum;
+    shared_sum_sq[thread_idx] = local_sum_sq;
+    workgroupBarrier();
+
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride >>= 1u) {
+        if (thread_idx < stride) {
+            shared_sum[thread_idx] += shared_sum[thread_idx + stride];
+            shared_sum_sq[thread_idx] += shared_sum_sq[thread_idx + stride];
+        }
+        workgroupBarrier();
+    }
+
+    let n = f32(elements_per_group);
+    let mean = shared_sum[0] / n;
+    let mean_sq = shared_sum_sq[0] / n;
+    let variance = max(0.0, mean_sq - mean * mean);
+    let inv_std = 1.0 / sqrt(variance + u.eps);
+
+    workgroupBarrier();
+
+    // -------------------------------------------------------------------------
+    // Phase 2: Compute Gradient Stats for the group
+    // -------------------------------------------------------------------------
+    var local_dot_w: f32 = 0.0;
+    var local_dot_w_norm: f32 = 0.0;
+
+    for (var i = thread_idx; i < elements_per_group; i += WORKGROUP_SIZE) {
+        let c_in_group = i / spatial;
+        let s_idx = i % spatial;
+        let c = group_c_start + c_in_group;
+        let real_idx = c * spatial + s_idx;
+        
+        let x = input[real_idx];
+        let dy = grad_output[real_idx];
+        let w = weight[c];
+        let dy_w = dy * w;
+        
+        local_dot_w += dy_w;
+        local_dot_w_norm += dy_w * (x - mean);
+    }
+
+    shared_dot_w[thread_idx] = local_dot_w;
+    shared_dot_w_norm[thread_idx] = local_dot_w_norm;
+    workgroupBarrier();
+
+    for (var stride: u32 = WORKGROUP_SIZE / 2u; stride > 0u; stride >>= 1u) {
+        if (thread_idx < stride) {
+            shared_dot_w[thread_idx] += shared_dot_w[thread_idx + stride];
+            shared_dot_w_norm[thread_idx] += shared_dot_w_norm[thread_idx + stride];
+        }
+        workgroupBarrier();
+    }
+
+    let sum_dy_w = shared_dot_w[0];
+    let sum_dy_w_norm = shared_dot_w_norm[0];
+    let term1 = sum_dy_w;
+    let term2 = sum_dy_w_norm * (inv_std * inv_std);
+
+    // -------------------------------------------------------------------------
+    // Phase 3: Write Gradient Input
+    // -------------------------------------------------------------------------
+    for (var i = thread_idx; i < elements_per_group; i += WORKGROUP_SIZE) {
+        let c_in_group = i / spatial;
+        let s_idx = i % spatial;
+        let c = group_c_start + c_in_group;
+        let real_idx = c * spatial + s_idx;
+        
+        let x = input[real_idx];
+        let dy = grad_output[real_idx];
+        let w = weight[c];
+
+        let dx = inv_std * (dy * w - (term1 + (x - mean) * term2) / n);
+        grad_input[real_idx] = dx;
+    }
+}