Skip to content

metal: deepseek v4 sparse prefill #26512

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tarruda wants to merge 2 commits into from
Closed

metal: deepseek v4 sparse prefill #26512

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bccc54b
common: optimize DeepSeek V4 sparse attention
tarruda Jul 18, 2026
b5892a9
metal: optimize DeepSeek V4 sparse attention
tarruda Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ggml/include/ggml-rpc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,10 @@ extern "C" {

#define RPC_PROTO_MAJOR_VERSION 5
#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
Expand Down
22 changes: 20 additions & 2 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -574,6 +574,7 @@ extern "C" {
GGML_OP_DSV4_HC_COMB,
GGML_OP_DSV4_HC_PRE,
GGML_OP_DSV4_HC_POST,
GGML_OP_DSV4_SPARSE_PACK,

GGML_OP_UNARY,

Expand Down Expand Up @@ -2434,8 +2435,14 @@ extern "C" {
const struct ggml_tensor * a);

GGML_API void ggml_flash_attn_ext_add_sinks(
struct ggml_tensor * a,
struct ggml_tensor * sinks);
struct ggml_tensor * a,
struct ggml_tensor * sinks);

// Like ggml_flash_attn_ext_add_sinks(), but indexes sinks by the query-row
// dimension instead of the attention-head dimension.
GGML_API void ggml_flash_attn_ext_add_sinks_rows(
struct ggml_tensor * a,
struct ggml_tensor * sinks);

// TODO: needs to be adapted to ggml_flash_attn_ext
GGML_API struct ggml_tensor * ggml_flash_attn_back(
Expand Down Expand Up @@ -2601,6 +2608,17 @@ extern "C" {
struct ggml_tensor * weights,
struct ggml_tensor * mask);

// 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(
struct ggml_context * ctx,
struct ggml_tensor * raw_k,
struct ggml_tensor * comp_k,
struct ggml_tensor * raw_mask,
struct ggml_tensor * comp_mask,
struct ggml_tensor * comp_idx,
int64_t n_raw);

Comment on lines +2611 to +2621

@ggerganov ggerganov Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Earlier I didn't notice that this involves a new ggml OP. So, likely will not get merged as I think we are already working on similar functionality in another PR.

@am17an am17an Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ggerganov are you talking about #25917?

@ggerganov ggerganov Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes

// DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
// In short these operations are replacements for the original residual connection (x = transformer(x) + x)
// using a richer representation through streams.
Expand Down
3 changes: 2 additions & 1 deletion ggml/src/ggml-backend-meta.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -986,7 +986,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_DSV4_SPARSE_PACK: {
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
} break;
case GGML_OP_UNARY: {
Expand Down
5 changes: 5 additions & 0 deletions ggml/src/ggml-cpu/ggml-cpu.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2076,6 +2076,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_DSV4_SPARSE_PACK:
{
ggml_compute_forward_dsv4_sparse_pack(params, tensor);
} break;
case GGML_OP_MAP_CUSTOM1:
{
ggml_compute_forward_map_custom1(params, tensor);
Expand Down Expand Up @@ -2259,6 +2263,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_DSV4_SPARSE_PACK:
{
n_tasks = n_threads;
} break;
Expand Down
70 changes: 65 additions & 5 deletions ggml/src/ggml-cpu/ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8528,6 +8528,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
memcpy(&scale, (float *) dst->op_params + 0, sizeof(float));
memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float));
memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float));
const bool sinks_rows = ggml_get_op_params_i32(dst, 4);

if (logit_softcap != 0) {
scale /= logit_softcap;
Expand Down Expand Up @@ -8572,7 +8573,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
memset(VKQ32, 0, DV*sizeof(float));
}

const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + iq1*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL;
const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + (iq1%mask->ne[1])*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL;

// k indices
const int ik3 = iq3 / rk3;
Expand Down Expand Up @@ -8664,7 +8665,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(

// sinks - apply only on the first kv-chunk
if (sinks && ic_start == 0) {
const float s = ((float *)((char *) sinks->data))[h];
const float s = ((float *)((char *) sinks->data))[sinks_rows ? iq1 : h];

float ms = 1.0f;
float vs = 1.0f;
Expand Down Expand Up @@ -8764,6 +8765,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
memcpy(&scale, (float *) dst->op_params + 0, sizeof(float));
memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float));
memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float));
const bool sinks_rows = ggml_get_op_params_i32(dst, 4);

