Skip to content

Prefill improvements: 119 -> 210 tok/s at 32k depth, sparse FA 8.16 -> 1.10 seconds #2

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

Merged
Nathanw1014 merged 5 commits into from
Aug 17, 2026
Merged

Prefill improvements: 119 -> 210 tok/s at 32k depth, sparse FA 8.16 -> 1.10 seconds #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
24c7ead
vulkan: accelerate DeepSeek V4 sparse prefill FA
Mushoz Aug 12, 2026
4bbe53e
vulkan: split sparse prefill attention
Mushoz Aug 13, 2026
80728f3
vulkan: tile sparse prefill scratch
Mushoz Aug 13, 2026
733fbbb
vulkan: reuse sparse FA probability fragments
Mushoz Aug 13, 2026
11714b1
vulkan: cache sparse FA masks per key block
Mushoz Aug 13, 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
587 changes: 587 additions & 0 deletions docs/development/DSV4-vulkan-sparse-prefill-progress.md
View file
Open in desktop

Large diffs are not rendered by default.

137 changes: 131 additions & 6 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1040,6 +1040,7 @@ struct vk_device_struct {
vk_pipeline pipeline_lightning_indexer_cm_f16;
vk_pipeline pipeline_lightning_indexer_decode_cm_f16;
vk_pipeline pipeline_flash_attn_top_k_f16;
vk_pipeline pipeline_flash_attn_top_k_cm_f16;
vk_pipeline pipeline_flash_attn_gather_f16;
vk_pipeline pipeline_dsv4_hc_pre_f32;
vk_pipeline pipeline_dsv4_hc_comb_f32;
Expand Down Expand Up @@ -1849,6 +1850,8 @@ struct vk_op_flash_attn_top_k_push_constants {
uint32_t nb1, nb2, nb3;
float scale;
uint32_t has_sinks;
uint32_t profile_stage;
uint32_t split_mode;
};
static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128);

