@simulatte/doppler 0.1.8 → 0.2.0

package/src/gpu/kernels/attention_decode.wgsl

@@ -1,233 +0,0 @@
-// Fused Multi-Head Attention Kernel
-//
-// Implements fused Q @ K^T → scale → mask → softmax → @ V
-// Uses tiled/blocked approach to avoid materializing full attention matrix.
-// Supports grouped query attention (GQA) where numKVHeads < numHeads.
-//
-// Based on Flash Attention principles adapted for WebGPU.
-
-const MAX_HEAD_DIM: u32 = 64u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-override DECODE_WORKGROUP_SIZE: u32 = 256u;  // Decode kernel workgroup size (must match MAX_WORKGROUP_SIZE)
-
-struct Uniforms {
-    num_heads: u32,       // Number of query heads
-    num_kv_heads: u32,    // Number of KV heads (for GQA)
-    head_dim: u32,        // Dimension per head
-    seq_len: u32,         // Current sequence length (for KV)
-    query_len: u32,       // Query length (1 for decode, seq_len for prefill)
-    scale: f32,           // 1/sqrt(head_dim)
-    is_causal: u32,       // Apply causal mask (1 = yes)
-    start_pos: u32,       // Absolute position offset for causal masking
-    attn_softcap: f32,    // Gemma 2: 50.0, 0 = disabled
-    sliding_window: u32,  // Sliding window size (0 = disabled, >0 = window size)
-    kv_len_source: u32,   // 0 = use uniform seq_len, 1 = use buffer
-    kv_start: u32,
-    page_size: u32,
-    kv_layout: u32,
-    _pad: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> Q: array<f32>;       // [query_len, num_heads, head_dim]
-@group(0) @binding(2) var<storage, read> K: array<f32>;       // [seq_len, num_kv_heads, head_dim]
-@group(0) @binding(3) var<storage, read> V: array<f32>;       // [seq_len, num_kv_heads, head_dim]
-@group(0) @binding(4) var<storage, read_write> output: array<f32>; // [query_len, num_heads, head_dim]
-@group(0) @binding(5) var<storage, read> kv_len_buffer: array<u32>;
-@group(0) @binding(6) var<storage, read> page_table: array<u32>;
-
-// Online softmax accumulators (per-thread)
-// Sized for 256 to support attention_decode workgroup size
-var<workgroup> row_max: array<f32, MAX_WORKGROUP_SIZE>;
-var<workgroup> row_sum: array<f32, MAX_WORKGROUP_SIZE>;
-
-// Get KV head index for grouped query attention
-fn get_kv_head_idx(query_head_idx: u32) -> u32 {
-    // GQA: multiple query heads share one KV head
-    let heads_per_kv = u.num_heads / u.num_kv_heads;
-    return query_head_idx / heads_per_kv;
-}
-
-// Check if position should be masked (causal + sliding window attention)
-fn is_masked(query_pos: u32, key_pos: u32) -> bool {
-    // Compute absolute positions
-    let abs_query = query_pos + u.start_pos;
-    let abs_key = u.kv_start + key_pos;
-
-    // Causal mask: query can only attend to keys at same or earlier positions
-    if (u.is_causal != 0u && abs_key > abs_query) {
-        return true;
-    }
-
-    // Sliding window mask: query can only attend to keys within the window
-    // Key must be >= (query - window + 1), i.e., within the last `window` positions
-    if (u.sliding_window > 0u && abs_query >= u.sliding_window) {
-        if (abs_key < abs_query - u.sliding_window + 1u) {
-            return true;  // Key is too far in the past
-        }
-    }
-
-    return false;
-}
-
-fn get_kv_pos(key_pos: u32) -> u32 {
-    let abs_key = u.kv_start + key_pos;
-    if (u.kv_layout == 1u && u.sliding_window > 0u) {
-        return abs_key % u.sliding_window;
-    }
-    if (u.kv_layout == 2u) {
-        let page_idx = abs_key / u.page_size;
-        let in_page = abs_key - (page_idx * u.page_size);
-        let phys_page = page_table[page_idx];
-        return phys_page * u.page_size + in_page;
-    }
-    return abs_key;
-}
-
-fn get_kv_len() -> u32 {
-    if (u.kv_len_source == 0u) {
-        return u.seq_len;
-    }
-    return kv_len_buffer[0];
-}
-
-// Main attention kernel - one workgroup per (query_block, head)
-// Workgroups are dispatched linearly as: numQueryBlocks * numHeads.
-// head_idx and query_block_idx are derived from workgroup_id.x.
-
-// Simplified single-query attention for decode step
-// More efficient when query_len == 1
-@compute @workgroup_size(DECODE_WORKGROUP_SIZE, 1, 1)
-fn attention_decode(
-    @builtin(global_invocation_id) global_id: vec3<u32>,
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) wg_id: vec3<u32>
-) {
-    let head_idx = wg_id.x;
-    let thread_idx = local_id.x;
-
-    if (DECODE_WORKGROUP_SIZE > MAX_WORKGROUP_SIZE) {
-        return;
-    }
-
-    let kv_head_idx = get_kv_head_idx(head_idx);
-    let head_dim = u.head_dim;
-    if (head_dim > MAX_HEAD_DIM) {
-        return;
-    }
-    let seq_len = get_kv_len();
-    let scale = u.scale;
-
-    // Each thread handles a subset of key positions
-    let keys_per_thread = (seq_len + DECODE_WORKGROUP_SIZE - 1u) / DECODE_WORKGROUP_SIZE;
-
-    // Load query (single position)
-    var q_local: array<f32, MAX_HEAD_DIM>;
-    let q_offset = head_idx * head_dim;
-    for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-        q_local[d] = Q[q_offset + d];
-    }
-
-    // Pass 1: local max
-    var local_max: f32 = -3.402823e+38;
-    let query_pos = 0u;
-
-    for (var i: u32 = 0u; i < keys_per_thread; i = i + 1u) {
-        let key_pos = thread_idx * keys_per_thread + i;
-        if (key_pos >= seq_len) { break; }
-        if (is_masked(query_pos, key_pos)) { continue; }
-
-        let k_idx = get_kv_pos(key_pos);
-        let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
-
-        var score: f32 = 0.0;
-        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-            score = score + q_local[d] * K[k_offset + d];
-        }
-        score = score * scale;
-
-        // Gemma 2 attention softcapping
-        if (u.attn_softcap > 0.0) {
-            score = tanh(score / u.attn_softcap) * u.attn_softcap;
-        }
-
-        local_max = max(local_max, score);
-    }
-
-    // Store local max for reduction
-    row_max[thread_idx] = local_max;
-    workgroupBarrier();
-
-    // Parallel reduction to find global max
-    for (var stride: u32 = DECODE_WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
-        if (thread_idx < stride && thread_idx + stride < DECODE_WORKGROUP_SIZE) {
-            row_max[thread_idx] = max(row_max[thread_idx], row_max[thread_idx + stride]);
-        }
-        workgroupBarrier();
-    }
-
-    let global_max = row_max[0];
-
-    // Pass 2: local sum and output accumulator
-    var local_sum: f32 = 0.0;
-    var local_out: array<f32, MAX_HEAD_DIM>;
-    for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-        local_out[d] = 0.0;
-    }
-
-    for (var i: u32 = 0u; i < keys_per_thread; i = i + 1u) {
-        let key_pos = thread_idx * keys_per_thread + i;
-        if (key_pos >= seq_len) { break; }
-        if (is_masked(query_pos, key_pos)) { continue; }
-
-        let k_offset = key_pos * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
-        let v_idx = get_kv_pos(key_pos);
-        let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
-
-        var score: f32 = 0.0;
-        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-            score = score + q_local[d] * K[k_offset + d];
-        }
-        score = score * scale;
-
-        if (u.attn_softcap > 0.0) {
-            score = tanh(score / u.attn_softcap) * u.attn_softcap;
-        }
-
-        let exp_score = exp(score - global_max);
-        local_sum = local_sum + exp_score;
-        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-            local_out[d] = local_out[d] + exp_score * V[v_offset + d];
-        }
-    }
-
-    row_sum[thread_idx] = local_sum;
-    workgroupBarrier();
-
-    // Parallel reduction for sum
-    for (var stride: u32 = DECODE_WORKGROUP_SIZE / 2u; stride > 0u; stride = stride >> 1u) {
-        if (thread_idx < stride && thread_idx + stride < DECODE_WORKGROUP_SIZE) {
-            row_sum[thread_idx] = row_sum[thread_idx] + row_sum[thread_idx + stride];
-        }
-        workgroupBarrier();
-    }
-
-    let global_sum = row_sum[0];
-    let inv_sum = select(0.0, 1.0 / global_sum, global_sum > 0.0);
-
-    // Reduction for output (per-dimension)
-    let out_offset = head_idx * head_dim;
-    for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
-        row_sum[thread_idx] = local_out[d];
-        workgroupBarrier();
-        if (thread_idx == 0u) {
-            var sum: f32 = 0.0;
-            for (var t: u32 = 0u; t < DECODE_WORKGROUP_SIZE; t = t + 1u) {
-                sum = sum + row_sum[t];
-            }
-            output[out_offset + d] = sum * inv_sum;
-        }
-        workgroupBarrier();
-    }
-}

