#ifdef MNN_SUPPORT_FP16
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#endif

// Kernel 1: Depthwise Conv1D + SiLU
// Each work-item processes one (batch*channel, seq_pos) element.
// Input:  qkv [B, D, L], conv_state [B, D, conv_state_size], conv_weight [D, 1, K]
// Output: conv_out [B, D, L]
// conv_state is read but NOT updated here (updated by separate kernel).
__kernel void linear_attn_conv_silu(
    __private const int global_dim0,
    __global const FLOAT* qkv,
    __global const FLOAT* conv_state,
    __global const FLOAT* conv_weight,
    __global FLOAT* conv_out,
    __private const int batch,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int kernel_size,
    __private const int conv_state_size)
{
    const int gid = get_global_id(0);
    if (gid >= global_dim0) return;

    const int L = seq_len;
    const int D = conv_dim;
    const int K = kernel_size;
    const int css = conv_state_size;

    // Decompose: gid -> (batch_chan, seq_pos)
    const int l = gid % L;
    const int bd = gid / L;
    const int b = bd / D;
    const int d = bd % D;

    // Compute valid convolution for position l
    // Padded input = [conv_state[b,d,:], qkv[b,d,:]]
    // padded[pos]: if pos < css -> conv_state[b*D*css + d*css + pos]
    //              else -> qkv[b*D*L + d*L + (pos - css)]
    float sum = 0.0f;
    for (int k = 0; k < K; ++k) {
        int pos = l + k;
        float input_val;
        if (pos < css) {
            input_val = (float)conv_state[b * D * css + d * css + pos];
        } else {
            input_val = (float)qkv[b * D * L + d * L + (pos - css)];
        }
        sum += input_val * (float)conv_weight[d * K + k];
    }

    // SiLU activation: x * sigmoid(x)
    float sigmoid_val = 1.0f / (1.0f + exp(-sum));
    conv_out[b * D * L + d * L + l] = (FLOAT)(sum * sigmoid_val);
}

// Kernel 2: Update conv state with last (K-1) elements of padded input
// new_state[i] = padded[L + i], where padded = cat(old_state[css], qkv[L])
// Must execute AFTER linear_attn_conv_silu (which reads conv_state).
__kernel void linear_attn_conv_state_update(
    __private const int global_dim0,
    __global const FLOAT* qkv,
    __global FLOAT* conv_state,
    __private const int batch,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int conv_state_size)
{
    const int gid = get_global_id(0);
    if (gid >= global_dim0) return;

    const int L = seq_len;
    const int D = conv_dim;
    const int css = conv_state_size;

    const int i = gid % css;
    const int bd = gid / css;
    const int b = bd / D;
    const int d = bd % D;

    // new_state[i] = padded[L + i]
    // padded = cat(old_state[css], qkv[L])
    // position (L + i) in padded: if (L+i) < css -> old_state, else -> qkv[(L+i) - css]
    int pos = L + i;
    FLOAT val;
    if (pos < css) {
        val = conv_state[b * D * css + d * css + pos];
    } else {
        val = qkv[b * D * L + d * L + (pos - css)];
    }
    // Safe: we write to index i, read from index (L+i) where L >= 1, so (L+i) > i always
    conv_state[b * D * css + d * css + i] = val;
}

