vulkan : selected-key flash attention for QSA block selection

ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h

@@ -729,6 +729,18 @@ struct vk_op_flash_attn_top_k_push_constants {
 };
 static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128);
 
+struct vk_op_flash_attn_sel_push_constants {
+    uint32_t n_kv, n_sel, n_sort;
+    uint32_t nbq1, nbq2, nbq3;
+    uint32_t nbk1, nbk2, nbk3;
+    uint32_t nbv1, nbv2, nbv3;
+    uint32_t nbi1, nbi3;
+    uint32_t nb1, nb2, nb3;
+    float scale;
+    uint32_t n_split, chunk, n_head, n_tok, n_stream;
+};
+static_assert(sizeof(vk_op_flash_attn_sel_push_constants) <= 128);
+
 struct vk_op_gated_delta_net_push_constants {
     uint32_t H;
     uint32_t n_tokens;

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -1009,6 +1009,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_lightning_indexer_cm_small_f16[LI_NH_COUNT];
     vk_pipeline pipeline_lightning_indexer_decode_cm_f16[LI_NH_COUNT];
     vk_pipeline pipeline_flash_attn_top_k_f16;
+    vk_pipeline pipeline_flash_attn_sel_f16[3]; // maskless selected-key FA, GQA 4, 8, 12
     vk_pipeline pipeline_flash_attn_top_k_cm_f16;
     vk_pipeline pipeline_flash_attn_gather_f16;
     vk_pipeline pipeline_flash_attn_gather_dq[GGML_TYPE_COUNT];

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -3925,6 +3925,17 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             "flash_attn_top_k_f16", flash_attn_top_k_f16_len, flash_attn_top_k_f16_data, "main", 6,
             sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true,
             device->subgroup_size);
+        // shared memory: the sorted row (16 KiB), Q (1 KiB per head), a K and a V block (16 KiB each): at most GQA 12
+        static const char * sel_names[3] = { "flash_attn_sel_f16_gqa4", "flash_attn_sel_f16_gqa8", "flash_attn_sel_f16_gqa12" };
+        for (uint32_t i = 0; i < 3; ++i) {
+            const uint32_t gqa = 4 * (i + 1);
+            if (device->subgroup_shuffle && 64 * gqa <= device->properties.limits.maxComputeWorkGroupInvocations &&
+                device->properties.limits.maxComputeSharedMemorySize >= 65536) {
+                ggml_vk_create_pipeline(device, device->pipeline_flash_attn_sel_f16[i],
+                    sel_names[i], flash_attn_sel_f16_len, flash_attn_sel_f16_data, "main", 5,
+                    sizeof(vk_op_flash_attn_sel_push_constants), {1, 1, 1}, {64 * gqa, 64, gqa}, 1, true, true, 64);
+            }
+        }
         ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_f16,
             "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5,
             sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true,
@@ -8478,6 +8489,113 @@ static void ggml_vk_perf_mark_subop(ggml_backend_vk_context * ctx, vk_context& s
     subctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
 }
 
