@simulatte/doppler 0.1.0

package/src/gpu/kernels/attention_decode_tiered_int8_f16kv.wgsl

@@ -0,0 +1,242 @@
+// Tiered Decode Attention Kernel (int8 cold, f16 hot, f32 output)
+
+enable f16;
+
+const MAX_KV_LEN: u32 = 2048u;
+const MAX_WORKGROUP_SIZE: u32 = 256u;
+const PACK_FACTOR: u32 = 4u;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    num_heads: u32,
+    num_kv_heads: u32,
+    head_dim: u32,
+    cold_len: u32,
+    hot_len: u32,
+    seq_len: u32,
+    scale: f32,
+    is_causal: u32,
+    start_pos: u32,
+    attn_softcap: f32,
+    sliding_window: u32,
+    hot_window: u32,
+    hot_start: u32,
+    packed_stride: 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_hot: array<f16>;
+@group(0) @binding(3) var<storage, read> V_hot: array<f16>;
+@group(0) @binding(4) var<storage, read> K_cold: array<u32>;
+@group(0) @binding(5) var<storage, read> V_cold: array<u32>;
+@group(0) @binding(6) var<storage, read> scales_k: array<f16>;
+@group(0) @binding(7) var<storage, read> scales_v: array<f16>;
+@group(0) @binding(8) var<storage, read_write> output: array<f32>;
+
+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(abs_key: u32) -> bool {
+    let abs_query = u.start_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_hot_pos(hot_pos: u32) -> u32 {
+    if (u.hot_window > 0u) {
+        let abs_key = u.hot_start + hot_pos;
+        return abs_key % u.hot_window;
+    }
+    return hot_pos;
+}
+
+fn read_cold_k(
+    key_pos: u32,
+    kv_head_idx: u32,
+    tid: u32
+) -> f32 {
+    let pack_idx = tid / PACK_FACTOR;
+    if (pack_idx >= u.packed_stride) {
+        return 0.0;
+    }
+    let lane = tid - pack_idx * PACK_FACTOR;
+    let base = (key_pos * u.num_kv_heads + kv_head_idx) * u.packed_stride + pack_idx;
+    let packed_val = K_cold[base];
+    let raw = (packed_val >> (8u * lane)) & 0xffu;
+    var q = i32(raw);
+    if (q >= 128) { q -= 256; }
+    let scale_idx = key_pos * u.num_kv_heads + kv_head_idx;
+    let scale = f32(scales_k[scale_idx]);
+    return f32(q) * scale;
+}
+
+fn read_cold_v(
+    key_pos: u32,
+    kv_head_idx: u32,
+    tid: u32
+) -> f32 {
+    let pack_idx = tid / PACK_FACTOR;
+    if (pack_idx >= u.packed_stride) {
+        return 0.0;
+    }
+    let lane = tid - pack_idx * PACK_FACTOR;
+    let base = (key_pos * u.num_kv_heads + kv_head_idx) * u.packed_stride + pack_idx;
+    let packed_val = V_cold[base];
+    let raw = (packed_val >> (8u * lane)) & 0xffu;
+    var q = i32(raw);
+    if (q >= 128) { q -= 256; }
+    let scale_idx = key_pos * u.num_kv_heads + kv_head_idx;
+    let scale = f32(scales_v[scale_idx]);
+    return f32(q) * scale;
+}
+
+@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 cold_len = u.cold_len;
+    let hot_len = u.hot_len;
+    let total_len = cold_len + hot_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 || total_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 = Q[q_offset];
+    }
+
+    if (valid) {
+        shared_acc[tid] = 0.0;
+    }
+
+    for (var k_pos: u32 = 0u; k_pos < cold_len; k_pos++) {
+        var k_val: f32 = 0.0;
+        if (valid) {
+            k_val = read_cold_k(k_pos, kv_head_idx, tid);
+        }
+
+        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();
+    }
+
+    for (var k_pos: u32 = 0u; k_pos < hot_len; k_pos++) {
+        var k_val: f32 = 0.0;
+        if (valid) {
+            let k_idx = get_hot_pos(k_pos);
+            let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
+            k_val = f32(K_hot[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;
+            }
+            let score_idx = cold_len + k_pos;
+            let abs_key = u.hot_start + k_pos;
+            if (is_masked(abs_key)) {
+                s = -3.402823e+38;
+            }
+            shared_scores[score_idx] = s;
+        }
+        workgroupBarrier();
+    }
+
+    if (tid == 0u) {
+        var max_score: f32 = -3.402823e+38;
+        for (var k: u32 = 0u; k < total_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 = 0u; k < total_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 score_idx: u32 = 0u; score_idx < total_len; score_idx++) {
+            var v_val: f32 = 0.0;
+            if (score_idx < cold_len) {
+                v_val = read_cold_v(score_idx, kv_head_idx, tid);
+            } else {
+                let hot_pos = score_idx - cold_len;
+                let v_idx = get_hot_pos(hot_pos);
+                let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + tid;
+                v_val = f32(V_hot[v_offset]);
+            }
+            acc += shared_scores[score_idx] * v_val;
+        }
+        shared_acc[tid] = acc * inv_sum;
+    }
+    workgroupBarrier();
+
+    if (valid) {
+        let out_offset = head_idx * head_dim + tid;
+        output[out_offset] = shared_acc[tid];
+    }
+}

