QVAC-18665: TurboQuant (Vulkan) rebased into 8189

[jesusmb1995]

Brings previously merged TurboQuant PR #115 to 8189

Rebase & Fixes Summary

• Critical Fixes: Fixes an f16 upstream P*V dropout regression and resolves divergent barriers in copy_to_quant.comp.
• Testing: Added new test-backend-ops test leg (f16 FLASH_ATTN_EXT).
• Performance & PPL Verified: Perplexity and throughput perfectly match PR 115 baselines across AMD Strix and NVIDIA 5090 tests. See numbers/details on task comments https://app.asana.com/1/45238840754660/project/1209610027650937/task/1214639981590135
• Complete test-turboquant.sh executed both on 5090 and Strix

AI Review

• Score: Production ready (no functional regressions).
• For the full detailed AI review, bug fix details, and performance test notes, please see this comment.

[jesusmb1995]

macOS-arm64 test-backend-ops failure is pre-existing on temp-8189 baseline, not caused by this PR.

The current failing job (run 25753678331, job 75636543172) reports:

36: [SOFT_MAX_BACK] ERR = 0.001221320 > 0.000000100   SOFT_MAX_BACK(type=f32,ne=[1023,15,1,1],scale=1.000000,max_bias=0.000000): FAIL
36: [SOFT_MAX_BACK] ERR = 0.000900075 > 0.000000100   SOFT_MAX_BACK(type=f32,ne=[1023,1023,1,1],scale=1.000000,max_bias=0.000000): FAIL
36: [SOFT_MAX_BACK] ERR = 0.002131379 > 0.000000100   SOFT_MAX_BACK(type=f32,ne=[1023,15,1,1],scale=0.100000,max_bias=0.000000): FAIL
36: [SOFT_MAX_BACK] ERR = 0.000943703 > 0.000000100   SOFT_MAX_BACK(type=f32,ne=[1023,1023,1,1],scale=0.100000,max_bias=0.000000): FAIL
36:   Backend MTL0: FAIL

The previous Asserting on type 40 (TBQ/PQ on Metal mul_mat) abort that motivated the recent Metal exclusion patch is gone — the test now runs to completion and surfaces this separate, unrelated Metal SOFT_MAX_BACK numerical-accuracy regression.

The exact same four SOFT_MAX_BACK failures, with the same magnitudes, already appear on PR #131 (run 25588996013, job 75123047326), which targets temp-8189 and contains no TurboQuant changes. Likely root cause is one of the upstream Metal commits pulled in during the temp-8189 rebase (e34743a9b Add SOFT_MAX_BACK metal kernel / 5fa6d7787 metal: port OUT_PROD, SILU_BACK, SOFT_MAX_BACK, RMS_NORM_BACK ops to split architecture).

So: the Metal patch in this PR fixed the type-40 abort it was meant to fix; the remaining macOS-arm64 failure is a temp-8189 baseline issue and should be addressed in a separate PR (or by the temp-8189 maintainer) rather than in this rebase.

[jesusmb1995]

ubuntu-24-cmake-webgpu failure is a pre-existing flake on temp-8189, unrelated to this PR.

The current failure (run 25753678331, job 75636543611) is:

The following tests FAILED:
     32 - test-thread-safety (Timeout)

test-thread-safety runs the WebGPU backend through llvmpipe (Mesa software rasterizer) and hangs mid-generation, hitting the 3600s ctest timeout. No assertion, no error — just a hang.

Same exact failure (test-thread-safety Timeout on ubuntu-24-cmake-webgpu) reproduces on other recent PRs targeting temp-8189 that contain no TurboQuant changes:

• PR vulkan: Force f32 src1 through the strided cpy path which uses correct strides #130 (job 75039398337) — same test-thread-safety timeout
• PR Add GGML_OP_DELTA_NET_AR compute shader and dispatch path so Vulkan #129 (job 75029636148) — same job failed

Other PRs on the same base pass (#131, #128), confirming flake rather than deterministic regression. Likely root cause is upstream WebGPU/llvmpipe concurrency in the temp-8189 rebase, not anything introduced here.

Recommend a re-run of this job; not a TBQ-side blocker.

[jesusmb1995]

AI Review (multiple passes)

Critical

None

High

None

Medium

None

Low

• [ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp ~ lines 935-947] Indentation inside if (!is_tq_fht) { ... } for the get_rows block is misleading: the if (tname == "f16") / else / final string_to_spv("get_rows_..._f32", ...) lines sit at the outer for-loop indentation, but the closing } on line 946 actually closes the if (!is_tq_fht) (because braces ignore indentation). Functionally correct today, but a strong footgun for future edits — adding a line "outside" the if (visually) would still land inside it. Worth re-indenting before merge. Okayish because this edit minimizes conflicts with upstream rather than style.
• [ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp lines 142-151] For the upstream same-type Q4_0/Q8_0/TBQ paths, both the legacy #if BLOCK_SIZE > 1 block and the new #if defined(K_BLOCK_SIZE) && K_BLOCK_SIZE > 1 block fire and call setTensorLayoutBlockSizeNV on the same tensor twice with the same value. Harmless (idempotent) but redundant. The legacy block is also dead code under PR 133 because K_BLOCK_SIZE/V_BLOCK_SIZE are always ≥ BLOCK_SIZE for any case where BLOCK_SIZE > 1. Could be cleaned up to a single source of truth. However, less conflicts this way with upstream because original block is preserved.

