@simulatte/doppler 0.1.5 → 0.1.6

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.js

@@ -29,7 +29,6 @@ async function runSummary(target, query, key, value, summaryBuffer, uniforms, va
 }
 
 async function runApply(target, query, summaryBuffer, outputBuffer, uniforms, variant) {
-  const outputSize = uniforms.num_tokens * uniforms.hidden_size;
   await unifiedKernelWrapper(
     'sana_linear_attention_apply',
     target,
@@ -45,7 +44,7 @@ async function runApply(target, query, summaryBuffer, outputBuffer, uniforms, va
       _pad1: 0,
       _pad2: 0,
     },
-    Math.ceil(outputSize / WORKGROUP_SIZES.DEFAULT)
+    [Math.ceil(uniforms.hidden_size / WORKGROUP_SIZES.DEFAULT), uniforms.num_tokens, 1]
   );
 }
 

package/src/gpu/kernels/sana_linear_attention_apply.wgsl

@@ -18,14 +18,13 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let total = u.num_tokens * u.hidden_size;
-    if (idx >= total) {
+    let hidden = gid.x;
+    let token = gid.y;
+    if (token >= u.num_tokens || hidden >= u.hidden_size) {
         return;
     }
 
-    let token = idx / u.hidden_size;
-    let hidden = idx - token * u.hidden_size;
+    let idx = token * u.hidden_size + hidden;
     let head = hidden / u.head_dim;
     let dim = hidden - head * u.head_dim;
     let rows_per_head = u.head_dim + 1u;

package/src/gpu/kernels/sana_linear_attention_apply_f16.wgsl

@@ -20,14 +20,13 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let total = u.num_tokens * u.hidden_size;
-    if (idx >= total) {
+    let hidden = gid.x;
+    let token = gid.y;
+    if (token >= u.num_tokens || hidden >= u.hidden_size) {
         return;
     }
 
-    let token = idx / u.hidden_size;
-    let hidden = idx - token * u.hidden_size;
+    let idx = token * u.hidden_size + hidden;
     let head = hidden / u.head_dim;
     let dim = hidden - head * u.head_dim;
     let rows_per_head = u.head_dim + 1u;

package/src/gpu/kernels/sana_linear_attention_summary.wgsl

@@ -33,6 +33,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
     var acc: f32 = 0.0;
     for (var token: u32 = 0u; token < u.num_tokens; token = token + 1u) {
+        let query_value = query[token * u.hidden_size + hidden_base + col];
         let key_idx = token * u.hidden_size + hidden_base + col;
         let key_value = max(key[key_idx], 0.0);
         let value_value = select(
@@ -40,6 +41,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
             1.0,
             row == u.head_dim
         );
+        if (u.hidden_size == 0u) {
+            acc = acc + query_value;
+        }
         acc = acc + value_value * key_value;
     }
 

package/src/gpu/kernels/sana_linear_attention_summary_f16.wgsl

@@ -35,6 +35,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
     var acc: f32 = 0.0;
     for (var token: u32 = 0u; token < u.num_tokens; token = token + 1u) {
+        let query_value = f32(query[token * u.hidden_size + hidden_base + col]);
         let key_idx = token * u.hidden_size + hidden_base + col;
         let key_value = max(f32(key[key_idx]), 0.0);
         let value_value = select(
@@ -42,6 +43,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
             1.0,
             row == u.head_dim
         );
+        if (u.hidden_size == 0u) {
+            acc = acc + query_value;
+        }
         acc = acc + value_value * key_value;
     }
 

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

@@ -16,6 +16,7 @@ export interface SiLUOptions extends OutputBufferOptions {
   size?: number | null;
   gate?: Tensor | null;
   gateActivation?: 'silu' | 'sigmoid';
+  inputActivation?: 'silu' | 'identity';
   useVec4?: boolean;
   biasOffset?: number;
   swigluLimit: number | null;

package/src/gpu/kernels/silu.js

@@ -47,6 +47,18 @@ function createSiLUBindGroupEntries(uniformBuffer, input, output, gate) {
   ];
 }
 
+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 +72,7 @@ export async function runSiLU(
     useVec4 = false,
     swigluLimit,
     gateActivation = 'silu',
+    inputActivation = 'silu',
   } = options;
   const resolvedSwigluLimit = resolveSwigluLimit(swigluLimit, 'SiLU');
 
@@ -74,14 +87,17 @@ 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 dispatchPlan = planSiLUDispatch(device, inferredSize, useVec4);
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -89,7 +105,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);
     },
@@ -106,8 +122,7 @@ export async function runSiLU(
     entries,
   });
 
-  const workgroups = Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'silu');
+  dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
 
   uniformBuffer.destroy();
 
@@ -215,7 +230,7 @@ export async function runSiLURowSplit(
     ],
   });
 
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+  const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
   dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
 
   uniformBuffer.destroy();
@@ -269,7 +284,7 @@ export async function recordSiLURowSplit(
     ],
   });
 
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+  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');
@@ -288,6 +303,7 @@ export async function recordSiLU(
     outputBuffer = null,
     swigluLimit,
     gateActivation = 'silu',
+    inputActivation = 'silu',
   } = options;
   const resolvedSwigluLimit = resolveSwigluLimit(swigluLimit, 'SiLU');
 
@@ -302,14 +318,17 @@ 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 dispatchPlan = planSiLUDispatch(device, inferredSize, false);
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -317,7 +336,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);
     },
@@ -333,8 +352,7 @@ export async function recordSiLU(
     entries,
   });
 
-  const workgroups = Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT);
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu');
+  recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
 
   return createTensor(output, input.dtype, [inferredSize], 'silu_output');
 }

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