Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts (#25956)

Browse Source 
* improve fa of quantized kv cache

* Fix some bugs and some comments.

* fix v type check and some comments

* Fix build error caused by rebasing

* editorconfig checking pass
This commit is contained in:
Masashi Yoshimura
2026-07-31 09:08:40 +03:00
committed by GitHub
parent 5f55650a78
commit 8f4646a63e
7 changed files with 264 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl
+48 -3
View File
@@ -9,6 +9,12 @@ fn get_byte_i32(value: u32, index: u32) -> i32 {
#endif
#ifdef U32_DEQUANT_HELPERS
fn f16_from_u16(bits: u32) -> f16 {
let packed = unpack2x16float(bits);
return f16(packed[0]);
}
#ifdef DECLARE_BYTE_LOADERS_SRC
fn load_u16_at_src(byte_offset: u32) -> u32 {
let word = src[byte_offset / 4u];
@@ -36,7 +42,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
let d_bits = (word >> shift) & 0xFFFFu;
return unpack2x16float(d_bits)[0];
}
#endif
#endif // DECLARE_BYTE_LOADERS_SRC
#ifdef DECLARE_BYTE_LOADERS_SRC0
fn load_u16_at_src0(byte_offset: u32) -> u32 {
@@ -72,8 +78,47 @@ fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
let d_bits = (word >> shift) & 0xFFFFu;
return unpack2x16float(d_bits)[0];
}
#endif
#endif
#endif // DECLARE_BYTE_LOADERS_SRC0
#ifdef LOADERS_QUANTIZED_K
fn load_k_u16_at(byte_offset: u32) -> u32 {
let word = K[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}
fn load_k_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = K[word_idx];
if (shift == 0u) {
return lo;
}
let hi = K[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif // LOADERS_QUANTIZED_K
#ifdef LOADERS_QUANTIZED_V
fn load_v_u16_at(byte_offset: u32) -> u32 {
let word = V[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}
fn load_v_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = V[word_idx];
if (shift == 0u) {
return lo;
}
let hi = V[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif // LOADERS_QUANTIZED_V
#endif // U32_DEQUANT_HELPERS
ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl
+13 -14
View File
@@ -138,7 +138,7 @@ const FLOAT_MIN: f32 = -1.0e9;
// The number of Q rows processed per workgroup
var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;
#ifndef KV_DIRECT
#if !defined(K_DIRECT) || !defined(V_DIRECT)
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
@@ -183,13 +183,12 @@ fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u3
return (*buf)[scalar_index >> 2u];
}
#ifndef KV_DIRECT
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f16
#include "quant_inner_loops.tmpl"
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_Q4_0) && !defined(K_Q8_0)
#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
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;
@@ -204,7 +203,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
}
#endif
#if !defined(V_Q4_0) && !defined(V_Q8_0)
#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
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;
@@ -296,7 +295,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}
// load k tile into shared memory
#ifndef KV_DIRECT
#ifndef K_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#endif
@@ -306,7 +305,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
// TODO: this loop seems to be the current largest bottleneck
// this bracket exists to scope the lifetime of variables, reducing register pressure
{
#ifdef KV_DIRECT
#ifdef K_DIRECT
let k_block_row = kv_tile + subgroup_id * SG_MAT_N;
var k_global_offset = k_head_offset + k_block_row * params.stride_k1;
#else
@@ -318,7 +317,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
var q_cur = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, 0u, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + 0u, true, params.stride_k1);
#else
var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + 0u, true, HEAD_DIM_QK);
@@ -328,7 +327,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
for (; t + 1u < HEAD_DIM_QK / SG_MAT_K; t += 2u) {
let h0 = t * SG_MAT_K;
var q0 = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h0, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h0, true, params.stride_k1);
#else
var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h0, true, HEAD_DIM_QK);
@@ -339,7 +338,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
let h1 = (t + 1u) * SG_MAT_K;
var q1g = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h1, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h1, true, params.stride_k1);
#else
var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h1, true, HEAD_DIM_QK);
@@ -353,7 +352,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
if (t < HEAD_DIM_QK / SG_MAT_K) {
let h = t * SG_MAT_K;
var qn = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h, true, params.stride_k1);
#else
var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h, true, HEAD_DIM_QK);
@@ -365,7 +364,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
acc = subgroupMatrixMultiplyAccumulate(q_cur, k_cur, acc);
#ifdef KV_DIRECT
#ifdef K_DIRECT
k_global_offset += num_subgroups * SG_MAT_N * params.stride_k1;
#else
k_block_offset += num_subgroups * SG_MAT_N * HEAD_DIM_QK;
@@ -436,7 +435,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}
// load v tile into shared memory
#ifndef KV_DIRECT
#ifndef V_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif
@@ -464,7 +463,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
);
// load V submatrix from global or shared memory
#ifdef KV_DIRECT
#ifdef V_DIRECT
let v_block_row = kv_tile + kv_block * SG_MAT_N;
let v_global_offset = v_head_offset + v_block_row * params.stride_v1 + head_dim_block;
var v_sg_mat: subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K> = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(
ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_quant_staging.tmpl
+2 -43
View File
@@ -1,3 +1,5 @@
#include "quant_inner_loops.tmpl"
#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)
@@ -26,49 +28,6 @@
#define V_BYTES_PER_INNER_LOOP 4u
#endif
#if defined(K_Q4_0) || defined(K_Q8_0)
fn load_k_u16_at(byte_offset: u32) -> u32 {
let word = K[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}
fn load_k_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = K[word_idx];
if (shift == 0u) {
return lo;
}
let hi = K[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif
#if defined(V_Q4_0) || defined(V_Q8_0)
fn load_v_u16_at(byte_offset: u32) -> u32 {
let word = V[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}
fn load_v_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = V[word_idx];
if (shift == 0u) {
return lo;
}
let hi = V[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif
fn f16_from_u16(bits: u32) -> f16 {
let packed = unpack2x16float(bits);
return f16(packed[0]);
}
#if defined(K_Q4_0) || defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_tile.wgsl
+6 -3
View File
@@ -153,7 +153,6 @@ var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f16
#include "quant_inner_loops.tmpl"
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_Q4_0) && !defined(K_Q8_0)
@@ -270,7 +269,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
local_scores[slot] = FLOAT_MIN;
}
#ifndef KV_DIRECT
// The tile path stages K/V in shared memory so each tile can be reused across
// Q_TILE query rows. It therefore does not use the direct path.
#ifndef K_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#endif
@@ -333,7 +334,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
workgroupBarrier();
#ifndef KV_DIRECT
// The tile path stages K/V in shared memory so each tile can be reused across
// Q_TILE query rows. It therefore does not use the direct path.
#ifndef V_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif
ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
+144 -42
View File
@@ -196,49 +196,35 @@ struct Params {
// Just a very small float value.
const FLOAT_MIN: f32 = -1.0e9;
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
var<workgroup> q_shmem: array<f32, HEAD_DIM_QK>;
#ifndef KV_DIRECT
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
#endif
var<workgroup> o_shmem: array<f32, HEAD_DIM_V>;
// note that we reuse the same storage for both since we only need one at a time
var<workgroup> inter_shmem: array<f32, KV_TILE>;
#ifdef MASK
// storage for mask values
var<workgroup> mask_shmem: array<f32, KV_TILE>;
#endif
// note that we reuse the same storage for both since we only need one at a time
var<workgroup> inter_shmem: array<f32, KV_TILE>;
// Storage for row max and exp sum during online softmax
fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
var v = select(FLOAT_MIN,
inter_shmem[kv_idx] * params.scale,
kv_idx < KV_TILE);
#ifdef LOGIT_SOFTCAP
v = params.logit_softcap * tanh(v);
#if defined(K_DIRECT) || defined(V_DIRECT)
// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
// so caching it is more efficient, even on the direct path.
var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
#endif
#ifdef MASK
if (apply_mask) {
var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
v += select(mask_val, slope * mask_val, has_bias);
}
#endif
return v;
}
#ifndef KV_DIRECT
// K/V shared memory handling
#if !defined(K_DIRECT) || !defined(V_DIRECT)
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f32
#include "quant_inner_loops.tmpl"
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_Q4_0) && !defined(K_Q8_0)
#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE * 4u) {
let k_row = elem_idx / HEAD_DIM_QK;
@@ -256,7 +242,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
}
#endif
#if !defined(V_Q4_0) && !defined(V_Q8_0)
#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE * 4u) {
let v_row = elem_idx / HEAD_DIM_V;
@@ -273,7 +259,24 @@ fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u
}
}
#endif
#endif // !defined(K_DIRECT) || !defined(V_DIRECT)
// Storage for row max and exp sum during online softmax
fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
var v = select(FLOAT_MIN,
inter_shmem[kv_idx] * params.scale,
kv_idx < KV_TILE);
#ifdef LOGIT_SOFTCAP
v = params.logit_softcap * tanh(v);
#endif
#ifdef MASK
if (apply_mask) {
var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
v += select(mask_val, slope * mask_val, has_bias);
}
#endif
return v;
}
@compute @workgroup_size(WG_SIZE)
fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
@@ -355,12 +358,31 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
inter_shmem[elem_idx] = 0.0;
}
// load k tile into shared memory
#ifndef KV_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#ifdef K_DIRECT
// load only the scale factor (d) from each quantized block into shared memory on the direct path.
#if defined(K_Q8_0)
for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
let kv_row = kv_tile + j / HEAD_DIM_QK;
let block_idx = (j % HEAD_DIM_QK) / 32;
let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
d_shmem[j / 32] = d;
}
#elif defined(K_Q4_0)
for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
let kv_row = kv_tile + j / HEAD_DIM_QK;
let block_idx = (j % HEAD_DIM_QK) / 32;
let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
d_shmem[j / 32] = d;
}
#endif
#else
// load k tile into shared memory
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#endif // defined(K_DIRECT)
workgroupBarrier();
workgroupBarrier();
// accumulate q block * k block into registers across the entire KV tile
if (!skip_tile) {
@@ -381,9 +403,40 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
q_shmem[q_off + 1u],
q_shmem[q_off + 2u],
q_shmem[q_off + 3u]);
#ifdef KV_DIRECT
#ifdef K_DIRECT
#if defined(K_Q8_0)
let kv_row = kv_tile + kv_idx;
let block_idx = (i * 4u) / 32;
let id_in_block = (i * 4u) % 32;
let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
let q_byte_base = block_byte_base + 2u;
let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
let q8u4 = load_k_u32_at(q_byte_base + id_in_block);
let kv = vec4<f32>(
d * f32(get_byte_i32(q8u4, 0)),
d * f32(get_byte_i32(q8u4, 1)),
d * f32(get_byte_i32(q8u4, 2)),
d * f32(get_byte_i32(q8u4, 3)),
);
#elif defined(K_Q4_0)
let kv_row = kv_tile + kv_idx;
let block_idx = (i * 4u) / 32;
let id_in_block = (i * 4u) % 32;
let phase = id_in_block / 16;
let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
let q_byte_base = block_byte_base + 2u;
let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
let q8u4 = load_k_u32_at(q_byte_base + (id_in_block - phase * 16u));
let kv = vec4<f32>(
d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
);
#else
let idx = k_head_offset + (kv_tile + kv_idx) * params.stride_k1 + (i * 4u);
let kv = vec4<f32>(K[idx >> 2u]);
#endif
#else
let idx = kv_idx * HEAD_DIM_QK + (i * 4u);
let kv = vec4<f32>(
@@ -391,7 +444,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
kv_shmem[idx + 1u],
kv_shmem[idx + 2u],
kv_shmem[idx + 3u]);
#endif
#endif // defined(K_DIRECT)
partial_sum += dot(qv, kv);
}
}
@@ -473,12 +526,32 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}
}
// load v tile into shared memory
#ifndef KV_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif
workgroupBarrier();
#ifdef V_DIRECT
// load only `d` of quantized block into shared memory in the direct path
#if defined(V_Q8_0)
for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
let v_row = kv_tile + j / HEAD_DIM_V;
let block_idx = (j % HEAD_DIM_V) / 32;
let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
d_shmem[j / 32] = d;
}
#elif defined(V_Q4_0)
for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
let v_row = kv_tile + j / HEAD_DIM_V;
let block_idx = (j % HEAD_DIM_V) / 32;
let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
d_shmem[j / 32] = d;
}
#endif
#else
// load v tile into shared memory
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif // V_DIRECT
workgroupBarrier();
if (!skip_tile) {
// we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
@@ -501,9 +574,38 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}
let p = inter_shmem[kv_idx];
#ifdef KV_DIRECT
#ifdef V_DIRECT
#if defined(V_Q8_0)
let block_idx = (vec_col * 4u) / 32;
let id_in_block = (vec_col * 4u) % 32;
let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
let q_byte_base = block_byte_base + 2u;
let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
let q8u4 = load_v_u32_at(q_byte_base + id_in_block);
let v4 = vec4<f32>(
d * f32(get_byte_i32(q8u4, 0)),
d * f32(get_byte_i32(q8u4, 1)),
d * f32(get_byte_i32(q8u4, 2)),
d * f32(get_byte_i32(q8u4, 3)),
);
#elif defined(V_Q4_0)
let block_idx = (vec_col * 4u) / 32;
let id_in_block = (vec_col * 4u) % 32;
let phase = id_in_block / 16;
let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
let q_byte_base = block_byte_base + 2u;
let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
let q8u4 = load_v_u32_at(q_byte_base + (id_in_block - phase * 16u));
let v4 = vec4<f32>(
d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
);
#else
let v_idx = v_head_offset + v_row * params.stride_v1 + vec_col * 4u;
let v4 = vec4<f32>(V[v_idx >> 2u]);
#endif
#else
let v_idx = kv_idx * HEAD_DIM_V + vec_col * 4u;
let v4 = vec4<f32>(
@@ -511,7 +613,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
kv_shmem[v_idx + 1u],
kv_shmem[v_idx + 2u],
kv_shmem[v_idx + 3u]);
#endif
#endif // defined(V_DIRECT)
lo += p * v4;
}