Skip to content

vulkan: top_k radix select for k >= 1024 for Qwen 3.8 Flash Next #28032

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
0cc4m merged 4 commits into from
Aug 31, 2026
+471 −10
Merged

vulkan: top_k radix select for k >= 1024 for Qwen 3.8 Flash Next #28032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4276f0e
vulkan: add top-k radix sort shader for k >= 1024
0cc4m Aug 30, 2026
7690925
add Qwen 3.8 Flash Next top-k tests
0cc4m Aug 30, 2026
fc95190
add top-k qsa fusion
0cc4m Aug 30, 2026
fcc4a22
clean up code
0cc4m Aug 30, 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
239 changes: 229 additions & 10 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 @@ -657,6 +657,21 @@ static constexpr std::initializer_list<ggml_op> snake_pattern { GGM
GGML_OP_SQR, GGML_OP_MUL,
GGML_OP_ADD };

// qwen4 QSA indexer: gather per-block scores to cells + add f16 mask (cast+reshape) + top-k,
// fused into one radix-select. The cast/reshape are elided; the raw f16 mask is read in-shader.
static constexpr std::initializer_list<ggml_op> topk_qsa_pattern { GGML_OP_GET_ROWS, GGML_OP_PERMUTE,
GGML_OP_CONT, GGML_OP_CPY,
GGML_OP_RESHAPE, GGML_OP_ADD,
GGML_OP_TOP_K };
static constexpr std::initializer_list<std::array<int, 3>> topk_qsa_edges {
{ 1, 0, 0 }, // permute->src[0] == get_rows
{ 2, 0, 1 }, // cont->src[0] == permute
{ 4, 0, 3 }, // reshape->src[0] == cpy (mask cast)
{ 5, 0, 2 }, // add->src[0] == cont
{ 5, 1, 4 }, // add->src[1] == reshape
{ 6, 0, 5 }, // top_k->src[0] == add
};

//node #978 ( SOFT_MAX): ffn_moe_probs-15 ( 0K) [Vulka ] use=2: ffn_moe_logits-15 ( 0K) [Vulka ]
//node #979 ( RESHAPE): ffn_moe_probs-15 (re ( 0K) [Vulka ] use=1: ffn_moe_probs-15 ( 0K) [Vulka ]
//node #980 ( ARGSORT): ffn_moe_argsort-15 ( 0K) [Vulka ] use=1: ffn_moe_probs-15 ( 0K) [Vulka ]
Expand Down Expand Up @@ -1056,6 +1071,8 @@ struct vk_device_struct {
vk_pipeline pipeline_argsort_f32[num_argsort_pipelines];
vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
vk_pipeline pipeline_topk_f32[num_topk_pipelines];
vk_pipeline pipeline_topk_radix_f32;
vk_pipeline pipeline_topk_radix_qsa; // qwen4 QSA indexer fusion (f16 mask)
vk_pipeline pipeline_sum_rows_f32;
vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
Expand Down Expand Up @@ -1748,6 +1765,15 @@ struct vk_op_topk_push_constants {
uint32_t last_pass;
};

struct vk_op_topk_radix_push_constants {
uint32_t ncols;
uint32_t k;
uint32_t nrows;
uint32_t n_tps; // QSA only
uint32_t n_blocks; // QSA only
uint32_t n_stream; // QSA only
};

struct vk_op_im2col_push_constants {
uint64_t dst_addr;
uint32_t batch_offset; uint32_t offset_delta;
Expand Down Expand Up @@ -2438,6 +2464,8 @@ struct ggml_backend_vk_context {
int fused_ops_write_mask {};
topk_moe_mode fused_topk_moe_mode {};
bool fused_topk_moe_scale {};
// QSA indexer gather+add+top_k fused into one radix-select
bool fused_topk_qsa {};

// for GGML_VK_PERF_LOGGER
std::unique_ptr<vk_perf_logger> perf_logger;
Expand Down Expand Up @@ -5812,6 +5840,14 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
}
}

// large-k fallback: one workgroup per row, radix-select instead of a full sort. The QSA
// variant (spec constant 1) additionally gathers the qwen4 indexer input on the fly.
{
const uint32_t BLOCK_SIZE = 1u << std::min(10u, device->max_workgroup_size_log2);
ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_f32, "topk_radix_f32", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 0}, 1, true);
ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_qsa, "topk_radix_qsa", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 1}, 1, true);
}

ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);

ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
Expand Down Expand Up @@ -13936,6 +13972,31 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons
uint32_t nrows = ggml_nrows(src0);
uint32_t k = dst->ne[0];