package/src/gpu/kernels/attention_decode_chunked_f16.wgsl

@@ -1,183 +0,0 @@
-// Chunked Decode Attention Kernel (f16 QKV + f16 output)
-//
-// Same as attention_decode_chunked_f16kv.wgsl but Q/K/V are f16 and output is f16.
-
-enable f16;
-
-const MAX_KV_LEN: u32 = 2048u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    num_heads: u32,
-    num_kv_heads: u32,
-    head_dim: u32,
-    kv_len: u32,
-    seq_len: u32,
-    scale: f32,
-    is_causal: u32,
-    start_pos: u32,
-    attn_softcap: f32,
-    sliding_window: u32,
-    kv_len_source: u32,
-    kv_start: u32,
-    page_size: u32,
-    kv_layout: u32,
-    _pad: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> Q: array<f16>;
-@group(0) @binding(2) var<storage, read> K: array<f16>;
-@group(0) @binding(3) var<storage, read> V: array<f16>;
-@group(0) @binding(4) var<storage, read_write> output: array<f16>;
-@group(0) @binding(5) var<storage, read> kv_len_buffer: array<u32>;
-@group(0) @binding(6) var<storage, read> page_table: array<u32>;
-
-var<workgroup> shared_scores: array<f32, MAX_KV_LEN>;
-var<workgroup> shared_partial: array<f32, MAX_WORKGROUP_SIZE>;
-var<workgroup> shared_max: f32;
-var<workgroup> shared_sum: f32;
-var<workgroup> shared_acc: array<f32, MAX_WORKGROUP_SIZE>;
-
-fn get_kv_head_idx(query_head_idx: u32) -> u32 {
-    let heads_per_kv = u.num_heads / u.num_kv_heads;
-    return query_head_idx / heads_per_kv;
-}
-
-fn is_masked(key_pos: u32) -> bool {
-    let abs_query = u.start_pos;
-    let abs_key = u.kv_start + key_pos;
-    if (u.is_causal != 0u && abs_key > abs_query) { return true; }
-    if (u.sliding_window > 0u && abs_query >= u.sliding_window) {
-        if (abs_key < abs_query - u.sliding_window + 1u) { return true; }
-    }
-    return false;
-}
-
-fn get_kv_pos(key_pos: u32) -> u32 {
-    let abs_key = u.kv_start + key_pos;
-    if (u.kv_layout == 1u && u.sliding_window > 0u) {
-        return abs_key % u.sliding_window;
-    }
-    if (u.kv_layout == 2u) {
-        let page_idx = abs_key / u.page_size;
-        let in_page = abs_key - (page_idx * u.page_size);
-        let phys_page = page_table[page_idx];
-        return phys_page * u.page_size + in_page;
-    }
-    return abs_key;
-}
-
-fn get_kv_len() -> u32 {
-    if (u.kv_len_source == 0u) {
-        return u.kv_len;
-    }
-    return kv_len_buffer[0];
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) workgroup_id: vec3<u32>,
-) {
-    let head_idx = workgroup_id.x;
-    let tid = local_id.x;
-    let head_dim = u.head_dim;
-    let kv_len = get_kv_len();
-    let scale = u.scale;
-    let kv_head_idx = get_kv_head_idx(head_idx);
-
-    if (WORKGROUP_SIZE > MAX_WORKGROUP_SIZE || head_dim > WORKGROUP_SIZE || kv_len > MAX_KV_LEN) {
-        return;
-    }
-
-    let valid = tid < head_dim;
-
-    var q_val: f32 = 0.0;
-    if (valid) {
-        let q_offset = head_idx * head_dim + tid;
-        q_val = f32(Q[q_offset]);
-    }
-
-    if (valid) {
-        shared_acc[tid] = 0.0;
-    }
-
-    var start_k: u32 = 0u;
-    if (u.sliding_window > 0u && kv_len > u.sliding_window) {
-        start_k = kv_len - u.sliding_window;
-    }
-
-    for (var k_pos: u32 = start_k; k_pos < kv_len; k_pos++) {
-        var k_val: f32 = 0.0;
-        if (valid) {
-            let k_idx = get_kv_pos(k_pos);
-            let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
-            k_val = f32(K[k_offset]);
-        }
-
-        let partial = q_val * k_val;
-        shared_partial[tid] = partial;
-        workgroupBarrier();
-
-        if (tid == 0u) {
-            var dot: f32 = 0.0;
-            for (var d: u32 = 0u; d < head_dim; d++) {
-                dot += shared_partial[d];
-            }
-            var s = dot * scale;
-            if (u.attn_softcap > 0.0) {
-                s = tanh(s / u.attn_softcap) * u.attn_softcap;
-            }
-            if (is_masked(k_pos)) {
-                s = -3.402823e+38;
-            }
-            shared_scores[k_pos] = s;
-        }
-        workgroupBarrier();
-    }
-
-    if (tid == 0u) {
-        var max_score: f32 = -3.402823e+38;
-        for (var k: u32 = start_k; k < kv_len; k++) {
-            max_score = max(max_score, shared_scores[k]);
-        }
-        shared_max = max_score;
-    }
-    workgroupBarrier();
-
-    let max_score = shared_max;
-
-    if (tid == 0u) {
-        var sum_exp: f32 = 0.0;
-        for (var k: u32 = start_k; k < kv_len; k++) {
-            let w = exp(shared_scores[k] - max_score);
-            shared_scores[k] = w;
-            sum_exp += w;
-        }
-        shared_sum = sum_exp;
-    }
-    workgroupBarrier();
-
-    let sum_exp = shared_sum;
-    let inv_sum = select(0.0, 1.0 / sum_exp, sum_exp > 0.0);
-
-    if (valid) {
-        var acc: f32 = 0.0;
-        for (var k_pos: u32 = start_k; k_pos < kv_len; k_pos++) {
-            let v_idx = get_kv_pos(k_pos);
-            let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
-            let v_val = f32(V[v_offset]);
-            acc += shared_scores[k_pos] * v_val;
-        }
-        shared_acc[tid] = acc * inv_sum;
-    }
-    workgroupBarrier();
-
-    if (valid) {
-        let out_offset = head_idx * head_dim + tid;
-        output[out_offset] = f16(shared_acc[tid]);
-    }
-}