Expand Down Expand Up @@ -6133,6 +6136,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
"lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5,
sizeof(vk_op_lightning_indexer_push_constants), {16, 1, 1}, {device->subgroup_size}, 1, true, true,
device->subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_cm_f16,
"flash_attn_top_k_cm_f16", flash_attn_top_k_cm_f16_len, flash_attn_top_k_cm_f16_data, "main", 6,
sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true,
device->subgroup_size);
}
#endif
ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_f16,
Expand Down Expand Up @@ -11193,7 +11200,9 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context &
const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v,
const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) {
const ggml_tensor * top_k = dst->src[5];
if (!top_k || !ctx->device->pipeline_flash_attn_top_k_f16 ||
static const char * top_k_env = getenv("GGML_VK_FA_TOPK");
if ((top_k_env && top_k_env[0] == '0') ||
!top_k || (!ctx->device->pipeline_flash_attn_top_k_f16 && !ctx->device->pipeline_flash_attn_top_k_cm_f16) ||
q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 ||
!mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 ||
q->ne[0] != 512 || q->ne[1] < 64 || k->ne[0] != 512 || v->ne[0] != 512 ||
Expand All @@ -11219,11 +11228,11 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context &
return false;
}
const int64_t n_kv_active = n_kv_raw + top_k->ne[0];
if (k->ne[1] < 3 * n_kv_active) {
if (k->ne[1] <= n_kv_active) {
return false;
}

const vk_op_flash_attn_top_k_push_constants pc = {
vk_op_flash_attn_top_k_push_constants pc = {
(uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw,
(uint32_t) top_k->ne[0], (uint32_t) q->ne[2],
(uint32_t) (q->nb[1] / sizeof(float)),
Expand All @@ -11238,17 +11247,133 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context &
(uint32_t) (dst->nb[1] / sizeof(float)),
(uint32_t) (dst->nb[2] / sizeof(float)),
(uint32_t) (dst->nb[3] / sizeof(float)),
scale, sinks != nullptr,
scale, sinks != nullptr, 0, 0,
};

const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q);
const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf;
vk_pipeline pipeline = ctx->device->pipeline_flash_attn_top_k_f16;
static const char * top_k_cm_env = getenv("GGML_VK_FA_TOPK_CM");
const bool use_cm = (!top_k_cm_env || top_k_cm_env[0] != '0') && ctx->device->pipeline_flash_attn_top_k_cm_f16;
static const char * top_k_profile_env = getenv("GGML_VK_FA_TOPK_PROFILE");
pc.profile_stage = use_cm && top_k_profile_env ? atoi(top_k_profile_env) : 0;
vk_pipeline pipeline = use_cm ? ctx->device->pipeline_flash_attn_top_k_cm_f16 : ctx->device->pipeline_flash_attn_top_k_f16;

static const char * top_k_split_env = getenv("GGML_VK_FA_TOPK_SPLIT");
const uint32_t mask_stride = (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t));
const bool try_split = use_cm && (!top_k_split_env || top_k_split_env[0] != '0') && n_kv_raw > 0 && top_k->ne[0] > 0;
if (try_split) {
const uint32_t N = (uint32_t) q->ne[1];
const uint32_t D = 512;
const uint32_t NH = 64;
const uint32_t NS = (uint32_t) q->ne[3];
const uint32_t raw_kv = (uint32_t) n_kv_raw;
const uint32_t partitions = 2;
const uint32_t tile_size = std::min(N, 256u);
const uint32_t n_tiles = CEIL_DIV(N, tile_size);
const bool f32acc = true;
vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc);

const uint32_t q_stride = (uint32_t) (q->nb[1] / sizeof(float));
const uint32_t k_stride = (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t));
const bool aligned = raw_kv % tuning.block_cols == 0 && (q_stride & 7) == 0 && (k_stride & 7) == 0;
const vk_fa_pipeline_state raw_state = get_fa_pipeline_state(ctx->device, tuning, D, D, aligned, f32acc,
true, false, false, GGML_TYPE_F16, GGML_TYPE_F16);
if (raw_state.path == FA_COOPMAT1 && ctx->device->pipeline_flash_attn_split_k_reduce) {
vk_pipeline raw_pipeline;
{
std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
auto & pipelines = ctx->device->pipeline_flash_attn_f32_f16;
auto it = pipelines.find(raw_state);
if (it != pipelines.end()) {
raw_pipeline = it->second;
} else {
pipelines[raw_state] = raw_pipeline = std::make_shared<vk_pipeline_struct>();
}
}
ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, n_tiles);
ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_tiles);
ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, n_tiles);

const uint64_t partition_size = ((uint64_t) D * NH + NH * 2) * sizeof(float) * tile_size * NS;
const uint64_t split_size = partition_size * partitions;
if (split_size <= ctx->device->properties.limits.maxStorageBufferRange) {
if (ctx->prealloc_size_split_k < split_size) {
ctx->prealloc_size_split_k = split_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_split_k_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}

const uint32_t n_head_log2 = 64;
const uint32_t mask_stride_in_split_kv = 1u << 31;
const uint32_t packed_gqa = mask_stride_in_split_kv | 1u;
const uint32_t packed_partitions = (partitions << 16) | 1;
const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k);
const vk_subbuffer mask_buf = ggml_vk_tensor_subbuffer(ctx, mask);
const vk_subbuffer top_buf = ggml_vk_tensor_subbuffer(ctx, top_k);
const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
const auto sliced = [](const vk_subbuffer & buf, uint64_t offset) {
return vk_subbuffer{buf.buffer, buf.offset + offset, buf.size - offset};
};

