// AUTO-GENERATED from src/gpu/kernels/attention_small.wgsl.
// Edit the source kernel and src/gpu/kernels/codegen/wgsl-variants.js, then run `npm run kernels:codegen:sync`.
// Fused Multi-Head Attention Kernel (small tiles, f16 KV)
//
// Same algorithm as attention_small.wgsl but K/V are stored as f16.
// Shared tiles also use f16 to minimize workgroup storage.

enable f16;

const BLOCK_SIZE: u32 = 32u;
const WORKGROUP_SIZE: u32 = BLOCK_SIZE;
const HEAD_TILE_STRIDE: u32 = 32u;
const MAX_HEAD_DIM: u32 = 256u;

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,  // 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>;
@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>;

var<workgroup> shared_K: array<f16, BLOCK_SIZE * HEAD_TILE_STRIDE>;
var<workgroup> shared_V: array<f16, BLOCK_SIZE * HEAD_TILE_STRIDE>;

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 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(local_invocation_id) local_id: vec3<u32>,
    @builtin(workgroup_id) wg_id: vec3<u32>
) {
    let linear = wg_id.x;
    let num_heads = u.num_heads;
    let head_idx = linear % num_heads;
    let query_block_idx = linear / 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;
    let abs_query = query_pos + u.start_pos;

    var q_local: array<f32, MAX_HEAD_DIM>;
    var acc: array<f32, MAX_HEAD_DIM>;

    if (valid_query) {
        let q_offset = query_pos * 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];
            acc[d] = 0.0;
        }
    }

    var m_i: f32 = -3.402823e+38;
    var l_i: f32 = 0.0;

    let num_kv_blocks = (seq_len + BLOCK_SIZE - 1u) / BLOCK_SIZE;
    let head_tile = select(32u, 16u, head_dim >= 256u);
    let num_head_tiles = (head_dim + head_tile - 1u) / head_tile;
    var min_key_pos: u32 = 0u;
    var max_key_pos: u32 = seq_len;

    if (valid_query) {
        if (u.is_causal != 0u) {
            if (abs_query < u.kv_start) {
                max_key_pos = 0u;
            } else {
                let causal_limit = abs_query - u.kv_start + 1u;
                max_key_pos = min(seq_len, causal_limit);
            }
        }
        if (u.sliding_window > 0u && abs_query >= u.sliding_window) {
            let min_abs_key = abs_query - u.sliding_window + 1u;
            if (min_abs_key > u.kv_start) {
                min_key_pos = min_abs_key - u.kv_start;
            }
            min_key_pos = min(min_key_pos, seq_len);
        }
    }

    for (var kv_block: u32 = 0u; kv_block < num_kv_blocks; kv_block = kv_block + 1u) {
        let kv_block_start = kv_block * BLOCK_SIZE;

        var scores: array<f32, BLOCK_SIZE>;
        var key_active: array<u32, BLOCK_SIZE>;
        var probs: array<f32, BLOCK_SIZE>;
        for (var k_init: u32 = 0u; k_init < BLOCK_SIZE; k_init = k_init + 1u) {
            scores[k_init] = 0.0;
            key_active[k_init] = 0u;
            probs[k_init] = 0.0;
        }

        if (valid_query) {
            for (var k_mask: u32 = 0u; k_mask < BLOCK_SIZE; k_mask = k_mask + 1u) {
                let key_pos = kv_block_start + k_mask;
                if (key_pos < seq_len && key_pos >= min_key_pos && key_pos < max_key_pos) {
                    key_active[k_mask] = 1u;
                }
            }
        }

        for (var ht: u32 = 0u; ht < num_head_tiles; ht = ht + 1u) {
            let d0 = ht * head_tile;
            let tile_len = min(head_tile, head_dim - d0);

            let key_pos_load = kv_block_start + thread_idx;
            if (key_pos_load < seq_len) {
                let k_idx = get_kv_pos(key_pos_load);
                let k_offset = k_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + d0;
                for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                    shared_K[thread_idx * HEAD_TILE_STRIDE + td] = K[k_offset + td];
                }
            } else {
                for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                    shared_K[thread_idx * HEAD_TILE_STRIDE + td] = f16(0.0);
                }
            }

            workgroupBarrier();

            if (valid_query) {
                for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
                    if (key_active[k] == 0u) { continue; }

                    var dot_partial: f32 = 0.0;
                    for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                        dot_partial = dot_partial + q_local[d0 + td] * f32(shared_K[k * HEAD_TILE_STRIDE + td]);
                    }
                    scores[k] = scores[k] + dot_partial;
                }
            }

            workgroupBarrier();
        }

        var m_new: f32 = m_i;
        if (valid_query) {
            var block_max: f32 = -3.402823e+38;
            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
                if (key_active[k] == 0u) { continue; }

                var s = scores[k] * scale;
                // Gemma 2 attention softcapping
                if (u.attn_softcap > 0.0) {
                    s = tanh(s / u.attn_softcap) * u.attn_softcap;
                }
                scores[k] = s;
                block_max = max(block_max, s);
            }

            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) {
                if (key_active[k] == 0u) { continue; }
                let p = exp(scores[k] - m_new);
                probs[k] = p;
                l_i = l_i + p;
            }
        }

        // Accumulate V contribution by tiling head_dim again.
        // Barriers must be in uniform control flow, so only the math is guarded.
        for (var ht: u32 = 0u; ht < num_head_tiles; ht = ht + 1u) {
            let d0 = ht * head_tile;
            let tile_len = min(head_tile, head_dim - d0);

            let key_pos_load = kv_block_start + thread_idx;
            if (key_pos_load < seq_len) {
                let v_idx = get_kv_pos(key_pos_load);
                let v_offset = v_idx * u.num_kv_heads * head_dim + kv_head_idx * head_dim + d0;
                for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                    shared_V[thread_idx * HEAD_TILE_STRIDE + td] = V[v_offset + td];
                }
            } else {
                for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                    shared_V[thread_idx * HEAD_TILE_STRIDE + td] = f16(0.0);
                }
            }

            workgroupBarrier();

            if (valid_query) {
                for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
                    let p = probs[k];
                    if (p == 0.0) { continue; }

                    for (var td: u32 = 0u; td < tile_len; td = td + 1u) {
                        acc[d0 + td] = acc[d0 + td] + p * f32(shared_V[k * HEAD_TILE_STRIDE + td]);
                    }
                }
            }

            workgroupBarrier();
        }

        if (valid_query) {
            m_i = m_new;
        }
    }

    // Always write output when valid_query - write zeros if l_i is 0 to avoid garbage
    if (valid_query) {
        let out_offset = query_pos * 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;
        }
    }
}