@simulatte/doppler 0.1.0

package/src/gpu/kernels/cast.js

@@ -0,0 +1,422 @@
+
+
+import { getDevice } from '../device.js';
+import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { createTensor } from '../tensor.js';
+import { dispatch, recordDispatch } from './dispatch.js';
+import { createPipeline, createUniformBufferWithView } from './utils.js';
+import { GPU_LIMITS, WORKGROUP_SIZES } from './constants.js';
+import { trace } from '../../debug/index.js';
+import { DTYPE_SIZES } from '../../config/schema/index.js';
+
+// =============================================================================
+// Dispatch Helpers
+// =============================================================================
+
+
+function calculate2DDispatch(workgroups) {
+  const maxWorkgroupsPerDim = GPU_LIMITS.MAX_WORKGROUPS;
+  return workgroups <= maxWorkgroupsPerDim
+    ? [workgroups, 1, 1]
+    : [maxWorkgroupsPerDim, Math.ceil(workgroups / maxWorkgroupsPerDim), 1];
+}
+
+
+function lcm(a, b) {
+  
+  const gcd = (x, y) => {
+    let a0 = x;
+    let b0 = y;
+    while (b0 !== 0) {
+      const t = b0;
+      b0 = a0 % b0;
+      a0 = t;
+    }
+    return a0;
+  };
+  return (a / gcd(a, b)) * b;
+}
+
+
+export async function castF32ToF16(
+  input,
+  options = {}
+) {
+  const device = getDevice();
+  const { outputBuffer = null } = options;
+  const numElements = input.shape.reduce((a, b) => a * b, 1);
+
+  const pipeline = await createPipeline('cast', 'f32_to_f16');
+
+  const outputSize = numElements * DTYPE_SIZES.f16;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f32_to_f16_output');
+
+  const uniformBuffer = createUniformBufferWithView(
+    'cast_f32_to_f16_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+    },
+    null,
+    device
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'cast_f32_to_f16_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+
+  // Use 2D dispatch for large tensors (like embeddings with 300M+ elements)
+  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+
+  uniformBuffer.destroy();
+
+  return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+}
+
+
+export async function castF16ToF32(
+  input,
+  options = {}
+) {
+  const device = getDevice();
+  const { outputBuffer = null } = options;
+  const numElements = input.shape.reduce((a, b) => a * b, 1);
+
+  const pipeline = await createPipeline('cast', 'f16_to_f32');
+
+  const outputSize = numElements * DTYPE_SIZES.f32;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f16_to_f32_output');
+
+  const uniformBuffer = createUniformBufferWithView(
+    'cast_f16_to_f32_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+    },
+    null,
+    device
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'cast_f16_to_f32_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+
+  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+
+  uniformBuffer.destroy();
+
+  return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+}
+
+
+export async function recordCastF32ToF16(
+  recorder,
+  input,
+  options = {}
+) {
+  const device = recorder.device;
+  const { outputBuffer = null } = options;
+  const numElements = input.shape.reduce((a, b) => a * b, 1);
+
+  const pipeline = await createPipeline('cast', 'f32_to_f16');
+
+  const outputSize = numElements * DTYPE_SIZES.f16;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f32_to_f16_output');
+
+  const uniformBuffer = createUniformBufferWithView(
+    'cast_f32_to_f16_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+    },
+    recorder
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'cast_f32_to_f16_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+
+  // Use 2D dispatch for large tensors
+  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+
+  return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+}
+
+
+export async function recordCastF16ToF32(
+  recorder,
+  input,
+  options = {}
+) {
+  const device = recorder.device;
+  const { outputBuffer = null } = options;
+  const numElements = input.shape.reduce((a, b) => a * b, 1);
+
+  const pipeline = await createPipeline('cast', 'f16_to_f32');
+
+  const outputSize = numElements * DTYPE_SIZES.f32;
+  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f16_to_f32_output');
+
+  const uniformBuffer = createUniformBufferWithView(
+    'cast_f16_to_f32_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+    },
+    recorder
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'cast_f16_to_f32_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input.buffer } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+
+  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+
+  return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+}
+
+
+export async function runBF16ToF32(
+  input,
+  shape,
+  name = 'bf16_to_f32_output'
+) {
+  const numElements = shape.reduce((a, b) => a * b, 1);
+  trace.kernels(`BF16ToF32: Entry numElements=${numElements}, name=${name}, inputSize=${input.size}`);
+  const device = getDevice();
+
+  // Check for size limits (handle chunking if needed)
+  const limits = device.limits;
+  const maxBufferSize = limits.maxBufferSize;
+  const maxBindingSize = limits.maxStorageBufferBindingSize;
+  const outputSize = numElements * DTYPE_SIZES.f32;
+  trace.kernels(`BF16ToF32: outputSize=${outputSize}, maxBufferSize=${maxBufferSize}, maxBindingSize=${maxBindingSize}`);
+
+  if (outputSize > maxBufferSize) {
+    throw new Error(
+      `BF16->F32 output (${outputSize} bytes) exceeds device maxBufferSize (${maxBufferSize}). ` +
+      `This often happens for large-vocab models when converting embeddings/LM head. ` +
+      `Enable F16 and use BF16->F16 weights, or run on a device with a higher maxBufferSize.`
+    );
+  }
+
+  if (outputSize > maxBindingSize) {
+    // Need to chunk - output buffer can exist, but must be bound in smaller ranges.
+    return runBF16ToF32Chunked(input, shape, name, maxBindingSize);
+  }
+
+  const pipeline = await createPipeline('bf16_to_f32', 'default');
+  trace.kernels('BF16ToF32: Pipeline created');
+
+  const output = acquireBuffer(outputSize, undefined, name);
+  trace.kernels(`BF16ToF32: Output buffer acquired, size=${output.size}`);
+
+  const uniformBuffer = createUniformBufferWithView(
+    'bf16_to_f32_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+    },
+    null,
+    device
+  );
+  trace.kernels(`BF16ToF32: Uniform numElements=${numElements}`);
+
+  const bindGroup = device.createBindGroup({
+    label: 'bf16_to_f32_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+  trace.kernels('BF16ToF32: BindGroup created');
+
+  // Each thread processes 2 BF16 values (1 u32), so divide by 2 for thread count
+  // Then divide by 256 for workgroup count
+  const numPairs = Math.ceil(numElements / 2);
+  const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  trace.kernels(`BF16ToF32: Dispatching ${dispatchSize[0]}x${dispatchSize[1]} workgroups for ${numPairs} pairs (${numElements} elements)`);
+  dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f32');
+
+  uniformBuffer.destroy();
+
+  return createTensor(output, 'f32', [...shape], name);
+}
+
+
+export async function runBF16ToF16(
+  input,
+  shape,
+  name = 'bf16_to_f16_output'
+) {
+  const numElements = shape.reduce((a, b) => a * b, 1);
+  const device = getDevice();
+  const pipeline = await createPipeline('bf16_to_f16', 'default');
+
+  const limits = device.limits;
+  const maxBufferSize = limits.maxBufferSize;
+  const maxBindingSize = limits.maxStorageBufferBindingSize;
+  const outputSize = numElements * DTYPE_SIZES.f16;
+
+  if (outputSize > maxBufferSize) {
+    throw new Error(
+      `BF16->F16 output (${outputSize} bytes) exceeds device maxBufferSize (${maxBufferSize}).`
+    );
+  }
+  if (outputSize > maxBindingSize) {
+    throw new Error(
+      `BF16->F16 output (${outputSize} bytes) exceeds device maxStorageBufferBindingSize (${maxBindingSize}).`
+    );
+  }
+
+  const output = acquireBuffer(outputSize, undefined, name);
+
+  const uniformBuffer = createUniformBufferWithView(
+    'bf16_to_f16_uniforms',
+    16,
+    (view) => {
+      view.setUint32(0, numElements, true);
+      view.setUint32(4, 0, true);
+      view.setUint32(8, 0, true);
+    },
+    null,
+    device
+  );
+
+  const bindGroup = device.createBindGroup({
+    label: 'bf16_to_f16_bind_group',
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuffer } },
+      { binding: 1, resource: { buffer: input } },
+      { binding: 2, resource: { buffer: output } },
+    ],
+  });
+
+  const numPairs = Math.ceil(numElements / 2);
+  const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
+  const dispatchSize = calculate2DDispatch(workgroups);
+
+  dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f16');
+
+  uniformBuffer.destroy();
+
+  return createTensor(output, 'f16', [...shape], name);
+}
+
+
+async function runBF16ToF32Chunked(
+  input,
+  shape,
+  name,
+  maxBindingSize
+) {
+  const numElements = shape.reduce((a, b) => a * b, 1);
+  const device = getDevice();
+  const pipeline = await createPipeline('bf16_to_f32', 'default');
+
+  // Calculate chunk size
+  const alignmentBytes = device.limits.minStorageBufferOffsetAlignment;
+
+  const inElemAlign = Math.max(1, Math.floor(alignmentBytes / DTYPE_SIZES.bf16)); // BF16 elements
+  const outElemAlign = Math.max(1, Math.floor(alignmentBytes / DTYPE_SIZES.f32)); // F32 elements
+  const elemAlign = lcm(inElemAlign, outElemAlign);
+
+  let maxElementsPerChunk = Math.floor(maxBindingSize / DTYPE_SIZES.f32); // F32 output bytes
+  maxElementsPerChunk -= maxElementsPerChunk % elemAlign;
+  if (maxElementsPerChunk <= 0) {
+    throw new Error(`BF16->F32 chunk size underflow (maxBindingSize=${maxBindingSize}, alignment=${alignmentBytes})`);
+  }
+  const numChunks = Math.ceil(numElements / maxElementsPerChunk);
+
+  // Create full output buffer
+  const outputSize = numElements * DTYPE_SIZES.f32;
+  const output = acquireBuffer(outputSize, undefined, name);
+
+  trace.kernels(`BF16ToF32: Chunking ${numElements} elements in ${numChunks} chunks`);
+
+  for (let chunkIdx = 0; chunkIdx < numChunks; chunkIdx++) {
+    const chunkStart = chunkIdx * maxElementsPerChunk;
+    const chunkEnd = Math.min((chunkIdx + 1) * maxElementsPerChunk, numElements);
+    const chunkSize = chunkEnd - chunkStart;
+
+    const uniformBuffer = createUniformBufferWithView(
+      `bf16_to_f32_chunk${chunkIdx}_uniforms`,
+      16,
+      (view) => {
+        view.setUint32(0, chunkSize, true);
+        view.setUint32(4, 0, true);
+        view.setUint32(8, 0, true);
+      },
+      null,
+      device
+    );
+
+    const inputOffsetBytes = chunkStart * DTYPE_SIZES.bf16;
+    const outputOffsetBytes = chunkStart * DTYPE_SIZES.f32;
+    const inputPairs = Math.ceil(chunkSize / 2);
+    const inputSizeBytes = inputPairs * DTYPE_SIZES.f32; // Pairs read as u32
+    const outputSizeBytes = chunkSize * DTYPE_SIZES.f32;
+
+    const bindGroup = device.createBindGroup({
+      label: `bf16_to_f32_chunk${chunkIdx}_bind_group`,
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input, offset: inputOffsetBytes, size: inputSizeBytes } },
+        { binding: 2, resource: { buffer: output, offset: outputOffsetBytes, size: outputSizeBytes } },
+      ],
+    });
+
+    // Each thread processes 2 BF16 values
+    const numPairs = Math.ceil(chunkSize / 2);
+    const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = calculate2DDispatch(workgroups);
+
+    dispatch(device, pipeline, bindGroup, dispatchSize, `bf16_to_f32_chunk${chunkIdx}`);
+
+    uniformBuffer.destroy();
+  }
+
+  return createTensor(output, 'f32', [...shape], name);
+}