// Kernel 3: Gated Delta Rule (Steps 2-5 fused)
// Each work-item processes one (batch, head) pair across all timesteps.
// Uses float32 exclusively for numerical stability of the recurrence.
__kernel void linear_attn_gated_delta_rule(
    __global const FLOAT* conv_out,
    __global const FLOAT* gate,
    __global const FLOAT* beta_in,
    __global FLOAT* recurrent_state,
    __global FLOAT* attn_out,
    __private const int batch,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_k_heads,
    __private const int num_v_heads,
    __private const int d_k,
    __private const int d_v,
    __private const int key_dim,
    __private const int val_dim,
    __private const int gqa_factor,
    __private const float q_scale)
{
    const int dv4 = (d_v + 3) / 4;
    const int total = dv4 * num_v_heads * batch;
    const int gid = get_global_id(1);
    if (gid >= total) return;
    const int x = (gid % dv4) << 2;
    const int bh = gid / dv4;
    const int h = bh % num_v_heads;
    const int b = bh / num_v_heads;
    const int k_head = h / gqa_factor; // GQA: corresponding K-head
    __local float4 __sum[LOCAL_SIZE];
    __local float4 rec_local[K_SIZE];
    __local float key_local[K_SIZE];
    const int lid = get_local_id(0);

    // State pointer: recurrent_state layout [B, H, d_k, d_v]
    const int state_offset = (b * num_v_heads + h) * K_SIZE * d_v + x;
    #ifdef DECODE_PHASE
    const int conv_base = b * conv_dim;
    float g_t    = (float)(gate[b * num_v_heads + h]);
    float4 beta_t = (float4)(beta_in[b * num_v_heads + h]);
    float4 decay_val = (float4)(exp(g_t));
    const int out_offset = b * num_v_heads * d_v + h * d_v;
    float4 s = (float4)(0);
    for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
        float4 rec = convert_float4(vload4(0, recurrent_state + state_offset + i * d_v))* decay_val;
        float key = (float)(conv_out[conv_base + key_dim + k_head * K_SIZE + i]);
        s += rec * (float4)(key);
        rec_local[i] = rec;
        key_local[i] = key;
    }
    __sum[lid] = s;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int i = LOCAL_SIZE/2; i > 0; i /= 2){
        if (lid < i)
            __sum[lid] += __sum[lid + i];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    s = __sum[0];
    float4 v_data = convert_float4(vload4(0, conv_out + conv_base + 2 * key_dim + h * d_v + x));
    float4 delta = beta_t * (v_data - s);
    s = (float4)(0);
    for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
        float4 recurrent_state_data = rec_local[i]  + (float4)(key_local[i]) * delta;
        s += recurrent_state_data * (float4)(conv_out[conv_base + (k_head * K_SIZE + i)]) * (float4)(q_scale);
        vstore4(CONVERT_FLOAT4(recurrent_state_data), 0, recurrent_state + state_offset + i * d_v);
    }
    __sum[lid] = s;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int i = LOCAL_SIZE/2; i > 0; i /= 2){
        if (lid < i)
            __sum[lid] += __sum[lid + i];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    s = __sum[0];
    if(lid == 0){
        vstore4(CONVERT_FLOAT4(s), 0, attn_out + out_offset + x);
    }
    #else
    for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
        rec_local[i] = convert_float4(vload4(0, recurrent_state + state_offset + i * d_v));
    }
    for (int t = 0; t < seq_len; ++t) {
        const int conv_base = b * seq_len * conv_dim;
        float g_t    = (float)(gate[b * seq_len * num_v_heads + t * num_v_heads + h]);
        float4 beta_t = (float4)(beta_in[b * seq_len * num_v_heads + t * num_v_heads + h]);
        float4 decay_val = (float4)(exp(g_t));
        const int out_offset = (b * seq_len + t) * num_v_heads * d_v + h * d_v;
        float4 s = (float4)(0);
        for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
            float4 rec = rec_local[i] * decay_val;
            float key = (float)(conv_out[conv_base + (key_dim + k_head * K_SIZE + i) * seq_len + t]);
            s += rec * (float4)(key);
            key_local[i] = key;
        }
        __sum[lid] = s;
        barrier(CLK_LOCAL_MEM_FENCE);
        for(int i = LOCAL_SIZE/2; i > 0; i /= 2){
            if (lid < i)
                __sum[lid] += __sum[lid + i];
            barrier(CLK_LOCAL_MEM_FENCE);
        }
        s = __sum[0];
        float4 v_data;
        v_data.x = (float)(conv_out[conv_base + (2 * key_dim + h * d_v + x) * seq_len + t]);
        v_data.y = (float)(conv_out[conv_base + (2 * key_dim + h * d_v + x + 1) * seq_len + t]);
        v_data.z = (float)(conv_out[conv_base + (2 * key_dim + h * d_v + x + 2) * seq_len + t]);
        v_data.w = (float)(conv_out[conv_base + (2 * key_dim + h * d_v + x + 3) * seq_len + t]);
        float4 delta = beta_t * (v_data - s);
        s = (float4)(0);
        for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
            float4 recurrent_state_data = rec_local[i] * decay_val + (float4)(key_local[i]) * delta;
            s += recurrent_state_data * (float4)(conv_out[conv_base + (k_head * K_SIZE + i) * seq_len + t]) * (float4)(q_scale);
            rec_local[i] = recurrent_state_data;
        }
        __sum[lid] = s;
        barrier(CLK_LOCAL_MEM_FENCE);
        for(int i = LOCAL_SIZE/2; i > 0; i /= 2){
            if (lid < i)
                __sum[lid] += __sum[lid + i];
            barrier(CLK_LOCAL_MEM_FENCE);
        }
        s = __sum[0];
        if(lid == 0){
            vstore4(CONVERT_FLOAT4(s), 0, attn_out + out_offset + x);
        }
    }
    for (int i = lid; i < K_SIZE; i+=LOCAL_SIZE) {
        vstore4(CONVERT_FLOAT4(rec_local[i]), 0, recurrent_state + state_offset + i * d_v);
    }
    #endif
}

