@simulatte/doppler 0.1.4 → 0.1.6

package/src/gpu/kernels/attention.js

@@ -780,6 +780,23 @@ function resolveAttentionExecution(recorder) {
   };
 }
 
+function assertAttentionBindGroupBuffer(kernelName, variant, bindingIndex, bindingLabel, buffer, details = []) {
+  const isGpuBuffer = buffer && (
+    typeof GPUBuffer === 'undefined'
+      ? true
+      : buffer instanceof GPUBuffer
+  );
+  if (isGpuBuffer) {
+    return;
+  }
+  const detailText = details.filter(Boolean).join(', ');
+  throw new Error(
+    `[${kernelName}] variant="${variant}" binding ${bindingIndex} "${bindingLabel}" requires a GPUBuffer` +
+    (detailText ? ` (${detailText})` : '') +
+    '.'
+  );
+}
+
 function releaseAttentionUniform(execution, uniformBuffer) {
   if (!execution.recorder) {
     releaseUniformBuffer(uniformBuffer);
@@ -867,6 +884,26 @@ async function executeAttentionBDPA(
     slidingWindow,
   });
 
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 0, 'uniforms', uniformBuffer);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 1, 'Q', Q?.buffer, [
+    `QLabel=${Q?.label ?? 'unknown'}`,
+    `QDtype=${Q?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 2, 'basisK', basisK?.buffer, [
+    `basisKLabel=${basisK?.label ?? 'unknown'}`,
+    `basisKDtype=${basisK?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 3, 'basisV', basisV?.buffer, [
+    `basisVLabel=${basisV?.label ?? 'unknown'}`,
+    `basisVDtype=${basisV?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 4, 'pagedK', pagedK);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 5, 'pagedV', pagedV);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 6, 'index', index);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 7, 'ropeCos', ropeCos);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 8, 'ropeSin', ropeSin);
+  assertAttentionBindGroupBuffer('attention_bdpa', variant, 9, 'output', outputBuf);
+
   const bindGroup = execution.device.createBindGroup({
     label: 'attention_bdpa_bind_group',
     layout: pipeline.getBindGroupLayout(0),
@@ -982,6 +1019,24 @@ async function executeAttention(
 
   const kvLenBinding = kvLenBuffer || getKvLenFallbackBuffer(execution.device);
   const pageTableBinding = kvPageTable || getPageTableFallbackBuffer(execution.device);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 0, 'uniforms', uniformBuffer);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 1, 'Q', Q?.buffer, [
+    `QLabel=${Q?.label ?? 'unknown'}`,
+    `QDtype=${Q?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 2, 'K', K?.buffer, [
+    `KLabel=${K?.label ?? 'unknown'}`,
+    `KDtype=${K?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 3, 'V', V?.buffer, [
+    `VLabel=${V?.label ?? 'unknown'}`,
+    `VDtype=${V?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 4, 'output', outputBuf);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 5, 'kvLen', kvLenBinding);
+  assertAttentionBindGroupBuffer('attention', plan.variant, 6, 'pageTable', pageTableBinding, [
+    `kvLayout=${kvLayout}`,
+  ]);
   const bindGroup = execution.device.createBindGroup({
     label: 'attention_bind_group',
     layout: pipeline.getBindGroupLayout(0),
@@ -1099,6 +1154,31 @@ async function executeAttentionTiered(
   });
 
   const pageTableBinding = coldPageTable || getPageTableFallbackBuffer(execution.device);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 0, 'uniforms', uniformBuffer);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 1, 'Q', Q?.buffer, [
+    `QLabel=${Q?.label ?? 'unknown'}`,
+    `QDtype=${Q?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 2, 'hotK', hotK?.buffer, [
+    `hotKLabel=${hotK?.label ?? 'unknown'}`,
+    `hotKDtype=${hotK?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 3, 'hotV', hotV?.buffer, [
+    `hotVLabel=${hotV?.label ?? 'unknown'}`,
+    `hotVDtype=${hotV?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 4, 'coldK', coldK?.buffer, [
+    `coldKLabel=${coldK?.label ?? 'unknown'}`,
+    `coldKDtype=${coldK?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 5, 'coldV', coldV?.buffer, [
+    `coldVLabel=${coldV?.label ?? 'unknown'}`,
+    `coldVDtype=${coldV?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 6, 'output', outputBuf);
+  assertAttentionBindGroupBuffer('attention_tiered', variant, 7, 'pageTable', pageTableBinding, [
+    `coldLayout=${coldLayout}`,
+  ]);
   const bindGroup = execution.device.createBindGroup({
     label: 'attention_tiered_bind_group',
     layout: pipeline.getBindGroupLayout(0),
@@ -1200,6 +1280,24 @@ async function executeAttentionTieredQuant(
     packedStride,
   });
 
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 0, 'uniforms', uniformBuffer);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 1, 'Q', Q?.buffer, [
+    `QLabel=${Q?.label ?? 'unknown'}`,
+    `QDtype=${Q?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 2, 'hotK', hotK?.buffer, [
+    `hotKLabel=${hotK?.label ?? 'unknown'}`,
+    `hotKDtype=${hotK?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 3, 'hotV', hotV?.buffer, [
+    `hotVLabel=${hotV?.label ?? 'unknown'}`,
+    `hotVDtype=${hotV?.dtype ?? 'unknown'}`,
+  ]);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 4, 'coldPackedK', coldPackedK);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 5, 'coldPackedV', coldPackedV);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 6, 'coldScalesK', coldScalesK);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 7, 'coldScalesV', coldScalesV);
+  assertAttentionBindGroupBuffer('attention_tiered_quant', variant, 8, 'output', outputBuf);
   const bindGroup = execution.device.createBindGroup({
     label: 'attention_tiered_quant_bind_group',
     layout: pipeline.getBindGroupLayout(0),

package/src/gpu/kernels/bias_add.wgsl

@@ -14,6 +14,10 @@ struct Uniforms {
     dim: u32,
     data_offset: u32,  // byte offset into data buffer (divide by 4 for F32)
     bias_offset: u32,  // byte offset into bias buffer (divide by 4 for F32)
+    token_stride: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 override WORKGROUP_SIZE: u32 = 256u;
@@ -24,17 +28,15 @@ override WORKGROUP_SIZE: u32 = 256u;
 
 @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) {
+    let d = gid.x;
+    let token = gid.z * max(u.token_stride, 1u) + gid.y;
+    if (token >= u.num_tokens || d >= u.dim) {
         return;
     }
 
     // Convert byte offsets to F32 indices
     let data_base = u.data_offset / 4u;
     let bias_base = u.bias_offset / 4u;
-
-    let d = idx % u.dim;
+    let idx = token * u.dim + d;
     data[data_base + idx] = data[data_base + idx] + bias[bias_base + d];
 }
-

package/src/gpu/kernels/bias_add_f16.wgsl

@@ -18,6 +18,10 @@ struct Uniforms {
     dim: u32,
     data_offset: u32,  // byte offset into data buffer (divide by 2 for F16)
     bias_offset: u32,  // byte offset into bias buffer (divide by 2 for F16)
+    token_stride: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 override WORKGROUP_SIZE: u32 = 256u;
@@ -28,17 +32,16 @@ override WORKGROUP_SIZE: u32 = 256u;
 
 @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) {
+    let d = gid.x;
+    let token = gid.z * max(u.token_stride, 1u) + gid.y;
+    if (token >= u.num_tokens || d >= u.dim) {
         return;
     }
 
     // Convert byte offsets to F16 indices
     let data_base = u.data_offset / 2u;
     let bias_base = u.bias_offset / 2u;
-
-    let d = idx % u.dim;
+    let idx = token * u.dim + d;
     let out = f32(data[data_base + idx]) + f32(bias[bias_base + d]);
     data[data_base + idx] = f16(out);
 }

package/src/gpu/kernels/conv2d.js

@@ -58,7 +58,7 @@ async function _conv2d(target, input, weight, bias, options = {}) {
       kernel_h: kernelH, kernel_w: kernelW,
       stride, pad, _pad0: 0, _pad1: 0,
     },
-    Math.ceil((outChannels * outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT)
+    [Math.ceil((outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
   );
 
   if (tempBias) {

package/src/gpu/kernels/conv2d.wgsl

@@ -27,19 +27,18 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let out_height = u.out_height;
     let out_width = u.out_width;
-    let out_size = u.out_channels * out_height * out_width;
-    if (idx >= out_size) {
+    let out_spatial = out_height * out_width;
+    let out_spatial_idx = gid.x;
+    let out_c = gid.y;
+    if (out_c >= u.out_channels || out_spatial_idx >= out_spatial) {
         return;
     }
 
-    let out_spatial = out_height * out_width;
-    let out_c = idx / out_spatial;
-    let rem = idx - out_c * out_spatial;
-    let out_y = rem / out_width;
-    let out_x = rem - out_y * out_width;
+    let out_y = out_spatial_idx / out_width;
+    let out_x = out_spatial_idx - out_y * out_width;
+    let idx = out_c * out_spatial + out_spatial_idx;
 
     var sum: f32 = bias[out_c];
 

package/src/gpu/kernels/conv2d_f16.wgsl

@@ -29,19 +29,18 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let out_height = u.out_height;
     let out_width = u.out_width;
-    let out_size = u.out_channels * out_height * out_width;
-    if (idx >= out_size) {
+    let out_spatial = out_height * out_width;
+    let out_spatial_idx = gid.x;
+    let out_c = gid.y;
+    if (out_c >= u.out_channels || out_spatial_idx >= out_spatial) {
         return;
     }
 
-    let out_spatial = out_height * out_width;
-    let out_c = idx / out_spatial;
-    let rem = idx - out_c * out_spatial;
-    let out_y = rem / out_width;
-    let out_x = rem - out_y * out_width;
+    let out_y = out_spatial_idx / out_width;
+    let out_x = out_spatial_idx - out_y * out_width;
+    let idx = out_c * out_spatial + out_spatial_idx;
 
     var sum: f32 = f32(bias[out_c]);
 

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

@@ -0,0 +1,29 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+import type { WeightBuffer } from '../weight-buffer.js';
+
+export interface DepthwiseConv2DOptions extends OutputBufferOptions {
+  channels: number;
+  height: number;
+  width: number;
+  kernelH: number;
+  kernelW: number;
+  stride?: number;
+  pad?: number;
+}
+
+export declare function runDepthwiseConv2D(
+  input: Tensor,
+  weight: GPUBuffer | WeightBuffer,
+  bias: GPUBuffer | WeightBuffer | null,
+  options: DepthwiseConv2DOptions
+): Promise<Tensor>;
+
+export declare function recordDepthwiseConv2D(
+  recorder: CommandRecorder,
+  input: Tensor,
+  weight: GPUBuffer | WeightBuffer,
+  bias: GPUBuffer | WeightBuffer | null,
+  options: DepthwiseConv2DOptions
+): Promise<Tensor>;

package/src/gpu/kernels/depthwise_conv2d.js

@@ -0,0 +1,99 @@
+import { getDevice } from '../device.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../tensor.js';
+import { getBuffer } from '../weight-buffer.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { selectRuleValue } from './rule-registry.js';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectDepthwiseConv2DVariant(isF16) {
+  return selectRuleValue('depthwiseConv2d', 'variant', { isF16 });
+}
+
+async function _depthwiseConv2D(target, input, weight, bias, options = {}) {
+  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
+  const device = target?.device || getDevice();
+  const {
+    channels,
+    height,
+    width,
+    kernelH,
+    kernelW,
+    stride = 1,
+    pad = 0,
+    outputBuffer = null,
+  } = options;
+
+  if (
+    !Number.isFinite(channels) ||
+    !Number.isFinite(height) ||
+    !Number.isFinite(width) ||
+    !Number.isFinite(kernelH) ||
+    !Number.isFinite(kernelW)
+  ) {
+    throw new Error('DepthwiseConv2D requires explicit dimensions.');
+  }
+
+  const outHeight = Math.floor((height + pad * 2 - kernelH) / stride) + 1;
+  const outWidth = Math.floor((width + pad * 2 - kernelW) / stride) + 1;
+  if (outHeight <= 0 || outWidth <= 0) {
+    throw new Error(`DepthwiseConv2D invalid output size: ${outHeight}x${outWidth}`);
+  }
+
+  const isF16 = input.dtype === 'f16';
+  const variant = selectDepthwiseConv2DVariant(isF16);
+  const bytesPerElement = dtypeBytes(input.dtype);
+  const outputSize = channels * outHeight * outWidth * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'depthwise_conv2d_output');
+  const outSpatial = outHeight * outWidth;
+
+  const weightBuffer = getBuffer(weight);
+  let biasBuffer = getBuffer(bias);
+  let tempBias = null;
+  if (!biasBuffer) {
+    const biasSize = channels * bytesPerElement;
+    tempBias = acquireBuffer(biasSize, undefined, 'depthwise_conv2d_bias_zero');
+    biasBuffer = tempBias;
+    const paddedSize = Math.ceil(biasSize / 4) * 4;
+    device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
+  }
+
+  await unifiedKernelWrapper(
+    'depthwise_conv2d',
+    target,
+    variant,
+    [input, weightBuffer, biasBuffer, output],
+    {
+      channels,
+      height,
+      width,
+      out_height: outHeight,
+      out_width: outWidth,
+      kernel_h: kernelH,
+      kernel_w: kernelW,
+      stride,
+      pad,
+      _pad0: 0,
+      _pad1: 0,
+    },
+    [Math.ceil(outSpatial / WORKGROUP_SIZES.DEFAULT), channels, 1]
+  );
+
+  if (tempBias) {
+    if (recorder) {
+      recorder.trackTemporaryBuffer(tempBias);
+    } else {
+      releaseBuffer(tempBias);
+    }
+  }
+
+  return createTensor(output, input.dtype, [channels, outHeight, outWidth], 'depthwise_conv2d_output');
+}
+
+export async function runDepthwiseConv2D(input, weight, bias, options = {}) {
+  return _depthwiseConv2D(null, input, weight, bias, options);
+}
+
+export async function recordDepthwiseConv2D(recorder, input, weight, bias, options = {}) {
+  return _depthwiseConv2D(recorder, input, weight, bias, options);
+}

package/src/gpu/kernels/depthwise_conv2d.wgsl

@@ -0,0 +1,55 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    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,
+    _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> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: 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 out_spatial = u.out_height * u.out_width;
+    let spatial_idx = gid.x;
+    let channel = gid.y;
+    if (spatial_idx >= out_spatial || channel >= u.channels) {
+        return;
+    }
+    let out_y = spatial_idx / u.out_width;
+    let out_x = spatial_idx - out_y * u.out_width;
+
+    var sum: f32 = bias[channel];
+    let pad = i32(u.pad);
+
+    for (var ky: u32 = 0u; ky < u.kernel_h; ky = ky + 1u) {
+        let in_y = i32(out_y * u.stride + ky) - pad;
+        if (in_y < 0 || in_y >= i32(u.height)) {
+            continue;
+        }
+        for (var kx: u32 = 0u; kx < u.kernel_w; kx = kx + 1u) {
+            let in_x = i32(out_x * u.stride + kx) - pad;
+            if (in_x < 0 || in_x >= i32(u.width)) {
+                continue;
+            }
+            let input_idx = (channel * u.height + u32(in_y)) * u.width + u32(in_x);
+            let weight_idx = ((channel * u.kernel_h + ky) * u.kernel_w + kx);
+            sum = sum + input[input_idx] * weight[weight_idx];
+        }
+    }
+
+    output[channel * out_spatial + spatial_idx] = sum;
+}

package/src/gpu/kernels/depthwise_conv2d_f16.wgsl

@@ -0,0 +1,59 @@
+// Depthwise Conv2D Kernel (NCHW, f16)
+
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    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,
+    _pad2: 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> weight: array<f16>;
+@group(0) @binding(3) var<storage, read> bias: array<f16>;
+@group(0) @binding(4) var<storage, read_write> output: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let out_spatial = u.out_height * u.out_width;
+    let spatial_idx = gid.x;
+    let channel = gid.y;
+    if (spatial_idx >= out_spatial || channel >= u.channels) {
+        return;
+    }
+    let out_y = spatial_idx / u.out_width;
+    let out_x = spatial_idx - out_y * u.out_width;
+
+    var sum: f32 = f32(bias[channel]);
+    let pad = i32(u.pad);
+
+    for (var ky: u32 = 0u; ky < u.kernel_h; ky = ky + 1u) {
+        let in_y = i32(out_y * u.stride + ky) - pad;
+        if (in_y < 0 || in_y >= i32(u.height)) {
+            continue;
+        }
+        for (var kx: u32 = 0u; kx < u.kernel_w; kx = kx + 1u) {
+            let in_x = i32(out_x * u.stride + kx) - pad;
+            if (in_x < 0 || in_x >= i32(u.width)) {
+                continue;
+            }
+            let input_idx = (channel * u.height + u32(in_y)) * u.width + u32(in_x);
+            let weight_idx = ((channel * u.kernel_h + ky) * u.kernel_w + kx);
+            sum = sum + f32(input[input_idx]) * f32(weight[weight_idx]);
+        }
+    }
+
+    output[channel * out_spatial + spatial_idx] = f16(sum);
+}

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

@@ -0,0 +1,27 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+import type { WeightBuffer } from '../weight-buffer.js';
+
+export interface GroupedPointwiseConv2DOptions extends OutputBufferOptions {
+  inChannels: number;
+  outChannels: number;
+  height: number;
+  width: number;
+  groups: number;
+}
+
+export declare function runGroupedPointwiseConv2D(
+  input: Tensor,
+  weight: GPUBuffer | WeightBuffer,
+  bias: GPUBuffer | WeightBuffer | null,
+  options: GroupedPointwiseConv2DOptions
+): Promise<Tensor>;
+
+export declare function recordGroupedPointwiseConv2D(
+  recorder: CommandRecorder,
+  input: Tensor,
+  weight: GPUBuffer | WeightBuffer,
+  bias: GPUBuffer | WeightBuffer | null,
+  options: GroupedPointwiseConv2DOptions
+): Promise<Tensor>;

package/src/gpu/kernels/grouped_pointwise_conv2d.js

@@ -0,0 +1,93 @@
+import { getDevice } from '../device.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../tensor.js';
+import { getBuffer } from '../weight-buffer.js';
+import { unifiedKernelWrapper } from './utils.js';
+import { selectRuleValue } from './rule-registry.js';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectGroupedPointwiseConv2DVariant(isF16) {
+  return selectRuleValue('groupedPointwiseConv2d', 'variant', { isF16 });
+}
+
+async function _groupedPointwiseConv2D(target, input, weight, bias, options = {}) {
+  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
+  const device = target?.device || getDevice();
+  const {
+    inChannels,
+    outChannels,
+    height,
+    width,
+    groups,
+    outputBuffer = null,
+  } = options;
+
+  if (
+    !Number.isFinite(inChannels) ||
+    !Number.isFinite(outChannels) ||
+    !Number.isFinite(height) ||
+    !Number.isFinite(width) ||
+    !Number.isFinite(groups)
+  ) {
+    throw new Error('GroupedPointwiseConv2D requires explicit dimensions.');
+  }
+  if (groups <= 0 || inChannels % groups !== 0 || outChannels % groups !== 0) {
+    throw new Error(
+      `GroupedPointwiseConv2D requires inChannels/outChannels divisible by groups. Got ${inChannels}/${outChannels}/${groups}.`
+    );
+  }
+
+  const isF16 = input.dtype === 'f16';
+  const variant = selectGroupedPointwiseConv2DVariant(isF16);
+  const bytesPerElement = dtypeBytes(input.dtype);
+  const outputSize = outChannels * height * width * bytesPerElement;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'grouped_pointwise_conv2d_output');
+  const spatial = height * width;
+
+  const weightBuffer = getBuffer(weight);
+  let biasBuffer = getBuffer(bias);
+  let tempBias = null;
+  if (!biasBuffer) {
+    const biasSize = outChannels * bytesPerElement;
+    tempBias = acquireBuffer(biasSize, undefined, 'grouped_pointwise_conv2d_bias_zero');
+    biasBuffer = tempBias;
+    const paddedSize = Math.ceil(biasSize / 4) * 4;
+    device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
+  }
+
+  await unifiedKernelWrapper(
+    'grouped_pointwise_conv2d',
+    target,
+    variant,
+    [input, weightBuffer, biasBuffer, output],
+    {
+      in_channels: inChannels,
+      out_channels: outChannels,
+      height,
+      width,
+      groups,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
+    [Math.ceil(spatial / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+  );
+
+  if (tempBias) {
+    if (recorder) {
+      recorder.trackTemporaryBuffer(tempBias);
+    } else {
+      releaseBuffer(tempBias);
+    }
+  }
+
+  return createTensor(output, input.dtype, [outChannels, height, width], 'grouped_pointwise_conv2d_output');
+}
+
+export async function runGroupedPointwiseConv2D(input, weight, bias, options = {}) {
+  return _groupedPointwiseConv2D(null, input, weight, bias, options);
+}
+
+export async function recordGroupedPointwiseConv2D(recorder, input, weight, bias, options = {}) {
+  return _groupedPointwiseConv2D(recorder, input, weight, bias, options);
+}

package/src/gpu/kernels/grouped_pointwise_conv2d.wgsl

@@ -0,0 +1,44 @@
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    in_channels: u32,
+    out_channels: u32,
+    height: u32,
+    width: u32,
+    groups: 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> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: 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 spatial = u.height * u.width;
+    let spatial_idx = gid.x;
+    let out_channel = gid.y;
+    if (spatial_idx >= spatial || out_channel >= u.out_channels) {
+        return;
+    }
+    let y = spatial_idx / u.width;
+    let x = spatial_idx - y * u.width;
+
+    let in_per_group = u.in_channels / u.groups;
+    let out_per_group = u.out_channels / u.groups;
+    let group_idx = out_channel / out_per_group;
+    let in_offset = group_idx * in_per_group;
+
+    var sum: f32 = bias[out_channel];
+    for (var i: u32 = 0u; i < in_per_group; i = i + 1u) {
+        let input_idx = ((in_offset + i) * u.height + y) * u.width + x;
+        let weight_idx = out_channel * in_per_group + i;
+        sum = sum + input[input_idx] * weight[weight_idx];
+    }
+
+    output[out_channel * spatial + spatial_idx] = sum;
+}