// Flash-Attention-Style Prefill Kernel (head_dim=256, f16 KV)
//
// Flash attention with KV-axis splitting to raise RDNA3 occupancy. The existing
// attention_head256_f16kv.wgsl dispatches (num_query_blocks × num_heads, 1, 1)
// workgroups, which is too few for RDNA3 at common prefill shapes (e.g. 16
// workgroups for prefill=64, 8 heads). This variant multiplies the workgroup
// count by num_kv_splits along the KV axis and produces partial (output, m, l)
// per split — a companion reduction kernel merges them.
//
// Algorithm: online-softmax flash attention per split.
//
// Pass 1 (this kernel): each workgroup processes one KV slice for its
// (query_block, head) pair; writes un-normalised acc + m + l to partial buffers.
// Pass 2 (flash_reduce): merges partials across splits for each (query, head).

enable f16;

const BLOCK_SIZE: u32 = 32u;
const WORKGROUP_SIZE: u32 = BLOCK_SIZE;
const HEAD_DIM: u32 = 256u;
const HEAD_DIM_VECS: 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,
    num_kv_splits: 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>;
// Per-split partials: [num_query_blocks, num_heads, num_kv_splits, BLOCK_SIZE, HEAD_DIM]
// When num_kv_splits == 1 (single-split fast path), this buffer is bound to the
// FINAL output buffer and we write normalised output directly — skipping the
// reduce pass entirely.
@group(0) @binding(4) var<storage, read_write> partial_acc: array<f32>;
// Per-split stats: [num_query_blocks, num_heads, num_kv_splits, BLOCK_SIZE]
@group(0) @binding(5) var<storage, read_write> partial_m: array<f32>;
@group(0) @binding(6) var<storage, read_write> partial_l: array<f32>;
@group(0) @binding(7) var<storage, read> kv_len_buffer: array<u32>;
@group(0) @binding(8) var<storage, read> page_table: array<u32>;

// Separate shared tiles for K and V so they can be loaded in parallel within
// the same loop iteration. Saves a workgroupBarrier per KV block vs the
// shared-tile-reuse pattern of the non-flash head256 kernel.
var<workgroup> shared_K: array<vec4<f16>, BLOCK_SIZE * HEAD_DIM_VECS>;
var<workgroup> shared_V: array<vec4<f16>, BLOCK_SIZE * HEAD_DIM_VECS>;