for (uint32_t tile = 0; tile < n_tiles; ++tile) {
if (tile != 0) {
ggml_vk_sync_buffers(ctx, subctx);
}
const uint32_t token_offset = tile * tile_size;
const uint32_t tile_n = std::min(tile_size, N - token_offset);
const vk_subbuffer tile_q = sliced(q_buf, (uint64_t) token_offset * q->nb[1]);
const vk_subbuffer tile_mask = sliced(mask_buf, (uint64_t) token_offset * mask->nb[1]);
const vk_subbuffer tile_top = sliced(top_buf, (uint64_t) token_offset * top_k->nb[1]);
const vk_subbuffer tile_dst = sliced(dst_buf, (uint64_t) token_offset * dst->nb[2]);
const vk_flash_attn_push_constants raw_pc = {
tile_n, raw_kv,
NH, tile_n, NS,
NH, NS,
1, NS,
1, NS,
(uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3],
q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3],
k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3],
k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3],
scale, 0.0f, 0.0f,
n_head_log2, 1.0f, 1.0f,
packed_gqa, mask_stride, packed_partitions,
};

ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline,
{tile_q, k_buf, k_buf, tile_mask, tile_q, split_buf, tile_q},
raw_pc, {tile_n, NH, NS});
ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)");

pc.n_batch = tile_n;
pc.split_mode = 1;
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{tile_q, k_buf, tile_mask, sinks_buf, tile_top, split_buf},
pc, {tile_n, (uint32_t) CEIL_DIV(q->ne[2], 32), NS});
ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_SELECTED (sub-op)");

ctx->prealloc_split_k_need_sync = true;
ggml_vk_sync_buffers(ctx, subctx);
const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, tile_n, NS, partitions, sinks != nullptr};
ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce,
{split_buf, sinks_buf, tile_dst}, reduce_pc, {NH, D, tile_n * NS});
ctx->prealloc_split_k_need_sync = true;
ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_REDUCE (sub-op)");
}
return true;
}
}
}

ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf,
ggml_vk_tensor_subbuffer(ctx, top_k), ggml_vk_tensor_subbuffer(ctx, dst)},
pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], 8), (uint32_t) q->ne[3]});
pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], use_cm ? 32 : 8), (uint32_t) q->ne[3]});
ggml_vk_perf_mark_subop(ctx, subctx, use_cm ? "FA_TOP_K_CM (sub-op)" : "FA_TOP_K_SPARSE (sub-op)");
return true;
}

Expand Down
40 changes: 23 additions & 17 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ ACC_TYPE perElemOpStoreCol0(const in uint32_t r, const in uint32_t c, const in A
// Load the slope matrix, indexed by Q's dimension 2.
ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2)
{
const uint32_t h = iq2 + (r % p.gqa_ratio);
const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff));

uint32_t n_head_log2 = p.mask_n_head_log2 & N_LOG2_MASK;

Expand All @@ -186,32 +186,40 @@ ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const i
// Load the sink value, indexed by Q's dimension 2.
ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2)
{
const uint32_t h = iq2 + (r % p.gqa_ratio);
const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff));

return ACC_TYPE(data_s[h]);
}

uint32_t i, N, KV, split_k_index, Tr, start_j, end_j,
gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3,
q_stride, k_stride, v_stride, m_stride;
q_stride, k_stride, v_stride, m_stride, gqa_ratio, split_k_num, output_k_num;
bool partial_output;