+// maskless selected-key attention (QSA block selection): the pipeline that takes this op, or nullptr
+static vk_pipeline ggml_vk_flash_attn_sel_pipeline(const vk_device & device, const ggml_tensor * op) {
+    const ggml_tensor * q   = op->src[0];
+    const ggml_tensor * k   = op->src[1];
+    const ggml_tensor * v   = op->src[2];
+    const ggml_tensor * ids = op->src[5];
+    if (ids == nullptr || op->src[3] != nullptr || op->src[4] != nullptr || ggml_get_op_params_i32(op, 4) != 0) {
+        return nullptr;
+    }
+    float max_bias = 0.0f;
+    float logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) op->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float));
+    if (max_bias != 0.0f || logit_softcap != 0.0f || op->type != GGML_TYPE_F32 || !ggml_is_contiguous(op) ||
+        q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || ids->type != GGML_TYPE_I32 ||
+        q->ne[0] != 256 || k->ne[0] != 256 || v->ne[0] != 256 ||
+        k->ne[1] != v->ne[1] || k->ne[2] != v->ne[2] || k->ne[3] != q->ne[3] || v->ne[3] != q->ne[3] ||
+        k->ne[2] < 1 || q->ne[2] % k->ne[2] != 0 ||
+        ids->ne[0] < 1 || ids->ne[0] > 4096 || ids->ne[1] < q->ne[1] || ids->ne[2] != 1 ||
+        (ids->ne[3] != 1 && ids->ne[3] != q->ne[3]) ||
+        q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t) ||
+        ids->nb[0] != sizeof(int32_t)) {
+        return nullptr;
+    }
+    for (int i = 1; i < 4; ++i) {
+        if (q->nb[i] % sizeof(float) || k->nb[i] % sizeof(ggml_fp16_t) || v->nb[i] % sizeof(ggml_fp16_t) ||
+            ids->nb[i] % sizeof(int32_t)) {
+            return nullptr;
+        }
+    }
+    // the shader indexes with 32-bit element offsets
+    if (ggml_nbytes(k) / sizeof(ggml_fp16_t) > UINT32_MAX || ggml_nbytes(v) / sizeof(ggml_fp16_t) > UINT32_MAX ||
+        ggml_nbytes(q) / sizeof(float) > UINT32_MAX || ggml_nbytes(op) / sizeof(float) > UINT32_MAX) {
+        return nullptr;
+    }
+    switch (q->ne[2] / k->ne[2]) {
+        case 4:  return device->pipeline_flash_attn_sel_f16[0];
+        case 8:  return device->pipeline_flash_attn_sel_f16[1];
+        case 12: return device->pipeline_flash_attn_sel_f16[2];
+        default: return nullptr;
+    }
+}
+
+static bool ggml_vk_flash_attn_sel(ggml_backend_vk_context * ctx, vk_context & subctx,
+        const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, ggml_tensor * dst) {
+    vk_pipeline pipeline = ggml_vk_flash_attn_sel_pipeline(ctx->device, dst);
+    if (!pipeline) {
+        return false;
+    }
+    const ggml_tensor * ids = dst->src[5];
+    float scale = 1.0f;
+    memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
+    uint32_t n_sort = 1;
+    while (n_sort < (uint32_t) ids->ne[0]) {
+        n_sort <<= 1;
+    }
+    // few queries (decode) give too few workgroups to fill the GPU: split each row, at least 64 keys per split
+    const uint32_t n_sel = (uint32_t) ids->ne[0];
+    const uint32_t n_wg  = (uint32_t) (q->ne[1] * k->ne[2] * q->ne[3]);
+    uint32_t n_split = n_wg < 64 ? std::min(CEIL_DIV(64u, n_wg), CEIL_DIV(n_sel, 64u)) : 1u;
+    const uint32_t chunk = ROUNDUP_POW2(CEIL_DIV(n_sel, std::max(n_split, 1u)), 32u); // whole key blocks
+    n_split = CEIL_DIV(n_sel, chunk);
+    const vk_op_flash_attn_sel_push_constants pc = {
+        (uint32_t) k->ne[1], (uint32_t) ids->ne[0], n_sort,
+        (uint32_t) (q->nb[1] / sizeof(float)), (uint32_t) (q->nb[2] / sizeof(float)), (uint32_t) (q->nb[3] / sizeof(float)),
+        (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)),
+        (uint32_t) (v->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[3] / sizeof(ggml_fp16_t)),
+        (uint32_t) (ids->nb[1] / sizeof(int32_t)), ids->ne[3] == 1 ? 0u : (uint32_t) (ids->nb[3] / sizeof(int32_t)),
+        (uint32_t) (dst->nb[1] / sizeof(float)), (uint32_t) (dst->nb[2] / sizeof(float)), (uint32_t) (dst->nb[3] / sizeof(float)),
+        scale, n_split, chunk, (uint32_t) q->ne[2], (uint32_t) q->ne[1], (uint32_t) q->ne[3],
+    };
+    const vk_subbuffer q_buf   = ggml_vk_tensor_subbuffer(ctx, q);
+    const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
+    const std::array<uint32_t, 3> grid = { (uint32_t) q->ne[1] * n_split, (uint32_t) k->ne[2], (uint32_t) q->ne[3] };
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    if (n_split == 1) {
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
+            { q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, v), ggml_vk_tensor_subbuffer(ctx, ids), dst_buf },
+            pc, grid);
+        return true;
+    }
+    const uint32_t D  = 256;
+    const uint32_t NH = (uint32_t) q->ne[2];
+    const uint32_t N  = (uint32_t) q->ne[1];
+    const uint32_t NS = (uint32_t) q->ne[3];
+    const size_t split_size = ((size_t) D * NH * N * NS * n_split + (size_t) NH * 2 * N * NS * n_split) * sizeof(float);
+    ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1);
+    if (ctx->prealloc_size_split_k < split_size) {
+        ctx->prealloc_size_split_k = split_size;
+        ggml_vk_preallocate_buffers(ctx, subctx);
+    }
+    if (ctx->prealloc_split_k_need_sync) {
+        ggml_vk_sync_buffers(ctx, subctx);
+    }
+    const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
+        { q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, v), ggml_vk_tensor_subbuffer(ctx, ids), split_buf },
+        pc, grid);
+    ctx->prealloc_split_k_need_sync = true;
+    ggml_vk_sync_buffers(ctx, subctx);
+    const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = { D, NH, N, N, NS, n_split, 0 };
+    ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce,
+        { split_buf, q_buf, dst_buf }, reduce_pc, { NH, D, N * NS });
+    ctx->prealloc_split_k_need_sync = true;
+    return true;
+}
+
 static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & subctx,
         const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v,
         const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) {
@@ -9223,6 +9341,12 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
 
     assert(dst->type == GGML_TYPE_F32);
     assert(q->type == GGML_TYPE_F32);
+    if (dst->src[5] != nullptr && mask == nullptr) {
+        // supports_op only takes a maskless selected-key op when flash_attn_sel does
+        const bool done = ggml_vk_flash_attn_sel(ctx, subctx, q, k, v, dst);
+        GGML_ASSERT(done);
+        return;
+    }
     if (ggml_vk_flash_attn_top_k(ctx, subctx, q, k, v, mask, sinks, dst)) {
         return;
     }
@@ -16777,6 +16901,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_FLASH_ATTN_EXT:
             {
                 bool coopmat2 = device->coopmat2;
+                // maskless selected-key attention (src[5] without a mask): only flash_attn_sel reads the list, the dense kernel would attend to all cells
+                if (op->src[5] != nullptr && op->src[3] == nullptr) {
+                    return ggml_vk_flash_attn_sel_pipeline(device, op) != nullptr;
+                }
                 uint32_t HSK = op->src[1]->ne[0];
                 uint32_t HSV = op->src[2]->ne[0];
                 if ((HSK % 8) != 0 || (HSV % 8) != 0) {

ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp

@@ -0,0 +1,209 @@
+#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_arithmetic : require
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_shuffle : require
+
+// Maskless selected-key flash attention (QSA complete-block selection): each query's ids row names the cells it
+// may attend to, -1 for none. One workgroup per (query, K/V head, stream), one subgroup per query head of that K/V
+// head, so each K/V row is read once per query. The row is sorted and duplicates are skipped, like the CPU path.
+// Keys go in blocks of BK: two lanes compute one key's dot product (half the dims each), so the softmax needs two
+// subgroup reductions per block instead of one per key. With n_split > 1 (few queries) each workgroup takes one
+// chunk of the sorted row and writes a partial result in the layout of flash_attn_split_k_reduce.comp.
+
+layout(constant_id = 0) const uint WORKGROUP_SIZE = 768;
+layout(constant_id = 1) const uint SUBGROUP_SIZE = 64;
+layout(constant_id = 2) const uint GQA = 12;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly  buffer QBuf   { float     data_q[];   };
+layout(binding = 1) readonly  buffer KBuf   { float16_t data_k[];   };
+layout(binding = 2) readonly  buffer VBuf   { float16_t data_v[];   };
+layout(binding = 3) readonly  buffer IdsBuf { int       data_ids[]; };
+layout(binding = 4) writeonly buffer DstBuf { float     data_dst[]; };
+
+layout(push_constant) uniform Parameters {
+    uint n_kv;
+    uint n_sel;
+    uint n_sort;
+    uint nbq1;
+    uint nbq2;
+    uint nbq3;
+    uint nbk1;
+    uint nbk2;
+    uint nbk3;
+    uint nbv1;
+    uint nbv2;
+    uint nbv3;
+    uint nbi1;
+    uint nbi3;
+    uint nb1;
+    uint nb2;
+    uint nb3;
+    float scale;
+    uint n_split;
+    uint chunk;
+    uint n_head;
+    uint n_tok;
+    uint n_stream;
+} p;
+
+// The pipeline is created with subgroup size 64 and head size 256 only (see ggml_vk_flash_attn_sel_pipeline).
+const uint HEAD_SIZE = 256;
+const uint LANES = 64;
+const uint PER_LANE = HEAD_SIZE / LANES;
+const uint BK = 32;
+const uint HALF = HEAD_SIZE / 2;
+const uint MAX_SORT = 4096;
+const uint NONE = 0xffffffffu;
+const float FLT_MAX = 3.402823466e38;
+
+shared uint sorted_key[MAX_SORT];
+shared uint block_key[BK];
+shared float q_sh[GQA * HEAD_SIZE];
+shared float16_t kt_sh[HEAD_SIZE * BK]; // K block, transposed: [dim][key], so the per-lane key reads hit different banks
+shared float16_t v_sh[BK * HEAD_SIZE];  // V block: [key][dim]
+shared float p_sh[GQA * BK];
+
+void main() {
+    const float NEG_INF = uintBitsToFloat(0xff800000);
+    const uint tid = gl_LocalInvocationIndex;
+    const uint lane = gl_SubgroupInvocationID;
+    const uint sg = gl_SubgroupID;
+    const uint split = gl_WorkGroupID.x % p.n_split;
+    const uint query = gl_WorkGroupID.x / p.n_split;
+    const uint kv_head = gl_WorkGroupID.y;
+    const uint stream = gl_WorkGroupID.z;
+    const uint head = kv_head * GQA + sg;
+
+    // No early return: every invocation must reach every barrier below.
+
+    // the query's row, invalid entries as NONE, sorted ascending (NONE last)
+    const uint ids_base = stream * p.nbi3 + query * p.nbi1;
+    for (uint i = tid; i < p.n_sort; i += WORKGROUP_SIZE) {
+        uint key = NONE;
+        if (i < p.n_sel) {
+            const int id = data_ids[ids_base + i];
+            if (id >= 0 && uint(id) < p.n_kv) {
+                key = uint(id);
+            }
+        }
+        sorted_key[i] = key;
+    }
+    const uint q_base = stream * p.nbq3 + head * p.nbq2 + query * p.nbq1;
+    [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+        q_sh[sg * HEAD_SIZE + lane + i * LANES] = data_q[q_base + lane + i * LANES];
+    }
+    barrier();
+    for (uint size = 2; size <= p.n_sort; size <<= 1) {
+        for (uint stride = size >> 1; stride > 0; stride >>= 1) {
+            for (uint i = tid; i < p.n_sort; i += WORKGROUP_SIZE) {
+                const uint j = i ^ stride;
+                if (j > i) {
+                    const uint a = sorted_key[i];
+                    const uint b = sorted_key[j];
+                    if ((a > b) == ((i & size) == 0)) {
+                        sorted_key[i] = b;
+                        sorted_key[j] = a;
+                    }
+                }
+            }
+            barrier();
+        }
+    }
+
+    float acc[PER_LANE];
+    [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+        acc[i] = 0.0;
+    }
+    float row_max = NEG_INF;
+    float row_sum = 0.0;
+
+    const uint k_base = stream * p.nbk3 + kv_head * p.nbk2;
+    const uint v_base = stream * p.nbv3 + kv_head * p.nbv2;
+    const uint kcol = lane % BK;   // the key this lane scores
+    const uint khalf = lane / BK;  // and which half of its dims
+
+    const uint kb_begin = split * p.chunk;
+    const uint kb_end = min(p.n_sel, kb_begin + p.chunk);
+    for (uint kb = kb_begin; kb < kb_end; kb += BK) {
+        // sorted: once the block starts at NONE, so does everything after it (uniform, read from shared memory)
+        if (sorted_key[kb] == NONE) {
+            break;
+        }
+        if (tid < BK) {
+            const uint s = kb + tid;
+            uint key = s < kb_end ? sorted_key[s] : NONE;
+            if (key != NONE && s > 0 && sorted_key[s - 1] == key) {
+                key = NONE; // duplicate
+            }
+            block_key[tid] = key;
+        }
+        barrier();
+
+        for (uint idx = tid; idx < BK * HEAD_SIZE; idx += WORKGROUP_SIZE) {
+            const uint col = idx / HEAD_SIZE;
+            const uint dim = idx % HEAD_SIZE;
+            const uint key = block_key[col];
+            kt_sh[dim * BK + col] = key != NONE ? data_k[k_base + key * p.nbk1 + dim] : float16_t(0.0);
+            v_sh[idx] = key != NONE ? data_v[v_base + key * p.nbv1 + dim] : float16_t(0.0);
+        }
+        barrier();
+
+        float s = 0.0;
+        const uint qoff = sg * HEAD_SIZE + khalf * HALF;
+        const uint koff = khalf * HALF * BK + kcol;
+        for (uint d = 0; d < HALF; ++d) {
+            s += q_sh[qoff + d] * float(kt_sh[koff + d * BK]);
+        }
+        s += subgroupShuffleXor(s, BK);
+        const bool valid = block_key[kcol] != NONE;
+        s = valid ? s * p.scale : NEG_INF;
+
+        const float new_max = max(row_max, subgroupMax(s));
+        const float old_scale = new_max == row_max ? 1.0 : exp(row_max - new_max);
+        const float weight = valid ? exp(s - new_max) : 0.0;
+        row_sum = row_sum * old_scale + subgroupAdd(khalf == 0 ? weight : 0.0);
+        row_max = new_max;
+        if (khalf == 0) {
+            p_sh[sg * BK + kcol] = weight;
+        }
+        subgroupMemoryBarrierShared();
+        subgroupBarrier();
+
+        [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+            acc[i] *= old_scale;
+        }
+        for (uint j = 0; j < BK; ++j) {
+            const float w = p_sh[sg * BK + j];
+            [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+                acc[i] += w * float(v_sh[j * HEAD_SIZE + lane + i * LANES]);
+            }
+        }
+        barrier();
+    }
+
+    if (p.n_split == 1) {
+        const uint dst_base = stream * p.nb3 + query * p.nb2 + head * p.nb1;
+        const float inv_sum = row_sum == 0.0 ? 0.0 : 1.0 / row_sum;
+        [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+            data_dst[dst_base + lane + i * LANES] = acc[i] * inv_sum;
+        }
+    } else {
+        // partial: unnormalized O, then L and M per head (an empty split writes M = -FLT_MAX, not -inf, so the
+        // reduce never computes exp(-inf - -inf))
+        const uint slot = split + p.n_split * (query + p.n_tok * stream);
+        const uint o_base = HEAD_SIZE * p.n_head * slot + HEAD_SIZE * head;
+        [[unroll]] for (uint i = 0; i < PER_LANE; ++i) {
+            data_dst[o_base + lane + i * LANES] = acc[i];
+        }
+        if (lane == 0) {
+            const uint lm = HEAD_SIZE * p.n_head * p.n_tok * p.n_stream * p.n_split + p.n_head * 2 * slot + head;
+            data_dst[lm] = row_sum;
+            data_dst[lm + p.n_head] = max(row_max, -FLT_MAX);
+        }
+    }
+}