// tournament path is faster where it fits; use radix-select only past its k limit
const uint32_t k_min_pipeline = std::max((uint32_t) log2f(float(k)) + 1, ctx->device->subgroup_size_log2);
if (k_min_pipeline >= num_topk_pipelines || ctx->device->pipeline_topk_f32[k_min_pipeline] == nullptr) {
vk_pipeline pipeline = ctx->device->pipeline_topk_radix_f32;
GGML_ASSERT(pipeline != nullptr);

if (ctx->prealloc_x_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}

vk_op_topk_radix_push_constants pc { ncols, k, nrows, 0, 0, 0 };
std::array<uint32_t, 3> elements {
pipeline->wg_denoms[0],
std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]),
1,
};
// the non-QSA path only uses bindings 0/1; bind valid buffers for the unused QSA slots
vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0);
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{ src0_buf, dst_buf, src0_buf, src0_buf, src0_buf }, pc, elements);
return;
}

vk_op_topk_push_constants pc { ncols, ncols, ncols, k, nrows, 0, 0 };

if (ctx->prealloc_x_need_sync) {
Expand Down Expand Up @@ -14039,6 +14100,55 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons
ctx->prealloc_x_need_sync = true;
}

static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_cgraph * cgraph, int node_idx) {
const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0];
const ggml_tensor * add = cgraph->nodes[node_idx + ctx->num_additional_fused_ops - 1];
ggml_tensor * top_k = cgraph->nodes[node_idx + ctx->num_additional_fused_ops];

const ggml_tensor * scores = get_rows->src[0]; // [n_tps, n_blocks, n_stream]
const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream]

// raw f16 mask: follow the reshape/cpy chain back to the materialized input
const ggml_tensor * mask = add->src[1];
while (mask->op == GGML_OP_RESHAPE || mask->op == GGML_OP_CPY) {
mask = mask->src[0];
}

const uint32_t n_tps = scores->ne[0];
const uint32_t n_blocks = scores->ne[1];
const uint32_t n_stream = scores->ne[2];
const uint32_t n_kv = cell_blk->ne[0];
const uint32_t width = top_k->ne[0];
const uint32_t nrows = n_tps * n_stream;

vk_pipeline pipeline = ctx->device->pipeline_topk_radix_qsa;
GGML_ASSERT(pipeline != nullptr);

// scratch holds the gathered+masked input, materialized once and reused across passes
const size_t scratch_size = size_t{ n_kv } * nrows * sizeof(float);
if (ctx->prealloc_size_x < scratch_size) {
ctx->prealloc_size_x = scratch_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_x_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}

vk_op_topk_radix_push_constants pc { n_kv, width, nrows, n_tps, n_blocks, n_stream };
std::array<uint32_t, 3> elements {
pipeline->wg_denoms[0],
std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]),
1,
};
vk_subbuffer scratch_buf { ctx->prealloc_x, 0, ctx->prealloc_x->size };
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{ ggml_vk_tensor_subbuffer(ctx, scores), ggml_vk_tensor_subbuffer(ctx, top_k),
ggml_vk_tensor_subbuffer(ctx, cell_blk), ggml_vk_tensor_subbuffer(ctx, mask),
scratch_buf }, pc, elements);
ctx->prealloc_x_need_sync = true;
}

static void ggml_vk_sum(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
vk_op_sum_rows_push_constants p = vk_op_sum_rows_push_constants_init(src0, dst, ggml_nelements(src0));
ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_SUM, p);
Expand Down Expand Up @@ -15700,7 +15810,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr

