From 1065050f59e0a8d2ddb5d95594830c9818a2b9fb Mon Sep 17 00:00:00 2001 From: shenron0101 Date: Fri, 21 Aug 2026 01:40:32 +0800 Subject: [PATCH] vulkan: updated lightning_indexer.comp and ggml-vulkan.cpp with 128-lane dot-product reduction moved from a shared-memory tree to subgroupAdd. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 44 +++-- .../ggml-vulkan/vulkan-shaders/fa_types.glsl | 5 + .../vulkan-shaders/flash_attn_dequant.glsl | 40 ++--- .../vulkan-shaders/lightning_indexer.comp | 79 ++++++--- .../vulkan-shaders/vulkan-shaders-gen.cpp | 6 +- tests/test-backend-ops.cpp | 151 +++++++++--------- 6 files changed, 181 insertions(+), 144 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c5ac3670b6..633f39df55 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3937,11 +3937,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 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, @@ -5869,10 +5874,15 @@ 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 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, false, li_subgroup); + } } { @@ -11710,10 +11720,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const } 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]; - } + // 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: @@ -18621,9 +18630,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm 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() @@ -18632,7 +18638,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } // the shader block size is hardcoded to head size 128 - if (q->ne[0] != 128 || k->ne[0] != 128) { + if (q->ne[0] != 128) { return false; } @@ -18643,6 +18649,12 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm (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; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl index a8ac5327fb..6f414ded12 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -1,3 +1,6 @@ +#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 @@ -48,3 +51,5 @@ bool fa_type_needs_shmem(uint ty) { default: return false; } } + +#endif // !defined(GGML_FA_TYPES_COMP) 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 ecf08af46e..746148d008 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl @@ -3,8 +3,13 @@ // switches on FaTypeK / FaTypeV. After spec-constant specialization the driver // folds away every path except the one matching the K/V type for this pipeline. // -// Included by flash_attn.comp and flash_attn_cm1.comp. Not included by -// flash_attn_cm2.comp, which has its own buffer_reference-based decode path. +// Included by flash_attn.comp, flash_attn_cm1.comp and lightning_indexer.comp. +// Not included by flash_attn_cm2.comp, which has its own buffer_reference-based +// decode path. +// +// Define FA_K_ONLY before including to drop the V views and the FaTypeV switch. +// Shaders that only read K do not declare FaTypeV, and binding 2 is theirs to +// use for something else. // // We use macros (rather than per-quant decode functions taking a struct) on // purpose: the FA shaders don't enable GL_EXT_shader_explicit_arithmetic_types_float16 @@ -15,37 +20,22 @@ // 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; +layout (binding = 1) readonly buffer K_PACKED_Q4_0 { block_q4_0_packed16 data[]; } k_packed_q4_0; +layout (binding = 1) readonly buffer K_PACKED_Q4_1 { block_q4_1_packed16 data[]; } k_packed_q4_1; +layout (binding = 1) readonly buffer K_PACKED_Q5_0 { block_q5_0_packed16 data[]; } k_packed_q5_0; +layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; } k_packed_q5_1; +layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0; +layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl; +layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16; + #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 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp index ddca015960..10680f6019 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -1,24 +1,29 @@ #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 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; +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; -#define FaTypeK K_TYPE #include "flash_attn_dequant.glsl" -layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; +// 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[]; }; @@ -48,15 +53,20 @@ layout(push_constant) uniform PushConstants { uint d_nb3; }; -shared float k_row[128]; -shared float partials[128]; +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(K_TYPE)) { + if (fa_type_needs_shmem(FaTypeK)) { init_iq_shmem(gl_WorkGroupSize); } @@ -70,18 +80,21 @@ void main() { 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) { + // 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 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); + 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; @@ -91,10 +104,29 @@ void main() { 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]; + const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_row[tid]; + +#if USE_SUBGROUP_ADD + const float sg_sum = subgroupAdd(prod); + if (gl_SubgroupInvocationID == 0) { + sg_partials[gl_SubgroupID] = sg_sum; + } barrier(); - for (uint stride = 64; stride > 0; stride >>= 1) { + 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]; } @@ -107,6 +139,7 @@ void main() { // the read of partials[0] above must complete before the next iteration // overwrites partials[tid] barrier(); +#endif } if (tid == 0) { 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 a743c7a015..1790197616 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1066,7 +1066,11 @@ 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", {}); + // 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 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"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 26f22ef719..9ea1dc41b0 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -186,7 +186,24 @@ static void init_tensor_kq_mask(ggml_tensor * tensor, float min = -1.0f, float m ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), ne0*ne1*ne2*ne3); - ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); + if (ggml_is_contiguous(tensor)) { + ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); + return; + } + + // scatter the rows to their strides, leaving any padding zeroed + GGML_TENSOR_LOCALS(size_t, nb, tensor, nb); + + std::vector data(ggml_nbytes(tensor), 0); + for (int32_t i3 = 0; i3 < ne3; ++i3) { + for (int32_t i2 = 0; i2 < ne2; ++i2) { + for (int32_t i1 = 0; i1 < ne1; ++i1) { + const size_t src = ((size_t)(i3*ne2 + i2)*ne1 + i1)*ne0; + memcpy(data.data() + i3*nb3 + i2*nb2 + i1*nb1, &data_f16[src], ne0*sizeof(ggml_fp16_t)); + } + } + } + ggml_backend_tensor_set(tensor, data.data(), 0, data.size()); } // generate a lower triangular matrix @@ -7625,18 +7642,25 @@ struct input_tensor { size_t view_offs = 0; }; +static std::array default_strides(const input_tensor & src) { + std::array nb; + nb[0] = ggml_type_size(src.type); + nb[1] = nb[0] * (src.ne[0] / ggml_blck_size(src.type)); + nb[2] = nb[1] * src.ne[1]; + nb[3] = nb[2] * src.ne[2]; + return nb; +} + +// strides to build the tensor with, resolving nb[0] == 0 to the contiguous defaults +static std::array effective_strides(const input_tensor & src) { + return src.nb[0] == 0 ? default_strides(src) : src.nb; +} + static bool is_non_contiguous(const input_tensor & src) { if (src.nb[0] == 0) { return false; } - const size_t default_nb0 = ggml_type_size(src.type); - const size_t default_nb1 = default_nb0 * (src.ne[0] / ggml_blck_size(src.type)); - const size_t default_nb2 = default_nb1 * src.ne[1]; - const size_t default_nb3 = default_nb2 * src.ne[2]; - return src.nb[0] != default_nb0 || - src.nb[1] != default_nb1 || - src.nb[2] != default_nb2 || - src.nb[3] != default_nb3; + return src.nb != default_strides(src); } static std::string var_to_str(const std::vector& sources) { @@ -7701,18 +7725,20 @@ struct test_generic_op : public test_case { for (size_t i = 0; i < source_count; ++i) { const input_tensor& src = sources[i]; - if (is_non_contiguous(src)) { + if (is_non_contiguous(src) || src.view_offs != 0) { + const std::array nb = effective_strides(src); + size_t total_size; const size_t blck_size = ggml_blck_size(src.type); if (blck_size == 1) { total_size = ggml_type_size(src.type); for (int d = 0; d < 4; d++) { - total_size += (src.ne[d] - 1) * src.nb[d]; + total_size += (src.ne[d] - 1) * nb[d]; } } else { - total_size = src.ne[0] * src.nb[0] / blck_size; + total_size = src.ne[0] * nb[0] / blck_size; for (int d = 1; d < 4; d++) { - total_size += (src.ne[d] - 1) * src.nb[d]; + total_size += (src.ne[d] - 1) * nb[d]; } } total_size += src.view_offs; @@ -7724,9 +7750,9 @@ 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], src.view_offs); + nb[1], nb[2], 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]; + source_tensors[i]->nb[0] = nb[0]; } else { source_tensors[i] = ggml_new_tensor_4d(ctx, src.type, src.ne[0], src.ne[1], src.ne[2], src.ne[3]); } @@ -7797,7 +7823,16 @@ struct test_generic_op : public test_case { break; } - if (op == GGML_OP_LIGHTNING_INDEXER && t->view_src != nullptr) { + // LIGHTNING_INDEXER: src[3] is the KQ mask + if (op == GGML_OP_LIGHTNING_INDEXER && i == 3) { + init_tensor_kq_mask(t); + continue; + } + + // init_tensor_uniform writes the elements contiguously, which does not + // reach the whole extent of the padded views these cases use. Fill the + // rows at their strides instead and leave the padding zeroed. + if (op == GGML_OP_LIGHTNING_INDEXER) { std::vector data(ggml_nbytes(t), 0); std::vector row(t->ne[0]); std::vector imatrix(t->ne[0], 1.0f); @@ -10139,69 +10174,27 @@ static std::vector> make_test_cases_eval() { 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")); + // padded strides and a nonzero view offset on every source, for each K type + // and both mask layouts (broadcast over streams, and one mask per stream) + auto padded_k = [](ggml_type type_K) { + const size_t ts = ggml_type_size(type_K); + const size_t row = ts * (128 / ggml_blck_size(type_K)); + const size_t nb1 = row + ts; + const size_t nb2 = nb1 + ts; + return input_tensor{ type_K, { 128, 1, 64, 4 }, { ts, nb1, nb2, 64*nb2 + nb1 }, 256 }; + }; + + 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}) { + for (int64_t nm : { 1, 4 }) { + 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 }, + padded_k(type_K), + { GGML_TYPE_F32, { 32, 3, 1, 4 }, { 4, 144, 448, 480 }, 256 }, + { GGML_TYPE_F16, { 64, 3, 1, nm }, { 2, 144, 448, 480 }, 256 }, + }, std::string("padded_") + ggml_type_name(type_K) + "_k_nm" + std::to_string(nm))); + } + } return test_cases; }