diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f6cbaecb7f..c5ac3670b6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -767,6 +767,21 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::array 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; @@ -1065,6 +1080,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; @@ -1845,6 +1861,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; @@ -5833,6 +5869,12 @@ 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); + for (ggml_type k_type : lightning_indexer_k_types) { + const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32"; + const uint32_t block_bytes = k_type == GGML_TYPE_F32 ? 4 * sizeof(float) : (uint32_t)ggml_type_size(k_type); + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), lightning_indexer_f32_len, lightning_indexer_f32_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, block_bytes}, 1); + } + { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; const char * gdn_names[][2] = { @@ -11667,6 +11709,13 @@ 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: + if (src0->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && dst->src[3]->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + 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]; @@ -12741,6 +12790,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]; @@ -15747,6 +15845,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); @@ -18512,6 +18615,37 @@ 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]; + if (!q || !k || !w || !m) { + return false; + } + + // 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 || k->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; + } + } + return true; + } case GGML_OP_GATED_DELTA_NET: { const uint32_t S_v = op->src[2]->ne[0]; @@ -19515,6 +19649,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], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl new file mode 100644 index 0000000000..a8ac5327fb --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -0,0 +1,50 @@ +// 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; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 3c64f91dad..0ce4503a88 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -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. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl index 8ba4725f33..ecf08af46e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl @@ -15,23 +15,39 @@ // F32 is fed as a vec4 "block" (4 floats), matching what dequant_funcs_cm2.glsl // does for F32 in the cm2 shader. FaBlockBytesK/V == 16 for F32. layout (binding = 1) readonly buffer K_PACKED_F32 { vec4 data[]; } k_packed_f32; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_F32 { vec4 data[]; } v_packed_f32; +#endif layout (binding = 1) readonly buffer K_PACKED_Q4_0 { block_q4_0_packed16 data[]; } k_packed_q4_0; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_Q4_0 { block_q4_0_packed16 data[]; } v_packed_q4_0; +#endif layout (binding = 1) readonly buffer K_PACKED_Q4_1 { block_q4_1_packed16 data[]; } k_packed_q4_1; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_Q4_1 { block_q4_1_packed16 data[]; } v_packed_q4_1; +#endif layout (binding = 1) readonly buffer K_PACKED_Q5_0 { block_q5_0_packed16 data[]; } k_packed_q5_0; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_Q5_0 { block_q5_0_packed16 data[]; } v_packed_q5_0; +#endif layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; } k_packed_q5_1; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1; +#endif layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0; +#endif layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl; +#endif layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16; +#ifndef FA_K_ONLY layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16; +#endif // Q4_1 and Q5_1 packed32 views: aliased to the same memory as the packed16 // views, used by the MMQ K-side hot path for fast 4-uint loads. @@ -130,7 +146,9 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl) case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16) } - } else { + } +#ifndef FA_K_ONLY + else { switch (FaTypeV) { case FA_TYPE_F32: FA_DEQUANT4_F32 (v_packed_f32) case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0) @@ -142,5 +160,6 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16) } } +#endif return FLOAT_TYPEV4(0); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp new file mode 100644 index 0000000000..ddca015960 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -0,0 +1,116 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +#define DATA_A_IQ4_NL +#define FLOAT_TYPE float +#define FLOAT_TYPEV4 vec4 +#define FA_K_ONLY +#define BINDING_IDX_K 0u + +#include "types.glsl" +#include "fa_types.glsl" + +layout(constant_id = 0) const uint K_TYPE = FA_TYPE_F32; +layout(constant_id = 1) const uint K_BLOCK_BYTES = 4; + +#define FaTypeK K_TYPE +#include "flash_attn_dequant.glsl" + +layout(local_size_x = 128, 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[128]; +shared float partials[128]; + +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(K_TYPE)) { + 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; + + if (K_TYPE == FA_TYPE_F16) { + k_row[tid] = float(k_f16[k_offset / 2 + tid]); + } else if (K_TYPE == FA_TYPE_F32) { + k_row[tid] = k_f32[k_offset / 4 + tid]; + } else if (K_TYPE == FA_TYPE_BF16) { + k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / 2 + tid])); + } else if (tid < 32) { + const uint coord = 4 * tid; + const uint block_size = 32u; + const uint ib = coord / block_size; + const uint iqs = coord % block_size; + const vec4 values = dequantize4(ib, iqs, k_offset / K_BLOCK_BYTES, 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(); + + float score = 0.0; + for (uint h = 0; h < n_heads; ++h) { + partials[tid] = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_row[tid]; + barrier(); + + for (uint stride = 64; 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(); + } + + 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]); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index caa0c889a4..a743c7a015 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1066,6 +1066,8 @@ void process_shaders() { string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", {}); + 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"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 17098825bc..26f22ef719 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7622,6 +7622,7 @@ struct input_tensor { ggml_type type; std::array ne; std::array nb; // strides (0 = use default contiguous strides) + size_t view_offs = 0; }; static bool is_non_contiguous(const input_tensor & src) { @@ -7647,6 +7648,9 @@ static std::string var_to_str(const std::vector& sources) { if (is_non_contiguous(src)) { oss << "nb[" << src.nb[0] << "," << src.nb[1] << "," << src.nb[2] << "," << src.nb[3] << "]"; } + if (src.view_offs != 0) { + oss << "offs[" << src.view_offs << "]"; + } first = false; } return oss.str(); @@ -7711,6 +7715,7 @@ struct test_generic_op : public test_case { total_size += (src.ne[d] - 1) * src.nb[d]; } } + total_size += src.view_offs; // Convert bytes to elements, padded to block size for quantized types const size_t type_size = ggml_type_size(src.type); @@ -7719,7 +7724,7 @@ struct test_generic_op : public test_case { ggml_tensor * backing = ggml_new_tensor_1d(ctx, src.type, backing_elements); source_tensors[i] = ggml_view_4d(ctx, backing, src.ne[0], src.ne[1], src.ne[2], src.ne[3], - src.nb[1], src.nb[2], src.nb[3], 0); + src.nb[1], src.nb[2], src.nb[3], src.view_offs); // nb[0] does not get set by view_4d, so set it manually source_tensors[i]->nb[0] = src.nb[0]; } else { @@ -7765,6 +7770,7 @@ struct test_generic_op : public test_case { case GGML_OP_CPY: return 5e-4; case GGML_OP_SOFT_MAX: + case GGML_OP_LIGHTNING_INDEXER: return 1e-6; case GGML_OP_RWKV_WKV7: return 5e-3; @@ -7791,6 +7797,37 @@ struct test_generic_op : public test_case { break; } + if (op == GGML_OP_LIGHTNING_INDEXER && t->view_src != nullptr) { + std::vector data(ggml_nbytes(t), 0); + std::vector row(t->ne[0]); + std::vector imatrix(t->ne[0], 1.0f); + + for (int64_t i3 = 0; i3 < t->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < t->ne[2]; ++i2) { + for (int64_t i1 = 0; i1 < t->ne[1]; ++i1) { + const size_t offset = i3*t->nb[3] + i2*t->nb[2] + i1*t->nb[1]; + const int64_t row_index = (i3*t->ne[2] + i2)*t->ne[1] + i1; + for (int64_t i0 = 0; i0 < t->ne[0]; ++i0) { + row[i0] = ((int) ((row_index*t->ne[0] + i0) % 31) - 15) / 16.0f; + } + + if (t->type == GGML_TYPE_F32) { + memcpy(data.data() + offset, row.data(), ggml_row_size(t->type, t->ne[0])); + } else if (t->type == GGML_TYPE_F16) { + ggml_fp32_to_fp16_row(row.data(), (ggml_fp16_t *) (data.data() + offset), t->ne[0]); + } else if (t->type == GGML_TYPE_BF16) { + ggml_fp32_to_bf16_row(row.data(), (ggml_bf16_t *) (data.data() + offset), t->ne[0]); + } else { + GGML_ASSERT(ggml_is_quantized(t->type)); + ggml_quantize_chunk(t->type, row.data(), data.data() + offset, 0, 1, t->ne[0], imatrix.data()); + } + } + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()); + continue; + } + // FLASH_ATTN_EXT: src[3] is the KQ mask if (op == GGML_OP_FLASH_ATTN_EXT && i == 3) { init_tensor_kq_mask(t); @@ -10098,10 +10135,73 @@ static std::vector> make_test_cases_eval() { } for (int kv : { 1, 7, 8, 63, 64, 65 }) { - for (ggml_type type_K : {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}) { + for (ggml_type type_K : {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}) { test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 32, 4, 1, type_K)); } } + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_F16, { 128, 1, 64, 4 }, { 2, 272, 288, 18464 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_f16_k_mask_broadcast")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_F32, { 128, 1, 64, 4 }, { 4, 528, 544, 34880 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 4 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_f32_k_per_stream_mask")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_BF16, { 128, 1, 64, 4 }, { 2, 274, 290, 18562 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_bf16_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_Q8_0, { 128, 1, 64, 4 }, { 34, 170, 204, 13158 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_q8_0_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_Q5_1, { 128, 1, 64, 4 }, { 24, 120, 144, 9336 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_q5_1_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_Q5_0, { 128, 1, 64, 4 }, { 22, 110, 132, 8558 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_q5_0_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_Q4_1, { 128, 1, 64, 4 }, { 20, 100, 120, 7780 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_q4_1_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_Q4_0, { 128, 1, 64, 4 }, { 18, 90, 108, 7002 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_q4_0_k")); + test_cases.emplace_back(new test_generic_op( + GGML_OP_LIGHTNING_INDEXER, GGML_TYPE_F32, { 64, 3, 1, 4 }, {}, { + { GGML_TYPE_F32, { 128, 32, 3, 4 }, { 4, 528, 16928, 50848 }, 256 }, + { GGML_TYPE_IQ4_NL, { 128, 1, 64, 4 }, { 18, 90, 108, 7002 }, 256 }, + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, 1 }, { 2, 144, 448, 480 }, 256 }, + }, "padded_iq4_nl_k")); return test_cases; }