mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-04 02:37:27 +02:00
Feature: Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend (#27453)
* vulkan: add LIGHTNING_INDEXER op
* vulkan: updated lightning_indexer.comp and ggml-vulkan.cpp with 128-lane dot-product reduction moved from a shared-memory tree to subgroupAdd.
* vulkan: cleanup; Skip bounds checks
* vulkan: cleanup FA_K_ONLY
* Revert "vulkan: cleanup FA_K_ONLY"
This reverts commit fdcbdd9151.
* vulkan: restore interleaved K/V buffer ordering
* vulkan: Remove FA_K_ONLY
* vulkan: Revert flash_attn_dequant
* vulkan: Revert tests in backend-ops.cpp
This commit is contained in:
@@ -767,6 +767,21 @@ static constexpr std::initializer_list<std::array<int, 3>> rms_norm_mul_rope_vie
|
||||
{ 4, 0, 3 }, // set_rows->src[0] == view
|
||||
};
|
||||
|
||||
static constexpr std::array<ggml_type, 9> lightning_indexer_k_types = {
|
||||
GGML_TYPE_F32,
|
||||
GGML_TYPE_F16,
|
||||
GGML_TYPE_BF16,
|
||||
GGML_TYPE_Q8_0,
|
||||
GGML_TYPE_Q5_1,
|
||||
GGML_TYPE_Q5_0,
|
||||
GGML_TYPE_Q4_1,
|
||||
GGML_TYPE_Q4_0,
|
||||
GGML_TYPE_IQ4_NL,
|
||||
};
|
||||
|
||||
static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) {
|
||||
return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end();
|
||||
}
|
||||
|
||||
struct vk_device_struct {
|
||||
std::recursive_mutex mutex;
|
||||
@@ -1068,6 +1083,7 @@ struct vk_device_struct {
|
||||
vk_pipeline pipeline_rwkv_wkv6_f32;
|
||||
vk_pipeline pipeline_rwkv_wkv7_f32;
|
||||
vk_pipeline pipeline_gated_linear_attn_f32;
|
||||
vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT];
|
||||
// [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128
|
||||
vk_pipeline pipeline_gated_delta_net[4][2];
|
||||
vk_pipeline pipeline_ssm_scan_f32_d128;
|
||||
@@ -1848,6 +1864,26 @@ struct vk_op_gated_linear_attn_push_constants {
|
||||
uint32_t H;
|
||||
float scale;
|
||||
};
|
||||
struct vk_op_lightning_indexer_push_constants {
|
||||
uint32_t n_kv;
|
||||
uint32_t n_heads;
|
||||
uint32_t n_tokens;
|
||||
uint32_t n_streams;
|
||||
uint32_t n_masks;
|
||||
uint32_t dispatch_x;
|
||||
uint32_t q_nb1;
|
||||
uint32_t q_nb2;
|
||||
uint32_t q_nb3;
|
||||
uint32_t k_nb2;
|
||||
uint32_t k_nb3;
|
||||
uint32_t w_nb1;
|
||||
uint32_t w_nb3;
|
||||
uint32_t m_nb1;
|
||||
uint32_t m_nb3;
|
||||
uint32_t d_nb1;
|
||||
uint32_t d_nb3;
|
||||
};
|
||||
static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128);
|
||||
struct vk_op_gated_delta_net_push_constants {
|
||||
uint32_t H;
|
||||
uint32_t n_tokens;
|
||||
@@ -3904,11 +3940,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const
|
||||
return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type};
|
||||
}
|
||||
|
||||
// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as
|
||||
// a vec4 "block" of 4 floats, everything else uses its ggml block size.
|
||||
static uint32_t fa_block_bytes(ggml_type t) {
|
||||
if (t == GGML_TYPE_F32) {
|
||||
return 16u;
|
||||
}
|
||||
return (uint32_t) ggml_type_size(t);
|
||||
}
|
||||
|
||||
static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& state) {
|
||||
const auto fa_block_bytes = [](ggml_type t) -> uint32_t {
|
||||
if (t == GGML_TYPE_F32) return 16u;
|
||||
return (uint32_t) ggml_type_size(t);
|
||||
};
|
||||
return {
|
||||
/* 0 WorkGroupSize */ state.workgroup_size,
|
||||
/* 1 Br */ state.Br,
|
||||
@@ -5847,6 +5888,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
|
||||
ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1);
|
||||
|
||||
{
|
||||
const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support;
|
||||
const size_t li_len = li_subgroup ? lightning_indexer_subgroup_f32_len : lightning_indexer_f32_len;
|
||||
const void * li_data = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data;
|
||||
|
||||
for (ggml_type k_type : lightning_indexer_k_types) {
|
||||
const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32";
|
||||
ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup);
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
const uint32_t gdn_sizes[] = {16, 32, 64, 128};
|
||||
const char * gdn_names[][2] = {
|
||||
@@ -11697,6 +11749,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
|
||||
return ctx->device->pipeline_gated_linear_attn_f32;
|
||||
}
|
||||
return nullptr;
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
// only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer()
|
||||
if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) {
|
||||
return ctx->device->pipeline_lightning_indexer_f32[src1->type];
|
||||
}
|
||||
return nullptr;
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
|
||||
const uint32_t S_v = dst->src[2]->ne[0];
|
||||
@@ -12772,6 +12830,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context&
|
||||
pc, { (uint32_t)(n_seqs * n_heads), 1, 1 });
|
||||
}
|
||||
|
||||
static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * q = dst->src[0];
|
||||
const ggml_tensor * k = dst->src[1];
|
||||
const ggml_tensor * w = dst->src[2];
|
||||
const ggml_tensor * m = dst->src[3];
|
||||
|
||||
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op);
|
||||
GGML_ASSERT(pipeline != nullptr);
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
|
||||
|
||||
const uint32_t n_kv = k->ne[2];
|
||||
const uint32_t n_heads = q->ne[1];
|
||||
const uint32_t n_tokens = q->ne[2];
|
||||
const uint32_t n_streams = q->ne[3];
|
||||
const uint32_t n_masks = m->ne[3];
|
||||
|
||||
const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]);
|
||||
const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
|
||||
const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x);
|
||||
|
||||
// q, w and dst are f32 and m is f16, so their strides are passed in elements;
|
||||
// k may be quantized, so its strides stay in bytes
|
||||
const uint32_t q_nb1 = q->nb[1] / sizeof(float);
|
||||
const uint32_t q_nb2 = q->nb[2] / sizeof(float);
|
||||
const uint32_t q_nb3 = q->nb[3] / sizeof(float);
|
||||
const uint32_t k_nb2 = k->nb[2];
|
||||
const uint32_t k_nb3 = k->nb[3];
|
||||
const uint32_t w_nb1 = w->nb[1] / sizeof(float);
|
||||
const uint32_t w_nb3 = w->nb[3] / sizeof(float);
|
||||
const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t);
|
||||
const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t);
|
||||
const uint32_t d_nb1 = dst->nb[1] / sizeof(float);
|
||||
const uint32_t d_nb3 = dst->nb[3] / sizeof(float);
|
||||
|
||||
const vk_op_lightning_indexer_push_constants pc = {
|
||||
n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x,
|
||||
q_nb1, q_nb2, q_nb3,
|
||||
k_nb2, k_nb3,
|
||||
w_nb1, w_nb3,
|
||||
m_nb1, m_nb3,
|
||||
d_nb1, d_nb3,
|
||||
};
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
|
||||
{ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)},
|
||||
pc, {dispatch_x, dispatch_y, 1});
|
||||
}
|
||||
|
||||
static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * src_q = dst->src[0];
|
||||
const ggml_tensor * src_v = dst->src[2];
|
||||
@@ -15898,6 +16005,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
|
||||
|
||||
break;
|
||||
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
ggml_vk_lightning_indexer(ctx, compute_ctx, node);
|
||||
|
||||
break;
|
||||
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
ggml_vk_gated_delta_net(ctx, compute_ctx, node);
|
||||
|
||||
@@ -18676,6 +18788,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
case GGML_OP_GATED_LINEAR_ATTN:
|
||||
// the shader block size is hardcoded to head_size 64
|
||||
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64;
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
{
|
||||
const ggml_tensor * q = op->src[0];
|
||||
const ggml_tensor * k = op->src[1];
|
||||
const ggml_tensor * w = op->src[2];
|
||||
const ggml_tensor * m = op->src[3];
|
||||
|
||||
// the q/w/m types and the shape relationships between q, k, w, m and dst
|
||||
// are already asserted in ggml_lightning_indexer()
|
||||
if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the shader block size is hardcoded to head size 128
|
||||
if (q->ne[0] != 128) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the shader indexes the buffers by element stride, and is dispatched
|
||||
// without allow_misalign
|
||||
for (const ggml_tensor * t : {q, k, w, m, op}) {
|
||||
if (t->nb[0] != ggml_type_size(t->type) ||
|
||||
(vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) {
|
||||
return false;
|
||||
}
|
||||
// the strides get scaled down from bytes, so the division must be exact
|
||||
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
|
||||
if (t->nb[i] % ggml_type_size(t->type) != 0) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
{
|
||||
const uint32_t S_v = op->src[2]->ne[0];
|
||||
@@ -19685,6 +19831,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
|
||||
const float * op_params = (const float *)tensor->op_params;
|
||||
tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1],
|
||||
src_clone[2], src_clone[3], src_clone[4], op_params[0]);
|
||||
} else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) {
|
||||
tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
|
||||
} else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
|
||||
tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1],
|
||||
src_clone[2], src_clone[3], src_clone[4], src_clone[5],
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
#if !defined(GGML_FA_TYPES_COMP)
|
||||
#define GGML_FA_TYPES_COMP
|
||||
|
||||
// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the
|
||||
// host can pass the type directly. Keep in sync with ggml.h.
|
||||
#define FA_TYPE_F32 0u
|
||||
#define FA_TYPE_F16 1u
|
||||
#define FA_TYPE_Q4_0 2u
|
||||
#define FA_TYPE_Q4_1 3u
|
||||
#define FA_TYPE_Q5_0 6u
|
||||
#define FA_TYPE_Q5_1 7u
|
||||
#define FA_TYPE_Q8_0 8u
|
||||
#define FA_TYPE_IQ4_NL 20u
|
||||
#define FA_TYPE_BF16 30u
|
||||
|
||||
// Number of matrix elements per buffer block, derived from the K/V type spec
|
||||
// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
|
||||
// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
|
||||
uint fa_block_elems(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_F32: return 4u;
|
||||
case FA_TYPE_F16: return 1u;
|
||||
case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
|
||||
case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
|
||||
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
|
||||
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
|
||||
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
|
||||
case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
|
||||
case FA_TYPE_BF16: return 1u;
|
||||
default: return 1u;
|
||||
}
|
||||
}
|
||||
|
||||
// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
|
||||
// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
|
||||
// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
|
||||
uint fa_quant_r_mmq(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
|
||||
case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
|
||||
case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
|
||||
case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
|
||||
case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
|
||||
default: return 1u;
|
||||
}
|
||||
}
|
||||
|
||||
bool fa_type_needs_shmem(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_IQ4_NL: return true;
|
||||
default: return false;
|
||||
}
|
||||
}
|
||||
|
||||
#endif // !defined(GGML_FA_TYPES_COMP)
|
||||
@@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
|
||||
#define BINDING_IDX_K 0
|
||||
#define BINDING_IDX_V 1
|
||||
|
||||
// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the
|
||||
// host can pass the type directly. Keep in sync with ggml.h.
|
||||
#define FA_TYPE_F32 0u
|
||||
#define FA_TYPE_F16 1u
|
||||
#define FA_TYPE_Q4_0 2u
|
||||
#define FA_TYPE_Q4_1 3u
|
||||
#define FA_TYPE_Q5_0 6u
|
||||
#define FA_TYPE_Q5_1 7u
|
||||
#define FA_TYPE_Q8_0 8u
|
||||
#define FA_TYPE_IQ4_NL 20u
|
||||
#define FA_TYPE_BF16 30u
|
||||
#include "fa_types.glsl"
|
||||
|
||||
#if defined(BFLOAT16)
|
||||
#define O_TYPE float
|
||||
@@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
|
||||
#define O_TYPEV4 FLOAT_TYPEV4
|
||||
#endif
|
||||
|
||||
// Number of matrix elements per buffer block, derived from the K/V type spec
|
||||
// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
|
||||
// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
|
||||
uint fa_block_elems(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_F32: return 4u;
|
||||
case FA_TYPE_F16: return 1u;
|
||||
case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
|
||||
case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
|
||||
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
|
||||
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
|
||||
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
|
||||
case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
|
||||
case FA_TYPE_BF16: return 1u;
|
||||
default: return 1u;
|
||||
}
|
||||
}
|
||||
|
||||
// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
|
||||
// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
|
||||
// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
|
||||
uint fa_quant_r_mmq(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
|
||||
case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
|
||||
case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
|
||||
case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
|
||||
case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
|
||||
default: return 1u;
|
||||
}
|
||||
}
|
||||
|
||||
bool fa_type_needs_shmem(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_IQ4_NL: return true;
|
||||
default: return false;
|
||||
}
|
||||
}
|
||||
|
||||
// These can't be `const` globals because GLSL forbids function calls in global
|
||||
// const initializers, even when the spec constants would let the driver fold
|
||||
// them. Macros expand at the use site and fold after specialization.
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#extension GL_EXT_shader_16bit_storage : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
||||
#extension GL_KHR_shader_subgroup_basic : enable
|
||||
#if USE_SUBGROUP_ADD
|
||||
#extension GL_KHR_shader_subgroup_arithmetic : enable
|
||||
#endif
|
||||
|
||||
#define BINDING_IDX_K 0u
|
||||
|
||||
#include "types.glsl"
|
||||
#include "fa_types.glsl"
|
||||
#define FaTypeV FA_TYPE_F32
|
||||
|
||||
layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32;
|
||||
layout(constant_id = 1) const uint FaBlockBytesK = 4;
|
||||
layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
|
||||
|
||||
#include "flash_attn_dequant.glsl"
|
||||
|
||||
// one workgroup computes one output element, one invocation per head element
|
||||
#define HEAD_SIZE 128
|
||||
|
||||
layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout(binding = 0) readonly buffer QBuf { float q[]; };
|
||||
layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
|
||||
layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
|
||||
layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
|
||||
layout(binding = 2) readonly buffer WBuf { float weights[]; };
|
||||
layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
|
||||
layout(binding = 4) writeonly buffer DstBuf { float dst[]; };
|
||||
|
||||
layout(push_constant) uniform PushConstants {
|
||||
uint n_kv;
|
||||
uint n_heads;
|
||||
uint n_tokens;
|
||||
uint n_streams;
|
||||
uint n_masks;
|
||||
uint dispatch_x;
|
||||
uint q_nb1;
|
||||
uint q_nb2;
|
||||
uint q_nb3;
|
||||
uint k_nb2;
|
||||
uint k_nb3;
|
||||
uint w_nb1;
|
||||
uint w_nb3;
|
||||
uint m_nb1;
|
||||
uint m_nb3;
|
||||
uint d_nb1;
|
||||
uint d_nb3;
|
||||
};
|
||||
|
||||
shared float k_row[HEAD_SIZE];
|
||||
|
||||
#if USE_SUBGROUP_ADD
|
||||
shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE];
|
||||
#else
|
||||
shared float partials[HEAD_SIZE];
|
||||
#endif
|
||||
|
||||
void main() {
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x;
|
||||
const uint n_outputs = n_kv * n_tokens * n_streams;
|
||||
|
||||
if (fa_type_needs_shmem(FaTypeK)) {
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
}
|
||||
|
||||
if (output_idx >= n_outputs) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint ik = output_idx % n_kv;
|
||||
const uint ts = output_idx / n_kv;
|
||||
const uint t = ts % n_tokens;
|
||||
const uint s = ts / n_tokens;
|
||||
const uint k_offset = ik * k_nb2 + s * k_nb3;
|
||||
|
||||
// k strides come in as bytes, so scale them down to the view being indexed
|
||||
const uint k_block_elems = fa_block_elems(FaTypeK);
|
||||
const uint k_elem_bytes = FaBlockBytesK / k_block_elems;
|
||||
|
||||
if (FaTypeK == FA_TYPE_F16) {
|
||||
k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
|
||||
} else if (FaTypeK == FA_TYPE_F32) {
|
||||
k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
|
||||
} else if (FaTypeK == FA_TYPE_BF16) {
|
||||
k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid]));
|
||||
} else if (4 * tid < HEAD_SIZE) {
|
||||
const uint coord = 4 * tid;
|
||||
const uint ib = coord / k_block_elems;
|
||||
const uint iqs = coord % k_block_elems;
|
||||
const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K);
|
||||
k_row[coord + 0] = values.x;
|
||||
k_row[coord + 1] = values.y;
|
||||
k_row[coord + 2] = values.z;
|
||||
k_row[coord + 3] = values.w;
|
||||
}
|
||||
barrier();
|
||||
|
||||
const float k_val = k_row[tid];
|
||||
|
||||
float score = 0.0;
|
||||
for (uint h = 0; h < n_heads; ++h) {
|
||||
const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val;
|
||||
|
||||
#if USE_SUBGROUP_ADD
|
||||
const float sg_sum = subgroupAdd(prod);
|
||||
if (gl_SubgroupInvocationID == 0) {
|
||||
sg_partials[gl_SubgroupID] = sg_sum;
|
||||
}
|
||||
barrier();
|
||||
|
||||
if (tid == 0) {
|
||||
float sum = 0.0;
|
||||
[[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) {
|
||||
sum += sg_partials[i];
|
||||
}
|
||||
score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3];
|
||||
}
|
||||
// the reads above must complete before the next iteration overwrites sg_partials
|
||||
barrier();
|
||||
#else
|
||||
partials[tid] = prod;
|
||||
barrier();
|
||||
|
||||
[[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
partials[tid] += partials[tid + stride];
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3];
|
||||
}
|
||||
// the read of partials[0] above must complete before the next iteration
|
||||
// overwrites partials[tid]
|
||||
barrier();
|
||||
#endif
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
|
||||
dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]);
|
||||
}
|
||||
}
|
||||
@@ -1069,6 +1069,12 @@ void process_shaders() {
|
||||
|
||||
string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
|
||||
|
||||
// Compile IQ4_NL support in so its shared LUT is available when K uses it.
|
||||
// K quant type is selected at runtime via the FaTypeK spec constant.
|
||||
std::map<std::string, std::string> li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}};
|
||||
string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict);
|
||||
string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}}));
|
||||
|
||||
string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
|
||||
|
||||
string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}}));
|
||||
|
||||
Reference in New Issue
Block a user