@simulatte/doppler 0.1.4 → 0.1.6

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
   );

package/src/gpu/kernels/rope.wgsl

@@ -26,8 +26,8 @@ struct Uniforms {
     start_pos: u32,        // Starting position (for decode)
     rope_base: f32,        // Base frequency (default 10000)
     rope_scale: f32,       // Scaling factor for extended context
-    _pad0: u32,
-    _pad1: u32,
+    rotary_dim: u32,       // Rotary slice within head_dim
+    interleaved: u32,      // 1 = adjacent pairs, 0 = rotate-half
 }
 
 @group(0) @binding(0) var<uniform> u: Uniforms;
@@ -46,7 +46,8 @@ fn main(
     let start_pos = u.start_pos;
 
     // Global thread index (one thread per complex pair)
-    let half_dim = head_dim / 2u;
+    let rotary_dim = u.rotary_dim;
+    let half_dim = rotary_dim / 2u;
     let total_pairs = seq_len * num_heads * half_dim;
     let idx = global_id.x;
 
@@ -68,16 +69,18 @@ fn main(
 
     // Apply "rotate-half" layout: pair (x[i], x[i + half_dim])
     let base_idx = pos * num_heads * head_dim + head_idx * head_dim;
-    let x0 = input[base_idx + pair_idx];
-    let x1 = input[base_idx + pair_idx + half_dim];
+    let first_idx = select(pair_idx, pair_idx * 2u, u.interleaved == 1u);
+    let second_idx = select(pair_idx + half_dim, pair_idx * 2u + 1u, u.interleaved == 1u);
+    let x0 = input[base_idx + first_idx];
+    let x1 = input[base_idx + second_idx];
 
     // Apply rotation
     let y0 = x0 * cos_val - x1 * sin_val;
     let y1 = x0 * sin_val + x1 * cos_val;
 
     // Write back
-    input[base_idx + pair_idx] = y0;
-    input[base_idx + pair_idx + half_dim] = y1;
+    input[base_idx + first_idx] = y0;
+    input[base_idx + second_idx] = y1;
 }
 
 // Compute frequencies on-the-fly (no precomputation needed)
@@ -91,9 +94,10 @@ fn rope_compute_freqs(
     let start_pos = u.start_pos;
     let rope_base = u.rope_base;
     let rope_scale = u.rope_scale;
+    let rotary_dim = u.rotary_dim;
 
     let idx = global_id.x;
-    let half_dim = head_dim / 2u;
+    let half_dim = rotary_dim / 2u;
     let total_pairs = seq_len * num_heads * half_dim;
 
     if (idx >= total_pairs) {
@@ -109,7 +113,7 @@ fn rope_compute_freqs(
     let actual_pos = f32(start_pos + pos) / rope_scale;
 
     // Compute frequency: 1 / (base^(2*pair_idx/head_dim))
-    let exponent = f32(pair_idx * 2u) / f32(head_dim);
+    let exponent = f32(pair_idx * 2u) / f32(rotary_dim);
     let freq = 1.0 / pow(rope_base, exponent);
     let theta = actual_pos * freq;
 
@@ -118,12 +122,14 @@ fn rope_compute_freqs(
 
     // Apply "rotate-half" layout: pair (x[i], x[i + half_dim])
     let base_idx = pos * num_heads * head_dim + head_idx * head_dim;
-    let x0 = input[base_idx + pair_idx];
-    let x1 = input[base_idx + pair_idx + half_dim];
+    let first_idx = select(pair_idx, pair_idx * 2u, u.interleaved == 1u);
+    let second_idx = select(pair_idx + half_dim, pair_idx * 2u + 1u, u.interleaved == 1u);
+    let x0 = input[base_idx + first_idx];
+    let x1 = input[base_idx + second_idx];
 
     // Apply rotation
-    input[base_idx + pair_idx] = x0 * cos_val - x1 * sin_val;
-    input[base_idx + pair_idx + half_dim] = x0 * sin_val + x1 * cos_val;
+    input[base_idx + first_idx] = x0 * cos_val - x1 * sin_val;
+    input[base_idx + second_idx] = x0 * sin_val + x1 * cos_val;
 }
 
 // Apply RoPE to both Q and K in one pass
@@ -138,10 +144,11 @@ fn rope_qk(
     let start_pos = u.start_pos;
     let rope_base = u.rope_base;
     let rope_scale = u.rope_scale;
+    let rotary_dim = u.rotary_dim;
 
     let idx = global_id.x;
     // Each thread handles one Q-K pair at one dimension pair
-    let half_dim = head_dim / 2u;
+    let half_dim = rotary_dim / 2u;
     let total_pairs = seq_len * num_heads * half_dim;
 
     if (idx >= total_pairs) {
@@ -156,7 +163,7 @@ fn rope_qk(
     let actual_pos = f32(start_pos + pos) / rope_scale;
 
     // Compute frequency
-    let exponent = f32(pair_idx * 2u) / f32(head_dim);
+    let exponent = f32(pair_idx * 2u) / f32(rotary_dim);
     let freq = 1.0 / pow(rope_base, exponent);
     let theta = actual_pos * freq;
 
@@ -168,16 +175,18 @@ fn rope_qk(
     let k_base_idx = q_base_idx + head_dim;  // K starts after Q
 
     // Process Q
-    let q0 = input[q_base_idx + pair_idx];
-    let q1 = input[q_base_idx + pair_idx + half_dim];
-    input[q_base_idx + pair_idx] = q0 * cos_val - q1 * sin_val;
-    input[q_base_idx + pair_idx + half_dim] = q0 * sin_val + q1 * cos_val;
+    let first_idx = select(pair_idx, pair_idx * 2u, u.interleaved == 1u);
+    let second_idx = select(pair_idx + half_dim, pair_idx * 2u + 1u, u.interleaved == 1u);
+    let q0 = input[q_base_idx + first_idx];
+    let q1 = input[q_base_idx + second_idx];
+    input[q_base_idx + first_idx] = q0 * cos_val - q1 * sin_val;
+    input[q_base_idx + second_idx] = q0 * sin_val + q1 * cos_val;
 
     // Process K
-    let k0 = input[k_base_idx + pair_idx];
-    let k1 = input[k_base_idx + pair_idx + half_dim];
-    input[k_base_idx + pair_idx] = k0 * cos_val - k1 * sin_val;
-    input[k_base_idx + pair_idx + half_dim] = k0 * sin_val + k1 * cos_val;
+    let k0 = input[k_base_idx + first_idx];
+    let k1 = input[k_base_idx + second_idx];
+    input[k_base_idx + first_idx] = k0 * cos_val - k1 * sin_val;
+    input[k_base_idx + second_idx] = k0 * sin_val + k1 * cos_val;
 }
 
 // Precompute frequency table (run once at init)
@@ -190,9 +199,10 @@ fn precompute_freqs(
     let seq_len = u.seq_len;  // maxSeqLen for precomputation
     let rope_base = u.rope_base;
     let rope_scale = u.rope_scale;
+    let rotary_dim = u.rotary_dim;
 
     let idx = global_id.x;
-    let half_dim = head_dim / 2u;
+    let half_dim = rotary_dim / 2u;
     let total_elements = seq_len * half_dim;
 
     if (idx >= total_elements) {
@@ -203,7 +213,7 @@ fn precompute_freqs(
     let dim_idx = idx % half_dim;
 
     let actual_pos = f32(pos) / rope_scale;
-    let exponent = f32(dim_idx * 2u) / f32(head_dim);
+    let exponent = f32(dim_idx * 2u) / f32(rotary_dim);
     let freq = 1.0 / pow(rope_base, exponent);
     let theta = actual_pos * freq;
 
@@ -218,6 +228,7 @@ fn rope_ntk_scaled(
     @builtin(global_invocation_id) global_id: vec3<u32>
 ) {
     let head_dim = u.head_dim;
+    let rotary_dim = u.rotary_dim;
     let num_heads = u.num_heads;
     let seq_len = u.seq_len;
     let start_pos = u.start_pos;
@@ -225,7 +236,7 @@ fn rope_ntk_scaled(
     let rope_scale = u.rope_scale;
 
     let idx = global_id.x;
-    let half_dim = head_dim / 2u;
+    let half_dim = rotary_dim / 2u;
     let total_pairs = seq_len * num_heads * half_dim;
 
     if (idx >= total_pairs) {
@@ -234,7 +245,7 @@ fn rope_ntk_scaled(
 
     // NTK scaling: increase base proportionally to scale factor
     // This preserves high-frequency components better than linear interpolation
-    rope_base = rope_base * pow(rope_scale, f32(head_dim) / (f32(head_dim) - 2.0));
+    rope_base = rope_base * pow(rope_scale, f32(rotary_dim) / (f32(rotary_dim) - 2.0));
 
     let pos = idx / (num_heads * half_dim);
     let remainder = idx % (num_heads * half_dim);
@@ -243,7 +254,7 @@ fn rope_ntk_scaled(
 
     let actual_pos = f32(start_pos + pos);
 
-    let exponent = f32(pair_idx * 2u) / f32(head_dim);
+    let exponent = f32(pair_idx * 2u) / f32(rotary_dim);
     let freq = 1.0 / pow(rope_base, exponent);
     let theta = actual_pos * freq;
 
@@ -251,11 +262,13 @@ fn rope_ntk_scaled(
     let sin_val = sin(theta);
 
     let base_idx = pos * num_heads * head_dim + head_idx * head_dim;
-    let x0 = input[base_idx + pair_idx];
-    let x1 = input[base_idx + pair_idx + half_dim];
+    let first_idx = select(pair_idx, pair_idx * 2u, u.interleaved == 1u);
+    let second_idx = select(pair_idx + half_dim, pair_idx * 2u + 1u, u.interleaved == 1u);
+    let x0 = input[base_idx + first_idx];
+    let x1 = input[base_idx + second_idx];
 
-    input[base_idx + pair_idx] = x0 * cos_val - x1 * sin_val;
-    input[base_idx + pair_idx + half_dim] = x0 * sin_val + x1 * cos_val;
+    input[base_idx + first_idx] = x0 * cos_val - x1 * sin_val;
+    input[base_idx + second_idx] = x0 * sin_val + x1 * cos_val;
 }
 
 // YaRN-style RoPE with attention scaling
@@ -265,6 +278,7 @@ fn rope_yarn(
     @builtin(global_invocation_id) global_id: vec3<u32>
 ) {
     let head_dim = u.head_dim;
+    let rotary_dim = u.rotary_dim;
     let num_heads = u.num_heads;
     let seq_len = u.seq_len;
     let start_pos = u.start_pos;
@@ -272,7 +286,7 @@ fn rope_yarn(
     let rope_scale = u.rope_scale;
 
     let idx = global_id.x;
-    let half_dim = head_dim / 2u;
+    let half_dim = rotary_dim / 2u;
     let total_pairs = seq_len * num_heads * half_dim;
 
     if (idx >= total_pairs) {
@@ -292,7 +306,7 @@ fn rope_yarn(
     let alpha: f32 = 1.0;
 
     // Compute original frequency
-    let exponent = f32(pair_idx * 2u) / f32(head_dim);
+    let exponent = f32(pair_idx * 2u) / f32(rotary_dim);
     let orig_freq = 1.0 / pow(rope_base, exponent);
 
     // Compute wavelength
@@ -300,8 +314,8 @@ fn rope_yarn(
 
     // Interpolation factor based on wavelength
     var ramp: f32;
-    let low_wavelength = f32(head_dim) / beta_fast;
-    let high_wavelength = f32(head_dim) / beta_slow;
+    let low_wavelength = f32(rotary_dim) / beta_fast;
+    let high_wavelength = f32(rotary_dim) / beta_slow;
 
     if (wavelength < low_wavelength) {
         ramp = 0.0;  // No interpolation for high frequencies
@@ -320,9 +334,11 @@ fn rope_yarn(
     let sin_val = sin(theta);
 
     let base_idx = pos * num_heads * head_dim + head_idx * head_dim;
-    let x0 = input[base_idx + pair_idx];
-    let x1 = input[base_idx + pair_idx + half_dim];
+    let first_idx = select(pair_idx, pair_idx * 2u, u.interleaved == 1u);
+    let second_idx = select(pair_idx + half_dim, pair_idx * 2u + 1u, u.interleaved == 1u);
+    let x0 = input[base_idx + first_idx];
+    let x1 = input[base_idx + second_idx];
 
-    input[base_idx + pair_idx] = x0 * cos_val - x1 * sin_val;
-    input[base_idx + pair_idx + half_dim] = x0 * sin_val + x1 * cos_val;
+    input[base_idx + first_idx] = x0 * cos_val - x1 * sin_val;
+    input[base_idx + second_idx] = x0 * sin_val + x1 * cos_val;
 }

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

@@ -0,0 +1,27 @@
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+import type { OutputBufferOptions } from './types.js';
+
+export interface SanaLinearAttentionOptions extends OutputBufferOptions {
+  numHeads: number;
+  headDim: number;
+  numTokens?: number;
+  hiddenSize?: number;
+  eps?: number;
+  summaryBuffer?: GPUBuffer | null;
+}
+
+export declare function runSanaLinearAttention(
+  query: Tensor,
+  key: Tensor,
+  value: Tensor,
+  options: SanaLinearAttentionOptions
+): Promise<Tensor>;
+
+export declare function recordSanaLinearAttention(
+  recorder: CommandRecorder,
+  query: Tensor,
+  key: Tensor,
+  value: Tensor,
+  options: SanaLinearAttentionOptions
+): Promise<Tensor>;

package/src/gpu/kernels/sana_linear_attention.js

@@ -0,0 +1,121 @@
+import { getDevice } from '../device.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';
+import { WORKGROUP_SIZES } from './constants.js';
+
+function selectSanaLinearAttentionVariant(isF16) {
+  return selectRuleValue('sanaLinearAttention', 'variant', { isF16 });
+}
+
+async function runSummary(target, query, key, value, summaryBuffer, uniforms, variant) {
+  const summarySize = uniforms.num_heads * (uniforms.head_dim + 1) * uniforms.head_dim;
+  await unifiedKernelWrapper(
+    'sana_linear_attention_summary',
+    target,
+    variant,
+    [query, key, value, summaryBuffer],
+    {
+      num_heads: uniforms.num_heads,
+      head_dim: uniforms.head_dim,
+      num_tokens: uniforms.num_tokens,
+      hidden_size: uniforms.hidden_size,
+      _pad0: 0,
+      _pad1: 0,
+    },
+    Math.ceil(summarySize / WORKGROUP_SIZES.DEFAULT)
+  );
+}
+
+async function runApply(target, query, summaryBuffer, outputBuffer, uniforms, variant) {
+  await unifiedKernelWrapper(
+    'sana_linear_attention_apply',
+    target,
+    variant,
+    [query, summaryBuffer, outputBuffer],
+    {
+      num_heads: uniforms.num_heads,
+      head_dim: uniforms.head_dim,
+      num_tokens: uniforms.num_tokens,
+      hidden_size: uniforms.hidden_size,
+      eps: uniforms.eps,
+      _pad0: 0,
+      _pad1: 0,
+      _pad2: 0,
+    },
+    [Math.ceil(uniforms.hidden_size / WORKGROUP_SIZES.DEFAULT), uniforms.num_tokens, 1]
+  );
+}
+
+async function _sanaLinearAttention(target, query, key, value, options = {}) {
+  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
+  const device = target?.device || getDevice();
+  if (!device) {
+    throw new Error('SanaLinearAttention requires a WebGPU device.');
+  }
+
+  const {
+    numHeads,
+    headDim,
+    numTokens = query.shape?.[0],
+    hiddenSize = query.shape?.[1],
+    eps = 1e-15,
+    outputBuffer = null,
+    summaryBuffer = null,
+  } = options;
+
+  if (
+    !Number.isFinite(numHeads) ||
+    !Number.isFinite(headDim) ||
+    !Number.isFinite(numTokens) ||
+    !Number.isFinite(hiddenSize)
+  ) {
+    throw new Error('SanaLinearAttention requires numHeads, headDim, numTokens, and hiddenSize.');
+  }
+  if (hiddenSize !== numHeads * headDim) {
+    throw new Error(`SanaLinearAttention hiddenSize mismatch: ${hiddenSize} != ${numHeads} * ${headDim}`);
+  }
+
+  const isF16 = query.dtype === 'f16';
+  const variant = selectSanaLinearAttentionVariant(isF16);
+  const temporarySummary = summaryBuffer || acquireBuffer(
+    numHeads * (headDim + 1) * headDim * Float32Array.BYTES_PER_ELEMENT,
+    undefined,
+    'sana_linear_attention_summary'
+  );
+  const output = outputBuffer || acquireBuffer(
+    numTokens * hiddenSize * dtypeBytes(query.dtype),
+    undefined,
+    'sana_linear_attention_output'
+  );
+
+  const uniforms = {
+    num_heads: numHeads,
+    head_dim: headDim,
+    num_tokens: numTokens,
+    hidden_size: hiddenSize,
+    eps,
+  };
+
+  await runSummary(target, query, key, value, temporarySummary, uniforms, variant);
+  await runApply(target, query, temporarySummary, output, uniforms, variant);
+
+  if (!summaryBuffer) {
+    if (recorder) {
+      recorder.trackTemporaryBuffer(temporarySummary);
+    } else {
+      releaseBuffer(temporarySummary);
+    }
+  }
+
+  return createTensor(output, query.dtype, [numTokens, hiddenSize], 'sana_linear_attention_output');
+}
+
+export async function runSanaLinearAttention(query, key, value, options = {}) {
+  return _sanaLinearAttention(null, query, key, value, options);
+}
+
+export async function recordSanaLinearAttention(recorder, query, key, value, options = {}) {
+  return _sanaLinearAttention(recorder, query, key, value, options);
+}