@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/moe.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice, getDeviceEpoch } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { WORKGROUP_SIZES, GPU_LIMITS } from './constants.js';
 import { dispatch, recordDispatch } from './dispatch.js';
@@ -61,6 +61,15 @@ function releaseUniformBuffer(execution, uniformBuffer) {
   }
 }
 
+function releaseTemporaryBuffer(execution, buffer) {
+  if (!buffer) return;
+  if (execution.recorder) {
+    execution.recorder.trackTemporaryBuffer(buffer);
+  } else {
+    releaseBuffer(buffer);
+  }
+}
+
 async function executeTopK(recorder, probs, numTokens, numExperts, topK, options = {}) {
   const execution = resolveExecution(recorder);
   const { normalize = true } = options;
@@ -79,21 +88,26 @@ async function executeTopK(recorder, probs, numTokens, numExperts, topK, options
     view.setUint32(12, normalize ? 1 : 0, true);
   });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'topk_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: probs } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-    ],
-  });
-
-  dispatchTopK(execution, pipeline, bindGroup, numTokens);
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return { indices, weights };
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'topk_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: probs } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+      ],
+    });
+    dispatchTopK(execution, pipeline, bindGroup, numTokens);
+    return { indices, weights };
+  } catch (error) {
+    releaseTemporaryBuffer(execution, indices);
+    releaseTemporaryBuffer(execution, weights);
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 export async function runTopK(probs, numTokens, numExperts, topK, options = {}) {
@@ -225,65 +239,72 @@ async function executeMoEGather(recorder, hiddenStates, expertIndices, numTokens
     view.setUint32(28, 0, true);
   });
 
-  const bindGroup = await createBindGroupWithValidation(execution.device, {
-    label: 'moe_gather_bind_group',
-    layout: explicitLayout,
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: hiddenStates.buffer } },
-      { binding: 2, resource: { buffer: expertIndices } },
-      { binding: 3, resource: { buffer: gatheredBuffer } },
-      { binding: 4, resource: { buffer: tokenCounts } },
-      { binding: 5, resource: { buffer: tokenMap } },
-    ],
-  }, `moe_gather:${dtypeLabel}`);
-
-  const countWorkgroups = Math.ceil((numTokens * topK) / WORKGROUP_SIZES.DEFAULT);
-
-  if (execution.recorder) {
-    const encoder = execution.recorder.getEncoder();
-    encoder.clearBuffer(tokenCounts);
-
-    const countPass = execution.recorder.beginComputePass('moe_gather_count');
-    countPass.setPipeline(countPipeline);
-    countPass.setBindGroup(0, bindGroup);
-    countPass.dispatchWorkgroups(countWorkgroups);
-    countPass.end();
-
-    const gatherPass = execution.recorder.beginComputePass('moe_gather_gather');
-    gatherPass.setPipeline(gatherPipeline);
-    gatherPass.setBindGroup(0, bindGroup);
-    gatherPass.dispatchWorkgroups(gatherDispatch.x, gatherDispatch.y, 1);
-    gatherPass.end();
-  } else {
-    const encoder = execution.device.createCommandEncoder({ label: 'moe_gather_encoder' });
-    encoder.clearBuffer(tokenCounts);
-
-    const countPass = encoder.beginComputePass({ label: 'moe_gather_count_pass' });
-    countPass.setPipeline(countPipeline);
-    countPass.setBindGroup(0, bindGroup);
-    countPass.dispatchWorkgroups(countWorkgroups);
-    countPass.end();
-
-    const gatherPass = encoder.beginComputePass({ label: 'moe_gather_gather_pass' });
-    gatherPass.setPipeline(gatherPipeline);
-    gatherPass.setBindGroup(0, bindGroup);
-    gatherPass.dispatchWorkgroups(gatherDispatch.x, gatherDispatch.y, 1);
-    gatherPass.end();
-
-    execution.device.queue.submit([encoder.finish()]);
+  try {
+    const bindGroup = await createBindGroupWithValidation(execution.device, {
+      label: 'moe_gather_bind_group',
+      layout: explicitLayout,
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: hiddenStates.buffer } },
+        { binding: 2, resource: { buffer: expertIndices } },
+        { binding: 3, resource: { buffer: gatheredBuffer } },
+        { binding: 4, resource: { buffer: tokenCounts } },
+        { binding: 5, resource: { buffer: tokenMap } },
+      ],
+    }, `moe_gather:${dtypeLabel}`);
+
+    const countWorkgroups = Math.ceil((numTokens * topK) / WORKGROUP_SIZES.DEFAULT);
+
+    if (execution.recorder) {
+      const encoder = execution.recorder.getEncoder();
+      encoder.clearBuffer(tokenCounts);
+
+      const countPass = execution.recorder.beginComputePass('moe_gather_count');
+      countPass.setPipeline(countPipeline);
+      countPass.setBindGroup(0, bindGroup);
+      countPass.dispatchWorkgroups(countWorkgroups);
+      countPass.end();
+
+      const gatherPass = execution.recorder.beginComputePass('moe_gather_gather');
+      gatherPass.setPipeline(gatherPipeline);
+      gatherPass.setBindGroup(0, bindGroup);
+      gatherPass.dispatchWorkgroups(gatherDispatch.x, gatherDispatch.y, 1);
+      gatherPass.end();
+    } else {
+      const encoder = execution.device.createCommandEncoder({ label: 'moe_gather_encoder' });
+      encoder.clearBuffer(tokenCounts);
+
+      const countPass = encoder.beginComputePass({ label: 'moe_gather_count_pass' });
+      countPass.setPipeline(countPipeline);
+      countPass.setBindGroup(0, bindGroup);
+      countPass.dispatchWorkgroups(countWorkgroups);
+      countPass.end();
+
+      const gatherPass = encoder.beginComputePass({ label: 'moe_gather_gather_pass' });
+      gatherPass.setPipeline(gatherPipeline);
+      gatherPass.setBindGroup(0, bindGroup);
+      gatherPass.dispatchWorkgroups(gatherDispatch.x, gatherDispatch.y, 1);
+      gatherPass.end();
+
+      execution.device.queue.submit([encoder.finish()]);
+    }
+
+    const gathered = createTensor(
+      gatheredBuffer,
+      hiddenStates.dtype,
+      [numExperts, maxTokensPerExpert, hiddenSize],
+      'moe_gathered'
+    );
+
+    return { gathered, tokenCounts, tokenMap, maxTokensPerExpert };
+  } catch (error) {
+    releaseTemporaryBuffer(execution, gatheredBuffer);
+    releaseTemporaryBuffer(execution, tokenCounts);
+    releaseTemporaryBuffer(execution, tokenMap);
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
   }
