@simulatte/doppler 0.1.5 → 0.1.6

package/src/gpu/kernels/repeat_channels.wgsl

@@ -14,16 +14,15 @@ struct Uniforms {
 
 @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_channels = u.in_channels * u.repeats;
-    let total = out_channels * spatial;
-    if (idx >= total) {
+    let spatial_idx = gid.x;
+    let out_channel = gid.y;
+    if (out_channel >= out_channels || spatial_idx >= spatial) {
         return;
     }
 
-    let out_channel = idx / spatial;
     let channel = out_channel / u.repeats;
-    let spatial_idx = idx - out_channel * spatial;
+    let idx = out_channel * spatial + spatial_idx;
     output[idx] = input[channel * spatial + spatial_idx];
 }

package/src/gpu/kernels/repeat_channels_f16.wgsl

@@ -16,16 +16,15 @@ struct Uniforms {
 
 @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_channels = u.in_channels * u.repeats;
-    let total = out_channels * spatial;
-    if (idx >= total) {
+    let spatial_idx = gid.x;
+    let out_channel = gid.y;
+    if (out_channel >= out_channels || spatial_idx >= spatial) {
         return;
     }
 
-    let out_channel = idx / spatial;
     let channel = out_channel / u.repeats;
-    let spatial_idx = idx - out_channel * spatial;
+    let idx = out_channel * spatial + spatial_idx;
     output[idx] = input[channel * spatial + spatial_idx];
 }

package/src/gpu/kernels/residual.js

@@ -63,6 +63,22 @@ function cleanupTemps(temps, recorder) {
   }
 }
 
+function planResidualDispatch(target, size, elementsPerWorkgroup) {
+  const device = target?.device;
+  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
+    ? device.limits.maxComputeWorkgroupsPerDimension
+    : 65535;
+  const dispatchStride = Math.min(size, maxPerDim * elementsPerWorkgroup);
+  return {
+    dispatchStride,
+    workgroups: [
+      Math.ceil(dispatchStride / elementsPerWorkgroup),
+      Math.ceil(size / dispatchStride),
+      1,
+    ],
+  };
+}
+
 async function _residualAdd(target, a, b, size, options = {}) {
   const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
   const { useVec4 = true, outputBuffer = null } = options;
@@ -75,15 +91,17 @@ async function _residualAdd(target, a, b, size, options = {}) {
   const outputSize = size * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'residual_output');
 
-  const workgroups = useVec4
-    ? Math.ceil(size / VEC4_ELEMENTS_PER_WG)
-    : Math.ceil(size / WORKGROUP_SIZES.DEFAULT);
+  const dispatchPlan = planResidualDispatch(
+    target,
+    size,
+    useVec4 ? VEC4_ELEMENTS_PER_WG : WORKGROUP_SIZES.DEFAULT
+  );
 
   await unifiedKernelWrapper(
     'residual', target, variant,
     [aAligned, bAligned, output],
-    { size },
-    workgroups
+    { size, scale: 1, _pad1: dispatchPlan.dispatchStride, _pad2: 0 },
+    dispatchPlan.workgroups
   );
 
   cleanupTemps(temps, recorder);
@@ -96,13 +114,31 @@ async function _biasAdd(target, data, bias, numTokens, dim, options = {}) {
 
   const { bias: biasAligned, temps } = await alignBiasTensor(data, bias, recorder);
   const variant = selectBiasAddVariant(data.dtype, biasAligned.dtype);
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+  const device = target?.device;
+  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
+    ? device.limits.maxComputeWorkgroupsPerDimension
+    : 65535;
+  const tokenStride = Math.min(numTokens, maxPerDim);
+
+  const workgroups = [
+    Math.ceil(dim / WORKGROUP_SIZES.DEFAULT),
+    tokenStride,
+    Math.ceil(numTokens / tokenStride),
+  ];
 
   await unifiedKernelWrapper(
     'bias_add', target, variant,
     [data, biasAligned],
-    { num_tokens: numTokens, dim, data_offset: dataOffset, bias_offset: biasOffset },
+    {
+      num_tokens: numTokens,
+      dim,
+      data_offset: dataOffset,
+      bias_offset: biasOffset,
+      token_stride: tokenStride,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
     workgroups
   );
 

package/src/gpu/kernels/residual.wgsl

@@ -23,7 +23,8 @@ 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 dispatch_stride = max(u._pad1, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }
@@ -35,7 +36,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 // This avoids requiring a different bind group layout with read_write on 'a'
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn add_inplace(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+    let dispatch_stride = max(u._pad1, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }
@@ -45,7 +47,8 @@ fn add_inplace(@builtin(global_invocation_id) gid: vec3<u32>) {
 // Fused residual + scale: output = a + scale * b
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn add_scaled(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+    let dispatch_stride = max(u._pad1, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }

package/src/gpu/kernels/residual_f16.wgsl

@@ -27,7 +27,8 @@ 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 dispatch_stride = max(u._pad1, 1u);
+    let idx = gid.y * dispatch_stride + gid.x;
     if (idx >= u.size) {
         return;
     }

package/src/gpu/kernels/residual_f16_vec4.wgsl

@@ -25,7 +25,8 @@ override WORKGROUP_SIZE_VEC4: u32 = 64u;
 // Vectorized version for better throughput
 @compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
 fn add_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x * 4u;
+    let dispatch_stride = max(u._pad1, 4u);
+    let idx = gid.y * dispatch_stride + gid.x * 4u;
     let size = u.size;
 
     if (idx >= size) {

package/src/gpu/kernels/residual_vec4.wgsl

@@ -23,7 +23,8 @@ override WORKGROUP_SIZE_VEC4: u32 = 64u;
 // Vectorized version for better throughput
 @compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
 fn add_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x * 4u;
+    let dispatch_stride = max(u._pad1, 4u);
+    let idx = gid.y * dispatch_stride + gid.x * 4u;
     let size = u.size;
 
     if (idx >= size) {

package/src/gpu/kernels/rmsnorm.js

@@ -58,6 +58,36 @@ function resolveNormWeightDtype(weight, hiddenSize) {
   return 'f32';
 }
 
+function assertRMSNormWeightBuffer(weight, weightBuffer, hiddenSize) {
+  const isGpuBuffer = weightBuffer && (
+    typeof GPUBuffer === 'undefined'
+      ? true
+      : weightBuffer instanceof GPUBuffer
+  );
+  if (isGpuBuffer) {
+    return;
+  }
+  const weightLabel = weight?.label ?? 'unknown';
+  const weightType = weight === null ? 'null' : weight === undefined ? 'undefined' : weight.constructor?.name || typeof weight;
+  const bufferType = weightBuffer === null ? 'null' : weightBuffer === undefined ? 'undefined' : weightBuffer.constructor?.name || typeof weightBuffer;
+  throw new Error(
+    `[rmsnorm] weight "${weightLabel}" requires a GPUBuffer ` +
+    `(weightType=${weightType}, bufferType=${bufferType}, hiddenSize=${hiddenSize ?? 'unknown'}).`
+  );
+}
+
+function planRMSNormDispatch(target, numTokens) {
+  const device = target?.device;
+  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
+    ? device.limits.maxComputeWorkgroupsPerDimension
+    : 65535;
+  const tokenStride = Math.min(numTokens, maxPerDim);
+  return {
+    tokenStride,
+    workgroups: [tokenStride, Math.ceil(numTokens / tokenStride), 1],
+  };
+}
+
 export function selectRMSNormKernel(options = {}, isF16 = false) {
   const { residual = null, hiddenSize = null } = options;
   const { smallThreshold } = getKernelThresholds().rmsnorm;
@@ -82,23 +112,34 @@ export async function runRMSNorm(
   const variant = selectRMSNormKernel(options, isF16);
   const inferredHiddenSize = inferHiddenSize(input, hiddenSize);
   const normWeightBuffer = getBuffer(weight);
+  assertRMSNormWeightBuffer(weight, normWeightBuffer, inferredHiddenSize);
   const normWeightDtype = resolveNormWeightDtype(weight, inferredHiddenSize);
 
   const bytesPerElement = isF16 ? 2 : 4;
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'rmsnorm_output');
+  const dispatchPlan = planRMSNormDispatch(null, batchSize);
 
   // Shader layout always includes the residual binding; when unused, bind a harmless placeholder.
-  const residualBuf = residual?.buffer || input.buffer;
+  const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
 
   await unifiedKernelWrapper(
     'rmsnorm',
     null,
     variant,
     [input, normWeightBuffer, outputBuf, residualBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps, has_residual: residual ? 1 : 0 },
-    batchSize,
+    {
+      hidden_size: inferredHiddenSize,
+      num_tokens: batchSize,
+      eps,
+      has_residual: residual ? 1 : 0,
+      token_stride: dispatchPlan.tokenStride,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
+    dispatchPlan.workgroups,
     { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
   );
 
@@ -117,22 +158,33 @@ export async function recordRMSNorm(
   const variant = selectRMSNormKernel(options, isF16);
   const inferredHiddenSize = inferHiddenSize(input, hiddenSize);
   const normWeightBuffer = getBuffer(weight);
+  assertRMSNormWeightBuffer(weight, normWeightBuffer, inferredHiddenSize);
   const normWeightDtype = resolveNormWeightDtype(weight, inferredHiddenSize);
 
   const bytesPerElement = isF16 ? 2 : 4;
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'rmsnorm_output');
+  const dispatchPlan = planRMSNormDispatch(recorder, batchSize);
 
-  const residualBuf = residual?.buffer || input.buffer;
+  const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
 
   await unifiedKernelWrapper(
     'rmsnorm',
     recorder,
     variant,
     [input, normWeightBuffer, outputBuf, residualBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps, has_residual: residual ? 1 : 0 },
-    batchSize,
+    {
+      hidden_size: inferredHiddenSize,
+      num_tokens: batchSize,
+      eps,
+      has_residual: residual ? 1 : 0,
+      token_stride: dispatchPlan.tokenStride,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
+    dispatchPlan.workgroups,
     { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
   );
 

package/src/gpu/kernels/rmsnorm.wgsl

@@ -39,6 +39,10 @@ struct Uniforms {
     num_tokens: u32,    // Number of tokens to process
     eps: f32,           // Epsilon for numerical stability (typically 1e-5 or 1e-6)
     has_residual: u32,  // Runtime flag: 1 = add residual after norm
+    token_stride: u32,  // Workgroup rows per dispatch row
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 @group(0) @binding(0) var<uniform> u: Uniforms;
@@ -82,6 +86,10 @@ fn should_add_residual() -> bool {
     return HAS_RESIDUAL || (u.has_residual != 0u);
 }
 
+fn token_index(wg_id: vec3<u32>) -> u32 {
+    return wg_id.y * max(u.token_stride, 1u) + wg_id.x;
+}
+
 // =============================================================================
 // Main Entry Point
 // =============================================================================
@@ -93,7 +101,7 @@ fn main(
     @builtin(local_invocation_id) local_id: vec3<u32>,
     @builtin(workgroup_id) wg_id: vec3<u32>
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -163,7 +171,7 @@ fn main_small(
     @builtin(local_invocation_id) local_id: vec3<u32>,
     @builtin(workgroup_id) wg_id: vec3<u32>
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -219,7 +227,7 @@ fn main_cached(
     @builtin(local_invocation_id) local_id: vec3<u32>,
     @builtin(workgroup_id) wg_id: vec3<u32>
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -288,7 +296,7 @@ fn main_subgroup(
     @builtin(subgroup_invocation_id) sg_lane: u32,
     @builtin(subgroup_size) sg_size: u32,
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -362,7 +370,7 @@ fn main_small_subgroup(
     @builtin(subgroup_invocation_id) sg_lane: u32,
     @builtin(subgroup_size) sg_size: u32,
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -414,4 +422,4 @@ fn main_small_subgroup(
         }
         output[base_offset + thread_idx] = result;
     }
-}
+}

package/src/gpu/kernels/rmsnorm_f16.wgsl

@@ -20,6 +20,10 @@ struct Uniforms {
     num_tokens: u32,    // Number of tokens to process
     eps: f32,           // Epsilon for numerical stability
     has_residual: u32,  // 1 if residual input provided, 0 otherwise
+    token_stride: u32,  // Workgroup rows per dispatch row
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 @group(0) @binding(0) var<uniform> u: Uniforms;
@@ -47,6 +51,10 @@ fn load_weight(idx: u32) -> f32 {
     return bitcast<f32>(weight[idx]);
 }
 
+fn token_index(wg_id: vec3<u32>) -> u32 {
+    return wg_id.y * max(u.token_stride, 1u) + wg_id.x;
+}
+
 // Main RMSNorm kernel - one workgroup per token
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(
@@ -54,7 +62,7 @@ fn main(
     @builtin(local_invocation_id) local_id: vec3<u32>,
     @builtin(workgroup_id) wg_id: vec3<u32>
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 
@@ -121,7 +129,7 @@ fn rmsnorm_small_f16(
     @builtin(local_invocation_id) local_id: vec3<u32>,
     @builtin(workgroup_id) wg_id: vec3<u32>
 ) {
-    let token_idx = wg_id.x;
+    let token_idx = token_index(wg_id);
     let thread_idx = local_id.x;
     let size = u.size;
 

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

@@ -15,6 +15,8 @@ import type { OutputBufferOptions } from './types.js';
 export interface RoPEOptions extends OutputBufferOptions {
   numHeads?: number;
   headDim?: number;
+  rotaryDim?: number;
+  interleaved?: boolean;
   ropeTheta?: number;
   startPos?: number;
 }

package/src/gpu/kernels/rope.js

@@ -13,18 +13,26 @@ async function _rope(target, input, freqsCos, freqsSin, seqLen, options = {}) {
   const {
     numHeads = 1,
     headDim = 64,
+    rotaryDim = headDim,
+    interleaved = false,
     ropeTheta = ropeDefaults.defaultTheta,
   } = options;
 
   if (headDim % 2 !== 0) {
     throw new Error(`RoPE headDim must be even, got ${headDim}`);
   }
+  if (rotaryDim % 2 !== 0) {
+    throw new Error(`RoPE rotaryDim must be even, got ${rotaryDim}`);
+  }
+  if (rotaryDim <= 0 || rotaryDim > headDim) {
+    throw new Error(`RoPE rotaryDim must be in (0, headDim]; got ${rotaryDim} for headDim ${headDim}`);
+  }
 
   const caps = getKernelCapabilities();
   const useF16 = input.dtype === 'f16' && caps.hasF16;
   const variant = selectRuleValue('rope', 'variant', { useF16 });
 
-  const halfDim = headDim / 2;
+  const halfDim = rotaryDim / 2;
   const workgroups = Math.ceil((seqLen * numHeads * halfDim) / WORKGROUP_SIZES.DEFAULT);
 
   await unifiedKernelWrapper(
@@ -34,9 +42,11 @@ async function _rope(target, input, freqsCos, freqsSin, seqLen, options = {}) {
       seq_len: seqLen,
       num_heads: numHeads,
       head_dim: headDim,
+      rotary_dim: rotaryDim,
       start_pos: options.startPos ?? ropeDefaults.defaultStartPos,
       rope_base: ropeTheta,
       rope_scale: 1.0,
+      interleaved: interleaved ? 1 : 0,
     },
     workgroups
   );