__kernel void l2_norm(__global const FLOAT* input,
                             __global FLOAT* output,
                             __private const int conv_dim,
                             __private const int head_k_dim,
                             __private const int gqa_factor,
                             __private const int key_dim,
                             __private const int seq_len)
{
    #ifdef USE_VEC
    const int hl = get_global_id(1);
    const int bk = get_global_id(2);
    const int seq_len4 = (seq_len + 3) / 4;
    const int h = hl / seq_len4;
    const int sq = hl % seq_len4;
    const int b = bk / 2;
    const int k = bk % 2;
    const int lid = get_local_id(0);
    const int k_head = h / gqa_factor; // GQA: corresponding K-head
    const int input_offset = (b * conv_dim + k * key_dim + k_head * head_k_dim) * seq_len + (sq << 2);
    __local float4 __sum[K_SIZE];
    float4 sum = 0.0f;
    for(int i = lid; i < K_SIZE; i += K_SIZE){
        float4 in = convert_float4(vload4(0, input + input_offset + i * seq_len));
        sum += in * in;
    }
    __sum[lid] = sum;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int i = K_SIZE/2; i > 0; i /= 2){
        if (lid < i)
            __sum[lid] += __sum[lid + i];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    sum = __sum[0];
    int remain = seq_len - (sq << 2);
    float4 invNorm = (float4)1.0f / sqrt(sum + (float4)1e-6f);
    if(remain == 1){
        for(int i = lid; i < K_SIZE; i += K_SIZE){
            float4 out = convert_float4(vload4(0, input + input_offset + i * seq_len)) * invNorm;
            output[input_offset + i * seq_len] = (FLOAT)out.s0;
        }
    }else if(remain == 2){
        for(int i = lid; i < K_SIZE; i += K_SIZE){
            float4 out = convert_float4(vload4(0, input + input_offset + i * seq_len)) * invNorm;
            vstore2(CONVERT_FLOAT2(out.s01), 0, output + input_offset + i * seq_len);
        }
    }else if(remain == 3){
        for(int i = lid; i < K_SIZE; i += K_SIZE){
            float4 out = convert_float4(vload4(0, input + input_offset + i * seq_len)) * invNorm;
            vstore3(CONVERT_FLOAT3(out.s012), 0, output + input_offset + i * seq_len);
        }
    }else{
        for(int i = lid; i < K_SIZE; i += K_SIZE){
            float4 out = convert_float4(vload4(0, input + input_offset + i * seq_len)) * invNorm;
            vstore4(CONVERT_FLOAT4(out), 0, output + input_offset + i * seq_len);
        }
    }
    #else
    const int h = get_global_id(1);
    const int bk = get_global_id(2);
    const int b = bk / 2;
    const int k = bk % 2;
    const int lid = get_local_id(0);
    const int k_head = h / gqa_factor; // GQA: corresponding K-head
    const int input_offset = b * conv_dim + k * key_dim + k_head * head_k_dim;
    __local float __sum[K_SIZE];
    float sum = 0.0f;
    for(int i = lid; i < K_SIZE; i += K_SIZE){
        float in = (float)(input[input_offset + i]);
        sum += in * in;
    }
    __sum[lid] = sum;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int i = K_SIZE/2; i > 0; i /= 2){
        if (lid < i)
            __sum[lid] += __sum[lid + i];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    sum = __sum[0];
    float invNorm = 1.0f / sqrt(sum + 1e-6f);
    for(int i = lid; i < K_SIZE; i += K_SIZE){
        float out = (float)(input[input_offset + i]) * invNorm;
        output[input_offset + i] = (FLOAT)out;
    }
    #endif
}

