@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/silu.js

@@ -1,13 +1,26 @@
 
 
 import { getDevice } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { WORKGROUP_SIZES } from './constants.js';
 import { dispatch, recordDispatch } from './dispatch.js';
 import { getPipelineFast, createUniformBufferWithView } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
 
+function destroyAfterSubmit(device, buffer) {
+  if (!buffer) {
+    return;
+  }
+  device.queue.onSubmittedWorkDone()
+    .then(() => {
+      buffer.destroy();
+    })
+    .catch(() => {
+      buffer.destroy();
+    });
+}
+
 function canUseF16(input) {
   return input.dtype === 'f16';
 }
@@ -47,6 +60,24 @@ function createSiLUBindGroupEntries(uniformBuffer, input, output, gate) {
   ];
 }
 
+function cleanupRunResources(uniformBuffer, ownedOutput) {
+  if (ownedOutput) {
+    releaseBuffer(ownedOutput);
+  }
+}
+
+function planSiLUDispatch(device, size, useVec4) {
+  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
+    ? device.limits.maxComputeWorkgroupsPerDimension
+    : 65535;
+  const laneWidth = useVec4 ? 4 : 1;
+  const chunkSize = maxPerDim * WORKGROUP_SIZES.DEFAULT * laneWidth;
+  const dispatchStride = Math.min(size, chunkSize);
+  const x = Math.min(maxPerDim, Math.ceil(dispatchStride / (WORKGROUP_SIZES.DEFAULT * laneWidth)));
+  const y = Math.max(1, Math.ceil(size / chunkSize));
+  return { dispatchStride, workgroups: [x, y, 1] };
+}
+
 
 export async function runSiLU(
   input,
@@ -60,6 +91,7 @@ export async function runSiLU(
     useVec4 = false,
     swigluLimit,
     gateActivation = 'silu',
+    inputActivation = 'silu',
   } = options;
   const resolvedSwigluLimit = resolveSwigluLimit(swigluLimit, 'SiLU');
 
@@ -74,14 +106,18 @@ export async function runSiLU(
     useSplit: false,
     useRowsplit: false,
   });
-  const constants = gate && gateActivation === 'sigmoid'
-    ? { ...(overrides || {}), GATE_USE_SIGMOID: true }
-    : overrides;
+  const constants = {
+    ...(overrides || {}),
+    ...(gate && gateActivation === 'sigmoid' ? { GATE_USE_SIGMOID: true } : {}),
+    ...(inputActivation === 'identity' ? { INPUT_USE_IDENTITY: true } : {}),
+  };
   const pipeline = await getPipelineFast('silu', variant, null, constants);
 
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
+  const dispatchPlan = planSiLUDispatch(device, inferredSize, useVec4);
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -89,7 +125,7 @@ export async function runSiLU(
     16,
     (view) => {
       view.setUint32(0, inferredSize, true);
-      view.setUint32(4, 0, true);
+      view.setUint32(4, dispatchPlan.dispatchStride, true);
       view.setFloat32(8, gate ? resolvedSwigluLimit : 0, true);
       view.setFloat32(12, 0, true);
     },
@@ -100,18 +136,21 @@ export async function runSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  const workgroups = Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'silu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -133,6 +172,7 @@ export async function runSwiGLURowsplitBias(
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'swiglu_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -149,23 +189,27 @@ export async function runSwiGLURowsplitBias(
   );
 
   // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'swiglu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: bias.buffer } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'swiglu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: bias.buffer } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+    dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
+    return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -187,6 +231,7 @@ export async function runSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -203,24 +248,28 @@ export async function runSiLURowSplit(
   );
 
   // Bind group: provide a dummy gate buffer to satisfy the fixed layout
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 
@@ -243,6 +292,7 @@ export async function recordSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -257,22 +307,28 @@ export async function recordSiLURowSplit(
     recorder
   );
 
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 
@@ -288,6 +344,7 @@ export async function recordSiLU(
     outputBuffer = null,
     swigluLimit,
     gateActivation = 'silu',
+    inputActivation = 'silu',
   } = options;
   const resolvedSwigluLimit = resolveSwigluLimit(swigluLimit, 'SiLU');
 
@@ -302,14 +359,18 @@ export async function recordSiLU(
     useSplit: false,
     useRowsplit: false,
   });
-  const constants = gate && gateActivation === 'sigmoid'
-    ? { ...(overrides || {}), GATE_USE_SIGMOID: true }
-    : overrides;
+  const constants = {
+    ...(overrides || {}),
+    ...(gate && gateActivation === 'sigmoid' ? { GATE_USE_SIGMOID: true } : {}),
+    ...(inputActivation === 'identity' ? { INPUT_USE_IDENTITY: true } : {}),
+  };
   const pipeline = await getPipelineFast('silu', variant, null, constants);
 
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
+  const dispatchPlan = planSiLUDispatch(device, inferredSize, false);
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -317,7 +378,7 @@ export async function recordSiLU(
     16,
     (view) => {
       view.setUint32(0, inferredSize, true);
-      view.setUint32(4, 0, true);
+      view.setUint32(4, dispatchPlan.dispatchStride, true);
       view.setFloat32(8, gate ? resolvedSwigluLimit : 0, true);
       view.setFloat32(12, 0, true);
     },
@@ -327,14 +388,19 @@ export async function recordSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  const workgroups = Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT);
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu');
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/silu.wgsl

