diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 482bd3666..ecf0a3a6d 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 8 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 0 +#define RPC_PROTO_PATCH_VERSION 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index c37ba9471..3c2e6734b 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -585,6 +585,7 @@ extern "C" { GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, GGML_OP_DSV4_HC_POST, + GGML_OP_TOP_K_MASK, GGML_OP_UNARY, @@ -2732,6 +2733,26 @@ extern "C" { struct ggml_tensor * post, struct ggml_tensor * comb); + // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf + // + // mask: [n_kv, m1, m2, m3] f16 or f32, broadcast to the rows of top_k + // top_k: [n_top, ne1, ne2, ne3] i32, picks of each row + // score: [n_pick, ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live + // blocks: [n_blk, n_pick] i32, can be NULL: the member cells of each pick, else a pick is a cell + // extra: [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick + // res: [n_kv, ne1, ne2, ne3] same type as mask + // + // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf + // picks and cells out of range are skipped + // + GGML_API struct ggml_tensor * ggml_top_k_mask( + struct ggml_context * ctx, + struct ggml_tensor * mask, + struct ggml_tensor * top_k, + struct ggml_tensor * score, + struct ggml_tensor * blocks, + struct ggml_tensor * extra); + // custom operators typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index edbe4da41..17f871c89 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1055,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: - case GGML_OP_DSV4_HC_POST: { + case GGML_OP_DSV4_HC_POST: + case GGML_OP_TOP_K_MASK: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); } break; case GGML_OP_UNARY: { diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index ab82e3d6c..4c33369ee 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_dsv4_hc_post(params, tensor); } break; + case GGML_OP_TOP_K_MASK: + { + ggml_compute_forward_top_k_mask(params, tensor); + } break; case GGML_OP_MAP_CUSTOM1: { ggml_compute_forward_map_custom1(params, tensor); @@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: + case GGML_OP_TOP_K_MASK: { n_tasks = n_threads; } break; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 7dff1d11b..02d5efffa 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11736,6 +11736,92 @@ void ggml_compute_forward_dsv4_hc_post( } } +// ggml_compute_forward_top_k_mask + +void ggml_compute_forward_top_k_mask( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * mask = dst->src[0]; + const ggml_tensor * top_k = dst->src[1]; + const ggml_tensor * score = dst->src[2]; + const ggml_tensor * blocks = dst->src[3]; + const ggml_tensor * extra = dst->src[4]; + + GGML_ASSERT(dst->type == mask->type); + GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type)); + + const int64_t n_kv = mask->ne[0]; + const int64_t ne1 = dst->ne[1]; + const int64_t ne2 = dst->ne[2]; + const int64_t n_rows = ggml_nrows(dst); + const int64_t n_top = top_k->ne[0]; + const int64_t n_blk = blocks ? blocks->ne[0] : 1; + const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv; + const int64_t n_extra = extra ? extra->ne[0] : 0; + + const size_t ts = ggml_type_size(dst->type); + + const int ith = params->ith; + const int nth = params->nth; + + for (int64_t ir = ith; ir < n_rows; ir += nth) { + const int64_t i1 = ir % ne1; + const int64_t i2 = (ir / ne1) % ne2; + const int64_t i3 = ir / (ne1*ne2); + + const char * mask_row = (const char *) mask->data + + (i1 % mask->ne[1])*mask->nb[1] + (i2 % mask->ne[2])*mask->nb[2] + (i3 % mask->ne[3])*mask->nb[3]; + char * dst_row = (char *) dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]; + + if (dst->type == GGML_TYPE_F16) { + const ggml_fp16_t ninf = GGML_CPU_FP32_TO_FP16(-INFINITY); + for (int64_t c = 0; c < n_kv; ++c) { + ((ggml_fp16_t *) dst_row)[c] = ninf; + } + } else { + for (int64_t c = 0; c < n_kv; ++c) { + ((float *) dst_row)[c] = -INFINITY; + } + } + + auto keep = [&](int32_t c) { + if (c >= 0 && c < n_kv) { + memcpy(dst_row + c*ts, mask_row + c*ts, ts); + } + }; + + const char * top_k_row = (const char *) top_k->data + i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3]; + + for (int64_t j = 0; j < n_top; ++j) { + const int32_t p = *(const int32_t *) (top_k_row + j*top_k->nb[0]); + if (p < 0 || p >= n_pick) { + continue; + } + if (score) { + const float s = *(const float *) ((const char *) score->data + + p*score->nb[0] + i1*score->nb[1] + i2*score->nb[2] + i3*score->nb[3]); + if (!(s > -INFINITY)) { + continue; + } + } + if (!blocks) { + keep(p); + continue; + } + for (int64_t k = 0; k < n_blk; ++k) { + keep(*(const int32_t *) ((const char *) blocks->data + k*blocks->nb[0] + p*blocks->nb[1])); + } + } + + if (extra) { + const char * extra_row = (const char *) extra->data + i1*extra->nb[1] + i2*extra->nb[2] + i3*extra->nb[3]; + for (int64_t j = 0; j < n_extra; ++j) { + keep(*(const int32_t *) (extra_row + j*extra->nb[0])); + } + } + } +} + // ggml_compute_forward_rwkv_wkv7 static void ggml_compute_forward_rwkv_wkv7_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 2728b08b6..00a2ea7dc 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f6ff9f893..e43807392 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -70,6 +70,7 @@ #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" #include "ggml-cuda/lightning-indexer.cuh" +#include "ggml-cuda/top-k-mask.cuh" #include "ggml.h" #include @@ -2570,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_LIGHTNING_INDEXER: ggml_cuda_lightning_indexer(ctx, dst); break; + case GGML_OP_TOP_K_MASK: + ggml_cuda_op_top_k_mask(ctx, dst); + break; default: return false; } @@ -6003,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return true; case GGML_OP_LIGHTNING_INDEXER: return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op); + case GGML_OP_TOP_K_MASK: + return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32; default: return false; diff --git a/ggml/src/ggml-cuda/top-k-mask.cu b/ggml/src/ggml-cuda/top-k-mask.cu new file mode 100644 index 000000000..9ff995240 --- /dev/null +++ b/ggml/src/ggml-cuda/top-k-mask.cu @@ -0,0 +1,141 @@ +#include "top-k-mask.cuh" +#include "convert.cuh" + +#define CUDA_TOP_K_MASK_BLOCK_SIZE 512 + +// one block per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells +template +static __global__ void top_k_mask_kernel( + const T * mask, const int32_t * top_k, const float * score, const int32_t * blocks, const int32_t * extra, T * dst, + const int64_t n_kv, const int64_t ne1, const int64_t ne2, + const int64_t m1, const int64_t m2, const int64_t m3, + const int64_t n_top, const int64_t n_pick, const int64_t n_blk, const int64_t n_extra, + const int64_t sm1, const int64_t sm2, const int64_t sm3, + const int64_t sd1, const int64_t sd2, const int64_t sd3, + const int64_t sk0, const int64_t sk1, const int64_t sk2, const int64_t sk3, + const int64_t ss0, const int64_t ss1, const int64_t ss2, const int64_t ss3, + const int64_t sb0, const int64_t sb1, + const int64_t se0, const int64_t se1, const int64_t se2, const int64_t se3, + const bool vec) { + ggml_cuda_pdl_lc(); + + const int64_t ir = blockIdx.x; + const int64_t i1 = ir % ne1; + const int64_t i2 = (ir / ne1) % ne2; + const int64_t i3 = ir / (ne1*ne2); + + const T * mask_row = mask + (i1 % m1)*sm1 + (i2 % m2)*sm2 + (i3 % m3)*sm3; + T * dst_row = dst + i1*sd1 + i2*sd2 + i3*sd3; + const int32_t * top_k_row = top_k + i1*sk1 + i2*sk2 + i3*sk3; + + ggml_cuda_pdl_sync(); + + const T ninf = ggml_cuda_cast(-INFINITY); + if (vec) { + constexpr int nv = sizeof(int4)/sizeof(T); + T fill[nv]; +#pragma unroll + for (int j = 0; j < nv; ++j) { + fill[j] = ninf; + } + const int4 fill4 = *(const int4 *) fill; + for (int64_t i = threadIdx.x; i < n_kv/nv; i += blockDim.x) { + ((int4 *) dst_row)[i] = fill4; + } + } else { + for (int64_t i = threadIdx.x; i < n_kv; i += blockDim.x) { + dst_row[i] = ninf; + } + } + + __syncthreads(); + + // one thread per pick or extra cell; cells kept twice get the same value + for (int64_t i = threadIdx.x; i < n_top + n_extra; i += blockDim.x) { + if (i >= n_top) { + const int32_t c = extra[(i - n_top)*se0 + i1*se1 + i2*se2 + i3*se3]; + if (c >= 0 && c < n_kv) { + dst_row[c] = mask_row[c]; + } + continue; + } + + const int32_t p = top_k_row[i*sk0]; + if (p < 0 || p >= n_pick) { + continue; + } + if (score && !(score[p*ss0 + i1*ss1 + i2*ss2 + i3*ss3] > -INFINITY)) { + continue; + } + if (!blocks) { + if (p < n_kv) { + dst_row[p] = mask_row[p]; + } + continue; + } + for (int64_t k = 0; k < n_blk; ++k) { + const int32_t c = blocks[k*sb0 + p*sb1]; + if (c >= 0 && c < n_kv) { + dst_row[c] = mask_row[c]; + } + } + } +} + +template +static void top_k_mask_cuda(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * mask = dst->src[0]; + const ggml_tensor * top_k = dst->src[1]; + const ggml_tensor * score = dst->src[2]; + const ggml_tensor * blocks = dst->src[3]; + const ggml_tensor * extra = dst->src[4]; + + const int64_t n_kv = mask->ne[0]; + + GGML_ASSERT(dst->nb[0] == sizeof(T) && mask->nb[0] == sizeof(T)); + GGML_ASSERT(dst->nb[1] % sizeof(T) == 0 && mask->nb[1] % sizeof(T) == 0); + + // 16-byte stores need every dst row aligned + const bool vec = (n_kv*sizeof(T)) % sizeof(int4) == 0 && + dst->nb[1] % sizeof(int4) == 0 && dst->nb[2] % sizeof(int4) == 0 && dst->nb[3] % sizeof(int4) == 0 && + (uintptr_t) dst->data % sizeof(int4) == 0; + + const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv; + + const size_t si = sizeof(int32_t); + const size_t sf = sizeof(float); + + const dim3 block_dims(CUDA_TOP_K_MASK_BLOCK_SIZE, 1, 1); + const dim3 grid_dims(ggml_nrows(dst), 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream()); + + ggml_cuda_kernel_launch(top_k_mask_kernel, launch_params, + (const T *) mask->data, (const int32_t *) top_k->data, + score ? (const float *) score->data : nullptr, + blocks ? (const int32_t *) blocks->data : nullptr, + extra ? (const int32_t *) extra->data : nullptr, + (T *) dst->data, + n_kv, dst->ne[1], dst->ne[2], + mask->ne[1], mask->ne[2], mask->ne[3], + top_k->ne[0], n_pick, blocks ? blocks->ne[0] : 1, extra ? extra->ne[0] : 0, + mask->nb[1]/sizeof(T), mask->nb[2]/sizeof(T), mask->nb[3]/sizeof(T), + dst->nb[1]/sizeof(T), dst->nb[2]/sizeof(T), dst->nb[3]/sizeof(T), + top_k->nb[0]/si, top_k->nb[1]/si, top_k->nb[2]/si, top_k->nb[3]/si, + score ? score->nb[0]/sf : 0, score ? score->nb[1]/sf : 0, score ? score->nb[2]/sf : 0, score ? score->nb[3]/sf : 0, + blocks ? blocks->nb[0]/si : 0, blocks ? blocks->nb[1]/si : 0, + extra ? extra->nb[0]/si : 0, extra ? extra->nb[1]/si : 0, extra ? extra->nb[2]/si : 0, extra ? extra->nb[3]/si : 0, + vec); +} + +void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + switch (dst->type) { + case GGML_TYPE_F16: + top_k_mask_cuda(ctx, dst); + break; + case GGML_TYPE_F32: + top_k_mask_cuda(ctx, dst); + break; + default: + GGML_ABORT("unsupported type"); + } +} diff --git a/ggml/src/ggml-cuda/top-k-mask.cuh b/ggml/src/ggml-cuda/top-k-mask.cuh new file mode 100644 index 000000000..930d6667b --- /dev/null +++ b/ggml/src/ggml-cuda/top-k-mask.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 01dc8d05c..004970db0 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1095,6 +1095,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "DSV4_HC_COMB", "DSV4_HC_PRE", "DSV4_HC_POST", + "TOP_K_MASK", "UNARY", @@ -1112,7 +1113,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1210,6 +1211,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "dsv4_hc_comb(mixes, scale, base)", "dsv4_hc_pre(x, weights)", "dsv4_hc_post(x, residual, post, comb)", + "top_k_mask(mask, top_k, score, blocks, extra)", "unary(x)", @@ -1227,7 +1229,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6630,6 +6632,48 @@ struct ggml_tensor * ggml_dsv4_hc_post( return result; } +// ggml_top_k_mask + +struct ggml_tensor * ggml_top_k_mask( + struct ggml_context * ctx, + struct ggml_tensor * mask, + struct ggml_tensor * top_k, + struct ggml_tensor * score, + struct ggml_tensor * blocks, + struct ggml_tensor * extra) { + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(mask->nb[0] == ggml_type_size(mask->type)); + GGML_ASSERT(top_k->type == GGML_TYPE_I32); + GGML_ASSERT(top_k->ne[1] % mask->ne[1] == 0 && top_k->ne[2] % mask->ne[2] == 0 && top_k->ne[3] % mask->ne[3] == 0); + + if (score) { + GGML_ASSERT(score->type == GGML_TYPE_F32); + GGML_ASSERT(score->ne[1] == top_k->ne[1] && score->ne[2] == top_k->ne[2] && score->ne[3] == top_k->ne[3]); + } + + if (blocks) { + GGML_ASSERT(blocks->type == GGML_TYPE_I32); + GGML_ASSERT(blocks->ne[2] == 1 && blocks->ne[3] == 1); + GGML_ASSERT(!score || blocks->ne[1] == score->ne[0]); + } + + if (extra) { + GGML_ASSERT(extra->type == GGML_TYPE_I32); + GGML_ASSERT(extra->ne[1] == top_k->ne[1] && extra->ne[2] == top_k->ne[2] && extra->ne[3] == top_k->ne[3]); + } + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, mask->type, mask->ne[0], top_k->ne[1], top_k->ne[2], top_k->ne[3]); + + result->op = GGML_OP_TOP_K_MASK; + result->src[0] = mask; + result->src[1] = top_k; + result->src[2] = score; + result->src[3] = blocks; + result->src[4] = extra; + + return result; +} + //////////////////////////////////////////////////////////////////////////////// struct ggml_hash_set ggml_hash_set_new(size_t size) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9ee64b566..ca84fc405 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -165,6 +165,12 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = { /*.n_tokens_per_seq =*/ 1, }; +static const llm_fused_op_probe llm_fused_op_top_k_mask_probe = { + /*.op =*/ LLM_FUSED_OP_TOP_K_MASK, + /*.name =*/ "fused top-k mask", + /*.n_tokens_per_seq =*/ 1, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -328,6 +334,10 @@ llama_context::llama_context( cparams.fused_dsv4_hc_post = true; cparams.auto_fhc = true; + // only the indexer models build this mask, skip the probe for the others + cparams.fused_top_k_mask = true; + cparams.auto_ftkm = hparams.indexer_top_k > 0; + // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -671,6 +681,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post); cparams.auto_fhc = false; } + + if (cparams.auto_ftkm) { + LLAMA_LOG_INFO("%s: resolving fused top-k mask support:\n", func); + resolve(llm_fused_op_top_k_mask_probe, cparams.fused_top_k_mask); + cparams.auto_ftkm = false; + } } static int llama_graph_n_input_tensors(ggml_cgraph * gf) { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index f4eb181a7..0c260f7b6 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -48,6 +48,8 @@ struct llama_cparams { bool fused_dsv4_hc_comb; bool fused_dsv4_hc_post; bool auto_fhc; + bool fused_top_k_mask; // use fused top-k sparse attention mask + bool auto_ftkm; bool no_perf; bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP] bool op_offload; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 59caffe87..8c31b5130 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3375,6 +3375,27 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +ggml_tensor * llm_graph_context::build_fused_top_k_mask( + ggml_tensor * kq_mask, + ggml_tensor * top_k, + ggml_tensor * score, + ggml_tensor * blocks, + ggml_tensor * extra, + int il) const { + const int64_t m1 = kq_mask->ne[1]; + const int64_t m2 = kq_mask->ne[2]; + const int64_t m3 = kq_mask->ne[3]; + + auto rows = [&](ggml_tensor * t) { + return t ? ggml_reshape_4d(ctx0, t, t->ne[0], m1, m2, m3) : nullptr; + }; + + ggml_tensor * cur = ggml_top_k_mask(ctx0, kq_mask, rows(top_k), rows(score), blocks, rows(extra)); + res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il}); + + return cur; +} + ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_dsa * inp, ggml_tensor * wo, @@ -3407,31 +3428,36 @@ ggml_tensor * llm_graph_context::build_attn( const auto & kq_mask = inp->get_kq_mask_mla(); - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + ggml_tensor * kq_mask_top_k = nullptr; + if (cparams.fused_top_k_mask) { + kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il); + } else { + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - // combine with the original kq mask - kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + } ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); diff --git a/src/llama-graph.h b/src/llama-graph.h index abe5133c9..9c5f3aea9 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -53,6 +53,7 @@ enum llm_fused_op { LLM_FUSED_OP_DSV4_HC_PRE, LLM_FUSED_OP_DSV4_HC_COMB, LLM_FUSED_OP_DSV4_HC_POST, + LLM_FUSED_OP_TOP_K_MASK, }; enum llm_ffn_op_type : int { @@ -1335,6 +1336,15 @@ struct llm_graph_context { float kq_scale, int il) const; + // ggml_top_k_mask with top_k, score and extra viewed in the row layout of kq_mask, see ggml.h + ggml_tensor * build_fused_top_k_mask( + ggml_tensor * kq_mask, // [n_kv, n_tokens/n_stream, 1, n_stream] + ggml_tensor * top_k, // [n_top, n_tokens] + ggml_tensor * score, // [n_pick, n_tokens], can be null + ggml_tensor * blocks, // [n_blk, n_pick], can be null + ggml_tensor * extra, // [n_extra, n_tokens], can be null + int il) const; + llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const; llm_graph_input_dsv4 * build_inp_dsv4() const; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index af04dddd3..184585a9e 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -683,6 +683,13 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask( GGML_ASSERT(kq_mask); GGML_ASSERT(top_k); + if (cparams.fused_top_k_mask) { + ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il); + cb(kq_mask_top_k, name, il); + + return kq_mask_top_k; + } + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 90870bd96..5056c901b 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -323,13 +323,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0); inp->n_sel = (uint32_t) n_sel; - // the scatter maps the dead slots of this mask to dump rows - inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens); - ggml_set_input(inp->sel_mask); - // set_input_kpool always fills the mask, so it stays allocated in every graph - ggml_build_forward_expand(gf, inp->sel_mask); - // one view for all layers: each view of a host input is a separate upload - inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens); + // the scatter maps the dead slots of this mask to dump rows, the fused mask op reads liveness from the scores instead + if (!cparams.fused_top_k_mask) { + inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens); + ggml_set_input(inp->sel_mask); + // set_input_kpool always fills the mask, so it stays allocated in every graph + ggml_build_forward_expand(gf, inp->sel_mask); + // one view for all layers: each view of a host input is a separate upload + inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens); + } } inp->n_new = n_new; @@ -973,6 +975,17 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED + // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere + if (cparams.fused_top_k_mask) { + ggml_tensor * tail = hparams.indexer_kpool_select_tail ? inp_kpool->tail_idxs : nullptr; + GGML_ASSERT(kpool*n_top_pool + (tail ? tail->ne[0] : 0) == (int64_t) inp_kpool->n_sel); + + ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, il); + cb(sel, "indexer_sel", il); + + return sel; + } + // The selection mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score. ggml_tensor * sel_score = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens] diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 22e715984..ec0b438fa 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -813,6 +813,16 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel( ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], unordered cb(top_k, "indexer_top_k", il); + // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere + if (cparams.fused_top_k_mask) { + GGML_ASSERT(kpool*n_top_pool + inp_kpool->tail_idxs->ne[0] == inp_kpool->n_sel); + + ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs, il); + cb(sel, "indexer_sel", il); + + return sel; + } + // the top blocks, then the incomplete tail with n_kv for missing cells ggml_tensor * sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs, ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens] diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 58f3275da..9f2d1f54d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9128,6 +9128,148 @@ struct test_lightning_indexer : public test_case { } }; +// GGML_OP_TOP_K_MASK +struct test_top_k_mask : public test_case { + const ggml_type type; + const int64_t n_kv; + const int64_t n_blk; // 0: no blocks, a pick is a cell + const int64_t n_pick; // number of blocks, ignored without blocks + const int64_t n_top; + const int64_t n_extra; // 0: no extra cells + const bool score; // false: every pick is live + const int64_t nb; // rows per stream + const int64_t nh; // rows of top_k per mask row, broadcast over ne2 + const int64_t ns; // streams + const bool real; // blocks of consecutive cells, live picks, as in a real cache + + std::string vars() override { + return VARS_TO_STR11(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real); + } + + // the result is a copy of the mask or -inf, the harness checks the -inf pattern, the finite values must match exactly + double max_err() override { + return 0.0; + } + + double err(const float * a, const float * b, size_t n) override { + double result = 0.0; + for (size_t i = 0; i < n; ++i) { + if (std::isfinite(a[i]) && std::isfinite(b[i])) { + result = std::max(result, (double) fabsf(a[i] - b[i])); + } + } + return result; + } + + size_t op_size(ggml_tensor * t) override { + const size_t ts = ggml_type_size(type); + const size_t nc = std::max(n_blk, 1); + return ggml_nbytes(t) + nb*nh*ns*(n_top*(sizeof(int32_t) + (score ? sizeof(float) : 0) + nc*(sizeof(int32_t) + 2*ts)) + n_extra*(sizeof(int32_t) + 2*ts)); + } + + test_top_k_mask(ggml_type type = GGML_TYPE_F16, int64_t n_kv = 256, int64_t n_blk = 4, int64_t n_pick = 64, + int64_t n_top = 16, int64_t n_extra = 3, bool score = true, int64_t nb = 32, int64_t nh = 1, int64_t ns = 1, bool real = false) + : type(type), n_kv(n_kv), n_blk(n_blk), n_pick(n_blk ? n_pick : n_kv), n_top(n_top), n_extra(n_extra), score(score), + nb(nb), nh(nh), ns(ns), real(real) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * m = ggml_new_tensor_4d(ctx, type, n_kv, nb, 1, ns); + ggml_set_name(m, "m"); + + ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top, nb, nh, ns); + ggml_set_name(top_k, "top_k"); + + ggml_tensor * s = nullptr; + if (score) { + s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_pick, nb, nh, ns); + ggml_set_name(s, "score"); + } + + ggml_tensor * blocks = nullptr; + if (n_blk > 0) { + blocks = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_blk, n_pick); + ggml_set_name(blocks, "blocks"); + } + + ggml_tensor * extra = nullptr; + if (n_extra > 0) { + extra = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_extra, nb, nh, ns); + ggml_set_name(extra, "extra"); + } + + ggml_tensor * out = ggml_top_k_mask(ctx, m, top_k, s, blocks, extra); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::mt19937 gen(0x4B50); + + // cells and picks sometimes fall out of range, the op skips them + auto init_cells = [&](ggml_tensor * t) { + std::vector data(ggml_nelements(t)); + if (real) { + // block p holds cells p*n_blk..p*n_blk + n_blk - 1, the extra cells are the last cells + const bool is_blocks = strcmp(t->name, "blocks") == 0; + for (int64_t i = 0; i < (int64_t) data.size(); ++i) { + const int64_t c = is_blocks ? i : n_kv - 1 - i % t->ne[0]; + data[i] = c < n_kv ? (int32_t) c : (int32_t) n_kv; + } + } else { + std::uniform_int_distribution cell(0, n_kv - 1); + for (auto & c : data) { + const uint32_t r = gen() % 16; + c = r == 0 ? (int32_t) n_kv : r == 1 ? -1 : cell(gen); + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + }; + + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + if (type == GGML_TYPE_F16) { + init_tensor_kq_mask(t); + } else { + std::vector data(ggml_nelements(t)); + std::uniform_real_distribution dis(-1.0f, 1.0f); + for (auto & v : data) { + v = gen() % 4 == 0 ? -INFINITY : dis(gen); + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); + } + } else if (strcmp(t->name, "score") == 0) { + // -inf marks the picks a row cannot see + std::vector data(ggml_nelements(t)); + std::uniform_real_distribution dis(-1.0f, 1.0f); + for (auto & v : data) { + v = !real && gen() % 4 == 0 ? -INFINITY : dis(gen); + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); + } else if (strcmp(t->name, "top_k") == 0) { + std::vector data(ggml_nelements(t)); + std::vector order(n_pick); + for (int64_t r = 0; r < ggml_nrows(t); ++r) { + for (int64_t i = 0; i < n_pick; ++i) { + order[i] = i; + } + std::shuffle(order.begin(), order.end(), gen); + for (int64_t j = 0; j < n_top; ++j) { + const bool out = !real && j == 0 && r % 5 == 0; + data[r*n_top + j] = out ? (r % 2 ? -1 : (int32_t) n_pick) : order[j % n_pick]; + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else if (strcmp(t->name, "blocks") == 0 || strcmp(t->name, "extra") == 0) { + init_cells(t); + } else { + // the harness sentinels and the output, the callback checks the sentinels for NaN + init_tensor_uniform(t); + } + } + } +}; + // Deserializable generic test case struct input_tensor { ggml_type type; @@ -9912,6 +10054,25 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true)); test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true)); + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { + // blocks with extra cells and scores (k-pool) + test_cases.emplace_back(new test_top_k_mask(type)); + test_cases.emplace_back(new test_top_k_mask(type, 1001, 4, 70, 16, 3, true, 7)); + test_cases.emplace_back(new test_top_k_mask(type, 256, 4, 64, 64, 0, true, 33)); + test_cases.emplace_back(new test_top_k_mask(type, 512, 8, 16, 32, 7, true, 5, 1, 4)); + test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 1)); + test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 64, 1, 2)); + // cells without scores (DSA) + test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 16, 0, false, 32)); + test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 64, 0, false, 7, 1, 3)); + test_cases.emplace_back(new test_top_k_mask(type, 4096, 0, 0, 2048, 0, false, 5)); + // cells with scores and extra cells + test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 32, 5, true, 9)); + // blocks picked per head, the mask broadcast over the heads + test_cases.emplace_back(new test_top_k_mask(type, 512, 16, 32, 4, 0, true, 6, 4, 1)); + test_cases.emplace_back(new test_top_k_mask(type, 1024, 8, 128, 16, 8, false, 3, 2, 2)); + } + // glu ops for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) { for (int v : {0, 1}) { @@ -12745,6 +12906,14 @@ static std::vector> make_test_cases_perf() { } } + // top_k_mask, Flash-Next QSA shapes: 512 pools of 4 + a 3-cell tail, and DSA shapes: 2048 cells + for (int64_t n_kv : { 18432, 131072 }) { + for (int64_t nb : { 1, 2048 }) { + test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 4, n_kv/4, 512, 3, true, nb, 1, 1, true)); + test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 0, 0, 2048, 0, false, nb, 1, 1, true)); + } + } + // launch-overhead isolation: single L2_NORM launch vs batched siblings at the GDN // production shape (strided qkv views) -- perf-mode only, the eval list has its own // 2/4-wide coverage