@simulatte/doppler 0.1.4 → 0.1.5

package/src/converter/parsers/diffusion.js

@@ -4,6 +4,13 @@ const SD3_LAYOUT = {
   id: 'sd3',
   requiredComponents: ['transformer', 'text_encoder', 'text_encoder_2', 'text_encoder_3', 'vae', 'scheduler'],
   weightedComponents: ['transformer', 'text_encoder', 'text_encoder_2', 'text_encoder_3', 'vae'],
+  matches(modelIndex, components) {
+    return (
+      components.has('text_encoder_2') &&
+      components.has('text_encoder_3') &&
+      getComponentClassName(modelIndex?.transformer) === 'SD3Transformer2DModel'
+    );
+  },
   tokenizerSpecs: [
     {
       modelIndexKey: 'tokenizer',
@@ -66,6 +73,10 @@ const FLUX_LAYOUT = {
   id: 'flux',
   requiredComponents: ['transformer', 'text_encoder', 'vae', 'scheduler'],
   weightedComponents: ['transformer', 'text_encoder', 'vae'],
+  matches(modelIndex) {
+    const transformerClass = getComponentClassName(modelIndex?.transformer);
+    return typeof transformerClass === 'string' && /^Flux/i.test(transformerClass);
+  },
   tokenizerSpecs: [
     {
       modelIndexKey: 'tokenizer',
@@ -91,7 +102,39 @@ const FLUX_LAYOUT = {
   ],
 };
 
-const LAYOUTS = [SD3_LAYOUT, FLUX_LAYOUT];
+const SANA_LAYOUT = {
+  id: 'sana',
+  requiredComponents: ['transformer', 'text_encoder', 'tokenizer', 'vae', 'scheduler'],
+  weightedComponents: ['transformer', 'text_encoder', 'vae'],
+  matches(modelIndex) {
+    return (
+      getComponentClassName(modelIndex?.transformer) === 'SanaTransformer2DModel' &&
+      getComponentClassName(modelIndex?.text_encoder) === 'Gemma2Model'
+    );
+  },
+  tokenizerSpecs: [
+    {
+      modelIndexKey: 'tokenizer',
+      componentId: 'text_encoder',
+      type: 'bundled',
+      assets: [
+        { suffix: 'tokenizer/tokenizer.json', targetName: 'tokenizer_tokenizer.json', kind: 'text', required: true },
+        { suffix: 'tokenizer/tokenizer_config.json', targetName: 'tokenizer_config.json', kind: 'text', required: false },
+        { suffix: 'tokenizer/special_tokens_map.json', targetName: 'tokenizer_special_tokens_map.json', kind: 'text', required: false },
+        { suffix: 'tokenizer/tokenizer.model', targetName: 'tokenizer_tokenizer.model', kind: 'binary', required: false },
+      ],
+      config: {
+        type: 'bundled',
+        tokenizerFile: 'tokenizer_tokenizer.json',
+        configFile: 'tokenizer_config.json',
+        specialTokensFile: 'tokenizer_special_tokens_map.json',
+        sentencePieceFile: 'tokenizer_tokenizer.model',
+      },
+    },
+  ],
+};
+
+const LAYOUTS = [SD3_LAYOUT, FLUX_LAYOUT, SANA_LAYOUT];
 
 function toAbortError(message = 'Cancelled') {
   if (typeof DOMException === 'function') {
@@ -112,12 +155,26 @@ function listModelComponents(modelIndex) {
   return Object.keys(modelIndex || {}).filter((key) => !key.startsWith('_'));
 }
 
+function getComponentClassName(componentEntry) {
+  if (Array.isArray(componentEntry) && componentEntry.length >= 2 && typeof componentEntry[1] === 'string') {
+    return componentEntry[1];
+  }
+  if (componentEntry && typeof componentEntry === 'object' && typeof componentEntry._class_name === 'string') {
+    return componentEntry._class_name;
+  }
+  return null;
+}
+
 export function detectDiffusionLayout(modelIndex) {
   const components = new Set(listModelComponents(modelIndex));
   for (const layout of LAYOUTS) {
-    if (layout.requiredComponents.every((component) => components.has(component))) {
-      return layout;
+    if (!layout.requiredComponents.every((component) => components.has(component))) {
+      continue;
     }
+    if (typeof layout.matches === 'function' && !layout.matches(modelIndex, components)) {
+      continue;
+    }
+    return layout;
   }
   const listed = [...components].sort().join(', ') || '(none)';
   const expected = LAYOUTS
@@ -199,6 +256,9 @@ export async function parseDiffusionModel(adapter) {
   const tensors = [];
 
   for (const componentId of layout.requiredComponents) {
+    if (componentId === 'tokenizer') {
+      continue;
+    }
     const configSuffix = defaultConfigPath(componentId);
     const config = await readJson(configSuffix, `${componentId} config`);
     if (componentId === 'transformer' && config && !config.weight_format) {

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,98 @@
+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 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((channels * outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT)
+  );
+
+  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,58 @@
+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 idx = gid.x;
+    let out_spatial = u.out_height * u.out_width;
+    let out_size = u.channels * out_spatial;
+    if (idx >= out_size) {
+        return;
+    }
+
+    let channel = idx / out_spatial;
+    let rem = idx - channel * out_spatial;
+    let out_y = rem / u.out_width;
+    let out_x = rem - 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[idx] = sum;
+}

package/src/gpu/kernels/depthwise_conv2d_f16.wgsl

@@ -0,0 +1,62 @@
+// 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 idx = gid.x;
+    let out_spatial = u.out_height * u.out_width;
+    let out_size = u.channels * out_spatial;
+    if (idx >= out_size) {
+        return;
+    }
+
+    let channel = idx / out_spatial;
+    let rem = idx - channel * out_spatial;
+    let out_y = rem / u.out_width;
+    let out_x = rem - 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[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,92 @@
+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 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((outChannels * height * width) / WORKGROUP_SIZES.DEFAULT)
+  );
+
+  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,47 @@
+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 idx = gid.x;
+    let spatial = u.height * u.width;
+    let out_size = u.out_channels * spatial;
+    if (idx >= out_size) {
+        return;
+    }
+
+    let out_channel = idx / spatial;
+    let rem = idx - out_channel * spatial;
+    let y = rem / u.width;
+    let x = rem - 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[idx] = sum;
+}

package/src/gpu/kernels/grouped_pointwise_conv2d_f16.wgsl

@@ -0,0 +1,51 @@
+// Grouped Pointwise Conv2D Kernel (NCHW, f16)
+
+enable f16;
+
+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<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 idx = gid.x;
+    let spatial = u.height * u.width;
+    let out_size = u.out_channels * spatial;
+    if (idx >= out_size) {
+        return;
+    }
+
+    let out_channel = idx / spatial;
+    let rem = idx - out_channel * spatial;
+    let y = rem / u.width;
+    let x = rem - 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 = 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 + f32(input[input_idx]) * f32(weight[weight_idx]);
+    }
+
+    output[idx] = f16(sum);
+}

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

@@ -174,6 +174,18 @@ export {
   type Conv2DOptions,
 } from './conv2d.js';
 
+export {
+  runDepthwiseConv2D,
+  recordDepthwiseConv2D,
+  type DepthwiseConv2DOptions,
+} from './depthwise_conv2d.js';
+
+export {
+  runGroupedPointwiseConv2D,
+  recordGroupedPointwiseConv2D,
+  type GroupedPointwiseConv2DOptions,
+} from './grouped_pointwise_conv2d.js';
+
 // Gather (Embedding Lookup)
 export {
   runGather,
@@ -250,6 +262,24 @@ export {
   type SampleResult,
 } from './sample.js';
 
+export {
+  runSanaLinearAttention,
+  recordSanaLinearAttention,
+  type SanaLinearAttentionOptions,
+} from './sana_linear_attention.js';
+
+export {
+  runRepeatChannels,
+  recordRepeatChannels,
+  type RepeatChannelsOptions,
+} from './repeat_channels.js';
+
+export {
+  runReLU,
+  recordReLU,
+  type ReLUOptions,
+} from './relu.js';
+
 // Fused FFN (Tier 2 P0)
 export {
   runFusedFFN,

package/src/gpu/kernels/index.js

@@ -139,6 +139,16 @@ export {
   recordConv2D,
 } from './conv2d.js';
 
+export {
+  runDepthwiseConv2D,
+  recordDepthwiseConv2D,
+} from './depthwise_conv2d.js';
+
+export {
+  runGroupedPointwiseConv2D,
+  recordGroupedPointwiseConv2D,
+} from './grouped_pointwise_conv2d.js';
+
 // Gather (Embedding Lookup)
 export {
   runGather,
@@ -205,6 +215,21 @@ export {
   isGPUSamplingAvailable,
 } from './sample.js';
 
+export {
+  runSanaLinearAttention,
+  recordSanaLinearAttention,
+} from './sana_linear_attention.js';
+
+export {
+  runRepeatChannels,
+  recordRepeatChannels,
+} from './repeat_channels.js';
+
+export {
+  runReLU,
+  recordReLU,
+} from './relu.js';
+
 // Fused FFN (Tier 2 P0)
 export {
   runFusedFFN,

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

@@ -0,0 +1,18 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+
+export interface ReLUOptions extends OutputBufferOptions {
+  count?: number | null;
+}
+
+export declare function runReLU(
+  input: Tensor,
+  options?: ReLUOptions
+): Promise<Tensor>;
+
+export declare function recordReLU(
+  recorder: CommandRecorder,
+  input: Tensor,
+  options?: ReLUOptions
+): Promise<Tensor>;