-
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  const gathered = createTensor(
-    gatheredBuffer,
-    hiddenStates.dtype,
-    [numExperts, maxTokensPerExpert, hiddenSize],
-    'moe_gathered'
-  );
-
-  return { gathered, tokenCounts, tokenMap, maxTokensPerExpert };
 }
 
 export async function runMoEGather(hiddenStates, expertIndices, numTokens, hiddenSize, numExperts, topK, options = {}) {
@@ -298,6 +319,7 @@ export async function recordMoEGather(recorder, hiddenStates, expertIndices, num
 async function executeScatterAdd(recorder, expertOutputs, indices, weights, numTokens, hiddenSize, numExperts, topK, options = {}) {
   const execution = resolveExecution(recorder);
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
 
   const pipeline = await createPipeline('scatter_add', 'default');
 
@@ -312,42 +334,49 @@ async function executeScatterAdd(recorder, expertOutputs, indices, weights, numT
     view.setUint32(12, numExperts, true);
   });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'scatter_add_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: expertOutputs.buffer } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-      { binding: 4, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT);
-
-  if (execution.recorder) {
-    execution.recorder.getEncoder().clearBuffer(outputBuf);
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'scatter_add_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: expertOutputs.buffer } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+        { binding: 4, resource: { buffer: outputBuf } },
+      ],
+    });
 
