qwen4-next-fixes: upstream perf fixes for qwen4exp (Vulkan TOP_K, kv-cells, GDN/LID, graph splits) - #229
Conversation
for_each_token_in tested all LLAMA_MAX_SEQ sequences for every used cell, while a cell almost always belongs to one. The scan now stops once the cell's own sequences have been seen. Same visit order, same callback arguments, so behaviour is unchanged. get_prev_tokens is the only caller, so this affects the n-gram path. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs: 55k context generation 56.3 -> 74.3 t/s 132k context generation 33.6 -> 50.9 t/s Prompt processing is unchanged, the scan is amortised over the ubatch there. The gain follows the number of used cells, so it grows with context and is invisible on short prompts.
| static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { | ||
| uint32_t ncols = src0->ne[0]; | ||
| uint32_t nrows = ggml_nrows(src0); | ||
| uint32_t k = dst->ne[0]; | ||
|
|
||
| if (k > (1u << ctx->device->max_workgroup_size_log2)) { |
There was a problem hiding this comment.
This enters the large-k path on k alone, but the scratch it goes on to allocate is sized from nrows (scratch_sz + perm_sz at L16290), and the only thing bounding nrows is TOPK_LARGE_K_MAX_ROWS over in supports_op at L21380.
The two sites agree today, but only by a chain that isn't visible from either one: pipeline_topk_f32[i] is created only for i <= max_workgroup_size_log2, so !pipeline_topk_f32[min_pipeline] happens to imply k > 1u << max_workgroup_size_log2. Nothing states that, and if topk pipeline creation is ever relaxed, or the row bound lowered, this branch will size prealloc_x from an unbounded nrows with no guard. At pp2048 / 64k context that is a ~1 GB scratch plus a ~512 MB permutation, and unlike the mul_mat paths there is no GGML_ABORT("Requested preallocation size is too large") backstop.
Could the predicate be factored into one helper used by both, so the coupling is explicit?
static bool ggml_vk_topk_use_large_k(const vk_device & device, int64_t k, int64_t nrows) {
return device->vulkan_memory_model &&
(uint32_t)k > (1u << device->max_workgroup_size_log2) &&
nrows <= TOPK_LARGE_K_MAX_ROWS;
}plus a GGML_ASSERT on the row count at the top of ggml_vk_topk_large_k. The "k too large but too many rows" case can't reach the backend at all, so asserting is the right shape.
There was a problem hiding this comment.
Done in 0543b4ba3. supports_op now goes through ggml_vk_topk_can_use_large_k(), and ggml_vk_topk_large_k asserts vulkan_memory_model and nrows <= TOPK_LARGE_K_MAX_ROWS on entry.
One deliberate deviation from the sketch: the runtime dispatch in ggml_vk_topk still routes on k alone. Routing on the full predicate there would send a hypothetical k-large/rows-large node into the regular path — where pipeline_topk_f32[min_pipeline] doesn't exist — instead of hitting the assert. Since supports_op makes that shape unreachable on the backend, k-only routing + the assert is the failure shape we want. Comment added at the dispatch site saying so.
Verified on gfx1151: test-backend-ops -b Vulkan0 -o TOP_K,ARGSORT 549/549, with the nrows=33 cases still correctly unsupported (CPU fallback). Will propose the same change upstream on ggml-org#28062.
| uint32_t ncols_pad_log2 = (uint32_t)ceilf(log2f(float(ncols))); | ||
| uint32_t ncolsp2 = 1 << ncols_pad_log2; | ||
|
|
||
| bool use_small = ncols_pad_log2 <= ctx->device->max_workgroup_size_log2 && |
There was a problem hiding this comment.
The split leaves use_small computed twice from the same inputs: here it decides whether to allocate the prealloc_x scratch, and at L16188 it decides whether to bind that scratch or out_buf as binding 1.
The expressions match, so this is correct as written. But if they ever drift so the caller says small and the callee says large, binding 1 becomes dst_buf and argsort_large writes 2 * ncolsp2 * nrows ints of tmp_idx into a destination tensor holding only ncols * nrows ints, i.e. an overflow into whatever else lives in that backend buffer. That is a silent, data-dependent corruption rather than a crash.
Passing use_small into ggml_vk_argsort_rows (or returning it) removes the possibility for free, and also lets the callee drop its own pipeline_idx recompute.
There was a problem hiding this comment.
Done in 0543b4ba3. ggml_vk_argsort_use_small() is now the single definition of the predicate: ggml_vk_argsort uses it to size the scratch and passes the result into ggml_vk_argsort_rows, which selects binding 1 from the parameter instead of recomputing. Caller/callee drift is no longer possible.
ggml_vk_topk_large_k passes the helper's result too; it's always false there (k > 2^wg_log2 and k <= ncols implies ncols_pad_log2 > max_workgroup_size_log2), so behavior is unchanged.
| @@ -250,10 +250,10 @@ llama_context::llama_context( | |||
|
|
|||
| cparams.fused_gdn_ar = true; | |||
| cparams.fused_gdn_ch = true; | |||
| cparams.auto_fgdn = true; | |||
| cparams.auto_fgdn = false; | |||
There was a problem hiding this comment.
This makes fused_gdn_ar, fused_gdn_ch and fused_lid unconditionally true on every backend, and the PR was validated on Vulkan, CUDA and Metal only. GGML_OP_LIGHTNING_INDEXER is implemented in CPU, CUDA, Metal and Vulkan and nowhere else; GGML_OP_GATED_DELTA_NET coverage varies beyond those.
With the probe gone, a backend that lacks the op no longer falls back to the unfused decomposition (build_delta_net_autoregressive / build_delta_net_chunking in delta-net-base.cpp). Instead the graph keeps the fused node and ggml-backend schedules it on the CPU, so every GDN/LID layer becomes a round trip per token, and the "%s not supported, set to disabled" warning that used to surface this is no longer emitted.
Given the OpenCL work on this fork, could we get a llama-bench run on OpenCL before this lands? If those regress, keeping the probe and opting the three validated backends out of it would get the same win without the blast radius.
There was a problem hiding this comment.
The mechanism you describe is right: with the probe gone, a backend lacking the fused op keeps the fused node and ggml-backend schedules it on the CPU, and the "not supported, set to disabled" warning is no longer emitted.
One correction on the blast radius for this PR's target model: qwen4exp never emits GGML_OP_LIGHTNING_INDEXER. Its indexer is built from primitive ops (the rectify-and-sum over head dot products, see the comment at qwen4exp.cpp:576) — which is exactly why the Vulkan TOP_K fix in this PR matters to it. qwen4exp.cpp doesn't reference fused_lid at all, and its only use of fused_gdn_ar/ch is a shape check at :884. So for qwen4-next itself:
- LID: no impact — the model never builds the fused op on any backend.
- GDN: OpenCL and SYCL both implement
GGML_OP_GATED_DELTA_NETin this tree (OpenCL: F32→F32, S_v ∈ {16,32,64,128}; qwen4exp runs S_v=128 since head_v_dim = ssm_d_state), so the fused node stays on-device there too.
The real LID exposure is deepseek4 / glm-dsa / deepseek32 on OpenCL and SYCL — those models do consult cparams.fused_lid, and both backends lack the op, so they'd take the CPU round trip per token where they previously got the unfused decomposition. That's upstream's accepted trade-off in ggml-org#27877 ("for backends that don't yet support these ops, there is CPU fallback"), but I'm happy to scope it: keep the probe and only skip it when every device in the context is CPU/CUDA/Metal/Vulkan. That keeps the init cost saving where we validated it without touching OpenCL/SYCL behavior for the DSA models.
On the OpenCL bench ask: I don't have an OpenCL-capable device on hand (validated here on Vulkan gfx1151, CUDA GB10, Metal M3 Ultra) — if you can point me at a box with the OpenCL backend running I'll do the before/after llama-bench run. Otherwise, want me to do the scoped-probe change instead?
…l-org#28032) * vulkan: add top-k radix sort shader for k >= 1024 * add Qwen 3.8 Flash Next top-k tests * add top-k qsa fusion * clean up code
0543b4b to
65c7c96
Compare
|
Update: swapped the Vulkan TOP_K commit from the argsort-based ggml-org#28062 to the version that landed upstream, ggml-org#28032 (radix select for k >= 1024 + fused QSA indexer variant), cherry-picked as 65c7c96 from the upstream squash merge daef7b6. This makes the two Vulkan review threads moot:
Re-validated on gfx1151 with the merged implementation: The llama-context.cpp thread (fused GDN/LID probe removal) is unaffected by this swap — still open on your call there. |
Review StatusCurrent Status: ❌ PENDING Pending reviews: Needs 1 Management or Team Lead. |
7 participants
Summary
Cherry-picks from ggml-org/llama.cpp for Qwen3.8-Flash-Next (qwen4exp), validated on Vulkan (Strix Halo), CUDA (GB10), and Metal (M3 Ultra):
build_inp_embdsplitdaef7b687(supersedes the argsort-based vulkan: support TOP_K when k exceeds the workgroup limit ggml-org/llama.cpp#28062 we had here initially)Benchmarks: Qwen3.8-Flash-Next UD-IQ1_S,
-p 2048 -b 2048 -ub 2048 -t 128 -fa on -ngl 999 --load-mode dioVulkan — qvac-dev-strix-1 (Radeon 8060S, RADV gfx1151)
(pp @ d0 dip matches the upstream author's own 8060S measurement; the depth gains dominate.)
Metal — qvac-dev-mac-arm64 (Apple M3 Ultra 96GB)
CUDA — NVIDIA GB10 (128GB unified)
Adjacent warm runs (machine shows ~10% thermal drift between runs, so only back-to-back runs are comparable):
tg slowdown with context largely gone on all three backends (d20000 tg retention: Vulkan 56→80%, Metal 82→92%, CUDA 80→86%).
Test plan
test-backend-opsTOP_K/TOPK_QSA/GATED_DELTA_NET/LIGHTNING_INDEXER: Vulkan 457/457 (gfx1151 + GB10), CUDA 738/738, Metal 644/644 (incl. new qwen4exp large-k and QSA fusion shapes)llama-completionsmoke test: identical coherent output on Vulkan, CUDA, Metal