@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/repeat_channels.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';
@@ -32,23 +32,31 @@ async function _repeatChannels(target, input, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = outChannels * height * width * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'repeat_channels_output');
+  const ownedOutput = outputBuffer ? null : output;
 
-  await unifiedKernelWrapper(
-    'repeat_channels',
-    target,
-    variant,
-    [input, output],
-    {
-      in_channels: inChannels,
-      height,
-      width,
-      repeats,
-      _pad0: 0,
-    },
-    Math.ceil((outChannels * height * width) / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
+  try {
+    await unifiedKernelWrapper(
+      'repeat_channels',
+      target,
+      variant,
+      [input, output],
+      {
+        in_channels: inChannels,
+        height,
+        width,
+        repeats,
+        _pad0: 0,
+      },
+      [Math.ceil((height * width) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
+
+    return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runRepeatChannels(input, options = {}) {

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,9 +63,26 @@ 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;
+  const ownsOutput = outputBuffer == null;
 
   const { a: aAligned, b: bAligned, temps } = await alignResidualInputs(a, b, recorder);
   const outputDtype = inferOutputDtype(aAligned, bAligned);
@@ -75,19 +92,28 @@ 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);
-
-  await unifiedKernelWrapper(
-    'residual', target, variant,
-    [aAligned, bAligned, output],
-    { size },
-    workgroups
+  const dispatchPlan = planResidualDispatch(
+    target,
+    size,
+    useVec4 ? VEC4_ELEMENTS_PER_WG : WORKGROUP_SIZES.DEFAULT
   );
 
-  cleanupTemps(temps, recorder);
-  return createTensor(output, outputDtype, [size], 'residual_output');
+  try {
+    await unifiedKernelWrapper(
+      'residual', target, variant,
+      [aAligned, bAligned, output],
+      { size, scale: 1, _pad1: dispatchPlan.dispatchStride, _pad2: 0 },
+      dispatchPlan.workgroups
+    );
+    return createTensor(output, outputDtype, [size], 'residual_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  } finally {
+    cleanupTemps(temps, recorder);
+  }
 }
 
 async function _biasAdd(target, data, bias, numTokens, dim, options = {}) {
@@ -96,18 +122,38 @@ 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);
-
-  await unifiedKernelWrapper(
-    'bias_add', target, variant,
-    [data, biasAligned],
-    { num_tokens: numTokens, dim, data_offset: dataOffset, bias_offset: biasOffset },
-    workgroups
-  );
-
-  cleanupTemps(temps, recorder);
-  return createTensor(data.buffer, data.dtype, [numTokens, dim], 'bias_add_output');
+  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),
+  ];
+
+  try {
+    await unifiedKernelWrapper(
+      'bias_add', target, variant,
+      [data, biasAligned],
+      {
+        num_tokens: numTokens,
+        dim,
+        data_offset: dataOffset,
+        bias_offset: biasOffset,
+        token_stride: tokenStride,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      workgroups
+    );
+    return createTensor(data.buffer, data.dtype, [numTokens, dim], 'bias_add_output');
+  } finally {
+    cleanupTemps(temps, recorder);
+  }
 }
 
 export async function runResidualAdd(a, b, size, options = {}) {

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

@@ -1,7 +1,7 @@
 
 
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer, getBufferRequestedSize } from '../../memory/buffer-pool.js';
+import { acquireBuffer, getBufferRequestedSize, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { getKernelThresholds, padToQ4KBlock } from '../../config/schema/index.js';
 import { selectRuleValue } from './rule-registry.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,27 +112,46 @@ 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 ownedOutput = outputBuffer ? null : outputBuf;
+  const dispatchPlan = planRMSNormDispatch(null, batchSize);
 
   // Shader layout always includes the residual binding; when unused, bind a harmless placeholder.
-  const residualBuf = residual?.buffer || input.buffer;
-
-  await unifiedKernelWrapper(
-    'rmsnorm',
-    null,
-    variant,
-    [input, normWeightBuffer, outputBuf, residualBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps, has_residual: residual ? 1 : 0 },
-    batchSize,
-    { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
-  );
-
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
+
+  try {
+    await unifiedKernelWrapper(
+      'rmsnorm',
+      null,
+      variant,
+      [input, normWeightBuffer, outputBuf, residualBuf],
+      {
+        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' }
+    );
+
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function recordRMSNorm(
@@ -117,24 +166,43 @@ 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 residualBuf = residual?.buffer || input.buffer;
-
-  await unifiedKernelWrapper(
-    'rmsnorm',
-    recorder,
-    variant,
-    [input, normWeightBuffer, outputBuf, residualBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps, has_residual: residual ? 1 : 0 },
-    batchSize,
-    { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
-  );
-
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
+  const dispatchPlan = planRMSNormDispatch(recorder, batchSize);
+
+  const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
+
+  try {
+    await unifiedKernelWrapper(
+      'rmsnorm',
+      recorder,
+      variant,
+      [input, normWeightBuffer, outputBuf, residualBuf],
+      {
+        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' }
+    );
+
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

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,29 @@ 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}`);
+  }
+  if (input.dtype === 'f16' && (rotaryDim !== headDim || interleaved)) {
+    throw new Error('RoPE f16 kernel requires rotaryDim === headDim and interleaved === false.');
+  }
 
   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 +45,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
   );