package/src/gpu/kernels/attention_decode_chunked_f16kv.wgsl

@@ -1,208 +0,0 @@
-// Chunked Decode Attention Kernel (f16 KV)
-//
-// Optimized for models with few heads and large headDim (e.g., Gemma 3: 4 heads × 256 dim)
-// Uses parallel reduction across headDim within each workgroup.
-//
-// Dispatch: numHeads workgroups, each with 256 threads
-// Each thread handles one dimension of the head.
-
-enable f16;
-
-const MAX_KV_LEN: u32 = 2048u;
-const MAX_WORKGROUP_SIZE: u32 = 256u;
-
-// Workgroup size for decode
-override WORKGROUP_SIZE: u32 = 256u;  // Must match MAX_WORKGROUP_SIZE
-
-struct Uniforms {
-    num_heads: u32,
-    num_kv_heads: u32,
-    head_dim: u32,
-    kv_len: u32,       // Current KV cache length
-    seq_len: u32,      // Query length (1 for decode)
-    scale: f32,
-    is_causal: u32,
-    start_pos: u32,
-    attn_softcap: f32,    // Gemma 2: 50.0, 0 = disabled
-    sliding_window: u32,  // Sliding window size (0 = disabled, >0 = window size)
-    kv_len_source: u32,   // 0 = use uniform kv_len, 1 = use buffer
-    kv_start: u32,
-    page_size: u32,
-    kv_layout: u32,
-    _pad: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> Q: array<f32>;
-@group(0) @binding(2) var<storage, read> K: array<f16>;
-@group(0) @binding(3) var<storage, read> V: array<f16>;
-@group(0) @binding(4) var<storage, read_write> output: array<f32>;
-@group(0) @binding(5) var<storage, read> kv_len_buffer: array<u32>;
-@group(0) @binding(6) var<storage, read> page_table: array<u32>;
-
-// Shared memory for parallel reductions
-var<workgroup> shared_scores: array<f32, MAX_KV_LEN>;
-var<workgroup> shared_partial: array<f32, MAX_WORKGROUP_SIZE>;
-var<workgroup> shared_max: f32;
-var<workgroup> shared_sum: f32;
-var<workgroup> shared_acc: array<f32, MAX_WORKGROUP_SIZE>;
-
-fn get_kv_head_idx(query_head_idx: u32) -> u32 {
-    let heads_per_kv = u.num_heads / u.num_kv_heads;
-    return query_head_idx / heads_per_kv;
-}
-
-// Check if position should be masked (sliding window attention)
-// For decode, query_pos = start_pos (we're at the latest position)
-fn is_masked(key_pos: u32) -> bool {
-    let abs_query = u.start_pos;
-    let abs_key = u.kv_start + key_pos;
-    // Causal mask (key must be <= query position)
-    if (u.is_causal != 0u && abs_key > abs_query) { return true; }
-    // Sliding window mask
-    if (u.sliding_window > 0u && abs_query >= u.sliding_window) {
-        if (abs_key < abs_query - u.sliding_window + 1u) { return true; }
-    }
-    return false;
-}
-
-fn get_kv_pos(key_pos: u32) -> u32 {
-    let abs_key = u.kv_start + key_pos;
-    if (u.kv_layout == 1u && u.sliding_window > 0u) {
-        return abs_key % u.sliding_window;
-    }
-    if (u.kv_layout == 2u) {
-        let page_idx = abs_key / u.page_size;
-        let in_page = abs_key - (page_idx * u.page_size);
-        let phys_page = page_table[page_idx];
-        return phys_page * u.page_size + in_page;
-    }
-    return abs_key;
-}
-
-fn get_kv_len() -> u32 {
-    if (u.kv_len_source == 0u) {
-        return u.kv_len;
-    }
-    return kv_len_buffer[0];
-}
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(local_invocation_id) local_id: vec3<u32>,
-    @builtin(workgroup_id) workgroup_id: vec3<u32>,
-) {
-    let head_idx = workgroup_id.x;
-    let tid = local_id.x;
-    let head_dim = u.head_dim;
-    let kv_len = get_kv_len();
-    let scale = u.scale;
-    let kv_head_idx = get_kv_head_idx(head_idx);
-
-    if (WORKGROUP_SIZE > MAX_WORKGROUP_SIZE || head_dim > WORKGROUP_SIZE || kv_len > MAX_KV_LEN) {
-        return;
-    }
-
-    // Only threads < headDim participate in computation
-    let valid = tid < head_dim;
-
-    // Load Q value for this thread's dimension
-    var q_val: f32 = 0.0;
-    if (valid) {
-        let q_offset = head_idx * head_dim + tid;
-        q_val = Q[q_offset];
-    }
-
-    // Initialize accumulator
-    if (valid) {
-        shared_acc[tid] = 0.0;
-    }
-
-    var start_k: u32 = 0u;
-    if (u.sliding_window > 0u && kv_len > u.sliding_window) {
-        start_k = kv_len - u.sliding_window;
-    }
-
-    // Phase 1: Compute attention scores (Q @ K^T) in parallel
-    // Each thread computes partial dot product for its dimension
-    for (var k_pos: u32 = start_k; k_pos < kv_len; k_pos++) {
-        var k_val: f32 = 0.0;
-        if (valid) {
-            let k_idx = get_kv_pos(k_pos);
-            let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
-            k_val = f32(K[k_offset]);
-        }
-
-        // Partial dot product
-        let partial = q_val * k_val;
-        shared_partial[tid] = partial;
-        workgroupBarrier();
-
-        // Parallel reduction for dot product (thread 0 accumulates)
-        // Using sequential reduction for correctness - could optimize with tree reduction
-        if (tid == 0u) {
-            var dot: f32 = 0.0;
-            for (var d: u32 = 0u; d < head_dim; d++) {
-                dot += shared_partial[d];
-            }
-            var s = dot * scale;
-            // Gemma 2 attention softcapping
-            if (u.attn_softcap > 0.0) {
-                s = tanh(s / u.attn_softcap) * u.attn_softcap;
-            }
-            // Apply masking (causal + sliding window)
-            if (is_masked(k_pos)) {
-                s = -3.402823e+38;  // -FLT_MAX
-            }
-            shared_scores[k_pos] = s;
-        }
-        workgroupBarrier();
-    }
-
-    // Phase 2: Find max score for numerical stability
-    if (tid == 0u) {
-        var max_score: f32 = -3.402823e+38;
-        for (var k: u32 = start_k; k < kv_len; k++) {
-            max_score = max(max_score, shared_scores[k]);
-        }
-        shared_max = max_score;
-    }
-    workgroupBarrier();
-
-    let max_score = shared_max;
-
-    // Phase 3: Compute softmax weights and accumulate V
-    if (tid == 0u) {
-        var sum_exp: f32 = 0.0;
-        for (var k: u32 = start_k; k < kv_len; k++) {
-            let w = exp(shared_scores[k] - max_score);
-            shared_scores[k] = w;  // Store weight for V accumulation
-            sum_exp += w;
-        }
-        shared_sum = sum_exp;
-    }
-    workgroupBarrier();
-
-    let sum_exp = shared_sum;
-    // Use select instead of early return to maintain uniform control flow
-    let inv_sum = select(0.0, 1.0 / sum_exp, sum_exp > 0.0);
-
-    // Phase 4: Weighted sum of V (parallel across dimensions)
-    if (valid) {
-        var acc: f32 = 0.0;
-        for (var k_pos: u32 = start_k; k_pos < kv_len; k_pos++) {
-            let v_idx = get_kv_pos(k_pos);
-            let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
-            let v_val = f32(V[v_offset]);
-            acc += shared_scores[k_pos] * v_val;
-        }
-        shared_acc[tid] = acc * inv_sum;
-    }
-    workgroupBarrier();
-
-    // Phase 5: Write output
-    if (valid) {
-        let out_offset = head_idx * head_dim + tid;
-        output[out_offset] = shared_acc[tid];
-    }
-}