diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 13b8e73db151..7d8fa430d55c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/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; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index ab50260ae6ac..b03930ab8232 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/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]; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 6d002063f15e..b8b9c5c4c273 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/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 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) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp new file mode 100644 index 000000000000..5a9ecd215ff4 --- /dev/null +++ b/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); + } + } +} 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 f3e45215a7ef..23f43fab49c5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -869,6 +869,7 @@ void process_shaders() { string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); + string_to_spv("flash_attn_sel_f16", "flash_attn_sel.comp", {}); #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) string_to_spv("flash_attn_top_k_cm_f16", "flash_attn_top_k_cm.comp", {}); #endif diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index c32477868eef..a976cfa1584f 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -20,6 +20,7 @@ static bool hybrid_idx_no_recr(const llama_memory_recurrent * r) { #include #include #include +#include #include #include #include @@ -28,6 +29,21 @@ static bool hybrid_idx_no_recr(const llama_memory_recurrent * r) { // llama_memory_hybrid_idx // +// can dev run maskless selected-key attention with the shapes of QSA layer il? (a probe op, as for weight buffer types) +static bool qsa_selected_key_supported(ggml_backend_dev_t dev, const llama_hparams & hp, int il) { + ggml_init_params params = { 8*ggml_tensor_overhead(), nullptr, true }; + ggml_context_ptr ctx { ggml_init(params) }; + const int64_t d = hp.n_embd_head_k(); + const int64_t n_sel = hp.indexer_top_k + hp.dsv4_compress_ratios[il] - 1; + ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, d, 1, hp.n_head(il), 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, d, 4*n_sel, hp.n_head_kv(il), 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, d, 4*n_sel, hp.n_head_kv(il), 1); + ggml_tensor * ids = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_I32, n_sel, 1, 1, 1); + ggml_tensor * fa = ggml_flash_attn_ext(ctx.get(), q, k, v, nullptr, 1.0f, 0.0f, 0.0f); + ggml_flash_attn_ext_add_top_k(fa, ids, 0); + return ggml_backend_dev_supports_op(dev, fa); +} + llama_memory_hybrid_idx::llama_memory_hybrid_idx( const llama_model & model, /* attn */ @@ -78,6 +94,22 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); }()) { + if (mem_idx) { + // the QSA layers must all run where a kernel reads the selection list: HIP (RDNA3.5), or a backend that takes the probe op + // CUDA is not asked: its supports_op does not look at the list + bool all = true, any = false; + for (int il = 0; il < (int) model.hparams.n_layer_all; ++il) { + if (!model.hparams.has_kv(il) || !filter_idx(il) || model.hparams.dsv4_compress_ratios[il] <= 0) { continue; } + const char * reg = ggml_backend_reg_name(ggml_backend_dev_backend_reg(model.dev_layer(il))); + any = true; + all = all && (std::strcmp(reg, "ROCm") == 0 || + (std::strcmp(reg, "CUDA") != 0 && qsa_selected_key_supported(model.dev_layer(il), model.hparams, il))); + } + selected_key_attn = any && all && offload; + if (const char * e = getenv("LLAMA_QSA_SELECTED_KEY")) { selected_key_attn = atoi(e) != 0; } + LLAMA_LOG_INFO("%s: QSA attention: %s\n", __func__, selected_key_attn ? + "selected-key kernels (maskless block selection)" : "masked top-k (no selected-key kernels on this backend)"); + } if (!mem_idx || !offload || n_swa != 0) { return; } qsa_prefix = qsa_prefix_state(kv_size); const int layers = model.hparams.n_layer_all; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index a28c00cdc061..f26209435dca 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -99,6 +99,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { void qsa_invalidate(); bool qsa_prefix_matches(const llama_ubatch & ubatch) const; bool qsa_fast(int il, const llama_ubatch & ubatch) const; + bool qsa_selected_key_attn() const { return selected_key_attn; } ggml_tensor * qsa_cache(ggml_context * ctx, int il, int64_t blocks) const; void qsa_fill_updates(ggml_tensor * members, ggml_tensor * positions, ggml_tensor * rows) const; void qsa_commit(int il) const; @@ -118,6 +119,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_tensor * tails, const llama_ubatch & ubatch, uint32_t ratio) const; bool incremental_qsa = false; + // the QSA layers run where a kernel reads the maskless block selection (HIP on RDNA3.5, Vulkan); otherwise the masked top-k + bool selected_key_attn = false; bool qsa_recover_pending = false; bool qsa_recover(llama_seq_id seq); qsa_prefix_state qsa_prefix; @@ -176,6 +179,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { bool qsa_prefix_matches(const llama_ubatch & u) const { return mem && mem->qsa_prefix_matches(u); } bool qsa_fast(int il, const llama_ubatch & u) const { return mem && mem->qsa_fast(il, u); } + bool qsa_selected_key_attn() const { return mem && mem->qsa_selected_key_attn(); } // cells the block metadata must cover: the KV view is sized by occupied cells, but a cache whose positions // run ahead of its cells (an MTP draft never receives the image cells an M-RoPE image pins to one position) // has blocks past that view. Never below get_n_kv(), padded like it so graph reuse keeps its cadence. diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index d071c5281acb..6112baca5909 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -827,12 +827,12 @@ static int64_t qwen4exp_query_strip(int64_t n_tokens, int64_t n_stream) { // append the query's own partial block as the tail (up to ratio-1 cells), -1 where nothing is visible. The // selected rows then carry the visibility themselves, so attention can run without a mask. Only the selected-key // attention kernels understand that layout (F16 K/V, head 256, single stream; qsa_decode takes 1..512 queries, -// qsa_prefill larger ubatches). They exist only in the HIP backend for RDNA3.5; the gate does not check the device, so -// on any other backend the resulting maskless op is not taken by a QSA kernel. +// qsa_prefill larger ubatches). The HIP backend (RDNA3.5) and Vulkan (flash_attn_sel) have them, so the selection needs +// selected_key_attn (all QSA layers run there). Dense kernels do not read the list: ggml-cuda aborts, Vulkan declines. static bool qwen4exp_use_block_selection(bool blk_bias, int64_t n_stream, int64_t ratio, int64_t n_kv, const llama_ubatch & ubatch, const llama_cparams & cparams, const llama_hparams & hparams, - ggml_type type_k, ggml_type type_v) { - return blk_bias && cparams.causal_attn && n_stream==1 && ratio>1 && hparams.indexer_top_k%ratio==0 && + ggml_type type_k, ggml_type type_v, bool selected_key_attn) { + return selected_key_attn && blk_bias && cparams.causal_attn && n_stream==1 && ratio>1 && hparams.indexer_top_k%ratio==0 && n_kv>hparams.indexer_top_k+ratio-1 && n_kv<=16777216 && ubatch.token && cparams.flash_attn && cparams.offload_kqv && hparams.f_max_alibi_bias==0.0f && !hparams.attn_soft_cap && hparams.n_embd_head_k()==256 && hparams.n_embd_head_v()==256 && @@ -982,7 +982,7 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { // the selection mode and the strip bounds are baked into the graph, so a reused graph must agree on them const auto * attn = mctx->get_attn(); const bool blocks = attn && qwen4exp_use_block_selection(blk_bias, n_stream, ratio, n_kv, - params.ubatch, params.cparams, params.hparams, attn->type_k(), attn->type_v()); + params.ubatch, params.cparams, params.hparams, attn->type_k(), attn->type_v(), mctx->qsa_selected_key_attn()); const bool scalar = blocks && params.hparams.n_swa == 0 && mctx->qsa_scalar_visibility(params.ubatch, ratio); res &= (tail_idxs != nullptr) == scalar; res &= compact == scalar; @@ -1145,7 +1145,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); const auto * attn_ctx = mctx_hyb->get_attn(); const bool blocks = attn_ctx && qwen4exp_use_block_selection(blk_bias, n_stream, r, n_kv, ubatch, cparams, hparams, - attn_ctx->type_k(), attn_ctx->type_v()); + attn_ctx->type_k(), attn_ctx->type_v(), mctx_hyb->qsa_selected_key_attn()); const bool scalar = blocks && hparams.n_swa == 0 && mctx_hyb->qsa_scalar_visibility(ubatch, (uint32_t) r); qsa->compact = scalar; qsa->maskless = scalar;