fn zero_vec4_f16() -> vec4<f16> {
    return vec4<f16>(f16(0.0), f16(0.0), f16(0.0), f16(0.0));
}

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>
) {
    if (u.head_dim != HEAD_DIM) {
        return;
    }

    // Dispatch layout: linear = (query_block_idx * num_heads + head_idx) * num_kv_splits + kv_split_idx
    let num_heads = u.num_heads;
    let num_kv_splits = u.num_kv_splits;
    let linear = wg_id.x;
    let kv_split_idx = linear % num_kv_splits;
    let linear_qh = linear / num_kv_splits;
    let head_idx = linear_qh % num_heads;
    let query_block_idx = linear_qh / num_heads;
    let thread_idx = local_id.x;

    let kv_head_idx = get_kv_head_idx(head_idx);
    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;

    // KV slice bounds for this split
    let kv_per_split = (seq_len + num_kv_splits - 1u) / num_kv_splits;
    let kv_slice_start = kv_split_idx * kv_per_split;
    let kv_slice_end = min(kv_slice_start + kv_per_split, seq_len);

    var q_local: array<vec4<f32>, HEAD_DIM_VECS>;
    var acc: array<vec4<f32>, HEAD_DIM_VECS>;

    if (valid_query) {
        let q_offset = query_pos * num_heads * HEAD_DIM + head_idx * HEAD_DIM;
        for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
            let base = q_offset + d4 * 4u;
            q_local[d4] = vec4<f32>(Q[base], Q[base + 1u], Q[base + 2u], Q[base + 3u]);
            acc[d4] = vec4<f32>(0.0);
        }
    }

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

    // Clamp KV slice to causal/sliding constraints for this query
    var min_key_pos: u32 = kv_slice_start;
    var max_key_pos: u32 = kv_slice_end;

    if (valid_query) {
        if (u.is_causal != 0u) {
            if (abs_query < u.kv_start) {
                max_key_pos = min(max_key_pos, 0u);
            } else {
                let causal_limit = abs_query - u.kv_start + 1u;
                max_key_pos = min(max_key_pos, 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) {
                let sw_min = min_abs_key - u.kv_start;
                if (sw_min > min_key_pos) {
                    min_key_pos = sw_min;
                }
            }
            min_key_pos = min(min_key_pos, seq_len);
        }
    }

    // Iterate KV blocks WITHIN our split range
    let num_kv_blocks_in_slice = (kv_slice_end - kv_slice_start + BLOCK_SIZE - 1u) / BLOCK_SIZE;

    for (var kv_block: u32 = 0u; kv_block < num_kv_blocks_in_slice; kv_block = kv_block + 1u) {
        let kv_block_start = kv_slice_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 < kv_slice_end && key_pos >= min_key_pos && key_pos < max_key_pos) {
                    key_active[k_mask] = 1u;
                }
            }
        }

        // Load K and V in parallel into separate shared tiles. A single barrier
        // synchronises both loads before we score K and accumulate V.
        let key_pos_load = kv_block_start + thread_idx;
        let shared_row = thread_idx * HEAD_DIM_VECS;
        if (key_pos_load < kv_slice_end) {
            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;
            let v_idx = k_idx;
            let v_offset = v_idx * u.num_kv_heads * HEAD_DIM + kv_head_idx * HEAD_DIM;
            for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                let kb = k_offset + d4 * 4u;
                let vb = v_offset + d4 * 4u;
                shared_K[shared_row + d4] = vec4<f16>(
                    K[kb], K[kb + 1u], K[kb + 2u], K[kb + 3u]
                );
                shared_V[shared_row + d4] = vec4<f16>(
                    V[vb], V[vb + 1u], V[vb + 2u], V[vb + 3u]
                );
            }
        } else {
            for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                shared_K[shared_row + d4] = zero_vec4_f16();
                shared_V[shared_row + d4] = zero_vec4_f16();
            }
        }

        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; }
                let key_row = k * HEAD_DIM_VECS;
                var dot_partial: f32 = 0.0;
                for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                    dot_partial = dot_partial + dot(
                        q_local[d4], vec4<f32>(shared_K[key_row + d4])
                    );
                }
                var s = dot_partial * scale;
                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 d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                acc[d4] = acc[d4] * 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;
            }

            // V is already in shared_V — apply immediately (no reload, no barrier).
            for (var k: u32 = 0u; k < BLOCK_SIZE; k = k + 1u) {
                let p = probs[k];
                if (p == 0.0) { continue; }
                let value_row = k * HEAD_DIM_VECS;
                for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                    acc[d4] = acc[d4] + p * vec4<f32>(shared_V[value_row + d4]);
                }
            }
            m_i = m_new;
        }

        workgroupBarrier();
    }

    if (valid_query) {
        let qh_idx = query_block_idx * num_heads + head_idx;
        if (num_kv_splits == 1u) {
            // Single-split fast path: skip the reduce pass entirely. Write the
            // FINAL normalised output directly to the output buffer (which is
            // bound in slot 4 under the partial_acc name). Layout matches
            // [query_len, num_heads, head_dim].
            let inv_l_i = select(0.0, 1.0 / l_i, l_i > 0.0);
            let out_offset = query_pos * num_heads * HEAD_DIM + head_idx * HEAD_DIM;
            for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                let base = out_offset + d4 * 4u;
                let normed = acc[d4] * inv_l_i;
                partial_acc[base] = normed.x;
                partial_acc[base + 1u] = normed.y;
                partial_acc[base + 2u] = normed.z;
                partial_acc[base + 3u] = normed.w;
            }
            // m/l buffers are still bound but unused on the single-split path.
            // Leave them untouched to avoid spurious writes.
        } else {
            // Multi-split: write UN-NORMALISED partial acc and per-query m/l for
            // this split. Layout:
            //   partial_acc[(qb*num_heads + h) * num_kv_splits * BLOCK_SIZE * HEAD_DIM
            //               + kv_split * BLOCK_SIZE * HEAD_DIM
            //               + thread_idx * HEAD_DIM + d]
            let partial_base = (qh_idx * num_kv_splits + kv_split_idx) * BLOCK_SIZE * HEAD_DIM
                              + thread_idx * HEAD_DIM;
            for (var d4: u32 = 0u; d4 < HEAD_DIM_VECS; d4 = d4 + 1u) {
                let base = partial_base + d4 * 4u;
                partial_acc[base] = acc[d4].x;
                partial_acc[base + 1u] = acc[d4].y;
                partial_acc[base + 2u] = acc[d4].z;
                partial_acc[base + 3u] = acc[d4].w;
            }
            let stats_idx = (qh_idx * num_kv_splits + kv_split_idx) * BLOCK_SIZE + thread_idx;
            partial_m[stats_idx] = m_i;
            partial_l[stats_idx] = l_i;
        }
    }
}