break;
case GGML_OP_GET_ROWS:
ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node);
if (ctx->fused_topk_qsa) {
ggml_vk_topk_qsa(ctx, compute_ctx, cgraph, node_idx);
} else {
ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node);
}

break;
case GGML_OP_GET_ROWS_BACK:
Expand Down Expand Up @@ -17111,6 +17225,92 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
return true;
}

// Manual op-sequence match (ggml_can_fuse_subgraph rejects the mask's external reshape/cpy).
static bool ggml_vk_match_ops(const struct ggml_cgraph * cgraph, int node_idx,
const std::initializer_list<ggml_op> & ops) {
if (node_idx + (int) ops.size() > cgraph->n_nodes) {
return false;
}
for (size_t j = 0; j < ops.size(); ++j) {
const ggml_tensor * node = cgraph->nodes[node_idx + j];
if (node->op != ops.begin()[j] ||
(node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 ||
(node->flags & GGML_TENSOR_FLAG_OUTPUT) != 0) {
return false;
}
}
return true;
}

// True if the qwen4 QSA indexer top-k can be fused at node_idx (the get_rows).
static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) {
if (ctx->device->disable_fusion || !ctx->device->pipeline_topk_radix_qsa) {
return false;
}

const int n_ops = topk_qsa_pattern.size();
if (!ggml_vk_match_ops(cgraph, node_idx, topk_qsa_pattern) ||
!ggml_check_edges(cgraph, node_idx, topk_qsa_edges)) {
return false;
}

// elided nodes must be single-use (cpy counts its own src[1] self-reference)
for (int j = 0; j < n_ops - 1; ++j) {
const ggml_tensor * node = cgraph->nodes[node_idx + j];
const int32_t want = node->op == GGML_OP_CPY ? 2 : 1;
if (ggml_node_get_use_count(cgraph, node_idx + j) != want) {
return false;
}
}

const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0];
const ggml_tensor * add = cgraph->nodes[node_idx + n_ops - 2];
const ggml_tensor * top_k = cgraph->nodes[node_idx + n_ops - 1];

const ggml_tensor * scores = get_rows->src[0]; // [n_tps, n_blocks, n_stream]
const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream]
const ggml_tensor * expanded = add->src[0]; // [n_kv, n_tps, n_stream]

// raw mask: follow the reshape/cpy chain back to the materialized f16 input
const ggml_tensor * mask = add->src[1];
while (mask && (mask->op == GGML_OP_RESHAPE || mask->op == GGML_OP_CPY)) {
mask = mask->src[0];
}
if (!mask || mask->type != GGML_TYPE_F16) {
return false;
}

if (scores->type != GGML_TYPE_F32 || cell_blk->type != GGML_TYPE_I32 || top_k->type != GGML_TYPE_I32) {
return false;
}
if (!ggml_is_contiguous(scores) || !ggml_is_contiguous(cell_blk) || !ggml_is_contiguous(mask) ||
!ggml_is_contiguous(expanded) || !ggml_is_contiguous(top_k)) {
return false;
}

const int64_t n_tps = scores->ne[0];
const int64_t n_blocks = scores->ne[1];
const int64_t n_stream = scores->ne[2];
const int64_t n_kv = cell_blk->ne[0];
const int64_t width = top_k->ne[0];

// pin the indexer layout the shader's addressing assumes
if (scores->ne[3] != 1 || cell_blk->ne[1] != n_stream || ggml_nrows(cell_blk) != n_stream ||
ggml_nelements(mask) != n_kv * n_tps * n_stream ||
expanded->ne[0] != n_kv || expanded->ne[1] != n_tps || expanded->ne[2] != n_stream ||
top_k->ne[1] != n_tps || top_k->ne[2] != n_stream || top_k->ne[3] != 1 ||
n_blocks <= 0 || n_kv <= 0 || width <= 0 || width > n_kv) {
return false;
}