if (logit_softcap != 0) {
scale /= logit_softcap;
Expand Down Expand Up @@ -8853,7 +8855,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
if (mask) {
bool can_skip = true;
for (int tq = 0; tq < tile_rows; tq++) {
const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + ((iq1 + tq)%mask->ne[1])*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
for (int tk = 0; tk < kv_tile; tk++) {
mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]);
if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) {
Expand Down Expand Up @@ -8956,9 +8958,8 @@ static void ggml_compute_forward_flash_attn_ext_tiled(

// sinks (apply only to valid rows in the tile)
if (sinks) {
const float s = ((float *)((char *) sinks->data))[h];

for (int tq = 0; tq < tile_rows; tq++) {
const float s = ((float *)((char *) sinks->data))[sinks_rows ? iq1 + tq : h];
float ms = 1.0f;
float vs = 1.0f;

Expand Down Expand Up @@ -11229,6 +11230,65 @@ void ggml_compute_forward_dsv4_hc_post(
}
}

// ggml_compute_forward_dsv4_sparse_pack

void ggml_compute_forward_dsv4_sparse_pack(
const ggml_compute_params * params,
ggml_tensor * dst) {
const ggml_tensor * raw_k = dst->src[0];
const ggml_tensor * comp_k = dst->src[1];
const ggml_tensor * raw_mask = dst->src[2];
const ggml_tensor * comp_mask = dst->src[3];
const ggml_tensor * comp_idx = dst->src[4];

const int64_t d = raw_k->ne[0];
const int64_t nq = raw_mask->ne[1];
const int64_t nt = dst->ne[1];
const int64_t nr = ggml_get_op_params_i32(dst, 0);
const int64_t nc = comp_idx->ne[0];
const int64_t nk = nr + nc;

GGML_ASSERT(dst->type == GGML_TYPE_F16);

for (int64_t it = params->ith; it < nt; it += params->nth) {
const int64_t iq = it % nq;
const int64_t is = it / nq;
ggml_fp16_t * out = (ggml_fp16_t *) ((char *) dst->data + it*dst->nb[1]);
ggml_fp16_t * out_k = out;
ggml_fp16_t * out_m = out + d*nk;

int64_t ir = 0;
for (int64_t idx = 0; idx < raw_k->ne[2] && ir < nr; ++idx) {
const ggml_fp16_t m = *(const ggml_fp16_t *) ((const char *) raw_mask->data +
idx*raw_mask->nb[0] + iq*raw_mask->nb[1] + is*raw_mask->nb[3]);
if (!std::isfinite(GGML_CPU_FP16_TO_FP32(m))) {
continue;
}
memcpy(out_k + ir*d, (const char *) raw_k->data + idx*raw_k->nb[2] + is*raw_k->nb[3],
d*sizeof(ggml_fp16_t));
out_m[ir] = m;
++ir;
}
for (; ir < nr; ++ir) {
memset(out_k + ir*d, 0, d*sizeof(ggml_fp16_t));
out_m[ir] = GGML_CPU_FP32_TO_FP16(-INFINITY);
}

for (int64_t i = 0; i < nc; ++i) {
const int64_t oi = nr + 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 < comp_k->ne[2]);
memcpy(out_k + oi*d, (const char *) comp_k->data + idx*comp_k->nb[2] + is*comp_k->nb[3],
d*sizeof(ggml_fp16_t));
const ggml_fp16_t m = *(const ggml_fp16_t *) ((const char *) comp_mask->data +
idx*comp_mask->nb[0] + iq*comp_mask->nb[1] + is*comp_mask->nb[3]);
out_m[oi] = m;
}

}
}

// ggml_compute_forward_rwkv_wkv7

static void ggml_compute_forward_rwkv_wkv7_f32(
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-cpu/ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,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_dsv4_sparse_pack(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);
Expand Down
15 changes: 14 additions & 1 deletion ggml/src/ggml-metal/ggml-metal-device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ 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_SPARSE_PACK: op_str = "dsv4_sparse_pack"; break;
default: GGML_ABORT("fatal error");
};

Expand Down Expand Up @@ -1330,6 +1331,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal
return res;
}

ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib) {
const char * name = "kernel_top_k_radix_f32_i32";
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
}
return res;
}

ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_TOP_K);

