ggml_type dst_type;
uint32_t head_dim_qk;
uint32_t head_dim_v;
- bool kv_direct;
+ bool k_direct;
+ bool v_direct;
bool kv_overlap;
bool has_mask;
bool has_sinks;
bool operator==(const ggml_webgpu_flash_attn_common_pipeline_key & other) const {
return q_type == other.q_type && k_type == other.k_type && v_type == other.v_type &&
dst_type == other.dst_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v &&
- kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask &&
- has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap;
+ k_direct == other.k_direct && v_direct == other.v_direct && kv_overlap == other.kv_overlap &&
+ has_mask == other.has_mask && has_sinks == other.has_sinks &&
+ uses_logit_softcap == other.uses_logit_softcap;
}
};
ggml_webgpu_hash_combine(seed, key.dst_type);
ggml_webgpu_hash_combine(seed, key.head_dim_qk);
ggml_webgpu_hash_combine(seed, key.head_dim_v);
- ggml_webgpu_hash_combine(seed, key.kv_direct);
+ ggml_webgpu_hash_combine(seed, key.k_direct);
+ ggml_webgpu_hash_combine(seed, key.v_direct);
ggml_webgpu_hash_combine(seed, key.kv_overlap);
ggml_webgpu_hash_combine(seed, key.has_mask);
ggml_webgpu_hash_combine(seed, key.has_sinks);
ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
}
-inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q,
- const ggml_tensor * K,
- const ggml_tensor * V,
- uint32_t kv_direct_align) {
- return K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && (Q->ne[0] % kv_direct_align == 0) &&
- (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+inline bool ggml_webgpu_flash_attn_k_direct(const ggml_tensor * Q, const ggml_tensor * K, uint32_t kv_direct_align) {
+ return (K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_Q4_0) &&
+ (Q->ne[0] % kv_direct_align == 0) && (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+}
+
+inline bool ggml_webgpu_flash_attn_v_direct(const ggml_tensor * Q, const ggml_tensor * V, uint32_t kv_direct_align) {
+ return ggml_webgpu_flash_attn_k_direct(Q, V, kv_direct_align);
}
inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key(
key.dst_type = context.dst->type;
key.head_dim_qk = (uint32_t) context.src0->ne[0];
key.head_dim_v = (uint32_t) context.src2->ne[0];
- key.kv_direct = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align);
- key.kv_overlap = kv_overlap;
- key.has_mask = context.src3 != nullptr;
- key.has_sinks = context.src4 != nullptr;
+ key.k_direct = ggml_webgpu_flash_attn_k_direct(context.src0, context.src1, kv_direct_align);
+ key.v_direct = ggml_webgpu_flash_attn_v_direct(context.src0, context.src2, kv_direct_align);
+ key.kv_overlap = kv_overlap;
+ key.has_mask = context.src3 != nullptr;
+ key.has_sinks = context.src4 != nullptr;
key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
return key;
}
defines.push_back("LOGIT_SOFTCAP");
variant += "_lgsc";
}
- if (key.kv_direct) {
- defines.push_back("KV_DIRECT");
- variant += "_kvdirect";
+ if (key.k_direct) {
+ defines.push_back("K_DIRECT");
+ variant += "_k_direct";
+ }
+ if (key.v_direct) {
+ defines.push_back("V_DIRECT");
+ variant += "_v_direct";
}
if (key.kv_overlap) {
defines.push_back("KV_OVERLAP");
if (ggml_is_quantized(key.k_type) || ggml_is_quantized(key.v_type)) {
defines.push_back("U32_DEQUANT_HELPERS");
+ if (ggml_is_quantized(key.k_type)) {
+ defines.push_back("LOADERS_QUANTIZED_K");
+ }
+ if (ggml_is_quantized(key.v_type)) {
+ defines.push_back("LOADERS_QUANTIZED_V");
+ }
}
return defines;
ggml_webgpu_flash_attn_pipeline_key key = {};
key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(
context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap);
- key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct;
- key.use_sg_matrix = decisions.use_sg_matrix;
+ key.common.k_direct &= decisions.use_sg_matrix && key.common.k_type == GGML_TYPE_F16;
+ key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16;
+ key.use_sg_matrix = decisions.use_sg_matrix;
const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u,
- key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+ key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask,
+ key.common.k_direct || key.common.v_direct);
GGML_ASSERT(max_kv_tile > 0);
decisions.kv_tile = decisions.use_sg_matrix ?
std::min(context.max_wg_size, std::max(GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE,
GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE * context.max_subgroup_size));
- if (key.common.kv_direct) {
+ if (key.common.k_direct || key.common.v_direct) {
decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD);
while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) {
decisions.kv_tile -= decisions.use_sg_matrix ? context.sg_mat_n : context.min_subgroup_size;
}
ggml_webgpu_flash_attn_vec_decisions decisions = {};
- decisions.kv_tile =
- ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
- key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+ decisions.kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
+ key.common.head_dim_v, key.common.has_mask,
+ key.common.k_direct || key.common.v_direct);
decisions.wg_size = context.max_subgroup_size;
std::string variant = "flash_attn_vec";
variant += "_mask_blk";
}
- uint32_t d_split = context.min_subgroup_size;
- if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16) {
- const uint32_t D = key.common.head_dim_qk | key.common.head_dim_v;
- const uint32_t D_lsb = D & (~(D - 1u));
- d_split = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
- }
+ uint32_t d_split = context.min_subgroup_size;
+ const uint32_t D = key.common.head_dim_qk | key.common.head_dim_v;
+ const uint32_t D_lsb = D & (~(D - 1u));
+ d_split = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
variant += "_dsplit" + std::to_string(d_split);
const auto & capabilities = ctx->webgpu_global_ctx->capabilities;
if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V)) {
const bool kv_direct =
- ggml_webgpu_flash_attn_kv_direct(Q, K, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
+ ggml_webgpu_flash_attn_k_direct(Q, K, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH) ||
+ ggml_webgpu_flash_attn_v_direct(Q, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
const uint32_t kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(
capabilities.limits.maxComputeWorkgroupStorageSize, (uint32_t) Q->ne[0], (uint32_t) V->ne[0],
mask != nullptr, kv_direct);
const uint32_t q_tile =
use_subgroup_matrix ? capabilities.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE;
const uint32_t kv_granularity = use_subgroup_matrix ? capabilities.sg_mat_n : 1u;
- const bool kv_direct = use_subgroup_matrix ?
- ggml_webgpu_flash_attn_kv_direct(src0, src1, src2, capabilities.sg_mat_k) :
- false;
+ const bool kv_direct = use_subgroup_matrix ?
+ ggml_webgpu_flash_attn_k_direct(src0, src1, capabilities.sg_mat_k) ||
+ ggml_webgpu_flash_attn_v_direct(src0, src2, capabilities.sg_mat_k) :
+ false;
const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
capabilities.limits.maxComputeWorkgroupStorageSize, q_tile, kv_granularity, (uint32_t) src0->ne[0],
(uint32_t) src2->ne[0], op->src[3] != nullptr, kv_direct);
#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];
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 {
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
// 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>;
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;
}
#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;
}
// 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
// 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
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);
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);
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);
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);
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;
}
// 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
);
// 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>>(
+#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)
#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) {
#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)
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
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
// Just a very small float value.
const FLOAT_MIN: f32 = -1.0e9;
-
-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> q_shmem: array<f32, HEAD_DIM_QK>;
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;
}
#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;
}
}
#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>,
inter_shmem[elem_idx] = 0.0;
}
+#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
-#ifndef KV_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
-#endif
+#endif // defined(K_DIRECT)
- workgroupBarrier();
+ workgroupBarrier();
// accumulate q block * k block into registers across the entire KV tile
if (!skip_tile) {
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>(
kv_shmem[idx + 1u],
kv_shmem[idx + 2u],
kv_shmem[idx + 3u]);
-#endif
+#endif // defined(K_DIRECT)
partial_sum += dot(qv, kv);
}
}
}
}
+
+#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
-#ifndef KV_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
-#endif
+#endif // V_DIRECT
- workgroupBarrier();
+ workgroupBarrier();
if (!skip_tile) {
// we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
}
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>(
kv_shmem[v_idx + 1u],
kv_shmem[v_idx + 2u],
kv_shmem[v_idx + 3u]);
-#endif
+#endif // defined(V_DIRECT)
lo += p * v4;
}
overview / pkg / ggml / sources / llama.cpp / commitdiff