-    const pass = execution.recorder.beginComputePass('scatter_add');
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-  } else {
-    const encoder = execution.device.createCommandEncoder({ label: 'scatter_add_encoder' });
-    encoder.clearBuffer(outputBuf);
-    const pass = encoder.beginComputePass({ label: 'scatter_add_pass' });
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-    execution.device.queue.submit([encoder.finish()]);
+    const workgroups = Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT);
+
+    if (execution.recorder) {
+      execution.recorder.getEncoder().clearBuffer(outputBuf);
+
+      const pass = execution.recorder.beginComputePass('scatter_add');
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+    } else {
+      const encoder = execution.device.createCommandEncoder({ label: 'scatter_add_encoder' });
+      encoder.clearBuffer(outputBuf);
+      const pass = encoder.beginComputePass({ label: 'scatter_add_pass' });
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+      execution.device.queue.submit([encoder.finish()]);
+    }
+
+    return createTensor(outputBuf, expertOutputs.dtype, [numTokens, hiddenSize], 'scatter_add_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseTemporaryBuffer(execution, outputBuf);
+    }
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
   }
-
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return createTensor(outputBuf, expertOutputs.dtype, [numTokens, hiddenSize], 'scatter_add_output');
 }
 
 export async function runScatterAdd(expertOutputs, indices, weights, numTokens, hiddenSize, numExperts, topK, options = {}) {
@@ -362,6 +391,7 @@ export async function recordScatterAdd(recorder, expertOutputs, indices, weights
 async function executeMoEBuildTokenOffsets(recorder, tokenCounts, tokenMap, numTokens, numExperts, topK, maxTokensPerExpert, options = {}) {
   const execution = resolveExecution(recorder);
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
 
   const explicitLayout = getMoEOffsetsBindGroupLayout(execution.device);
   const pipeline = await createPipeline('moe_offsets', 'default', explicitLayout);
@@ -380,38 +410,46 @@ async function executeMoEBuildTokenOffsets(recorder, tokenCounts, tokenMap, numT
     view.setUint32(28, 0, true);
   });
 
-  const bindGroup = await createBindGroupWithValidation(execution.device, {
-    label: 'moe_offsets_bind_group',
-    layout: explicitLayout,
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: tokenCounts } },
-      { binding: 2, resource: { buffer: tokenMap } },
-      { binding: 3, resource: { buffer: tokenOffsets } },
-    ],
-  }, 'moe_offsets');
-
-  const totalSlots = numExperts * maxTokensPerExpert;
-  const workgroups = Math.ceil(totalSlots / WORKGROUP_SIZES.DEFAULT);
-
-  if (execution.recorder) {
-    const pass = execution.recorder.beginComputePass('moe_offsets');
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-  } else {
-    const encoder = execution.device.createCommandEncoder({ label: 'moe_offsets_encoder' });
-    const pass = encoder.beginComputePass({ label: 'moe_offsets_pass' });
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-    execution.device.queue.submit([encoder.finish()]);
+  try {
+    const bindGroup = await createBindGroupWithValidation(execution.device, {
+      label: 'moe_offsets_bind_group',
+      layout: explicitLayout,
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: tokenCounts } },
+        { binding: 2, resource: { buffer: tokenMap } },
+        { binding: 3, resource: { buffer: tokenOffsets } },
+      ],
+    }, 'moe_offsets');
+
+    const totalSlots = numExperts * maxTokensPerExpert;
+    const workgroups = Math.ceil(totalSlots / WORKGROUP_SIZES.DEFAULT);
+
+    if (execution.recorder) {
+      const pass = execution.recorder.beginComputePass('moe_offsets');
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+    } else {
+      const encoder = execution.device.createCommandEncoder({ label: 'moe_offsets_encoder' });
+      const pass = encoder.beginComputePass({ label: 'moe_offsets_pass' });
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+      execution.device.queue.submit([encoder.finish()]);
+    }
+
+    return tokenOffsets;
+  } catch (error) {
+    if (ownsOutput) {
+      releaseTemporaryBuffer(execution, tokenOffsets);
+    }
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
   }