#ifdef CHUNK_PREFILL

// ======================== Chunked Prefill Kernels ========================
// Implements torch_chunk_gated_delta_rule from modeling_qwen3_5.py
// Decomposes sequential recurrence into chunk-parallel operations.
// CHUNK_SIZE must be defined at compile time (e.g. -DCHUNK_SIZE=64).

// Kernel C1: Cumulative sum of gate values within each chunk
// GWS: {num_v_heads, num_chunks, batch}
__kernel void chunk_g_cumsum(
    __global const FLOAT* gate,        // [B, L, H]
    __global float* g_cumsum,           // [B, H, num_chunks, CHUNK_SIZE] (float32 for precision)
    __private const int num_v_heads,
    __private const int seq_len,
    __private const int num_chunks)
{
    const int h = get_global_id(0);
    const int c = get_global_id(1);
    const int b = get_global_id(2);
    if (h >= num_v_heads || c >= num_chunks) return;

    const int H = num_v_heads;
    const int L = seq_len;
    const int C = CHUNK_SIZE;

    float cumsum = 0.0f;
    int out_base = ((b * H + h) * num_chunks + c) * C;
    for (int p = 0; p < C; ++p) {
        int l = c * C + p;
        float g_val = (l < L) ? (float)gate[b * L * H + l * H + h] : 0.0f;
        cumsum += g_val;
        g_cumsum[out_base + p] = cumsum;
    }
}

// Kernel C2: Build Neumann-corrected intra-chunk attention matrix
// attn = I + neumann_correction(-(k_beta @ k^T) * decay_mask)
// GWS: {CHUNK_SIZE, num_v_heads * num_chunks, batch}, LWS: {CHUNK_SIZE, 1, 1}