package/src/gpu/kernels/attention_f16.wgsl

@@ -0,0 +1,214 @@
+// AUTO-GENERATED from src/gpu/kernels/attention.wgsl.
+// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
+// Fused Multi-Head Attention Kernel (f16 QKV + f16 output)
+//
+// Same algorithm as attention.wgsl but Q/K/V are f16 and output is f16.
+// Computation is performed in f32 for stability.
+
+enable f16;
+
+// Tile sizes for blocked attention
+const MAX_BLOCK_SIZE: u32 = 64u;
+const MAX_HEAD_TILE: u32 = 64u;
+const MAX_HEAD_DIM: u32 = 64u;
+
+override BLOCK_SIZE: u32 = 64u;
+override HEAD_TILE: u32 = 64u;
+override WORKGROUP_SIZE: u32 = 64u;
+
+struct Uniforms {
+    num_heads: u32,
+    num_kv_heads: u32,
+    head_dim: u32,
+    seq_len: u32,
+    query_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_K: array<f32, MAX_BLOCK_SIZE * MAX_HEAD_TILE>;
+var<workgroup> shared_V: array<f32, MAX_BLOCK_SIZE * MAX_HEAD_TILE>;
+var<workgroup> shared_scores: array<f32, MAX_BLOCK_SIZE * MAX_BLOCK_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(query_pos: u32, key_pos: u32) -> bool {
+    let abs_query = query_pos + 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.seq_len;
+    }
+    return kv_len_buffer[0];
+}
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) wg_id: vec3<u32>
+) {
+    let linear = wg_id.x;
+    let head_idx = linear % u.num_heads;
+    let query_block_idx = linear / u.num_heads;
+    let thread_idx = local_id.x;
+
+    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 query_len = u.query_len;
+    let scale = u.scale;
+
+    let query_pos = query_block_idx * BLOCK_SIZE + thread_idx;
+    let valid_query = query_pos < query_len;
+
+    var m_i: f32 = -3.402823e+38;
+    var l_i: f32 = 0.0;
+    var acc: array<f32, MAX_HEAD_DIM>;
+
+    for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+        acc[d] = 0.0;
+    }
+
+    var q_local: array<f32, MAX_HEAD_DIM>;
+    if (valid_query) {
+        let q_offset = query_pos * u.num_heads * head_dim + head_idx * head_dim;
+        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+            q_local[d] = f32(Q[q_offset + d]);
+        }
+    }
+
+    let num_kv_blocks = (seq_len + BLOCK_SIZE - 1u) / BLOCK_SIZE;
+
+    for (var kv_block: u32 = 0u; kv_block < num_kv_blocks; kv_block = kv_block + 1u) {
+        let kv_block_start = kv_block * BLOCK_SIZE;
+
+        let k_load_idx = kv_block_start + thread_idx;
+        if (k_load_idx < seq_len) {
+            let k_idx = get_kv_pos(k_load_idx);
+            let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_K[thread_idx * head_dim + d] = f32(K[k_offset + d]);
+            }
+        } else {
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_K[thread_idx * head_dim + d] = 0.0;
+            }
+        }
+
+        let v_load_idx = kv_block_start + thread_idx;
+        if (v_load_idx < seq_len) {
+            let v_idx = get_kv_pos(v_load_idx);
+            let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_V[thread_idx * head_dim + d] = f32(V[v_offset + d]);
+            }
+        } else {
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_V[thread_idx * head_dim + d] = 0.0;
+            }
+        }
+
+        workgroupBarrier();
+
+        if (valid_query) {
+            var block_max: f32 = -3.402823e+38;
+
+            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
+                let key_pos = kv_block_start + k;
+                if (key_pos >= seq_len) { continue; }
+                if (is_masked(query_pos, key_pos)) { continue; }
+
+                var score: f32 = 0.0;
+                for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                    score = score + q_local[d] * shared_K[k * head_dim + d];
+                }
+                score = score * scale;
+
+                if (u.attn_softcap > 0.0) {
+                    score = tanh(score / u.attn_softcap) * u.attn_softcap;
+                }
+
+                block_max = max(block_max, score);
+                shared_scores[thread_idx * BLOCK_SIZE + k] = score;
+            }
+
+            let m_new = max(m_i, block_max);
+            let correction = exp(m_i - m_new);
+
+            l_i = l_i * correction;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                acc[d] = acc[d] * correction;
+            }
+
+            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
+                let key_pos = kv_block_start + k;
+                if (key_pos >= seq_len) { continue; }
+                if (is_masked(query_pos, key_pos)) { continue; }
+
+                let score = shared_scores[thread_idx * BLOCK_SIZE + k];
+                let p = exp(score - m_new);
+                l_i = l_i + p;
+
+                for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                    acc[d] = acc[d] + p * shared_V[k * head_dim + d];
+                }
+            }
+
+            m_i = m_new;
+        }
+
+        workgroupBarrier();
+    }
+
+    if (valid_query) {
+        let out_offset = query_pos * u.num_heads * head_dim + head_idx * head_dim;
+        let inv_l_i = select(0.0, 1.0 / l_i, l_i > 0.0);
+        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+            output[out_offset + d] = f16(acc[d] * inv_l_i);
+        }
+    }
+}