Notes

• [src/llama-graph.cpp] The upstream assert GGML_ASSERT(v_mla == nullptr); was relaxed to GGML_ASSERT(!(inp->self_rotk || inp->self_rotv) || v_mla == nullptr); (currently at line 2078) during the rebase to accommodate MLA. This is correct and handled properly. The other build_attn overload around line 2245 has no equivalent assert and unconditionally calls build_attn_mha(..., v_mla, ...), but this asymmetry is pre-existing and identical between PR 115 and PR 133, so not a regression.
• [ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp] The centroid trick implementation is identical to PR 115. It appeared in the diff comparison solely because surrounding variable names and macros (like K_BLOCK_SIZE) were updated during the rebase. The underlying mathematical behavior remains unchanged and correct.
• [ggml/src/ggml-vulkan/ggml-vulkan.cpp] The previous pass claimed the is_coopmat1_fa_type(kv_type) type_ok filter was added in PR 133. That was incorrect — the filter is byte-identical between PR 115 (commit 3840e11d5, lines 3094-3110) and PR 133. No change here.
• [ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp] The if (norm < 1e-15) return; early-out in PR 115 was restructured in PR 133 into an if (skip_quant) {} else {} plus if (!skip_quant) { /* TQ_NORM_CORRECTION + QJL */ } block, with skip_quant = subgroupBroadcast(norm < 1e-15, 0u). Because norm = sqrt(tq_wg_add(local_sq)) is already workgroup-uniform, the subgroupBroadcast is functionally redundant but provides a compiler hint that makes the inner barrier() calls provably subgroup-uniform on backends that lack good uniformity analysis. Same barrier()-inside-conditional pattern existed in PR 115 (after the early return); not a regression.
• [ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp] PR 133 deliberately stops emitting get_rows_tbq3_0 / get_rows_tbq4_0 / get_rows_pq3_0 / get_rows_pq4_0 (the if (!is_tq_fht) guard at line 936). The comment explains: element-wise dequant doesn't apply for TBQ/PQ since they need an FHT pass. PR 115 emitted those shader variants but they were never referenced from ggml-vulkan.cpp (verified with grep — no get_rows_tbq* / get_rows_pq* references exist). Intentional and safe behavior change.
• [ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl lines 95-97] PR 133 added #ifndef BLOCK_SIZE → #define BLOCK_SIZE 1 near the top of the file, in addition to the existing #ifndef K_BLOCK_SIZE → 1 and #ifndef V_BLOCK_SIZE → 1 fallbacks lower down. This is the actual fix for the upstream same-type f16 regression that motivated the new test-leg in test-turboquant.sh — without these defaults, #if BLOCK_SIZE == 1 / #if K_BLOCK_SIZE == 1 style guards in flash_attn_cm1.comp (e.g. the V_BLOCK_SIZE == 1 guard around the f16 PV coopmat path) silently evaluate as false when the macro is undefined and drop the entire PV section from the SPIR-V.
• [ggml/src/ggml-vulkan/ggml-vulkan.cpp lines 425-447 and 10169] vk_fa_pipeline_state gained a v_type field (default-initialised to GGML_TYPE_F16 in the constructor and overwritten with v->type at dispatch time). It participates in operator< so that mixed-K/V variants live as separate map entries under the same K-keyed pipeline bucket. Verified: the CREATE_FA and CREATE_FA_MIXED macros correctly filter by fa.first.v_type, and same-type Q4_0/Q8_0/TBQ paths still resolve to their own v_type bucket since v->type is assigned before pipelines.find(...). No issues found here.

Summary

The rebased PR 133 introduces critical fixes over PR 115 (specifically the #ifndef BLOCK_SIZE/K_BLOCK_SIZE/V_BLOCK_SIZE → 1 fallbacks in flash_attn_base.glsl that fix the f16 upstream P*V dropout regression, and the subgroupBroadcast for divergent barriers in copy_to_quant.comp). Re-review found no Critical/High/Medium regressions. The two remaining Low-severity stylistic / code-hygiene issues (misleading indentation in vulkan-shaders-gen.cpp, redundant setTensorLayoutBlockSizeNV call in flash_attn_cm2.comp) are intentionally preserved, as fixing them would increase merge conflicts with upstream without providing functional benefits. The decorative #ifdef D_TYPEV4 guard in flash_attn_base.glsl was dropped in this PR (unconditional OV4 declaration), since every generator call site already passes D_TYPEV4=vec4 and every consumer is unconditional.

Score

4/4 perfect (no functional regressions; 2 minor Low-severity hygiene items left intentionally to minimize upstream conflicts)