__kernel void chunk_build_neumann_attn_step0(
    __global const FLOAT* conv_out,    // [B, D, L]
    __global const FLOAT* beta_in,     // [B, L, H]
    __global const float* g_cumsum,    // [B, H, num_chunks, C] (float32)
    __global float* attn_matrix,       // [B, H, num_chunks, C, C] (float32)
    __private const int batch,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int head_k_dim,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks)
{
    const int rc = get_global_id(0);
    const int hc = get_global_id(1);
    const int b = get_global_id(2);
    if (hc >= num_v_heads * num_chunks || b >= batch || rc >= CHUNK_SIZE * CHUNK_SIZE) return;

    const int col = rc % CHUNK_SIZE;
    const int row = rc / CHUNK_SIZE;
    const int h = hc / num_chunks;
    const int c = hc % num_chunks;


    const int C = CHUNK_SIZE;
    const int dk = head_k_dim;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int k_head = h / gqa_factor;

    const int g_base = ((b * H + h) * num_chunks + c) * C;
    const int attn_out_base = ((b * H + h) * num_chunks + c) * C * C;
    const int conv_base = b * D * L;
    
    // Phase 1: Compute strictly-lower-triangular attn
    // attn[i, j] = -(k[i]*beta[i]) dot k[j] * exp(g[i]-g[j])
    float val = 0.0f;
    if (col < row) {
        const int l_i = c * C + row;
        const int l_j = c * C + col;
        if (l_i < L && l_j < L) {
            const float g_i = (float)g_cumsum[g_base + row];
            const float beta_i = (float)beta_in[b * L * H + l_i * H + h];
            const float g_j = (float)g_cumsum[g_base + col];
            const float decay = exp(g_i - g_j);
            float dot = 0.0f;
            for (int d = 0; d < dk; ++d) {
                float k_i_d = (float)conv_out[conv_base + (key_dim + k_head * dk + d) * L + l_i];
                float k_j_d = (float)conv_out[conv_base + (key_dim + k_head * dk + d) * L + l_j];
                dot += k_i_d * k_j_d;
            }
            val = -(beta_i * dot) * decay;
        }
    }
    attn_matrix[attn_out_base + row * C + col] = val;
}