Expand Down Expand Up @@ -1464,21 +1474,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(

// do bounds checks for the mask?
const bool bc_mask = op->src[3] && (op->src[3]->ne[1] % 8 != 0);
const bool scan_mask = has_mask && op->src[3]->ne[1] != 1;

snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d",
"flash_attn_ext",
ggml_type_name(op->src[1]->type),
dk,
dv);

snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d",
snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_scanm=%d_ns10=%d_ns20=%d_nsg=%d",
base,
has_mask,
has_sinks,
has_bias,
has_scap,
has_kvpad,
bc_mask,
scan_mask,
ns10,
ns20,
nsg);
Expand All @@ -1494,6 +1506,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT + 4);

ggml_metal_cv_set_bool(cv, bc_mask, FC_FLASH_ATTN_EXT + 10);
ggml_metal_cv_set_bool(cv, scan_mask, FC_FLASH_ATTN_EXT + 11);

ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT + 20);
ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT + 21);
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-metal/ggml-metal-device.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op);
Expand Down
15 changes: 15 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1375,6 +1375,21 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
ggml_is_contiguous_rows(op->src[1]) &&
ggml_is_contiguous_rows(op->src[2]) &&
ggml_is_contiguous_rows(op->src[3]);
case GGML_OP_DSV4_SPARSE_PACK:
return op->src[0]->type == GGML_TYPE_F16 &&
op->src[1]->type == GGML_TYPE_F16 &&
op->src[2]->type == GGML_TYPE_F16 &&
op->src[3]->type == GGML_TYPE_F16 &&
op->src[4]->type == GGML_TYPE_I32 &&
op->type == GGML_TYPE_F16 &&
op->src[0]->ne[0] == 512 &&
ggml_get_op_params_i32(op, 0) <= 128 &&
ggml_get_op_params_i32(op, 0) + op->src[4]->ne[0] <= 128 + 512 &&
ggml_is_contiguous_rows(op->src[0]) &&
ggml_is_contiguous_rows(op->src[1]) &&
ggml_is_contiguous_rows(op->src[2]) &&
ggml_is_contiguous_rows(op->src[3]) &&
ggml_is_contiguous_rows(op->src[4]);
case GGML_OP_SSM_CONV:
case GGML_OP_SSM_SCAN:
return has_simdgroup_reduction;
Expand Down
24 changes: 24 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,7 @@ typedef struct {
float m0;
float m1;
int32_t n_head_log2;
int32_t sinks_rows;
float logit_softcap;
} ggml_metal_kargs_flash_attn_ext;

Expand Down Expand Up @@ -436,6 +437,7 @@ typedef struct {
float m0;
float m1;
int32_t n_head_log2;
int32_t sinks_rows;
float logit_softcap;
} ggml_metal_kargs_flash_attn_ext_vec;

Expand Down Expand Up @@ -1238,6 +1240,28 @@ typedef struct {
uint64_t nb_d2;
} ggml_metal_kargs_dsv4_hc_post;

typedef struct {
int32_t n_embd;
int32_t n_batch;
int32_t n_raw;
int32_t n_raw_k;
int32_t n_comp;
uint64_t nb_rk2;
uint64_t nb_rk3;
uint64_t nb_ck2;
uint64_t nb_ck3;
uint64_t nb_rm0;
uint64_t nb_rm1;
uint64_t nb_rm3;
uint64_t nb_cm0;
uint64_t nb_cm1;
uint64_t nb_cm3;
uint64_t nb_ci0;
uint64_t nb_ci1;
uint64_t nb_ci3;
uint64_t nb_d1;
} ggml_metal_kargs_dsv4_sparse_pack;

typedef struct {
int32_t ne00;
int32_t ne01;
Expand Down
Loading