package/src/gpu/kernels/cast_f16_to_f32.wgsl

@@ -0,0 +1,31 @@
+// cast_f16_to_f32.wgsl
+
+/**
+ * Cast F16 to F32 Kernel
+ *
+ * Converts a buffer of f16 values to f32.
+ */
+
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    count: 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_write> output: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= u.count) {
+        return;
+    }
+    output[idx] = f32(input[idx]);
+}

package/src/gpu/kernels/cast_f32_to_f16.wgsl

@@ -0,0 +1,36 @@
+// Cast Kernel - f32 to f16
+//
+// Converts a flat f32 buffer to f16.
+// Used for KV-cache compression and future f16 pipelines.
+// Supports 2D dispatch for large tensors (>65535 workgroups).
+
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+const MAX_WG_X: u32 = 65535u;
+
+struct Uniforms {
+    size: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_wg: vec3<u32>
+) {
+    // Support 2D dispatch for large tensors
+    // Global index = x + y * (numWorkgroupsX * workgroupSize)
+    let i = gid.x + gid.y * num_wg.x * WORKGROUP_SIZE;
+    if (i >= u.size) {
+        return;
+    }
+    output[i] = f16(input[i]);
+}
+

package/src/gpu/kernels/check-finiteness.d.ts

@@ -0,0 +1,15 @@
+export declare const DEFAULT_FINITENESS_ABS_THRESHOLD: number;
+
+export declare function resolveFinitenessAbsThreshold(value: number | null | undefined): number;
+
+export declare function shouldTriggerFinitenessValue(value: number, absThreshold?: number): boolean;
+
+export declare function recordCheckFiniteness(
+  target: unknown,
+  inputBuffer: GPUBuffer,
+  size: number,
+  statusBuffer: GPUBuffer,
+  layerIdx?: number,
+  step?: number,
+  absThreshold?: number
+): void;