// only worth it in the radix regime; small k uses the faster tournament unfused
const uint32_t k_min_pipeline = std::max((uint32_t) log2f(float(width)) + 1, ctx->device->subgroup_size_log2);
if (k_min_pipeline < num_topk_pipelines && ctx->device->pipeline_topk_f32[k_min_pipeline]) {
return false;
}
return true;
}

static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph,
int node_idx) {
GGML_UNUSED(ctx);
Expand Down Expand Up @@ -17490,6 +17690,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg

ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
ctx->fused_topk_moe_scale = false;
ctx->fused_topk_qsa = false;
const char *fusion_string {};
if (!ctx->device->disable_fusion) {
uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i);
Expand Down Expand Up @@ -17579,6 +17780,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
// with a data dependency on that register. The overlap check still
// rejects partial overlaps (different base or size).
std::fill_n(op_srcs_fused_elementwise, 5, true);
} else if (ggml_vk_can_fuse_topk_qsa(ctx, cgraph, i)) {
ctx->num_additional_fused_ops = topk_qsa_pattern.size() - 1;
ctx->fused_topk_qsa = true;
fusion_string = "TOPK_QSA";
std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
} else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax_norm, { i + 3, i + 9 }) &&
ggml_check_edges(cgraph, i, topk_moe_early_softmax_norm_edges) &&
ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX_NORM)) {
Expand Down Expand Up @@ -17695,6 +17901,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
ctx->fused_ops_write_mask = 1;
ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
ctx->fused_topk_moe_scale = false;
ctx->fused_topk_qsa = false;
}
}

Expand Down Expand Up @@ -17891,6 +18098,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
if (keep_pattern(snake_pattern)) {
continue;
}
if (keep_pattern(topk_qsa_pattern)) {
continue;
}

// First, grab the next unused node.
current_set.push_back(first_unused);
Expand All @@ -17909,13 +18119,23 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
if (is_empty(graph->nodes[j])) {
continue;
}
// Don't pull forward nodes from fusion patterns
// Protect every interior QSA node (not just the start): the mask branch is
// independent, so it gets pulled out and breaks keep_pattern otherwise.
auto const &in_qsa_pattern = [&](int n) -> bool {

@jeffbolznv jeffbolznv Aug 30, 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.

I had to have codex explain this to me, but it makes sense now.

for (int o = 0; o < (int) topk_qsa_pattern.size(); ++o) {
if (n - o >= 0 && match_pattern(topk_qsa_pattern, n - o)) {
return true;
}
}
return false;
};
if (match_pattern(topk_moe_early_softmax_norm, j) ||
match_pattern(topk_moe_sigmoid_norm_bias, j) ||
match_pattern(topk_moe_sqrt_softplus_norm_bias, j) ||
match_pattern(topk_moe_early_softmax, j) ||
match_pattern(topk_moe_late_softmax, j) ||
match_pattern(snake_pattern, j)) {
match_pattern(snake_pattern, j) ||
in_qsa_pattern(j)) {
continue;
}
bool ok = true;
Expand Down Expand Up @@ -18717,15 +18937,14 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
if (!ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0])) {
return false;
}
// We could potentially support larger, using argsort to sort the
// whole thing. Not clear if this is needed.
uint32_t min_pipeline = (uint32_t)log2f(float(op->ne[0])) + 1;
if (min_pipeline >= num_topk_pipelines ||
!device->pipeline_topk_f32[min_pipeline]) {
return false;
// large k falls back to radix-select
const uint32_t min_pipeline =
std::max((uint32_t) log2f(float(op->ne[0])) + 1, device->subgroup_size_log2);
if (min_pipeline < num_topk_pipelines && device->pipeline_topk_f32[min_pipeline]) {
return true;
}
return device->pipeline_topk_radix_f32 != nullptr;
}
return true;
case GGML_OP_UPSCALE:
if (op->op_params[0] & GGML_SCALE_FLAG_ANTIALIAS) {
if ((op->op_params[0] & 0xFF) != GGML_SCALE_MODE_BILINEAR) {
Expand Down
Loading
Loading