diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index a88eab440113..57aab90d592e 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 5 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 1 +#define RPC_PROTO_PATCH_VERSION 3 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 104, "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 83de623960d0..f371002d0904 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -571,6 +571,8 @@ extern "C" { GGML_OP_SOLVE_TRI, GGML_OP_GATED_DELTA_NET, GGML_OP_LIGHTNING_INDEXER, + GGML_OP_DSV4_COMPRESS, + GGML_OP_DSV4_TOP_K_MASK, GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, GGML_OP_DSV4_HC_POST, @@ -2608,6 +2610,29 @@ extern "C" { struct ggml_tensor * weights, struct ggml_tensor * mask); + // DeepSeek V4 compressor weighted reduction. + // + // kv_state, score_state: [overlap ? 2*n_embd : n_embd, n_rows] + // read_idxs: [(overlap ? 2 : 1)*ratio*n_blocks] + // res: [n_embd, n_blocks] + GGML_API struct ggml_tensor * ggml_dsv4_compress( + struct ggml_context * ctx, + struct ggml_tensor * kv_state, + struct ggml_tensor * score_state, + struct ggml_tensor * read_idxs, + int32_t ratio, + bool overlap); + + // Builds the raw + selected-compressed F16 attention mask in one pass. + // raw_mask: [n_raw, n_query, 1, n_stream] + // comp_mask: [n_comp, n_query, 1, n_stream] + // comp_idx: [n_select, n_query, 1, n_stream] + GGML_API struct ggml_tensor * ggml_dsv4_top_k_mask( + struct ggml_context * ctx, + struct ggml_tensor * raw_mask, + struct ggml_tensor * comp_mask, + struct ggml_tensor * comp_idx); + // Packs per-token raw-window and Lightning-Indexer selections into the // strided K + mask storage consumed by DeepSeek V4 sparse flash attention. GGML_API struct ggml_tensor * ggml_dsv4_sparse_pack( diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index c8a345cf846f..ce0521b6b7b0 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -984,6 +984,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_DSV4_COMPRESS: + case GGML_OP_DSV4_TOP_K_MASK: case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index b471a24fbdd9..3e284329a884 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2064,6 +2064,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_lightning_indexer(params, tensor); } break; + case GGML_OP_DSV4_COMPRESS: + { + ggml_compute_forward_dsv4_compress(params, tensor); + } break; + case GGML_OP_DSV4_TOP_K_MASK: + { + ggml_compute_forward_dsv4_top_k_mask(params, tensor); + } break; case GGML_OP_DSV4_HC_COMB: { ggml_compute_forward_dsv4_hc_comb(params, tensor); @@ -2260,6 +2268,8 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_COUNT_EQUAL: case GGML_OP_SOLVE_TRI: case GGML_OP_GATED_DELTA_NET: + case GGML_OP_DSV4_COMPRESS: + case GGML_OP_DSV4_TOP_K_MASK: case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index df042d23ac12..800dd1487c8a 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10946,6 +10946,132 @@ void ggml_compute_forward_gated_delta_net( } +// ggml_compute_forward_dsv4_compress + +void ggml_compute_forward_dsv4_compress( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * kv_state = dst->src[0]; + const ggml_tensor * score_state = dst->src[1]; + const ggml_tensor * read_idxs = dst->src[2]; + + GGML_ASSERT(kv_state->type == GGML_TYPE_F32); + GGML_ASSERT(score_state->type == GGML_TYPE_F32); + GGML_ASSERT(read_idxs->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const int32_t ratio = ggml_get_op_params_i32(dst, 0); + const bool overlap = ggml_get_op_params_i32(dst, 1) != 0; + const int64_t n_embd = dst->ne[0]; + const int64_t n_blocks = dst->ne[1]; + const int64_t n_rows = kv_state->ne[1]; + const int64_t n_read = (overlap ? 2 : 1)*ratio; + + GGML_ASSERT(ratio > 0 && n_blocks > 0); + GGML_ASSERT(kv_state->ne[0] == (overlap ? 2 : 1)*n_embd); + GGML_ASSERT(kv_state->ne[0] == score_state->ne[0]); + GGML_ASSERT(kv_state->ne[1] == score_state->ne[1]); + GGML_ASSERT(read_idxs->ne[0] == n_read*n_blocks); + + GGML_TENSOR_LOCALS(size_t, nbk, kv_state, nb); + GGML_TENSOR_LOCALS(size_t, nbs, score_state, nb); + GGML_TENSOR_LOCALS(size_t, nbi, read_idxs, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t nr = n_embd*n_blocks; + const int64_t dr = (nr + params->nth - 1)/params->nth; + const int64_t ir0 = dr*params->ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i0 = ir % n_embd; + const int64_t ib = ir / n_embd; + + float score_max = -INFINITY; + for (int64_t j = 0; j < n_read; ++j) { + const bool cur_half = overlap && j >= ratio; + const int64_t jr = cur_half ? j - ratio : j; + const int64_t idx_pos = (cur_half ? ratio*n_blocks : 0) + ib*ratio + jr; + const int32_t idx = *(const int32_t *) ((const char *) read_idxs->data + idx_pos*nbi0); + + GGML_ASSERT(idx >= 0 && idx <= n_rows); + if (idx == n_rows) { + continue; + } + + const int64_t i_src = (cur_half ? n_embd : 0) + i0; + const float score = *(const float *) ((const char *) score_state->data + i_src*nbs0 + idx*nbs1); + score_max = MAX(score_max, score); + } + + float sum_v = 0.0f; + float sum_w = 0.0f; + if (score_max != -INFINITY) { + for (int64_t j = 0; j < n_read; ++j) { + const bool cur_half = overlap && j >= ratio; + const int64_t jr = cur_half ? j - ratio : j; + const int64_t idx_pos = (cur_half ? ratio*n_blocks : 0) + ib*ratio + jr; + const int32_t idx = *(const int32_t *) ((const char *) read_idxs->data + idx_pos*nbi0); + + if (idx == n_rows) { + continue; + } + + const int64_t i_src = (cur_half ? n_embd : 0) + i0; + const float score = *(const float *) ((const char *) score_state->data + i_src*nbs0 + idx*nbs1); + const float weight = expf(score - score_max); + const float value = *(const float *) ((const char *) kv_state->data + i_src*nbk0 + idx*nbk1); + sum_v += value*weight; + sum_w += weight; + } + } + + *(float *) ((char *) dst->data + i0*nbd0 + ib*nbd1) = sum_w > 0.0f ? sum_v/sum_w : 0.0f; + } +} + +// ggml_compute_forward_dsv4_top_k_mask + +void ggml_compute_forward_dsv4_top_k_mask( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * raw_mask = dst->src[0]; + const ggml_tensor * comp_mask = dst->src[1]; + const ggml_tensor * comp_idx = dst->src[2]; + + GGML_ASSERT(raw_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_idx->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F16); + + const int64_t n_raw = raw_mask->ne[0]; + const int64_t n_comp = comp_mask->ne[0]; + const int64_t n_sel = comp_idx->ne[0]; + const int64_t nq = raw_mask->ne[1]; + const int64_t nrows = nq*raw_mask->ne[3]; + const ggml_fp16_t neg_inf = GGML_CPU_FP32_TO_FP16(-INFINITY); + + for (int64_t row = params->ith; row < nrows; row += params->nth) { + const int64_t iq = row % nq; + const int64_t is = row / nq; + ggml_fp16_t * out = (ggml_fp16_t *) ((char *) dst->data + iq*dst->nb[1] + is*dst->nb[3]); + + for (int64_t i = 0; i < n_raw; ++i) { + out[i] = *(const ggml_fp16_t *) ((const char *) raw_mask->data + + i*raw_mask->nb[0] + iq*raw_mask->nb[1] + is*raw_mask->nb[3]); + } + std::fill(out + n_raw, out + n_raw + n_comp, neg_inf); + + for (int64_t i = 0; i < n_sel; ++i) { + const int32_t idx = *(const int32_t *) ((const char *) comp_idx->data + + i*comp_idx->nb[0] + iq*comp_idx->nb[1] + is*comp_idx->nb[3]); + GGML_ASSERT(idx >= 0 && idx < n_comp); + out[n_raw + idx] = *(const ggml_fp16_t *) ((const char *) comp_mask->data + + idx*comp_mask->nb[0] + iq*comp_mask->nb[1] + is*comp_mask->nb[3]); + } + } +} + // ggml_compute_forward_dsv4_hc_comb static void ggml_dsv4_hc_comb_norm_cols(float * comb, float eps) { diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index cb1629e0bfc0..bcc8815c9e5b 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -106,6 +106,8 @@ void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, s void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_dsv4_compress(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_dsv4_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst); 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); diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 8fe3a391f103..d813064d9c6c 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -66,6 +66,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_base(ggml const char * op_str = "undefined"; switch (op) { case GGML_OP_ADD_ID: op_str = "add_id"; break; + case GGML_OP_DSV4_COMPRESS: op_str = "dsv4_compress"; break; + case GGML_OP_DSV4_TOP_K_MASK: op_str = "dsv4_top_k_mask"; break; case GGML_OP_DSV4_SPARSE_PACK: op_str = "dsv4_sparse_pack"; break; default: GGML_ABORT("fatal error"); }; @@ -1541,19 +1543,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[0]; const int32_t ns20 = op->src[2]->nb[1]/op->src[2]->nb[0]; + // DSV4 decode has a single 512-wide KV head, 64 query heads and a + // one-row top-k mask. Specialize per-row mask skipping to this signature + // so the extra checks cannot affect ordinary causal-attention kernels. + const bool sparse_mask = has_mask && has_sinks && dk == 512 && dv == 512 && + op->src[0]->ne[1] == 1 && op->src[0]->ne[2] == 64 && + op->src[1]->ne[2] == 1 && op->src[3]->ne[1] == 1 && op->src[4]->ne[0] == 64; + snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", "flash_attn_ext_vec", ggml_type_name(op->src[1]->type), dk, dv); - snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", + snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_mskip=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", base, has_mask, has_sinks, has_bias, has_scap, has_kvpad, + sparse_mask, ns10, ns20, nsg, nwg); @@ -1567,6 +1577,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v ggml_metal_cv_set_bool(cv, has_bias, FC_FLASH_ATTN_EXT_VEC + 2); ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT_VEC + 3); ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4); + ggml_metal_cv_set_bool(cv, sparse_mask, FC_FLASH_ATTN_EXT_VEC + 5); ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20); ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c8933dc408ac..d27b655f2aca 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1351,6 +1351,25 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]); + case GGML_OP_DSV4_COMPRESS: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_I32 && + op->type == GGML_TYPE_F32 && + (ggml_get_op_params_i32(op, 1) ? 2 : 1)*ggml_get_op_params_i32(op, 0) <= 128 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous(op->src[2]); + case GGML_OP_DSV4_TOP_K_MASK: + return op->src[0]->type == GGML_TYPE_F16 && + op->src[1]->type == GGML_TYPE_F16 && + op->src[2]->type == GGML_TYPE_I32 && + op->type == GGML_TYPE_F16 && + op->src[2]->ne[0] <= op->src[1]->ne[0] && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op->src[2]) && + ggml_is_contiguous(op); case GGML_OP_DSV4_HC_PRE: return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 561826387e3f..aac5418bbe93 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1210,6 +1210,36 @@ typedef struct { float eps; } ggml_metal_kargs_dsv4_hc_comb; +typedef struct { + int32_t n_embd; + int32_t n_blocks; + int32_t n_rows; + int32_t ratio; + int32_t overlap; + uint64_t nb_k0; + uint64_t nb_k1; + uint64_t nb_s0; + uint64_t nb_s1; + uint64_t nb_i0; + uint64_t nb_d0; + uint64_t nb_d1; +} ggml_metal_kargs_dsv4_compress; + +typedef struct { + int32_t n_raw; + int32_t n_comp; + int32_t n_select; + int32_t n_query; + uint64_t nb_rm1; + uint64_t nb_rm3; + uint64_t nb_cm1; + uint64_t nb_cm3; + uint64_t nb_ci1; + uint64_t nb_ci3; + uint64_t nb_d1; + uint64_t nb_d3; +} ggml_metal_kargs_dsv4_top_k_mask; + typedef struct { int32_t n_embd; int32_t n_tokens; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 007e4282ca21..dcdd99e493da 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -324,6 +324,14 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_lightning_indexer(ctx, idx); } break; + case GGML_OP_DSV4_COMPRESS: + { + n_fuse = ggml_metal_op_dsv4_compress(ctx, idx); + } break; + case GGML_OP_DSV4_TOP_K_MASK: + { + n_fuse = ggml_metal_op_dsv4_top_k_mask(ctx, idx); + } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: @@ -1381,6 +1389,97 @@ int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_dsv4_compress(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + GGML_ASSERT(op->op == GGML_OP_DSV4_COMPRESS); + + const ggml_tensor * kv_state = op->src[0]; + const ggml_tensor * score_state = op->src[1]; + const ggml_tensor * read_idxs = op->src[2]; + + const int32_t ratio = ggml_get_op_params_i32(op, 0); + const int32_t overlap = ggml_get_op_params_i32(op, 1); + const int32_t n_read = (overlap ? 2 : 1)*ratio; + + GGML_ASSERT(kv_state->type == GGML_TYPE_F32); + GGML_ASSERT(score_state->type == GGML_TYPE_F32); + GGML_ASSERT(read_idxs->type == GGML_TYPE_I32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(n_read <= 128); + + ggml_metal_kargs_dsv4_compress args = { + /*.n_embd =*/ (int32_t) op->ne[0], + /*.n_blocks =*/ (int32_t) op->ne[1], + /*.n_rows =*/ (int32_t) kv_state->ne[1], + /*.ratio =*/ ratio, + /*.overlap =*/ overlap, + /*.nb_k0 =*/ kv_state->nb[0], + /*.nb_k1 =*/ kv_state->nb[1], + /*.nb_s0 =*/ score_state->nb[0], + /*.nb_s1 =*/ score_state->nb[1], + /*.nb_i0 =*/ read_idxs->nb[0], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + }; + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_base(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(kv_state), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(score_state), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(read_idxs), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + ggml_metal_encoder_set_threadgroup_memory_size(enc, n_read*sizeof(int32_t), 0); + + const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + ggml_metal_encoder_dispatch_threadgroups( + enc, (args.n_embd + nth - 1)/nth, args.n_blocks, 1, nth, 1, 1); + + return 1; +} + +int ggml_metal_op_dsv4_top_k_mask(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + GGML_ASSERT(op->op == GGML_OP_DSV4_TOP_K_MASK); + + const ggml_tensor * raw_mask = op->src[0]; + const ggml_tensor * comp_mask = op->src[1]; + const ggml_tensor * comp_idx = op->src[2]; + + ggml_metal_kargs_dsv4_top_k_mask args = { + /*.n_raw =*/ (int32_t) raw_mask->ne[0], + /*.n_comp =*/ (int32_t) comp_mask->ne[0], + /*.n_select =*/ (int32_t) comp_idx->ne[0], + /*.n_query =*/ (int32_t) raw_mask->ne[1], + /*.nb_rm1 =*/ raw_mask->nb[1], + /*.nb_rm3 =*/ raw_mask->nb[3], + /*.nb_cm1 =*/ comp_mask->nb[1], + /*.nb_cm3 =*/ comp_mask->nb[3], + /*.nb_ci1 =*/ comp_idx->nb[1], + /*.nb_ci3 =*/ comp_idx->nb[3], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d3 =*/ op->nb[3], + }; + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_base(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(raw_mask), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comp_mask), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comp_idx), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + + const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + ggml_metal_encoder_dispatch_threadgroups( + enc, args.n_query*raw_mask->ne[3], 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -1557,7 +1656,13 @@ int ggml_metal_op_dsv4_sparse_pack(ggml_metal_op_t ctx, int idx) { // Four 256-thread packers can reside per core; a 512-thread group only // allows two and exposes the random selected-row reads to more latency. const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); - ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + if (args.n_raw == 0) { + // Decode gathers independent compressed rows. Giving each row its own + // threadgroup exposes enough parallelism to hide the random-read latency. + ggml_metal_encoder_dispatch_threadgroups(enc, args.n_comp, op->ne[1], 1, nth, 1, 1); + } else { + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + } return 1; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 0f1c5d893baa..b5aeb63b3c25 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -54,6 +54,8 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx); int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_compress (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_top_k_mask (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_sparse_pack (ggml_metal_op_t ctx, int idx); int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index c904429be527..e62ad0d83058 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -7378,6 +7378,7 @@ constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_ constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]]; constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]]; constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]]; +constant bool FC_flash_attn_ext_vec_sparse_mask [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; //constant float FC_flash_attn_ext_vec_scale [[function_constant(FC_FLASH_ATTN_EXT_VEC + 10)]]; //constant float FC_flash_attn_ext_vec_max_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 11)]]; @@ -7576,6 +7577,21 @@ kernel void kernel_flash_attn_ext_vec( // each simdgroup processes 1 query and NE (NW/NL) cache elements FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short ii = 0; ii < NE; ++ii) { + active |= sm[cc*NE + ii] > -MAXHALF; + } + } + + // The mask decision is uniform across the simdgroup. Avoid + // loading a full K row when all NE logits in this group are + // masked (notably DSV4's 512-of-N compressed selection). + if (!active) { + continue; + } + if (is_same::value) { FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { mqk[cc] += dot((float4) pk4[cc*NE*NS10/4 + ii*NL], (float4) pq4[ii*NL]); @@ -7686,12 +7702,34 @@ kernel void kernel_flash_attn_ext_vec( const auto sst = ss + ty; FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short jj = 0; jj < NE; ++jj) { + active |= ss[cc*NE + jj] != 0.0f; + } + } + if (!active) { + continue; + } + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { lo[ii] += o4_t(float4(pv4[cc*NE*NS20/4 + ii*NL])*float4(sst[cc*NE])); } } } else { FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short jj = 0; jj < NE; ++jj) { + active |= ss[cc*NE + jj] != 0.0f; + } + } + if (!active) { + continue; + } + device const vd4_t * pv4 = (device const vd4_t *) (v + ((ic + NE*cc + ty)*args.nb21)); FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { @@ -11544,6 +11582,107 @@ kernel void kernel_lightning_indexer( } } +// Fuse the get_rows, softmax, multiply and reduction sequence used by the DSV4 +// compressors. One threadgroup owns a contiguous embedding tile for one output +// block, so its row plan is loaded once into threadgroup memory and all state +// reads remain coalesced across embedding lanes. +kernel void kernel_dsv4_compress( + constant ggml_metal_kargs_dsv4_compress & args, + device const char * kv_state, + device const char * score_state, + device const char * read_idxs, + device char * dst, + threadgroup int32_t * idxs [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int n_read = (args.overlap ? 2 : 1)*args.ratio; + const int ib = tgpig.y; + + for (int j = tiitg; j < n_read; j += ntg.x) { + const bool cur_half = args.overlap && j >= args.ratio; + const int jr = cur_half ? j - args.ratio : j; + const int idx_pos = (cur_half ? args.ratio*args.n_blocks : 0) + ib*args.ratio + jr; + idxs[j] = *(device const int32_t *) (read_idxs + idx_pos*args.nb_i0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int i0 = tgpig.x*ntg.x + tiitg; + if (i0 >= args.n_embd) { + return; + } + + float score_max = -INFINITY; + for (int j = 0; j < n_read; ++j) { + const int idx = idxs[j]; + if (idx < 0 || idx >= args.n_rows) { + continue; + } + + const bool cur_half = args.overlap && j >= args.ratio; + const int i_src = (cur_half ? args.n_embd : 0) + i0; + const float score = *(device const float *) (score_state + i_src*args.nb_s0 + idx*args.nb_s1); + score_max = max(score_max, score); + } + + float sum_v = 0.0f; + float sum_w = 0.0f; + if (score_max != -INFINITY) { + for (int j = 0; j < n_read; ++j) { + const int idx = idxs[j]; + if (idx < 0 || idx >= args.n_rows) { + continue; + } + + const bool cur_half = args.overlap && j >= args.ratio; + const int i_src = (cur_half ? args.n_embd : 0) + i0; + const float score = *(device const float *) (score_state + i_src*args.nb_s0 + idx*args.nb_s1); + const float weight = exp(score - score_max); + const float value = *(device const float *) (kv_state + i_src*args.nb_k0 + idx*args.nb_k1); + sum_v += value*weight; + sum_w += weight; + } + } + + *(device float *) (dst + i0*args.nb_d0 + ib*args.nb_d1) = sum_w > 0.0f ? sum_v/sum_w : 0.0f; +} + +// Materialize raw visibility and the Lightning-Indexer selection in one pass. +// TOP_K indices are unique, so selected rows can overwrite the initial -INF +// fill without atomics after a single threadgroup barrier. +kernel void kernel_dsv4_top_k_mask( + constant ggml_metal_kargs_dsv4_top_k_mask & args, + device const char * raw_mask, + device const char * comp_mask, + device const char * comp_idx, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort ntg[[threads_per_threadgroup]]) { + const int iq = tgpig % args.n_query; + const int is = tgpig / args.n_query; + + device const half * rm = (device const half *) (raw_mask + (uint64_t) iq*args.nb_rm1 + (uint64_t) is*args.nb_rm3); + device const half * cm = (device const half *) (comp_mask + (uint64_t) iq*args.nb_cm1 + (uint64_t) is*args.nb_cm3); + device const int * ci = (device const int *) (comp_idx + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); + device half * out = (device half *) (dst + (uint64_t) iq*args.nb_d1 + (uint64_t) is*args.nb_d3); + + for (int i = tiitg; i < args.n_raw; i += ntg) { + out[i] = rm[i]; + } + for (int i = tiitg; i < args.n_comp; i += ntg) { + out[args.n_raw + i] = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_device); + + for (int i = tiitg; i < args.n_select; i += ntg) { + const int idx = ci[i]; + if (idx >= 0 && idx < args.n_comp) { + out[args.n_raw + idx] = cm[idx]; + } + } +} + typedef decltype(kernel_lightning_indexer) kernel_lightning_indexer_t; template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; @@ -11717,10 +11856,9 @@ kernel void kernel_dsv4_hc_post_f32( } } -// One threadgroup owns one token. The Lightning-Indexer selection is shared by -// all 64 attention heads, so every cache row is fetched once and packed next to -// a head-broadcast mask. This layout lets the regular tiled FA kernel treat the -// heads as query rows without materializing any dense top-k mask. +// Prefill assigns one threadgroup per token. Decode's compressed-only mode +// assigns one threadgroup per selected row to expose enough random-read +// parallelism. The packed masks remain adjacent to the F16 key storage. kernel void kernel_dsv4_sparse_pack( constant ggml_metal_kargs_dsv4_sparse_pack & args, device const char * raw_k, @@ -11729,18 +11867,40 @@ kernel void kernel_dsv4_sparse_pack( device const char * comp_mask, device const char * comp_idx, device char * dst, - uint tgpig[[threadgroup_position_in_grid]], + uint3 tgpig[[threadgroup_position_in_grid]], ushort tiitg[[thread_index_in_threadgroup]], - ushort ntg[[threads_per_threadgroup]]) { + ushort3 ntg[[threads_per_threadgroup]]) { constexpr int max_selected = 128 + 512; threadgroup int selected_idx[max_selected]; threadgroup half selected_mask[max_selected]; - const int iq = tgpig % args.n_batch; - const int is = tgpig / args.n_batch; + if (args.n_raw == 0) { + const int i = tgpig.x; + const int it = tgpig.y; + const int iq = it % args.n_batch; + const int is = it / args.n_batch; + const int idx = *(device const int *) (comp_idx + + (uint64_t) i*args.nb_ci0 + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); + + device const half * src = (device const half *) (comp_k + + (uint64_t) idx*args.nb_ck2 + (uint64_t) is*args.nb_ck3); + device half * out = (device half *) (dst + (uint64_t) it*args.nb_d1); + for (int e = tiitg; e < args.n_embd; e += ntg.x) { + out[i*args.n_embd + e] = src[e]; + } + if (tiitg == 0) { + out[args.n_embd*args.n_comp + i] = *(device const half *) (comp_mask + + (uint64_t) idx*args.nb_cm0 + (uint64_t) iq*args.nb_cm1 + (uint64_t) is*args.nb_cm3); + } + return; + } + + const int it = tgpig.x; + const int iq = it % args.n_batch; + const int is = it / args.n_batch; const int nk = args.n_raw + args.n_comp; - device half * out = (device half *) (dst + (uint64_t) tgpig*args.nb_d1); + device half * out = (device half *) (dst + (uint64_t) it*args.nb_d1); device half * out_k = out; device half * out_m = out + args.n_embd*nk; @@ -11764,7 +11924,7 @@ kernel void kernel_dsv4_sparse_pack( selected_mask[n] = -INFINITY; } } - for (int i = tiitg; i < args.n_comp; i += ntg) { + for (int i = tiitg; i < args.n_comp; i += ntg.x) { const int oi = args.n_raw + i; const int idx = *(device const int *) (comp_idx + (uint64_t) i*args.nb_ci0 + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); @@ -11779,11 +11939,11 @@ kernel void kernel_dsv4_sparse_pack( if (idx >= 0) { device const half * src = (device const half *) (raw_k + (uint64_t) idx*args.nb_rk2 + (uint64_t) is*args.nb_rk3); - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[i*args.n_embd + e] = src[e]; } } else { - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[i*args.n_embd + e] = 0.0h; } } @@ -11797,7 +11957,7 @@ kernel void kernel_dsv4_sparse_pack( const int idx = selected_idx[oi]; device const half * src = (device const half *) (comp_k + (uint64_t) idx*args.nb_ck2 + (uint64_t) is*args.nb_ck3); - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[oi*args.n_embd + e] = src[e]; } if (tiitg == 0) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index ae904c6a1ad8..be7bad8b0a57 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1080,6 +1080,8 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "SOLVE_TRI", "GATED_DELTA_NET", "LIGHTNING_INDEXER", + "DSV4_COMPRESS", + "DSV4_TOP_K_MASK", "DSV4_HC_COMB", "DSV4_HC_PRE", "DSV4_HC_POST", @@ -1101,7 +1103,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1196,6 +1198,8 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "A X = B, A triangular, solve X", "gated_delta_net(q, k, v, g, beta, s)", "lightning_indexer(q, k, weights, mask)", + "dsv4_compress(kv_state, score_state, read_idxs)", + "dsv4_top_k_mask(raw_mask, comp_mask, comp_idx)", "dsv4_hc_comb(mixes, scale, base)", "dsv4_hc_pre(x, weights)", "dsv4_hc_post(x, residual, post, comb)", @@ -1217,7 +1221,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6358,6 +6362,75 @@ struct ggml_tensor * ggml_lightning_indexer( return result; } +// ggml_dsv4_compress + +struct ggml_tensor * ggml_dsv4_compress( + struct ggml_context * ctx, + struct ggml_tensor * kv_state, + struct ggml_tensor * score_state, + struct ggml_tensor * read_idxs, + int32_t ratio, + bool overlap) { + GGML_ASSERT(kv_state->type == GGML_TYPE_F32); + GGML_ASSERT(score_state->type == GGML_TYPE_F32); + GGML_ASSERT(read_idxs->type == GGML_TYPE_I32); + GGML_ASSERT(ratio > 0); + GGML_ASSERT(kv_state->ne[0] == score_state->ne[0]); + GGML_ASSERT(kv_state->ne[1] == score_state->ne[1]); + GGML_ASSERT(kv_state->ne[2] == 1 && kv_state->ne[3] == 1); + GGML_ASSERT(score_state->ne[2] == 1 && score_state->ne[3] == 1); + GGML_ASSERT(read_idxs->ne[1] == 1 && read_idxs->ne[2] == 1 && read_idxs->ne[3] == 1); + + const int64_t n_read_per_block = (overlap ? 2 : 1)*ratio; + GGML_ASSERT(read_idxs->ne[0] % n_read_per_block == 0); + + const int64_t n_blocks = read_idxs->ne[0]/n_read_per_block; + const int64_t n_embd = overlap ? kv_state->ne[0]/2 : kv_state->ne[0]; + + GGML_ASSERT(n_blocks > 0 && n_embd > 0); + GGML_ASSERT(kv_state->ne[0] == (overlap ? 2 : 1)*n_embd); + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_blocks); + + ggml_set_op_params_i32(result, 0, ratio); + ggml_set_op_params_i32(result, 1, overlap ? 1 : 0); + + result->op = GGML_OP_DSV4_COMPRESS; + result->src[0] = kv_state; + result->src[1] = score_state; + result->src[2] = read_idxs; + + return result; +} + +// ggml_dsv4_top_k_mask + +struct ggml_tensor * ggml_dsv4_top_k_mask( + struct ggml_context * ctx, + struct ggml_tensor * raw_mask, + struct ggml_tensor * comp_mask, + struct ggml_tensor * comp_idx) { + GGML_ASSERT(raw_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_idx->type == GGML_TYPE_I32); + GGML_ASSERT(raw_mask->ne[2] == 1 && comp_mask->ne[2] == 1 && comp_idx->ne[2] == 1); + GGML_ASSERT(raw_mask->ne[1] == comp_mask->ne[1]); + GGML_ASSERT(raw_mask->ne[1] == comp_idx->ne[1]); + GGML_ASSERT(raw_mask->ne[3] == comp_mask->ne[3]); + GGML_ASSERT(raw_mask->ne[3] == comp_idx->ne[3]); + GGML_ASSERT(comp_idx->ne[0] <= comp_mask->ne[0]); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, + raw_mask->ne[0] + comp_mask->ne[0], raw_mask->ne[1], 1, raw_mask->ne[3]); + + result->op = GGML_OP_DSV4_TOP_K_MASK; + result->src[0] = raw_mask; + result->src[1] = comp_mask; + result->src[2] = comp_idx; + + return result; +} + // ggml_dsv4_sparse_pack struct ggml_tensor * ggml_dsv4_sparse_pack( @@ -6379,7 +6452,8 @@ struct ggml_tensor * ggml_dsv4_sparse_pack( const int64_t nq = raw_mask->ne[1]; const int64_t nk = n_raw + comp_idx->ne[0]; - GGML_ASSERT(n_raw > 0 && n_raw <= raw_k->ne[2]); + GGML_ASSERT(n_raw >= 0 && n_raw <= raw_k->ne[2]); + GGML_ASSERT(nk > 0); GGML_ASSERT(comp_k->ne[0] == d); GGML_ASSERT(raw_k->ne[1] == 1 && comp_k->ne[1] == 1); GGML_ASSERT(comp_k->ne[3] == n_stream); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c53e66c35c6a..6e87f7d45acc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -61,6 +61,18 @@ static const llm_fused_op_probe llm_fused_op_lid_probe = { /*.n_tokens_per_seq =*/ 1, }; +static const llm_fused_op_probe llm_fused_op_dsv4_compress_probe = { + /*.op =*/ LLM_FUSED_OP_DSV4_COMPRESS, + /*.name =*/ "fused DeepSeek V4 compressor", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_dsv4_top_k_mask_probe = { + /*.op =*/ LLM_FUSED_OP_DSV4_TOP_K_MASK, + /*.name =*/ "fused DeepSeek V4 top-k mask", + /*.n_tokens_per_seq =*/ 1, +}; + static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = { /*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE, /*.name =*/ "fused DeepSeek V4 HC pre", @@ -260,6 +272,10 @@ llama_context::llama_context( cparams.fused_lid = true; cparams.auto_flid = true; + cparams.fused_dsv4_compress = true; + cparams.fused_dsv4_top_k_mask = true; + cparams.auto_fdsv4_aux = true; + cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; cparams.fused_dsv4_hc_post = true; @@ -587,6 +603,13 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_dsv4_sparse_probe, cparams.fused_dsv4_sparse); cparams.auto_fdsv4_sparse = false; } + + if (cparams.auto_fdsv4_aux) { + LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 auxiliary ops support:\n", func); + resolve(llm_fused_op_dsv4_compress_probe, cparams.fused_dsv4_compress); + resolve(llm_fused_op_dsv4_top_k_mask_probe, cparams.fused_dsv4_top_k_mask); + cparams.auto_fdsv4_aux = false; + } } void llama_context::sched_reserve() { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 933ba2c86123..122d303d2070 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -43,6 +43,9 @@ struct llama_cparams { bool auto_fgdn; bool fused_lid; // use fused lightning indexer bool auto_flid; + bool fused_dsv4_compress; + bool fused_dsv4_top_k_mask; + bool auto_fdsv4_aux; bool fused_dsv4_hc_pre; bool fused_dsv4_hc_comb; bool fused_dsv4_hc_post; diff --git a/src/llama-graph.h b/src/llama-graph.h index e15d822a2567..b68087db3f3d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -43,6 +43,8 @@ enum llm_fused_op { LLM_FUSED_OP_GDN_AR, LLM_FUSED_OP_GDN_CH, LLM_FUSED_OP_LIGHTNING_INDEXER, + LLM_FUSED_OP_DSV4_COMPRESS, + LLM_FUSED_OP_DSV4_TOP_K_MASK, LLM_FUSED_OP_DSV4_HC_PRE, LLM_FUSED_OP_DSV4_HC_COMB, LLM_FUSED_OP_DSV4_HC_POST, diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 2eae0105aa82..86c1b393080c 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -479,21 +479,28 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state( GGML_ASSERT(state_read_idxs->ne[0] == DSV4_HCA_RATIO*n_blocks); GGML_ASSERT(n_embd_head >= n_embd_head_rope); - ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs); - kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks); - cb(kv, name, il); + ggml_tensor * comp = nullptr; + if (cparams.fused_dsv4_compress) { + comp = ggml_dsv4_compress( + ctx0, kv_state, score_state, state_read_idxs, DSV4_HCA_RATIO, false); + res->add_fused_node({LLM_FUSED_OP_DSV4_COMPRESS, comp, il}); + } else { + ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(kv, name, il); - ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs); - score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks); - cb(score, name, il); + ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs); + score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(score, name, il); - ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3)); - ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3)); + ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - ggml_tensor * weights = ggml_soft_max(ctx0, scores); - ggml_tensor * comp = ggml_mul(ctx0, values, weights); - comp = ggml_sum_rows(ctx0, comp); - comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + } cb(comp, name, il); comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); @@ -540,44 +547,49 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta GGML_ASSERT(score_state->ne[0] == 2*n_embd_head); GGML_ASSERT(n_embd_head >= n_embd_head_rope); - kv_state = dsv4_append_zero_row(ctx0, kv_state, false); - score_state = dsv4_append_zero_row(ctx0, score_state, true); - - const int64_t n_read = ratio*n_blocks; - - ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs); - ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs); - - ggml_tensor * kv_prev = ggml_cont(ctx0, - ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0)); - kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks); - cb(kv_prev, name, il); - - ggml_tensor * score_prev = ggml_cont(ctx0, - ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0)); - score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks); - cb(score_prev, name, il); - - ggml_tensor * kv_cur = ggml_cont(ctx0, - ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], - n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head))); - kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks); - - ggml_tensor * score_cur = ggml_cont(ctx0, - ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], - n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head))); - score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks); - - ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1); - ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1); - - values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3)); - scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)); - - ggml_tensor * weights = ggml_soft_max(ctx0, scores); - ggml_tensor * comp = ggml_mul(ctx0, values, weights); - comp = ggml_sum_rows(ctx0, comp); - comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + ggml_tensor * comp = nullptr; + if (cparams.fused_dsv4_compress) { + comp = ggml_dsv4_compress( + ctx0, kv_state, score_state, state_read_idxs, (int32_t) ratio, true); + res->add_fused_node({LLM_FUSED_OP_DSV4_COMPRESS, comp, il}); + } else { + kv_state = dsv4_append_zero_row(ctx0, kv_state, false); + score_state = dsv4_append_zero_row(ctx0, score_state, true); + + const int64_t n_read = ratio*n_blocks; + ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs); + ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs); + + ggml_tensor * kv_prev = ggml_cont(ctx0, + ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0)); + kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks); + cb(kv_prev, name, il); + + ggml_tensor * score_prev = ggml_cont(ctx0, + ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0)); + score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks); + cb(score_prev, name, il); + + ggml_tensor * kv_cur = ggml_cont(ctx0, + ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], + n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head))); + kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * score_cur = ggml_cont(ctx0, + ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], + n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head))); + score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1); + ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1); + values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3)); + scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)); + + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + } cb(comp, name, il); comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); @@ -744,8 +756,6 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( const auto & inp_csa = inp_dsv4->get_csa(); GGML_ASSERT(inp_csa.kq_mask); - ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); - ggml_tensor * k_rot = inp_attn->self_k_rot; if (k_rot) { q = llama_mul_mat_hadamard(ctx0, q, k_rot); @@ -777,16 +787,66 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + // Selecting every compressed row is equivalent to using the visibility + // mask directly. Avoid building the Lightning Indexer and TOP_K graph until + // the compressed cache grows beyond the selection size. Auto probes are + // the only exception because they must materialize their fused ops. + ggml_tensor * top_k = nullptr; + if (cparams.auto_fdsv4_aux || cparams.auto_fdsv4_sparse || + n_csa > (int64_t) hparams.indexer_top_k) { + top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); + } + + // At sufficiently deep single-token decode, gather only the Lightning + // Indexer's selected compressed keys. This preserves the regular head layout + // used by Metal's vector Flash Attention while making its KV work fixed-size. + const bool gather_decode = cparams.fused_dsv4_sparse && cparams.flash_attn && + raw_k->type == GGML_TYPE_F16 && csa_k->type == GGML_TYPE_F16 && + q->ne[2] == 1 && csa_k->ne[3] == 1 && top_k && + top_k->ne[1] == 1 && top_k->ne[3] == 1 && + n_csa >= 2*(int64_t) hparams.indexer_top_k; + if (gather_decode) { + ggml_tensor * packed = ggml_dsv4_sparse_pack( + ctx0, raw_k, csa_k, raw_mask, inp_csa.kq_mask, top_k, 0); + cb(packed, "csa_gathered_pack", il); + res->add_fused_node({LLM_FUSED_OP_DSV4_SPARSE_PACK, packed, il}); + + const int64_t nk = top_k->ne[0]; + ggml_tensor * gathered = ggml_view_4d(ctx0, packed, csa_k->ne[0], 1, nk, 1, + csa_k->ne[0]*sizeof(ggml_fp16_t), csa_k->ne[0]*sizeof(ggml_fp16_t), packed->nb[1], 0); + cb(gathered, "csa_gathered_k", il); + + ggml_tensor * k_sel = ggml_concat(ctx0, raw_k, gathered, 2); + cb(k_sel, "csa_k_selected", il); + + ggml_tensor * comp_mask = ggml_view_4d(ctx0, packed, nk, 1, 1, 1, + nk*sizeof(ggml_fp16_t), nk*sizeof(ggml_fp16_t), packed->nb[1], + csa_k->ne[0]*nk*sizeof(ggml_fp16_t)); + cb(comp_mask, "csa_gathered_mask", il); + + ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, comp_mask, 0); + cb(kq_mask, "csa_lid_kq_mask", il); + + ggml_tensor * out = build_attn_mha( + q, k_sel, k_sel, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_csa_lid_gathered", il); + return out; + } + // The Lightning Indexer selects a different compressed working set for every - // token. Reinterpret the 64 attention heads as query rows of one MQA problem - // per token so Metal's tiled flash-attention kernel consumes only those keys. - // Dense FA remains faster for decode/small caches; auto probing forces this - // branch once during setup so unsupported layer devices can disable it. + // prefill token. Reinterpret the 64 attention heads as query rows of one MQA + // problem per token so Metal's tiled flash-attention kernel consumes only + // those keys. Auto probing forces this branch once during setup so unsupported + // layer devices can disable it. const bool sparse_probe = cparams.auto_fdsv4_sparse; const bool sparse_prefill = q->ne[2] >= 8 && n_csa > (int64_t) hparams.indexer_top_k; if (cparams.fused_dsv4_sparse && cparams.flash_attn && raw_k->type == GGML_TYPE_F16 && csa_k->type == GGML_TYPE_F16 && (sparse_probe || sparse_prefill)) { + GGML_ASSERT(top_k); const int64_t n_stream = csa_k->ne[3]; const int64_t nq = q->ne[2]/n_stream; const int64_t nt = q->ne[2]; @@ -829,9 +889,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( return out; } - ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); + ggml_tensor * csa_mask = inp_csa.kq_mask; + ggml_tensor * kq_mask = nullptr; + if (top_k && cparams.fused_dsv4_top_k_mask) { + kq_mask = ggml_dsv4_top_k_mask(ctx0, raw_mask, csa_mask, top_k); + cb(kq_mask, "csa_top_k_mask", il); + res->add_fused_node({LLM_FUSED_OP_DSV4_TOP_K_MASK, kq_mask, il}); + } else { + if (top_k) { + csa_mask = build_top_k_mask(csa_mask, top_k, "csa_top_k_mask", il); + } else { + cb(csa_mask, "csa_top_k_mask", il); + } + kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); + } - ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a5831a9f45f6..61d33fdd6b15 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3760,6 +3760,99 @@ struct test_snake_fuse : public test_case { }; +struct test_dsv4_compress : public test_case { + const int64_t n_embd; + const int64_t n_rows; + const int64_t n_blocks; + const int32_t ratio; + const bool overlap; + + std::string vars() override { + return VARS_TO_STR5(n_embd, n_rows, n_blocks, ratio, overlap); + } + + double max_nmse_err() override { return 1e-6; } + + test_dsv4_compress(int64_t n_embd, int64_t n_rows, int64_t n_blocks, int32_t ratio, bool overlap) + : n_embd(n_embd), n_rows(n_rows), n_blocks(n_blocks), ratio(ratio), overlap(overlap) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t n_state = (overlap ? 2 : 1)*n_embd; + const int64_t n_read = (overlap ? 2 : 1)*ratio*n_blocks; + + ggml_tensor * kv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_state, n_rows); + ggml_tensor * score = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_state, n_rows); + ggml_tensor * idx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_read); + ggml_set_name(kv, "kv"); + ggml_set_name(score, "score"); + ggml_set_name(idx, "idx"); + + ggml_tensor * out = ggml_dsv4_compress(ctx, kv, score, idx, ratio, overlap); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + for (size_t i = 0; i < data.size(); ++i) { + // Include the synthetic zero/-inf row sentinel used by the + // production read plans without materializing that row. + data[i] = (i % 11 == 0) ? (int32_t) n_rows : (int32_t) ((7*i + 3) % n_rows); + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else { + init_tensor_uniform(t, -2.0f, 2.0f); + } + } + } +}; + +struct test_dsv4_top_k_mask : public test_case { + const int64_t n_raw; + const int64_t n_comp; + const int64_t n_select; + const int64_t n_query; + const int64_t n_stream; + + std::string vars() override { + return VARS_TO_STR5(n_raw, n_comp, n_select, n_query, n_stream); + } + + test_dsv4_top_k_mask(int64_t n_raw, int64_t n_comp, int64_t n_select, int64_t n_query, int64_t n_stream) + : n_raw(n_raw), n_comp(n_comp), n_select(n_select), n_query(n_query), n_stream(n_stream) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * raw = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_raw, n_query, 1, n_stream); + ggml_tensor * comp = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_comp, n_query, 1, n_stream); + ggml_tensor * idx = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_select, n_query, 1, n_stream); + ggml_set_name(raw, "raw"); + ggml_set_name(comp, "comp"); + ggml_set_name(idx, "idx"); + + ggml_tensor * out = ggml_dsv4_top_k_mask(ctx, raw, comp, idx); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + for (int64_t row = 0; row < n_query*n_stream; ++row) { + for (int64_t i = 0; i < n_select; ++i) { + data[row*n_select + i] = (int32_t) ((7*i + 11*row + 3) % n_comp); + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else { + init_tensor_uniform(t, -2.0f, 2.0f); + } + } + } +}; + struct test_dsv4_hc : public test_case { static constexpr int64_t hc = 4; @@ -7364,17 +7457,17 @@ struct test_lightning_indexer : public test_case { struct test_dsv4_sparse_pack : public test_case { const int64_t nb; const int64_t ns; + const int64_t kr; - std::string vars() override { return VARS_TO_STR2(nb, ns); } + std::string vars() override { return VARS_TO_STR3(nb, ns, kr); } - test_dsv4_sparse_pack(int64_t nb = 3, int64_t ns = 2) : nb(nb), ns(ns) {} + test_dsv4_sparse_pack(int64_t nb = 3, int64_t ns = 2, int64_t kr = 7) : nb(nb), ns(ns), kr(kr) {} ggml_tensor * build_graph(ggml_context * ctx) override { constexpr int64_t d = 512; constexpr int64_t nr = 11; - constexpr int64_t nc = 17; - constexpr int64_t kr = 7; - constexpr int64_t kc = 13; + const int64_t nc = kr == 0 ? 517 : 17; + const int64_t kc = kr == 0 ? 512 : 13; ggml_tensor * raw_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nr, ns); ggml_tensor * cmp_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nc, ns); @@ -7396,7 +7489,7 @@ struct test_dsv4_sparse_pack : public test_case { void initialize_tensors(ggml_context * ctx) override { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { - const int32_t limit = 17; + const int32_t limit = kr == 0 ? 517 : 17; std::vector data(ggml_nelements(t)); for (size_t i = 0; i < data.size(); ++i) { data[i] = (int32_t) ((i*7 + 3) % limit); @@ -8127,6 +8220,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_snake_fuse(type, { 64, 32, 2, 3})); // ne[2] > 1 and ne[3] > 1 } + test_cases.emplace_back(new test_dsv4_compress(512, 17, 3, 4, true)); + test_cases.emplace_back(new test_dsv4_compress(512, 137, 2, 128, false)); + test_cases.emplace_back(new test_dsv4_top_k_mask(17, 137, 64, 3, 2)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 2500, 512, 1, 1)); + test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); @@ -9658,6 +9756,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 640, 64, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 641, 8, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); + for (int kv : { 640, 2628, 5128, }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 1}, kv, 1, true, true, + 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). @@ -9802,6 +9904,7 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_dsv4_sparse_pack()); + test_cases.emplace_back(new test_dsv4_sparse_pack(1, 1, 0)); return test_cases; } @@ -9813,6 +9916,15 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // DeepSeek V4 token-generation shapes. + test_cases.emplace_back(new test_dsv4_hc_comb(1, 4)); + test_cases.emplace_back(new test_dsv4_hc_pre(4096, 1)); + test_cases.emplace_back(new test_dsv4_hc_post(4096, 1)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2628, 1, 1, 1}, 512)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {5128, 1, 1, 1}, 512)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 2500, 512, 1, 1)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 5000, 512, 1, 1)); + // Conv2d: K=CRS=NPQ=4096 matmul performance uint32_t iwh_idx = 0; uint32_t kwh_idx = 1; @@ -10019,6 +10131,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // DeepSeek V4 compressed sparse attention decode: one 512-wide KV head is + // shared by 64 query heads and the visibility mask is broadcast per row. + for (int kv : { 2628, 5128, }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 1}, kv, 1, true, true, + 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + for (int kv : { 4096, 8192, 16384, }) { for (int hs : { 64, 128, }) { for (int nr : { 1, 4, }) {