__kernel void chunk_build_neumann_attn_step1(
    __global float* attn_matrix,       // [B, H, num_chunks, C, C] (float32)
    __private const int batch,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int head_k_dim,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks)
{
    const int col = get_global_id(0);
    const int hc = get_global_id(1);
    const int b = get_global_id(2);
    if (hc >= num_v_heads * num_chunks || b >= batch || col >= CHUNK_SIZE) return;

    const int lid = get_local_id(0);
    const int h = hc / num_chunks;
    const int c = hc % num_chunks;


    const int C = CHUNK_SIZE;
    const int dk = head_k_dim;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int k_head = h / gqa_factor;

    const int attn_out_base = ((b * H + h) * num_chunks + c) * C * C;
    
    // Phase 2: Neumann series correction (row by row)
    for (int r = 1; r < C; ++r) {
        float orig = attn_matrix[attn_out_base + r * C + lid];
        float correction = 0.0f;
        if (lid < r) {
            for (int k = 0; k < r; ++k) {
                correction += attn_matrix[attn_out_base + r * C + k] * attn_matrix[attn_out_base + k * C + lid];
            }
        }
        if (lid < r) {
            attn_matrix[attn_out_base + r * C + lid] = orig + correction;
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }

    // Phase 3: Add identity and write to global (float32)
    attn_matrix[attn_out_base + lid * C + lid] = attn_matrix[attn_out_base + lid * C + lid] + 1.0f;
}

// Kernel C3: v_corrected = attn_matrix @ (V * beta)
// GWS: {UP_DIV(dv, 4), CHUNK_SIZE * num_chunks, B * H}
__kernel void chunk_correct_v(
    __global const float* attn_matrix,  // [B, H, num_chunks, C, C] (float32)
    __global const FLOAT* conv_out,     // [B, D, L]
    __global const FLOAT* beta_in,      // [B, L, H]
    __global const float* g_cumsum,     // [B, H, num_chunks, C] (float32)
    __global float* v_corrected,        // [B, H, num_chunks, C, dv] (float32)
    __global float* k_cumdecay,         // [B, H, num_chunks, C, dk] (float32)
    __private const int global_dim0,
    __private const int global_dim1,
    __private const int global_dim2,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int head_k_dim,
    __private const int head_v_dim,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks)
{
    const int x = get_global_id(0);
    const int y = get_global_id(1);
    const int z = get_global_id(2);
    if (x >= global_dim0 || y >= global_dim1 || z >= global_dim2) return;

    const int j = x << 2;
    const int c = y / CHUNK_SIZE;
    const int row = y % CHUNK_SIZE;
    const int b = z / num_v_heads;
    const int h = z % num_v_heads;

    const int C = CHUNK_SIZE;
    const int dk = head_k_dim;
    const int dv = head_v_dim;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int k_head = h / gqa_factor;

    const int attn_base = ((b * H + h) * num_chunks + c) * C * C + row * C;
    const int g_base = ((b * H + h) * num_chunks + c) * C;
    const int conv_base = b * D * L;

    float4 result0 = (float4)(0.0f);
    float4 result1 = (float4)(0.0f);
    int l_p = c * C;
    for (int p = 0; p < C && l_p < L; ++p, ++l_p) {
        float a = attn_matrix[attn_base + p];
        float beta_p = (float)beta_in[b * L * H + l_p * H + h];
        float ab = a * beta_p;
        float g_p = g_cumsum[g_base + p];
        float coeff = ab * exp(g_p);
        float4 v_p;
        v_p.x = (float)conv_out[conv_base + (2 * key_dim + h * dv + j) * L + l_p];
        v_p.y = (float)conv_out[conv_base + (2 * key_dim + h * dv + j + 1) * L + l_p];
        v_p.z = (float)conv_out[conv_base + (2 * key_dim + h * dv + j + 2) * L + l_p];
        v_p.w = (float)conv_out[conv_base + (2 * key_dim + h * dv + j + 3) * L + l_p];
        result0 += ab * v_p;
        float4 k_p;
        k_p.x = (float)conv_out[conv_base + (key_dim + k_head * dk + j) * L + l_p];
        k_p.y = (float)conv_out[conv_base + (key_dim + k_head * dk + j + 1) * L + l_p];
        k_p.z = (float)conv_out[conv_base + (key_dim + k_head * dk + j + 2) * L + l_p];
        k_p.w = (float)conv_out[conv_base + (key_dim + k_head * dk + j + 3) * L + l_p];
        result1 += coeff * k_p;
    }

    int out_idx = ((b * H + h) * num_chunks + c) * C * dv + row * dv + j;
    vstore4((float4)(result0), 0, v_corrected + out_idx);
    vstore4((float4)(result1), 0, k_cumdecay + out_idx);
}

// Kernel C5: Intra-chunk QK attention with decay (lower triangular including diagonal)
// qk[i,j] = (q[i] dot k[j]) * q_scale * exp(g[i]-g[j]) for j <= i
// GWS: {CHUNK_SIZE, CHUNK_SIZE * num_chunks, B * H}
__kernel void chunk_qk_attn(
    __global const FLOAT* conv_out,    // [B, D, L]
    __global const float* g_cumsum,    // [B, H, num_chunks, C] (float32)
    __global float* qk_attn_out,       // [B, H, num_chunks, C, C] (float32, reuses attn_matrix buf)
    __private const int global_dim0,
    __private const int global_dim1,
    __private const int global_dim2,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int head_k_dim,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks,
    __private const float q_scale)
{
    const int col = get_global_id(0);
    const int y = get_global_id(1);
    const int z = get_global_id(2);
    if (col >= global_dim0 || y >= global_dim1 || z >= global_dim2) return;

    const int c = y / CHUNK_SIZE;
    const int row = y % CHUNK_SIZE;
    const int b = z / num_v_heads;
    const int h = z % num_v_heads;

    const int C = CHUNK_SIZE;
    const int dk = head_k_dim;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int k_head = h / gqa_factor;

    float val = 0.0f;
    if (col <= row) {
        const int l_row = c * C + row;
        const int l_col = c * C + col;
        if (l_row < L && l_col < L) {
            const int g_base = ((b * H + h) * num_chunks + c) * C;
            const float decay = exp(g_cumsum[g_base + row] - g_cumsum[g_base + col]);
            const int conv_base = b * D * L;
            float dot = 0.0f;
            for (int d = 0; d < dk; ++d) {
                float q_d = (float)conv_out[conv_base + (k_head * dk + d) * L + l_row];
                float k_d = (float)conv_out[conv_base + (key_dim + k_head * dk + d) * L + l_col];
                dot += q_d * k_d;
            }
            val = dot * q_scale * decay;
        }
    }

    int out_idx = ((b * H + h) * num_chunks + c) * C * C + row * C + col;
    qk_attn_out[out_idx] = val;
}

// Kernel C6: v_new = v_corrected - k_cumdecay @ state (per chunk)
// GWS: {UP_DIV(dv, 4), CHUNK_SIZE, B * H}
__kernel void chunk_compute_vnew(
    __global const float* v_corrected,    // [B, H, num_chunks, C, dv] (float32)
    __global const float* k_cumdecay,     // [B, H, num_chunks, C, dk] (float32)
    __global const FLOAT* recurrent_state,// [B, H, dk, dv]
    __global float* v_new,                // [B, H, C, dv] (float32)
    __private const int global_dim0,
    __private const int global_dim1,
    __private const int global_dim2,
    __private const int head_k_dim,
    __private const int head_v_dim,
    __private const int num_v_heads,
    __private const int num_chunks,
    __private const int chunk_idx)
{
    const int x = get_global_id(0);
    const int p = get_global_id(1);
    const int z = get_global_id(2);
    if (x >= global_dim0 || p >= global_dim1 || z >= global_dim2) return;

    const int j = x << 2;
    const int b = z / num_v_heads;
    const int h = z % num_v_heads;

    const int C = CHUNK_SIZE;
    const int dk = head_k_dim;
    const int dv = head_v_dim;
    const int H = num_v_heads;
    const int c = chunk_idx;

    const int kc_base = ((b * H + h) * num_chunks + c) * C * dk + p * dk;
    const int state_base = (b * H + h) * dk * dv;

    float4 v_prime = (float4)(0.0f);
    for (int d = 0; d < dk; ++d) {
        float kc_d = k_cumdecay[kc_base + d];
        float4 s = convert_float4(vload4(0, recurrent_state + state_base + d * dv + j));
        v_prime += kc_d * s;
    }

    int vc_idx = ((b * H + h) * num_chunks + c) * C * dv + p * dv + j;
    float4 vc = vload4(0, v_corrected + vc_idx);
    float4 vn = vc - v_prime;

    int vn_idx = (b * H + h) * C * dv + p * dv + j;
    vstore4((float4)(vn), 0, v_new + vn_idx);
}

// Kernel C7: Compute output and update recurrent state (per chunk)
// GWS: {dk * UP_DIV(dv, 4), H, batch}
__kernel void chunk_output_state_update(
    __global const FLOAT* conv_out,        // [B, D, L]
    __global const float* qk_attn_matrix,  // [B, H, num_chunks, C, C] (float32)
    __global const float* v_new,           // [B, H, C, dv] (float32)
    __global const float* g_cumsum,        // [B, H, num_chunks, C] (float32)
    __global FLOAT* recurrent_state,       // [B, H, dk, dv]
    __global FLOAT* attn_out,              // [B, L, H*dv]
    __private const int global_dim0,
    __private const int global_dim1,
    __private const int global_dim2,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int head_k_dim,
    __private const int head_v_dim,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks,
    __private const int chunk_idx,
    __private const float q_scale)
{
    const int x = get_global_id(0);
    const int h = get_global_id(1);
    const int b = get_global_id(2);
    if (x >= global_dim0 || h >= global_dim1 || b >= global_dim2) return;
    
    const int dk = head_k_dim;
    const int dv = head_v_dim;
    const int dv4 = (dv + 3) / 4;
    const int j = (x % dv4) << 2;
    const int d = x / dv4;
    const int C = CHUNK_SIZE;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int c = chunk_idx;
    const int k_head = h / gqa_factor;

    const int state_base = (b * H + h) * dk * dv;
    const int g_base = ((b * H + h) * num_chunks + c) * C;
    const int qk_base = ((b * H + h) * num_chunks + c) * C * C;
    const int vn_base = (b * H + h) * C * dv;
    const int conv_base = b * D * L;

    int last_p = min(C, L - c * C);
    if (last_p <= 0) return;

    // Phase 2: Update recurrent state
    // state[d,j] = state[d,j]*exp(g_last) + sum_p k[p,d]*exp(g_last-g[p])*v_new[p,j]
    float g_last = g_cumsum[g_base + last_p - 1];
    float decay_state = exp(g_last);

    float4 s = convert_float4(vload4(0, recurrent_state + state_base + d * dv + j)) * (float4)(decay_state);
    for (int p = 0; p < last_p; ++p) {
        int l_p = c * C + p;
        float g_p = g_cumsum[g_base + p];
        float k_d = (float)conv_out[conv_base + (key_dim + k_head * dk + d) * L + l_p] * exp(g_last - g_p);
        float4 vn = vload4(0, v_new + vn_base + p * dv + j);
        s += k_d * vn;
    }
    vstore4(CONVERT_FLOAT4(s), 0, recurrent_state + state_base + d * dv + j);
}

__kernel void chunk_output(
    __global const FLOAT* conv_out,        // [B, D, L]
    __global const float* qk_attn_matrix,  // [B, H, num_chunks, C, C] (float32)
    __global const float* v_new,           // [B, H, C, dv] (float32)
    __global const float* g_cumsum,        // [B, H, num_chunks, C] (float32)
    __global const FLOAT* recurrent_state, // [B, H, dk, dv]
    __global FLOAT* attn_out,              // [B, L, H*dv]
    __private const int global_dim0,
    __private const int global_dim1,
    __private const int global_dim2,
    __private const int conv_dim,
    __private const int seq_len,
    __private const int num_v_heads,
    __private const int dk,
    __private const int dv,
    __private const int key_dim,
    __private const int gqa_factor,
    __private const int num_chunks,
    __private const int chunk_idx,
    __private const float q_scale)
{
    const int x = get_global_id(0);
    const int h = get_global_id(1);
    const int b = get_global_id(2);
    if (x >= global_dim0 || h >= global_dim1 || b >= global_dim2) return;
    
    const int dv4 = (dv + 3) / 4;
    const int j = (x % dv4) << 2;
    const int p = x / dv4;
    const int C = CHUNK_SIZE;
    const int D = conv_dim;
    const int L = seq_len;
    const int H = num_v_heads;
    const int c = chunk_idx;
    const int k_head = h / gqa_factor;

    const int state_base = (b * H + h) * dk * dv;
    const int g_base = ((b * H + h) * num_chunks + c) * C;
    const int qk_base = ((b * H + h) * num_chunks + c) * C * C;
    const int vn_base = (b * H + h) * C * dv;
    const int conv_base = b * D * L;

    int last_p = min(C, L - c * C);
    if (last_p <= 0 || p >= last_p) return;

    // Phase 1: Compute output for each position p
    int l_p = c * C + p;
    float g_p = g_cumsum[g_base + p];
    float decay_q = exp(g_p) * q_scale;

    // Cross-chunk: q * scale * exp(g) @ state
    float4 inter = (float4)(0.0f);
    for (int d = 0; d < dk; d++) {
        float q_d = (float)conv_out[conv_base + (k_head * dk + d) * L + l_p] * decay_q;
        float4 s = convert_float4(vload4(0, recurrent_state + state_base + d * dv + j));
        inter += q_d * s;
    }
    
    // Intra-chunk: qk_attn @ v_new
    float4 intra = (float4)(0.0f);
    for (int r = 0; r <= p; ++r) {
        float qk = qk_attn_matrix[qk_base + p * C + r];
        float4 vn = vload4(0, v_new + vn_base + r * dv + j);
        intra += qk * vn;
    }

    const int out_offset = (b * L + l_p) * H * dv + h * dv + j;
    vstore4(CONVERT_FLOAT4(inter + intra), 0, attn_out + out_offset);
}

#endif // CHUNK_PREFILL