package/src/gpu/kernels/check-finiteness.js

@@ -0,0 +1,149 @@
+import { getDevice, getDeviceEpoch } from '../device.js';
+import { dispatchKernel } from './dispatch.js';
+import { createUniformBufferWithView, getOrCreateBindGroupLayout, getOrCreatePipelineLayout } from './utils.js';
+
+export const DEFAULT_FINITENESS_ABS_THRESHOLD = 65500;
+
+export function resolveFinitenessAbsThreshold(value) {
+  if (Number.isFinite(value) && value > 0) {
+    return value;
+  }
+  return DEFAULT_FINITENESS_ABS_THRESHOLD;
+}
+
+export function shouldTriggerFinitenessValue(value, absThreshold = DEFAULT_FINITENESS_ABS_THRESHOLD) {
+  if (!Number.isFinite(value)) {
+    return true;
+  }
+  return Math.abs(value) > resolveFinitenessAbsThreshold(absThreshold);
+}
+
+let checkFinitenessPipeline = null;
+let checkFinitenessPipelineEpoch = -1;
+
+const SHADER = /* wgsl */ `
+enable f16;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    size: u32,
+    layer: u32,
+    step: u32,
+    abs_threshold: f32,
+}
+
+@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_write> status: array<atomic<u32>>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) global_id: vec3<u32>
+) {
+    let idx = global_id.x;
+    
+    if (idx >= u.size) {
+        return;
+    }
+
+    let val = f32(input[idx]);
+    let bits = bitcast<u32>(val);
+    
+    // Check for NaN or Infinity (exponent bits all 1)
+    let non_finite = (bits & 0x7F800000u) == 0x7F800000u;
+    let exceeds_abs_threshold = abs(val) > u.abs_threshold;
+    if (non_finite || exceeds_abs_threshold) {
+        let old = atomicCompareExchangeWeak(&status[0], 0u, 1u);
+        if (old.exchanged) {
+            atomicStore(&status[1], u.layer);
+            atomicStore(&status[2], u.step);
+        }
+    }
+}
+`;
+
+function getCheckFinitenessBindGroupLayout(device) {
+    return getOrCreateBindGroupLayout(
+        'check_finiteness_bind_group_layout',
+        [
+            { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
+            { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
+            { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
+        ],
+        device
+    );
+}
+
+function getCheckFinitenessPipeline() {
+    const epoch = getDeviceEpoch();
+    if (checkFinitenessPipeline && checkFinitenessPipelineEpoch === epoch) return checkFinitenessPipeline;
+    const device = getDevice();
+    const shaderModule = device.createShaderModule({ code: SHADER });
+    const bindGroupLayout = getCheckFinitenessBindGroupLayout(device);
+
+    checkFinitenessPipeline = device.createComputePipeline({
+        layout: getOrCreatePipelineLayout('check_finiteness_pipeline_layout', [bindGroupLayout], device),
+        compute: {
+            module: shaderModule,
+            entryPoint: 'main',
+            constants: { WORKGROUP_SIZE: 256 },
+        },
+    });
+    checkFinitenessPipelineEpoch = epoch;
+
+    return checkFinitenessPipeline;
+}
+
+export function recordCheckFiniteness(
+    target,
+    inputBuffer,
+    size,
+    statusBuffer,
+    layerIdx = 0,
+    step = 0,
+    absThreshold = DEFAULT_FINITENESS_ABS_THRESHOLD
+) {
+    const isRecorder = target && typeof target.beginComputePass === 'function';
+    const device = isRecorder ? target.device : getDevice();
+    const pipeline = getCheckFinitenessPipeline();
+    const resolvedAbsThreshold = resolveFinitenessAbsThreshold(absThreshold);
+
+    const uniformBuffer = createUniformBufferWithView(
+        'check_finiteness_uniforms',
+        16,
+        (view) => {
+            view.setUint32(0, size, true);
+            view.setUint32(4, layerIdx, true);
+            view.setUint32(8, step, true);
+            view.setFloat32(12, resolvedAbsThreshold, true);
+        },
+        isRecorder ? target : null,
+        device
+    );
+
+    // Create bind group
+    const bindGroup = device.createBindGroup({
+        layout: getCheckFinitenessBindGroupLayout(device),
+        entries: [
+            { binding: 0, resource: { buffer: uniformBuffer } },
+            { binding: 1, resource: { buffer: inputBuffer } },
+            { binding: 2, resource: { buffer: statusBuffer } },
+        ],
+    });
+
+    const workgroups = Math.ceil(size / 256);
+    dispatchKernel(target, pipeline, bindGroup, workgroups, 'check_finiteness');
+
+    // Recorder-created uniform buffers come from the uniform cache and must not
+    // be destroyed as temporaries. Non-recorder path uses direct allocations.
+    if (!isRecorder) {
+        device.queue.onSubmittedWorkDone()
+          .then(() => {
+            uniformBuffer.destroy();
+          })
+          .catch(() => {
+            uniformBuffer.destroy();
+          });
+    }
+}

package/src/gpu/kernels/check-stop.d.ts

@@ -0,0 +1,31 @@
+/**
+ * check-stop.ts - GPU Stop Condition Kernel
+ *
+ * Checks if generation should stop based on EOS token or max length.
+ * Used in GPU-only decode loop to eliminate CPU roundtrips.
+ */
+
+import type { CommandRecorder } from '../command-recorder.js';
+
+export interface CheckStopParams {
+  sampledTokenBuffer: GPUBuffer;  // u32 storage buffer with sampled tokens
+  tokenIndex?: number;
+  shouldStopBuffer?: GPUBuffer;
+  eosTokenId: number;
+  maxTokens: number;
+  currentPos: number;
+}
+
+/**
+ * Record a stop condition check into the command buffer.
+ * Returns the shouldStop buffer (1 if should stop, 0 otherwise).
+ */
+export declare function recordCheckStop(
+  recorder: CommandRecorder,
+  params: CheckStopParams
+): GPUBuffer;
+
+/**
+ * Standalone check stop (submits immediately, for testing).
+ */
+export declare function checkStop(params: CheckStopParams): Promise<boolean>;