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

@@ -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<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,
@@ -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;
             }

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)

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
 

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) {

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<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"}}));
 

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<uint8_t> 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<size_t, 4> default_strides(const input_tensor & src) {
+    std::array<size_t, 4> 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<size_t, 4> 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<input_tensor>& 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<size_t, 4> 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<uint8_t> data(ggml_nbytes(t), 0);
                 std::vector<float> row(t->ne[0]);
                 std::vector<float> imatrix(t->ne[0], 1.0f);
@@ -10139,69 +10174,27 @@ static std::vector<std::unique_ptr<test_case>> 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;
 }