package/src/gpu/kernels/attention_f16kv.wgsl

@@ -0,0 +1,242 @@
+// AUTO-GENERATED from src/gpu/kernels/attention.wgsl.
+// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
+// Fused Multi-Head Attention Kernel (f16 KV)
+//
+// Same as attention.wgsl but K/V are stored as f16 for KV-cache compression.
+// Q and output remain f32; K/V are cast to f32 on load.
+
+enable f16;
+
+// Tile sizes for blocked attention
+const MAX_BLOCK_SIZE: u32 = 64u;
+const MAX_HEAD_TILE: u32 = 64u;
+const MAX_HEAD_DIM: u32 = 64u;
+
+override BLOCK_SIZE: u32 = 64u;       // Sequence tile size (must match WORKGROUP_SIZE)
+override HEAD_TILE: u32 = 64u;        // Head dimension tile
+override WORKGROUP_SIZE: u32 = 64u;   // One thread per query position
+
+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<f16>;       // [seq_len, num_kv_heads, head_dim]
+@group(0) @binding(3) var<storage, read> V: array<f16>;       // [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>;
+
+// Shared memory for tiled computation
+var<workgroup> shared_K: array<f32, MAX_BLOCK_SIZE * MAX_HEAD_TILE>;
+var<workgroup> shared_V: array<f32, MAX_BLOCK_SIZE * MAX_HEAD_TILE>;
+var<workgroup> shared_scores: array<f32, MAX_BLOCK_SIZE * MAX_BLOCK_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 {
+    let abs_query = query_pos + u.start_pos;
+    let abs_key = u.kv_start + key_pos;
+    // Causal mask
+    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.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.
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(
+    @builtin(global_invocation_id) global_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>,
+    @builtin(workgroup_id) wg_id: vec3<u32>
+) {
+    let linear = wg_id.x;
+    let head_idx = linear % u.num_heads;
+    let query_block_idx = linear / u.num_heads;
+    let thread_idx = local_id.x;
+
+    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 query_len = u.query_len;
+    let scale = u.scale;
+
+    // Query position this thread handles
+    let query_pos = query_block_idx * BLOCK_SIZE + thread_idx;
+    let valid_query = query_pos < query_len;
+
+    // Initialize online softmax accumulators
+    var m_i: f32 = -3.402823e+38;  // -inf for max tracking
+    var l_i: f32 = 0.0;            // Sum of exp(x - max)
+    var acc: array<f32, MAX_HEAD_DIM>;
+
+    // Initialize accumulator
+    for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+        acc[d] = 0.0;
+    }
+
+    // Load query for this thread into registers
+    var q_local: array<f32, MAX_HEAD_DIM>;
+    if (valid_query) {
+        let q_offset = query_pos * u.num_heads * head_dim + head_idx * head_dim;
+        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+            q_local[d] = Q[q_offset + d];
+        }
+    }
+
+    // Process key-value blocks
+    let num_kv_blocks = (seq_len + BLOCK_SIZE - 1u) / BLOCK_SIZE;
+
+    for (var kv_block: u32 = 0u; kv_block < num_kv_blocks; kv_block = kv_block + 1u) {
+        let kv_block_start = kv_block * BLOCK_SIZE;
+
+        // Collaborative load of K block into shared memory
+        let k_load_idx = kv_block_start + thread_idx;
+        if (k_load_idx < seq_len) {
+            let k_idx = get_kv_pos(k_load_idx);
+            let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_K[thread_idx * head_dim + d] = f32(K[k_offset + d]);
+            }
+        } else {
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_K[thread_idx * head_dim + d] = 0.0;
+            }
+        }
+
+        // Load V block
+        let v_load_idx = kv_block_start + thread_idx;
+        if (v_load_idx < seq_len) {
+            let v_idx = get_kv_pos(v_load_idx);
+            let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_V[thread_idx * head_dim + d] = f32(V[v_offset + d]);
+            }
+        } else {
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                shared_V[thread_idx * head_dim + d] = 0.0;
+            }
+        }
+
+        workgroupBarrier();
+
+        // Compute attention scores for this block
+        if (valid_query) {
+            // Find max in this block (for numerical stability)
+            var block_max: f32 = -3.402823e+38;
+
+            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
+                let key_pos = kv_block_start + k;
+                if (key_pos >= seq_len) { continue; }
+
+                // Check causal mask
+                if (is_masked(query_pos, key_pos)) { continue; }
+
+                // Compute Q @ K^T for this position
+                var score: f32 = 0.0;
+                for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                    score = score + q_local[d] * shared_K[k * head_dim + d];
+                }
+                score = score * scale;
+
+                // Gemma 2 attention softcapping
+                if (u.attn_softcap > 0.0) {
+                    score = tanh(score / u.attn_softcap) * u.attn_softcap;
+                }
+
+                block_max = max(block_max, score);
+                shared_scores[thread_idx * BLOCK_SIZE + k] = score;
+            }
+
+            // Online softmax update
+            let m_new = max(m_i, block_max);
+            let correction = exp(m_i - m_new);
+
+            // Rescale previous accumulator
+            l_i = l_i * correction;
+            for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                acc[d] = acc[d] * correction;
+            }
+
+            // Add contribution from this block
+            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
+                let key_pos = kv_block_start + k;
+                if (key_pos >= seq_len) { continue; }
+                if (is_masked(query_pos, key_pos)) { continue; }
+
+                let score = shared_scores[thread_idx * BLOCK_SIZE + k];
+                let p = exp(score - m_new);
+                l_i = l_i + p;
+
+                // Accumulate V contribution
+                for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+                    acc[d] = acc[d] + p * shared_V[k * head_dim + d];
+                }
+            }
+
+            m_i = m_new;
+        }
+
+        workgroupBarrier();
+    }
+
+    // Normalize by sum and write output
+    // Always write output when valid_query - write zeros if l_i is 0 to avoid garbage
+    if (valid_query) {
+        let out_offset = query_pos * u.num_heads * head_dim + head_idx * head_dim;
+        let inv_l_i = select(0.0, 1.0 / l_i, l_i > 0.0);
+        for (var d: u32 = 0u; d < head_dim; d = d + 1u) {
+            output[out_offset + d] = acc[d] * inv_l_i;
+        }
+    }
+}