void init_indices()
{
N = p.N;
KV = p.KV;
gqa_ratio = p.gqa_ratio & 0xffff;
split_k_num = p.k_num & 0xffff;
output_k_num = p.k_num >> 16;
partial_output = output_k_num != 0;
if (!partial_output) {
output_k_num = split_k_num;
}

if (p.k_num > 1) {
if (p.gqa_ratio > 1) {
if (split_k_num > 1) {
if (gqa_ratio > 1) {
i = 0;
// batch and split_k share gl_WorkGroupID.x
gqa_iq1 = gl_WorkGroupID.x / p.k_num;
split_k_index = gl_WorkGroupID.x % p.k_num;
gqa_iq1 = gl_WorkGroupID.x / split_k_num;
split_k_index = gl_WorkGroupID.x % split_k_num;
} else {
gqa_iq1 = 0;
split_k_index = gl_WorkGroupID.x % p.k_num;
i = gl_WorkGroupID.x / p.k_num;
split_k_index = gl_WorkGroupID.x % split_k_num;
i = gl_WorkGroupID.x / split_k_num;
}
} else if (p.gqa_ratio > 1) {
} else if (gqa_ratio > 1) {
i = 0;
gqa_iq1 = gl_WorkGroupID.x;
split_k_index = 0;
Expand All @@ -228,7 +236,7 @@ void init_indices()

// When not using grouped query attention, all rows share the same iq2, equal to gl_WorkGroupID.y.
// When using grouped query attention, each workgroup does gqa_ratio consecutive values of iq2.
iq2 = gl_WorkGroupID.y * p.gqa_ratio;
iq2 = gl_WorkGroupID.y * gqa_ratio;
iq3 = gl_WorkGroupID.z;

// broadcast factors
Expand All @@ -249,14 +257,12 @@ void init_indices()
// nb?1 are already divided by the type size and are in units of elements.
// When using grouped query attention, Q is indexed by iq2, so the stride
// should be nb02 (which is in bytes).
q_stride = p.gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01;
q_stride = gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01;
k_stride = p.nb11;
v_stride = p.nb21;
// When using grouped query attention, all rows use the same mask (stride 0).
// "p.gqa_ratio >> 16" is just a roundabout way of writing zero
// that prevents the compiler from folding the "&" through the select
// and breaking the alignment detection.
m_stride = (p.gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV;
const bool mask_stride_in_split_kv = (p.gqa_ratio & 0x80000000u) != 0;
const uint32_t mask_stride_override = p.gqa_ratio >> 16;
m_stride = mask_stride_in_split_kv ? p.split_kv : (mask_stride_override != 0 ? mask_stride_override : (gqa_ratio > 1 ? 0 : KV));
}

// Bias applied to softmax to stay in fp16 range.
Expand Down
16 changes: 8 additions & 8 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ void main() {
}
// Only load if the block is not all zeros
if (mask_opt_bits != MASK_OPT_ALL_ZERO) {
bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0;
bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0;

float max_mask = NEG_FLT_MAX_OVER_2;
[[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) {
Expand Down Expand Up @@ -545,10 +545,10 @@ void main() {

// If there is split_k, then the split_k resolve shader does the final
// division by L. Store the intermediate O value and per-row m and L values.
if (p.k_num > 1) {
if (p.gqa_ratio > 1) {
if (partial_output || split_k_num > 1) {
if (gqa_ratio > 1) {
// note: O and Q have swapped coord 1,2.
uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4;
uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)) / 4;

[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
if (tile_row(r) < N) {
Expand All @@ -561,7 +561,7 @@ void main() {
}
}

o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3));
o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3));
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
if (tile_row(r) < N) {
perElemOpStoreCol0(tile_row(r), 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N);
Expand All @@ -574,7 +574,7 @@ void main() {
const uint global_row = i * Br + row;

if (global_row < N) {
uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4;
uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)) / 4;

[[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) {
const uint d = d0 + col_tid;
Expand All @@ -584,7 +584,7 @@ void main() {
}

if (global_row < N && col_tid == 0) {
uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3));
uint32_t lm_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3));
data_o[lm_offset + iq2] = D_TYPE(Lf[r]);
data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]);
}
Expand Down Expand Up @@ -633,7 +633,7 @@ void main() {

uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4;

if (p.gqa_ratio > 1) {
if (gqa_ratio > 1) {
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
if (tile_row(r) < N) {
[[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) {
Expand Down
2 changes: 2 additions & 0 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ layout(push_constant) uniform Parameters {
uint nb3;
float scale;
uint has_sinks;
uint profile_stage;
uint split_mode;
} p;

// Shape constants pinned by the dispatch gate in ggml_vk_flash_attn_top_k: DeepSeek V4
Expand Down
Loading