Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 09:15:04 +02:00
Code Issues Packages Projects Releases Wiki Activity

vulkan: add LIGHTNING_INDEXER op

Browse Source
This commit is contained in:
shenron0101
2026-08-12 22:27:12 +08:00
parent 521a64cd01
commit 9283284175
7 changed files with 427 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+136
View File
@@ -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;
@@ -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],
ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl
+50
View File
@@ -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;
}
}
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
+1 -50
View File
@@ -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.
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl
+20 -1
View File
@@ -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);
}
ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp
+116
View File
@@ -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]);
}
}
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+2
View File
@@ -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"}}));
tests/test-backend-ops.cpp
+102 -2
View File
@@ -7622,6 +7622,7 @@ struct input_tensor {
ggml_type type;
std::array<int64_t, 4> ne;
std::array<size_t, 4> 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<input_tensor>& 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<uint8_t> data(ggml_nbytes(t), 0);
std::vector<float> row(t->ne[0]);
std::vector<float> 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<std::unique_ptr<test_case>> 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;
}