vulkan : add mul_mat_vec_id for the TurboQuant weight types - #264
Conversation
MoE decode for a TQ3_1S/TQ4_1S model ran entirely on the CPU: ggml-org#259 rejected MUL_MAT_ID for both types because no mul_mat_vec_id pipeline existed, and reaching ggml_vk_get_dequantize_mul_mat_vec_id() with a TQ src0 asserted on a null pipeline. No new shader source is needed. All of the expert indirection lives in mul_mat_vec_base.glsl (get_offsets(), reduce_result()), which both TQ mat-vec shaders already include, and the expert id arrives via gl_WorkGroupID.y, which neither shader touches. Compiling the existing sources with MUL_MAT_ID adds exactly one binding (the ids buffer, 5 -> 6, matching mul_mat_vec_id_num_bindings) and introduces no subgroup capability, so the 32-thread workgroup pin and the shared-memory butterfly that make these kernels wave64-safe are unaffected. spirv-val passes on both. The blanket reject is replaced by a size gate rather than removed. There is still no TQ mul_mm_id, so prompt processing takes the generic path, where ggml_vk_get_mul_mat_mat_id_pipeline() returns nullptr, qx_needs_dequant goes true, and the ENTIRE expert tensor is staged as f16 -- x_ne is ggml_nelements(src0), across all experts. Reject when that staging buffer would exceed maxStorageBufferRange, or ggml_vk_mul_mat_id_q_f16() reaches GGML_ABORT("Requested preallocation size is too large"). The gate is deliberately independent of src2->ne[1]. An n-dependent gate is worse than no support: weight_buft_supported() (llama-model-loader.cpp) probes with a fixed ids->ne[1] = 512, so a gate that answers differently at load and at decode parks the experts in one backend's buffer and then runs the op in the other, copying every expert tensor across the bus on every token. Worked example of what the gate excludes: DeepSeek-V4-Flash Config-I has ffn_{gate,up,down}_exps of 2048 x 4096 x 256 = 2,147,483,648 elements, so the f16 staging buffer is exactly 4,294,967,296 bytes -- one byte over RADV's UINT32_MAX maxStorageBufferRange. Its experts therefore stay on the CPU, as before this change. Lifting that needs a real TQ mul_mm_id, which reads the quantized data through dequant_funcs.glsl and never allocates the f16 buffer. test-backend-ops gains a TurboQuant MUL_MAT_ID sweep. TQ3_1S/TQ4_1S are in all_types but not base_types, so their only prior coverage was two cases; this covers both sides of the n <= 8 threshold ggml_vk_use_mul_mat_vec_id() splits on, three n_used counts, and broadcast, which exercises the expert-index wrap. Assisted-by: Claude
|
Tangential, but it came out of reading this area and it affects the branch this PR targets: I have filed #265 — the CUDA TQ3_1S correctness fixes from #256 are merged but not reachable from Relevant to this PR in two ways.
And |
|
Verified on wave32 and merging. This closes the gap #259 left, and closes it better than I would have. GB10, NVIDIA, warp size 32
The count arithmetic is a nice confirmation on its own: the tree was 24496 before this PR, your sweep adds 76 TQ cases, and the run lands on exactly 24572 with nothing else moved. And the case that mattered most: That is the decode shape that asserted on a null pipeline in #259 — the one my reject was standing in for. It now runs on the GPU and matches the reference. Two wave widths, your gfx1151 and this GB10, same conclusion. On the parts that are better than my rejectThe Compiling the existing sources with Running a real-model PPL A/B because synthetic kernel tests let the CUDA fused-TQ3_1S bug through is the correct lesson to draw from that incident, and 6.5037 vs 6.5144 inside the error band is the evidence that matters. The 7.86 to 397 pp512 and 4.30 to 20.89 tg128 numbers are the payoff — MoE decode for these types was running entirely on CPU before this. Filing the RADV Standing offer, renewedYou have now found and fixed two real gaps in this backend, each time with the diagnosis rather than just the patch. If you want the |
0463c8e
into
TheTom:feature/turboquant-kv-cache
|
Closing the loop on the "I have that in progress" note above, since it changes what this PR's size gate costs in practice. The TQ
13.4x on the pass, and the perplexity improves rather than degrades. The improvement is worth explaining rather than just reporting, because a shift that size is also what a subtle kernel bug looks like. On Qwen3.6-35B-A3B, where the comparison is GPU-versus-GPU rather than GPU-versus-q8_0-CPU, chunk[1] moves 4.5768 -> 4.5769 against a CPU-expert control — which is the agreement I would expect when no lossy intermediate is being removed. None of this needs to land in this PR. The size gate here is still the right conservative behaviour for a tree without a TQ |
|
Thank you — and for taking the time to verify it on wave32 rather than just taking my numbers. Two wave widths agreeing is worth more than either run alone, and the count arithmetic landing on 24572 is a better check than I could do from one machine. On I have already put what I found on #260 — short version, since it postdates your comment there and changes the shape of it a little:
I got one thing wrong on the way there and it is worth flagging: I first concluded the bounds-check change was a harmless refactor, because Happy to test a candidate on wave64 whenever he has one, or to take a pass at the fix myself if he would rather hand it over — his call, not mine. Separately: the TQ |
Wires the pieces together, so MoE prompt processing for a TQ3_1S/TQ4_1S model runs the rotated matmul instead of staging every expert tensor as f16. The getter cases and the rotate orchestration land in one commit on purpose: either alone runs the rotated matmul against an unrotated activation and returns silently wrong numbers. Flow when tq_rotate is set: 1. force staging, so d_Y is ctx->prealloc_y and never aliases src1 2. contiguous f32 copy src1 -> prealloc_y 3. barrier, then tq_rotate_act in place on prealloc_y 4. mul_mm_id reads the rotated f32 copy against centroid*scale weights Details that have to agree or the result is quietly wrong: - `tq_rotate` requires `mmp != nullptr && !x_non_contig`, i.e. that the TQ pipeline was really selected. On the f16 dequant fallback the weights are true weights and must NOT see a rotated activation. - Staging is forced even for a contiguous f32 src1. Otherwise y_f32_kernel is true, qy_needs_dequant is false, and d_Y aliases the graph's own tensor -- the hazard Metal's in-place rotate/un-rotate has to work around. The "not implemented" assert is relaxed for exactly this case. - effective_src1_type and y_sz both switch to f32, matching what the rotate writes. The staging pipeline is an f32 cpy, not the usual f16 one: the butterfly is 5 rounds of adds and f16 would lose precision quantization did not. - The prealloc_y reuse cache is keyed on the ROTATE pipeline pointer rather than the cpy pipeline. prealloc_y now holds rotated data, so a later non-rotated consumer of the same src1 must not reuse it; keying this way forces a re-stage for anyone else. - Only the ne11*ne12*ne13 rows the copy wrote are rotated; the padded_n rows above ne11 are left alone, as in the f16 staging path. - ggml_vk_get_mul_mat_mat_id_pipeline() ends in GGML_ASSERT(support_fp32acc), so listing TQ in its switch would ABORT on a coopmat2 device, where these pipelines are deliberately not created. It now returns nullptr when both acc variants are empty, falling back to the f16 dequant path. Not yet done: the maxStorageBufferRange gate in supports_op still rejects large expert tensors on the old premise that prompt processing must stage them as f16. With this path that premise no longer holds, but relaxing it is a separate commit pending review of #264. Assisted-by: Claude
2 participants
MoE decode for a TQ3_1S/TQ4_1S model ran entirely on the CPU. #259 rejected MUL_MAT_ID for both types because no
mul_mat_vec_idpipeline existed, and reachingggml_vk_get_dequantize_mul_mat_vec_id()with a TQ src0 asserted on a null pipeline.No new shader source is needed. All of the expert indirection already lives in
mul_mm_vec_base.glsl(get_offsets(),reduce_result()), which both TQ mat-vec shaders include, and the expert id arrives viagl_WorkGroupID.y, which neither shader touches. Compiling the existing sources withMUL_MAT_IDadds exactly one binding — the ids buffer, 5 → 6, matchingmul_mat_vec_id_num_bindings— and introduces no subgroup capability, so the 32-thread workgroup pin and shared-memory butterfly that make these kernels wave64-safe are unaffected.spirv-valpasses on both.The reject became a size gate rather than being removed
There is still no TQ
mul_mm_id, so prompt processing takes the generic path:ggml_vk_get_mul_mat_mat_id_pipeline()returns nullptr,qx_needs_dequantgoes true, and the entire expert tensor is staged as f16 (x_ne = ggml_nelements(src0), across all experts). This rejects the cases where that staging buffer would not fit, otherwiseggml_vk_mul_mat_id_q_f16()reachesGGML_ABORT("Requested preallocation size is too large").The gate is deliberately independent of
src2->ne[1]. An n-dependent gate is worse than no support:weight_buft_supported()(llama-model-loader.cpp) probes with a fixedids->ne[1] = 512, so a gate that answers differently at load and at decode parks the experts in one backend's buffer and then runs the op in the other, copying every expert tensor across the bus on every token.Worked example of what it excludes: DeepSeek-V4-Flash Config-I has
ffn_{gate,up,down}_expsof 2048 × 4096 × 256 = 2,147,483,648 elements, so the f16 staging buffer is exactly 4,294,967,296 bytes — one byte over RADV'sUINT32_MAXmaxStorageBufferRange. Its experts stay on the CPU, as before this change. Lifting that needs a real TQmul_mm_idthat reads the quantized data throughdequant_funcs.glsland never allocates the f16 buffer. I have that in progress.Validation on gfx1151 (Radeon 8060S, RADV, wave64)
test-backend-ops -o MUL_MAT_ID:891 + 4 newly-supported + 72 new sweep cases = 967, which reconciles exactly. That arithmetic matters here because of #242 — a fully skipped op still reports OK, so I counted executed cases rather than trusting the summary line.
Full sweep on an idle GPU: 24572/24572, 2/2 backends. An earlier run showed one failure, which turned out to be pre-existing: the same sweep on an image without any of this fails identically, on a
TOPK_MOEcase unrelated to TurboQuant. Filed separately as #261.Real-model A/B on
thetom-ai/Qwen3.6-35B-A3B-ConfigI(431 TQ3_1S tensors, 256 experts, sha256-verified), since synthetic kernel tests are exactly what let the CUDA fused-TQ3_1S bug survive:0.16% apart, far inside the error band.
llama-bench -ngl 99 -p 512 -n 128 -r 2:To be precise about what those numbers mean: both arms report
backend: Vulkanand the control is already on the GPU. What moved is the MoE expert matmuls, which the blanket reject forced onto the CPU per-op. The gain is not decode-only either — replacing the reject with the size gate enables both prefill (viamul_mm_id+ f16 dequant) and decode (viamul_mat_vec_id).Tests
test-backend-opsgains a TurboQuant MUL_MAT_ID sweep. TQ3_1S/TQ4_1S are inall_typesbut notbase_types, so their only prior coverage was two cases. This covers both sides of then <= 8threshold thatggml_vk_use_mul_mat_vec_id()splits on, threen_usedcounts, and broadcast, which exercises the expert-index wrap.AI usage disclosure: written with Claude Code; I reviewed every change and ran all the measurements above on my own hardware.