-
-  releaseUniformBuffer(execution, uniformBuffer);
-  return tokenOffsets;
 }
 
 export async function runMoEBuildTokenOffsets(tokenCounts, tokenMap, numTokens, numExperts, topK, maxTokensPerExpert, options = {}) {
@@ -425,6 +463,7 @@ export async function recordMoEBuildTokenOffsets(recorder, tokenCounts, tokenMap
 async function executeScatterAddDynamic(recorder, expertOutputs, indices, weights, tokenOffsets, numTokens, hiddenSize, topK, options = {}) {
   const execution = resolveExecution(recorder);
   const { outputBuffer = null, weightsDtype = 'f32' } = options;
+  const ownsOutput = outputBuffer == null;
 
   if (weightsDtype === 'f16' && expertOutputs.dtype !== 'f16') {
     throw new Error('ScatterAddDynamic f16 weights require f16 expert outputs');
@@ -447,43 +486,50 @@ async function executeScatterAddDynamic(recorder, expertOutputs, indices, weight
     view.setUint32(8, topK, true);
   });
 
-  const bindGroup = await createBindGroupWithValidation(execution.device, {
-    label: 'scatter_add_dynamic_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: expertOutputs.buffer } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-      { binding: 4, resource: { buffer: tokenOffsets } },
-      { binding: 5, resource: { buffer: outputBuf } },
-    ],
-  }, `scatter_add_dynamic:${variant}`);
-
-  const workgroups = Math.ceil((numTokens * topK * hiddenSize) / WORKGROUP_SIZES.DEFAULT);
-
-  if (execution.recorder) {
-    execution.recorder.getEncoder().clearBuffer(outputBuf);
-
-    const pass = execution.recorder.beginComputePass('scatter_add_dynamic');
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-  } else {
-    const encoder = execution.device.createCommandEncoder({ label: 'scatter_add_dynamic_encoder' });
-    encoder.clearBuffer(outputBuf);
-    const pass = encoder.beginComputePass({ label: 'scatter_add_dynamic_pass' });
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    pass.dispatchWorkgroups(workgroups);
-    pass.end();
-    execution.device.queue.submit([encoder.finish()]);
+  try {
+    const bindGroup = await createBindGroupWithValidation(execution.device, {
+      label: 'scatter_add_dynamic_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: expertOutputs.buffer } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+        { binding: 4, resource: { buffer: tokenOffsets } },
+        { binding: 5, resource: { buffer: outputBuf } },
+      ],
+    }, `scatter_add_dynamic:${variant}`);
+
+    const workgroups = Math.ceil((numTokens * topK * hiddenSize) / WORKGROUP_SIZES.DEFAULT);
+
+    if (execution.recorder) {
+      execution.recorder.getEncoder().clearBuffer(outputBuf);
+
+      const pass = execution.recorder.beginComputePass('scatter_add_dynamic');
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+    } else {
+      const encoder = execution.device.createCommandEncoder({ label: 'scatter_add_dynamic_encoder' });
+      encoder.clearBuffer(outputBuf);
+      const pass = encoder.beginComputePass({ label: 'scatter_add_dynamic_pass' });
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      pass.dispatchWorkgroups(workgroups);
+      pass.end();
+      execution.device.queue.submit([encoder.finish()]);
+    }
+
+    return createTensor(outputBuf, expertOutputs.dtype, [numTokens, hiddenSize], 'scatter_add_dynamic_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseTemporaryBuffer(execution, outputBuf);
+    }
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
   }
-
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return createTensor(outputBuf, expertOutputs.dtype, [numTokens, hiddenSize], 'scatter_add_dynamic_output');
 }
 
 export async function runScatterAddDynamic(expertOutputs, indices, weights, tokenOffsets, numTokens, hiddenSize, topK, options = {}) {

package/src/gpu/kernels/pixel_shuffle.js

@@ -1,4 +1,4 @@
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -25,19 +25,27 @@ async function _pixelShuffle(target, input, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = outChannels * outHeight * outWidth * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'pixel_shuffle_output');
-
-  await unifiedKernelWrapper(
-    'pixel_shuffle', target, variant,
-    [input, output],
-    {
-      out_channels: outChannels, out_height: outHeight, out_width: outWidth,
-      grid_width: gridWidth, grid_height: gridHeight, patch_size: patchSize,
-      patch_channels: inferredPatchChannels, _pad0: 0,
-    },
-    Math.ceil((outChannels * outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  return createTensor(output, input.dtype, [outChannels, outHeight, outWidth], 'pixel_shuffle_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'pixel_shuffle', target, variant,
+      [input, output],
+      {
+        out_channels: outChannels, out_height: outHeight, out_width: outWidth,
+        grid_width: gridWidth, grid_height: gridHeight, patch_size: patchSize,
+        patch_channels: inferredPatchChannels, _pad0: 0,
+      },
+      [Math.ceil((outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
+
+    return createTensor(output, input.dtype, [outChannels, outHeight, outWidth], 'pixel_shuffle_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runPixelShuffle(input, options = {}) {

package/src/gpu/kernels/pixel_shuffle.wgsl

@@ -19,17 +19,16 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let spatial_size = u.out_height * u.out_width;
-    let total = u.out_channels * spatial_size;
-    if (idx >= total) {
+    let spatial = gid.x;
+    let c = gid.y;
+    if (c >= u.out_channels || spatial >= spatial_size) {
         return;
     }
 
-    let c = idx / spatial_size;
-    let spatial = idx - c * spatial_size;
     let y = spatial / u.out_width;
     let x = spatial - y * u.out_width;
+    let idx = c * spatial_size + spatial;
 
     let grid_y = y / u.patch_size;
     let grid_x = x / u.patch_size;

package/src/gpu/kernels/pixel_shuffle_f16.wgsl

@@ -22,17 +22,16 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let spatial_size = u.out_height * u.out_width;
-    let total = u.out_channels * spatial_size;
-    if (idx >= total) {
+    let spatial = gid.x;
+    let c = gid.y;
+    if (c >= u.out_channels || spatial >= spatial_size) {
         return;
     }
 
-    let c = idx / spatial_size;
-    let spatial = idx - c * spatial_size;
     let y = spatial / u.out_width;
     let x = spatial - y * u.out_width;
+    let idx = c * spatial_size + spatial;
 
     let grid_y = y / u.patch_size;
     let grid_x = x / u.patch_size;

package/src/gpu/kernels/relu.js

@@ -1,4 +1,4 @@
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -18,22 +18,43 @@ function resolveCount(input, countOverride) {
   return Math.floor(input.buffer.size / dtypeBytes(input.dtype));
 }
 
+function planReluDispatch(target, size) {
+  const device = target?.device;
+  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
+    ? device.limits.maxComputeWorkgroupsPerDimension
+    : 65535;
+  const dispatchStride = Math.min(size, maxPerDim * WORKGROUP_SIZES.DEFAULT);
+  return {
+    dispatchStride,
+    workgroups: [Math.ceil(dispatchStride / WORKGROUP_SIZES.DEFAULT), 1, 1],
+  };
+}
+
 async function _relu(target, input, options = {}) {
   const { count = null, outputBuffer = null } = options;
   const size = resolveCount(input, count);
   const variant = selectReluVariant(input.dtype);
   const output = outputBuffer || acquireBuffer(size * dtypeBytes(input.dtype), undefined, 'relu_output');
+  const ownedOutput = outputBuffer ? null : output;
+  const dispatchPlan = planReluDispatch(target, size);
 
-  await unifiedKernelWrapper(
-    'relu',
-    target,
-    variant,
-    [input, output],
-    { size, _pad0: 0, _pad1: 0, _pad2: 0 },
-    Math.ceil(size / WORKGROUP_SIZES.DEFAULT)
-  );
+  try {
+    await unifiedKernelWrapper(
+      'relu',
+      target,
+      variant,
+      [input, output],
+      { size, _pad0: dispatchPlan.dispatchStride, _pad1: 0, _pad2: 0 },
+      dispatchPlan.workgroups
+    );
 
-  return createTensor(output, input.dtype, [...input.shape], 'relu_output');
+    return createTensor(output, input.dtype, [...input.shape], 'relu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runReLU(input, options = {}) {

package/src/gpu/kernels/relu.wgsl

@@ -13,7 +13,8 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+    let dispatch_stride = max(u._pad0, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }

package/src/gpu/kernels/relu_f16.wgsl

@@ -15,7 +15,8 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+    let dispatch_stride = max(u._pad0, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }