vulkan : add rotated mul_mm_id for the TurboQuant weight types - #266
Merged
TheTom merged 8 commits intoAug 7, 2026
Merged
Conversation
…on rotate) Groundwork for a real TQ mul_mm_id, so MoE prompt processing stops staging the entire expert tensor as f16. Not yet wired: no generation entries and no host pipelines, so this is inert until those land (tq3_1s/tq4_1s are not in type_names, so no matmul variant defines DATA_A_TQ*_1S today). mul_mm_funcs.glsl gains TQ3_1S and TQ4_1S A-side loads that produce centroid*scale and deliberately do NOT apply the inverse WHT. That is only correct against a pre-rotated activation, via <S.H.y.k, x> == k.<y, H.S.x>. LOAD_VEC_A is 8 for these types, which lines up exactly: one invocation covers one 3-byte packing group = 8 contiguous elements in k, so buf_idx uses the LOAD_VEC_A/2 contiguous form rather than Q4_0's split-nibble form. tq_rotate_act.comp applies signs -> butterfly -> 1/sqrt(32) to the activation. k lives here because the A-side loads omit it; mul_mat_vec_tq3_1s.comp folds k into the weight instead. Exactly one side must apply it. Deliberately NOT in place, unlike Metal. Metal rotates src1 in place, matmuls, then un-rotates, needing ggml_metal_op_mutates_tq_src1 plus barriers -- and its rotate indexes src1 flat while the matmul reads with nb11/nb12/nb13, so a strided or padded src1 rotates different elements than the matmul consumes. Writing to scratch costs n_tokens*k*4 bytes and removes that class entirely. Verified before any host code: - glslc + spirv-val clean for mul_mm and mul_mm_id on both types; rotate shader VALID at LocalSize 32 with zero subgroup capabilities, so the wave64 pin that makes the butterfly correct still holds. - The GLSL (tid & step)==0 pairing is bit-identical to the (j, j+step) loop nest of tq3_0_rht_forward() in ggml-turbo-quant.c (0.0 diff over random inputs). - <x, rht_inverse(c)> == <rotate(x), c> over 20000 random blocks: median relative error 1.9e-16, p99 1.1e-14, max 1.2e-12 (float64 rounding over 32 terms). Centroid tables copied from dequant_funcs.glsl, TQ3 asymmetric. Assisted-by: Claude
Emits matmul{,_id,_id_subgroup}_tq{3,4}_1s_{f32,f16} from mul_mm.comp plus the
type-independent tq_rotate_act pipeline. Still inert on the host side: no
pipelines are created and no dispatch path selects them yet.
Generated explicitly rather than by adding the types to type_names, matching
the mat-vec precedent in ggml-org#259 -- that loop would also emit q8_1 mmq variants,
for which these types have no integer-dot path.
coopmat2 is excluded: TQ has no dequant_funcs_cm2.glsl entry, and the target
(gfx1151) exposes KHR_coopmat only.
LOAD_VEC_A is pinned to 8 rather than taken from load_vec_quant, because the
A-side block in mul_mm_funcs.glsl indexes idx/4 and idx&3 to map one invocation
onto exactly one 3-byte packing group of 8 contiguous elements. Any other value
would silently misindex.
One rotate pipeline serves both types: TQ3 and TQ4 share the same 32-element
sign pattern and butterfly, and the shader takes no DATA_A_* define because it
only touches the activation.
Note on shared memory: ggml_vk_matmul_shmem_support() has no lut_size case for
these types, and that is now correct rather than accidental -- the A-side block
uses a function-local const centroid array, not a shared-memory LUT, so its
true LUT cost is zero.
Assisted-by: Claude
Declares and creates pipeline_tq_rotate_act. Still inert: no dispatch path selects it, and the TQ matmul pipelines are not registered yet. Landing this alone is deliberate. Registering TQ in ggml_vk_get_mul_mat_mat_id_pipeline() before the rotate orchestration exists would run the rotated matmul against an UNROTATED activation and return silently wrong numbers, so the getter cases and the orchestration have to land together. This commit only proves the shader builds and the pipeline is creatable. Workgroup fixed at 32 for the same reason as the TQ mat-vec pipelines: the butterfly pairs lanes through a 32-entry shared array, so it is correct only when the workgroup equals the block size. gfx1151 reports warp size 64. Assisted-by: Claude
The TQ generation block passed LOAD_VEC_B = load_vec, which is 8 on the fp16 path, but its float_type_dict only carried FLOAT_TYPE/V2/V4. load_b_to_shmem() references FLOAT_TYPEV8 whenever LOAD_VEC_B is 8, so every fp16 TQ matmul variant failed to compile with "'FLOAT_TYPEV8' : undeclared identifier" at mul_mm_funcs.glsl:638. The dict now matches the one the type_names loop builds. Caught by a real build, not by the local glslc pre-check: that check passed hand-specified defines and so could not reproduce the generator's dict. Running the actual vulkan-shaders-gen locally does reproduce it and takes about a minute, which is the right loop for this class of error. Assisted-by: Claude
…buffer
Simplification of the rotate design. The hazard being avoided is mutating
src1, a graph tensor -- not in-place work as such. ctx->prealloc_y is
context-owned scratch, so rotating THAT in place is safe and needs no extra
allocation: each workgroup reads, transforms and writes back exactly its own 32
contiguous elements, with no cross-workgroup dependency.
Two consequences the host side must honour, or this is silently wrong:
1. Staging must be FORCED for these types. If src1 is already contiguous and
the right dtype, d_Y would otherwise alias the real tensor and the rotate
would mutate the graph's activation -- exactly Metal's problem.
2. The staged copy must be f32, and the f32 B variant of the matmul selected
to match. The butterfly is 5 rounds of adds; doing it in f16 would lose
precision the quantization did not.
Rotating a packed scratch copy also sidesteps the stride mismatch in Metal's
version, where the rotate indexes src1 flat while the matmul reads it with
nb11/nb12/nb13: the staging pass has already made the copy contiguous.
Shader drops to a single read-write binding and 3 push constants
(k, nrows, stride); the pipeline creation is updated to match. Verified by
running the real vulkan-shaders-gen: exits clean, tq_rotate_act symbols emitted.
Assisted-by: Claude
Creates pipeline_dequant_mul_mat_mat_id[TQ3_1S/TQ4_1S] at all five non-coopmat2 sites (coopmat1, and the scalar subgroup / non-subgroup variants). Still inert: ggml_vk_get_mul_mat_mat_id_pipeline()'s type switch does not include these types, so nothing selects them yet. That split is deliberate. Adding the getter cases before the rotate orchestration exists would run the rotated matmul against an UNROTATED activation and return silently wrong numbers, so the getter and the orchestration land together in the next commit. coopmat2 is skipped: TQ has no dequant_funcs_cm2.glsl entry, and gfx1151 is KHR_coopmat only. Note for whoever adds the getter cases: ggml_vk_get_mul_mat_mat_id_pipeline() ends in GGML_ASSERT(support_fp32acc), so a type listed in its switch whose pipelines were NOT created on the running device ABORTS rather than falling back. Because coopmat2 is skipped here, the getter must explicitly return nullptr for TQ when both acc variants are empty, letting it degrade to the f16 dequant path instead of asserting on a coopmat2 device. Assisted-by: Claude
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 ggml-org#264. Assisted-by: Claude
…llback
The gate rejected MUL_MAT_ID for TQ weights whose f16 staging buffer would
exceed maxStorageBufferRange. That premise no longer holds everywhere: the
rotated mul_mm_id path reads the quantized weights directly through
mul_mm_funcs.glsl and never materialises the expert tensor as f16, so there is
no staging buffer to overflow.
The gate now applies only where the fallback is actually taken -- when the TQ
matmul pipelines do not exist on this device (they are not created for
coopmat2) or src0 is not dim01-contiguous. Those are exactly the two conditions
ggml_vk_mul_mat_id_q_f16() gates tq_rotate on, so the two stay in agreement; if
they ever diverge the fallback would hit
GGML_ABORT("Requested preallocation size is too large") again.
Still independent of src2->ne[1], for the same reason as before:
weight_buft_supported() probes with a fixed ids->ne[1] = 512, so an n-dependent
answer parks the experts in one backend's buffer and runs the op in the other.
Effect: DeepSeek-V4-Flash Config-I becomes eligible. Its
ffn_{gate,up,down}_exps are 2048 x 4096 x 256, whose f16 staging would have been
exactly 4,294,967,296 bytes against a 4,294,967,295 limit -- excluded by one
byte. Nothing is staged now, so the limit is not consulted.
Assisted-by: Claude
Owner
|
Verified on wave32 and merging. This is the other half of #264 and it closes the gap that PR deliberately left. GB10, NVIDIA, warp size 32, head
Same totals as #264 left the tree at, so the new The gate change is the part I looked at hardest, because it is the piece most likely to be wrong in a subtle way. Narrowing the f16 staging limit so it applies only to the dequant fallback rather than to Worth noting what you have built over three PRs: #259 wired the weight types, #264 put MoE decode on the GPU, and this puts prefill there too and removes the model-size ceiling that #264 needed. That is a coherent arc, each step verified on real models rather than synthetic shapes alone, and each one landing with the gap it does not yet close stated plainly rather than papered over. The stating-the-remaining-gap habit is why these have been quick to review. Merging. |
TheTom
merged commit Aug 7, 2026
1371e37
into
TheTom:feature/turboquant-kv-cache
9 of 25 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #264. That PR added
mul_mat_vec_idso MoE decode runs on the GPU for the TurboQuant weight types, but left prompt processing on the f16 dequant path — and gated out any model whose expert tensor could not be staged as f16. This adds a real TQmul_mm_id, which removes the staging entirely.The approach
The A side is loaded as
centroid * scalewith no inverse WHT, which is only correct against a pre-rotated activation. That is the identitywith
Hthe symmetric 32x32 Hadamard,S = diag(TQ_SIGNS)andk = 1/sqrt(32): rotate the activation once per tile instead of un-rotating every weight block.dequant_funcs.glslalready produced exactly this form, so the A-side loads inmul_mm_funcs.glslare the only new dequant code.LOAD_VEC_Ais pinned to 8 for these types, which lines up so one invocation covers one 3-byte packing group = 8 contiguous elements in k.kis applied in the rotate shader because the A-side loads omit it —mul_mat_vec_tq3_1s.compdoes the opposite and folds it into the weight. Exactly one side must.Rotating scratch, not src1
Metal's rotated
mul_mm_idrotatessrc1in place, matmuls, then un-rotates, which needs a hazard flag (ggml_metal_op_mutates_tq_src1) and barriers on both sides. Its rotate also indexessrc1flat while the matmul reads it withnb11/nb12/nb13, so a strided or paddedsrc1rotates different elements than the matmul consumes.This rotates
ctx->prealloc_y— the staging copy — instead, and never touches the graph's tensor. Staging is forced for these types even whensrc1is already contiguous f32, because otherwisey_f32_kernelis true,qy_needs_dequantis false, andd_Yaliases the realsrc1. The staging pass has already made the copy contiguous, so the stride mismatch cannot arise. Staged in f32: the butterfly is five rounds of adds and f16 would lose precision the quantization did not.The
prealloc_yreuse cache is keyed on the rotate pipeline rather than the copy pipeline, so a later non-rotated consumer of the samesrc1cannot reuse rotated data.The size gate from #264
It now applies only where the f16 dequant fallback is actually taken — when the TQ pipelines do not exist on the device (they are not created for coopmat2) or
src0is not dim01-contiguous. Those are the two conditionstq_rotateis gated on, so the two must agree; if they diverge the fallback hitsGGML_ABORT("Requested preallocation size is too large")again.ggml_vk_get_mul_mat_mat_id_pipeline()ends inGGML_ASSERT(support_fp32acc), so it now returns nullptr for TQ when both acc variants are empty rather than aborting on a coopmat2 device.Validation — gfx1151, RADV, wave64
test-backend-ops -o MUL_MAT_ID: 967/967, 2/2 backends on the rebased head. Cases at n=9/17/32 exercise the rotatedmul_mm_id; n=1/4/8 still takemul_mat_vec_id.Before any host code, the identity was checked directly: the GLSL
(tid & step)==0pairing is bit-identical to the(j, j+step)loop nest oftq3_0_rht_forward(), and<x, rht_inverse(c)> == <rotate(x), c>over 20000 random blocks gives median relative error 1.9e-16, max 1.2e-12.Batched perplexity, Qwen3.6-35B-A3B-ConfigI (431 TQ3_1S tensors, 256 experts),
-c 512 --chunks 2, three arms:mul_mm_idmul_mm_idMeasuring the middle arm matters: against experts-on-CPU this looks like 89x, but 44.8x of that is what #264 already delivers. This PR's own contribution is 2.0x over the dequant path, plus removing the staging buffer.
DeepSeek-V4-Flash Config-I,
-c 512 --chunks 2 -ub 64:Its
ffn_{gate,up,down}_expsare 2048 x 4096 x 256, whose f16 staging would be exactly 4,294,967,296 bytes against a 4,294,967,295 limit — excluded by one byte under #264's gate. Nothing is staged now.Two things I would rather state than have found
The perplexity improves, which is also what a subtle bug looks like. The reason is that
TQ3_1Shas.vec_dot_type = GGML_TYPE_Q8_0, so the CPU fallback quantizes the activation to q8_0 before every dot; the rotated path keeps f32. The three DeepSeek runs order monotonically by how much intermediate quantization each avoids (10.4794 CPU q8_0, 9.7819 mat-vec q8_0, 8.9692 rotated f32), and on Qwen — where the comparison is GPU-versus-GPU — chunk[1] moves 4.5768 to 4.5769.The rotated path is marginally less accurate than the dequant path (+0.46% versus -0.14% against the CPU control). Both are far inside the band. This is consistent with
buf_aholdingcentroid*scalein f16 on the fp16 path: the rotated domain has a different dynamic range than true weights.Not included
mul_mmfor plain MUL_MAT is untouched — it still uses the dequant path, which works and is validated. Onlymul_mm_idis wired here. coopmat2 is excluded throughout: nodequant_funcs_cm2.glslentry, and gfx1151 is KHR_coopmat only, so I cannot test it.AI usage disclosure: written with Claude Code; I reviewed every change and ran all measurements above on my own hardware.