@@ -10,13 +10,14 @@
 override WORKGROUP_SIZE: u32 = 256u;
 override HAS_GATE: bool = false;
 override GATE_USE_SIGMOID: bool = false;
+override INPUT_USE_IDENTITY: bool = false;
 override USE_SPLIT: bool = false;
 override USE_VEC4: bool = false;
 override USE_ROWSPLIT: bool = false;
 
 struct Uniforms {
     size: u32,          // Total output elements
-    rowsplit_dim: u32,  // Dim for rowsplit variants (0 when unused)
+    rowsplit_dim: u32,  // Row-split dim or dispatch stride for non-row-split variants
     clamp_max: f32,     // SwiGLU clamp (0 = disabled)
     _pad1: f32,
 }
@@ -35,6 +36,10 @@ fn silu(x: f32) -> f32 {
     return x * sigmoid(x);
 }
 
+fn apply_input_activation(x: f32) -> f32 {
+    return select(silu(x), x, INPUT_USE_IDENTITY);
+}
+
 fn clamp_swiglu(x: f32) -> f32 {
     if (u.clamp_max <= 0.0) {
         return x;
@@ -46,8 +51,9 @@ fn clamp_swiglu(x: f32) -> f32 {
 fn main(
     @builtin(global_invocation_id) global_id: vec3<u32>
 ) {
+    let dispatch_stride = max(u.rowsplit_dim, 1u);
     if (USE_VEC4) {
-        let base_idx = global_id.x * 4u;
+        let base_idx = global_id.y * dispatch_stride + global_id.x * 4u;
         if (base_idx >= u.size) {
             return;
         }
@@ -55,12 +61,12 @@ fn main(
         let remaining = min(4u, u.size - base_idx);
         for (var i: u32 = 0u; i < remaining; i = i + 1u) {
             let x = input[base_idx + i];
-            output[base_idx + i] = silu(x);
+            output[base_idx + i] = apply_input_activation(x);
         }
         return;
     }
 
-    let idx = global_id.x;
+    let idx = global_id.y * dispatch_stride + global_id.x;
     if (idx >= u.size) {
         return;
     }
@@ -70,12 +76,16 @@ fn main(
             return;
         }
         let dim = u.rowsplit_dim;
-        let token_idx = idx / dim;
-        let dim_idx = idx % dim;
+        let num_tokens = u.size / dim;
+        let token_idx = global_id.y;
+        let dim_idx = global_id.x;
+        if (token_idx >= num_tokens || dim_idx >= dim) {
+            return;
+        }
         let row_base = token_idx * dim * 2u;
         let g = input[row_base + dim_idx];
         let up = input[row_base + dim + dim_idx];
-        output[idx] = clamp_swiglu(silu(g) * up);
+        output[token_idx * dim + dim_idx] = clamp_swiglu(silu(g) * up);
         return;
     }
 
@@ -83,7 +93,7 @@ fn main(
         let up = input[idx];
         let g = gate[idx];
         let gateAct = select(silu(g), sigmoid(g), GATE_USE_SIGMOID);
-        output[idx] = clamp_swiglu(gateAct * up);
+        output[idx] = clamp_swiglu(gateAct * apply_input_activation(up));
         return;
     }
 
@@ -95,5 +105,5 @@ fn main(
     }
 
     let x = input[idx];
-    output[idx] = silu(x);
+    output[idx] = apply_input_activation(x);
 }

package/src/gpu/kernels/silu_f16.wgsl

@@ -9,13 +9,14 @@ enable f16;
 override WORKGROUP_SIZE: u32 = 256u;
 override HAS_GATE: bool = false;
 override GATE_USE_SIGMOID: bool = false;
+override INPUT_USE_IDENTITY: bool = false;
 override USE_SPLIT: bool = false;
 override USE_VEC4: bool = false;
 override USE_ROWSPLIT: bool = false;
 
 struct Uniforms {
     size: u32,          // Total output elements
-    rowsplit_dim: u32,  // Dim for rowsplit variants (0 when unused)
+    rowsplit_dim: u32,  // Row-split dim or dispatch stride for non-row-split variants
     clamp_max: f32,     // SwiGLU clamp (0 = disabled)
     _pad1: f32,
 }
@@ -34,6 +35,10 @@ fn silu(x: f32) -> f32 {
     return x * sigmoid(x);
 }
 
+fn apply_input_activation(x: f32) -> f32 {
+    return select(silu(x), x, INPUT_USE_IDENTITY);
+}
+
 fn clamp_swiglu(x: f32) -> f32 {
     if (u.clamp_max <= 0.0) {
         return x;
@@ -45,8 +50,9 @@ fn clamp_swiglu(x: f32) -> f32 {
 fn main(
     @builtin(global_invocation_id) global_id: vec3<u32>
 ) {
+    let dispatch_stride = max(u.rowsplit_dim, 1u);
     if (USE_VEC4) {
-        let base_idx = global_id.x * 4u;
+        let base_idx = global_id.y * dispatch_stride + global_id.x * 4u;
         if (base_idx >= u.size) {
             return;
         }
@@ -54,12 +60,12 @@ fn main(
         let remaining = min(4u, u.size - base_idx);
         for (var i: u32 = 0u; i < remaining; i = i + 1u) {
             let x = f32(input[base_idx + i]);
-            output[base_idx + i] = f16(silu(x));
+            output[base_idx + i] = f16(apply_input_activation(x));
         }
         return;
     }
 
-    let idx = global_id.x;
+    let idx = global_id.y * dispatch_stride + global_id.x;
     if (idx >= u.size) {
         return;
     }
@@ -69,12 +75,16 @@ fn main(
             return;
         }
         let dim = u.rowsplit_dim;
-        let token_idx = idx / dim;
-        let dim_idx = idx % dim;
+        let num_tokens = u.size / dim;
+        let token_idx = global_id.y;
+        let dim_idx = global_id.x;
+        if (token_idx >= num_tokens || dim_idx >= dim) {
+            return;
+        }
         let row_base = token_idx * dim * 2u;
         let g = f32(input[row_base + dim_idx]);
         let up = f32(input[row_base + dim + dim_idx]);
-        output[idx] = f16(clamp_swiglu(silu(g) * up));
+        output[token_idx * dim + dim_idx] = f16(clamp_swiglu(silu(g) * up));
         return;
     }
 
@@ -82,7 +92,7 @@ fn main(
         let up = f32(input[idx]);
         let g = f32(gate[idx]);
         let gateAct = select(silu(g), sigmoid(g), GATE_USE_SIGMOID);
-        output[idx] = f16(clamp_swiglu(gateAct * up));
+        output[idx] = f16(clamp_swiglu(gateAct * apply_input_activation(up)));
         return;
     }
 
@@ -94,5 +104,5 @@ fn main(
     }
 
     let x = f32(input[idx]);
-    output[idx] = f16(silu(x));
+    output[idx] = f16(apply_input_activation(x));
 }

package/src/gpu/kernels/softmax.js

@@ -1,6 +1,6 @@
 
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { createPipeline, createUniformBufferWithView, createBindGroupWithValidation } from './utils.js';
@@ -20,23 +20,34 @@ function selectSoftmaxVariant(innerSize) {
 
 async function _softmax(target, input, axis, options = {}) {
   const { batchSize = 1, size, seqLen, temperature = 1.0, outputBuffer = null } = options;
+  if (input.dtype !== 'f32') {
+    throw new Error(`Softmax requires f32 input, got ${input.dtype}.`);
+  }
 
-  const bytesPerElement = input.dtype === 'f16' ? 2 : 4;
+  const bytesPerElement = 4;
   const inferredSize = size || seqLen || (input.buffer.size / (batchSize * bytesPerElement));
   const variant = selectSoftmaxVariant(inferredSize);
   trace.kernels(`Softmax: size=${inferredSize}, variant=${variant}`);
 
   const outputSize = batchSize * inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'softmax_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'softmax', target, variant,
+      [input, output],
+      { inner_size: inferredSize, outer_size: batchSize, temperature },
+      batchSize
+    );
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
-  await unifiedKernelWrapper(
-    'softmax', target, variant,
-    [input, output],
-    { inner_size: inferredSize, outer_size: batchSize, temperature },
-    batchSize
-  );
-
-  return createTensor(output, input.dtype, [batchSize, inferredSize], 'softmax_output');
+  return createTensor(output, 'f32', [batchSize, inferredSize], 'softmax_output');
 }
 
 export async function runSoftmax(input, axis, options = {}) {
@@ -76,6 +87,7 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
 
   const indices = acquireBuffer(indicesSize, undefined, 'softmax_topk_indices');
   const weights = acquireBuffer(weightsSize, undefined, 'softmax_topk_weights');
+  let completed = false;
 
   const uniformBuffer = createUniformBufferWithView(
     'softmax_topk_uniforms', 16,
@@ -88,19 +100,26 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
     null, device
   );
 
-  const bindGroup = await createBindGroupWithValidation(device, {
-    label: 'softmax_topk_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: logits } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-    ],
-  }, `topk:${variant}`);
-
-  dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
-  uniformBuffer.destroy();
-
-  return { indices, weights };
+  try {
+    const bindGroup = await createBindGroupWithValidation(device, {
+      label: 'softmax_topk_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: logits } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+      ],
+    }, `topk:${variant}`);
+
+    dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
+    completed = true;
+    return { indices, weights };
+  } finally {
+    uniformBuffer.destroy();
+    if (!completed) {
+      releaseBuffer(indices);
+      releaseBuffer(weights);
+    }
+  }
 }