#if defined(K_Q4_0) || defined(K_Q8_0) || defined(V_Q4_0) || defined(V_Q8_0)
#define QUANT_SHMEM STAGING_SHMEM
#define QUANT_OUT_TYPE STAGING_OUT_TYPE
#include "quant_inner_loops.tmpl"
#undef QUANT_SHMEM
#undef QUANT_OUT_TYPE
#define BLOCK_SIZE 32
#define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
#define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
#endif

#if defined(K_Q4_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 18u
#define K_BYTES_PER_THREAD 8u
#define K_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_K_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
#elif defined(K_Q8_0)
#define K_NQ 16
#define K_BLOCK_SIZE_BYTES 34u
#define K_BYTES_PER_THREAD 16u
#define K_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_K_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
#endif

#if defined(V_Q4_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 18u
#define V_BYTES_PER_THREAD 8u
#define V_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_V_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
#elif defined(V_Q8_0)
#define V_NQ 16
#define V_BLOCK_SIZE_BYTES 34u
#define V_BYTES_PER_THREAD 16u
#define V_BYTES_PER_INNER_LOOP 4u
#define DEQUANT_V_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
#endif

#ifndef K_DIRECT
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
#if defined(K_Q4_0) || defined(K_Q8_0)
    for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
        let blck_idx = elem_idx / BLOCK_SIZE;
        let block_offset = (elem_idx % BLOCK_SIZE) / K_NQ;
        let k_row = blck_idx / BLOCKS_K;
        let global_k_row = kv_tile + k_row;
        let block_k = blck_idx % BLOCKS_K;
        let row_offset = k_row * HEAD_DIM_QK;
        let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k;
        let block_byte_base = global_block_idx * K_BLOCK_SIZE_BYTES;
        let d = f16_from_u16(load_k_u16_at(block_byte_base));
        let thread_byte_offset = block_offset * K_BYTES_PER_THREAD;
        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
        for (var j = 0u; j < K_BYTES_PER_THREAD / K_BYTES_PER_INNER_LOOP; j += 1u) {
            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * K_BYTES_PER_INNER_LOOP;
            let q_packed = load_k_u32_at(q_byte_offset);
            DEQUANT_K_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
        }
    }
#elif defined(FLASH_ATTN_SCALAR_KV)
    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
        let k_row = elem_idx / HEAD_DIM_QK;
        let k_col = elem_idx % HEAD_DIM_QK;
        let global_k_row = kv_tile + k_row;
        let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
        STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
            0.0,
            K[global_k_row_offset + k_col],
            global_k_row < params.seq_len_kv && k_col < HEAD_DIM_QK));
    }
#else
    for (var vec_idx_local = local_x; vec_idx_local < kv_count * Q_CHUNKS; vec_idx_local += WG_SIZE) {
        let kv_local = vec_idx_local / Q_CHUNKS;
        let chunk = vec_idx_local % Q_CHUNKS;
        let global_k_row = kv_tile + kv_local;
        let k_vec_index = (k_head_offset + global_k_row * params.stride_k1 + chunk * 4u) >> 2u;
        let k4 = K[k_vec_index];
        let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
        STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(k4.x);
        STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(k4.y);
        STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(k4.z);
        STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(k4.w);
    }
#endif
}
#endif // !defined(K_DIRECT)

#ifndef V_DIRECT
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
#if defined(V_Q4_0) || defined(V_Q8_0)
    for (var elem_idx = local_x * V_NQ; elem_idx < kv_count * HEAD_DIM_V; elem_idx += WG_SIZE * V_NQ) {
        let blck_idx = elem_idx / BLOCK_SIZE;
        let block_offset = (elem_idx % BLOCK_SIZE) / V_NQ;
        let v_row = blck_idx / BLOCKS_V;
        let global_v_row = kv_tile + v_row;
        let block_k = blck_idx % BLOCKS_V;
        let row_offset = v_row * HEAD_DIM_V;
        let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k;
        let block_byte_base = global_block_idx * V_BLOCK_SIZE_BYTES;
        let d = f16_from_u16(load_v_u16_at(block_byte_base));
        let thread_byte_offset = block_offset * V_BYTES_PER_THREAD;
        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
        for (var j = 0u; j < V_BYTES_PER_THREAD / V_BYTES_PER_INNER_LOOP; j += 1u) {
            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * V_BYTES_PER_INNER_LOOP;
            let q_packed = load_v_u32_at(q_byte_offset);
            DEQUANT_V_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
        }
    }
#elif defined(FLASH_ATTN_SCALAR_KV)
    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
        let v_row = elem_idx / HEAD_DIM_V;
        let v_col = elem_idx % HEAD_DIM_V;
        let global_v_row = kv_tile + v_row;
        let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
        STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
            0.0,
            V[global_v_row_offset + v_col],
            global_v_row < params.seq_len_kv && v_col < HEAD_DIM_V));
    }
#else
    for (var vec_idx_local = local_x; vec_idx_local < kv_count * V_CHUNKS; vec_idx_local += WG_SIZE) {
        let kv_local = vec_idx_local / V_CHUNKS;
        let chunk = vec_idx_local % V_CHUNKS;
        let global_v_row = kv_tile + kv_local;
        let v_vec_index = (v_head_offset + global_v_row * params.stride_v1 + chunk * 4u) >> 2u;
        let v4 = V[v_vec_index];
        let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
        STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(v4.x);
        STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(v4.y);
        STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(v4.z);
        STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(v4.w);
    }
#endif
}
#endif // !defined(V_DIRECT)
