feat(shaders): quantized matvec + subgroup-attention shader library (upstream PR2) - #88
Merged
ericcurtin merged 5 commits intoAug 22, 2026
Conversation
…upstream PR2)
Second PR of the upstreaming stack. Adds the quantized matvec kernels that every
model path dispatches, plus the two subgroup paged-attention variants:
- mlx4 family (4-bit): plain, _cols, _w8, _w16, w8sg, and the repack /
repack_batched f32 variants
- mlx2 / mlx6 / mlx8 repack variants
- nvfp4 family: plain, _e4m3, and the e4m3repack / repack f32 variants
- fp8: plain, _fast, and fp8repack
- q8_0_cols and f16_cols (column-batched dequant matvec)
- paged_attn_decode_f32_sg / _f16_sg (subgroup decode attention)
- relu2
These are shared by every model, so they land once here rather than being
duplicated across the per-model PRs that follow. The Rust-side push-constant
helpers for them already shipped with the compute-engine foundation (PR1), where
they were dead code; this activates them.
scripts/compile_shaders.sh gains the full compile map plus a skip-when-absent
guard, so one script drives any feature slice: it compiles the shaders present
in the tree and cleanly skips the rest. build.rs derives the shader registry
from the compiled output, so a skipped shader is simply absent from the registry
and its dispatch path is never reached in that slice. This lets each following
model PR add its own .comp files and have them compile with no script edit.
Because a skipped shader is not an error, two guards keep it from going
unnoticed:
- `quant_matvec_kernels_are_registered` names every kernel this PR ships and
asserts each survives into the runtime shader map. A count/self-consistency
check cannot catch a loss (registry and map shrink together and stay 1:1),
so the test names kernels rather than counting them: adding a shader never
breaks it, losing one always does. Verified by hiding a .comp and
confirming the failure names the exact missing kernel.
- `required_kernel_spirv_is_wellformed` checks each kernel's SPIR-V is
word-aligned, non-trivial, and carries the correct magic number, so a
truncated or empty module cannot reach the registry.
- compile_shaders.sh now counts skips and reports them in its summary line,
so an accidentally-missing source is visible in the build log.
Gates: cargo build --no-default-features --features multiple-pymethods (117
SPIR-V compiled); cargo test on the same slice 119 passed / 0 failed.
kdunn926
force-pushed
the
upstream/quant-matvec-shaders
branch
from
August 21, 2026 15:33
1ceb8d3 to
c82cadd
Compare
kdunn926
changed the title
… kernels
The registry tests prove the 23 shipped kernels are present and that their
SPIR-V is well-formed. That is a packaging check — it says nothing about
whether a kernel computes the right numbers. This adds a harness that runs
the real compiled SPIR-V on whatever Vulkan device the machine exposes and
compares each dispatch against the CPU quantizer/dequantizer reference that
already lives in `model` (quantize_mlx_affine_4bit / dequantize_mlx_affine /
dequantize_mlx_affine_bits / dequantize_nvfp4 / dequantize_fp8 /
quantize_q8_0), so the packing is never re-invented here.
20 kernels are covered: the mlx4 affine family (scalar, _cols, w8, w16,
w8sg, repack), mlx2/mlx6/mlx8, the two batched MoE twins (expert on grid.y
through meta[]), nvfp4 (folded-scale and raw-e4m3, plain and repack), fp8
(plain, fast, repack) and the column-batched q8_0/f16 kernels. The three
remaining registry entries (both paged_attn_decode_*_sg and relu2_f32) are
not matvecs and carry a written excuse; a structural test fails if a new
kernel ever lands with neither a plan nor an excuse.
CI shape:
* no ICD -> the tests no-op and PASS, so the default job stays green;
* a kernel the PipelineCache declined to compile (wave64-baked shaders on
a subgroup!=64 device) is reported SKIPPED, not failed;
* f16-typed kernels are skipped when the device lacks 16-bit storage +
shader_float16.
Comparison is a relative forward-error bound (max |gpu-ref| normalized by
that row's sum of absolute products), not bit-equality: dequant is exact on
both sides but the kernels split k across threads and recombine via an LDS
tree or subgroupAdd, several folding the affine term with fma. TOL=1e-5 sits
just above the k*eps worst case for k=128.
`harness_detects_a_corrupted_reference` runs a genuine dispatch against a
deliberately perturbed reference and asserts the comparator rejects it, so
this file cannot decay into a test that only knows how to pass.
Validated on Apple M2 Max via MoltenVK (subgroup_size=32, fp16=true): all
20 kernels ran, 0 skipped, worst observed relative error 9.5e-8.
Descriptor-binding and push-constant mismatches between a .comp and its
Rust dispatch site are silent host-side and only surface on device, as
garbage output or a hang -- and CI has no GPU. This adds a small
hand-rolled SPIR-V word-stream parser (no new dependency) and asserts the
invariants the dispatch path actually relies on, statically:
* every registered module is a GLCompute shader with a `main` entry
point and SPIR-V 1.6 (pipeline.rs hardcodes c"main"; device.rs asks
for API_VERSION_1_3, compile_shaders.sh targets vulkan1.3);
* descriptor bindings live in set 0, below MAX_BINDINGS, and are
contiguous from 0 -- compute.rs::record_to maps buffers[i] to
dst_binding(i) positionally, so a hole misroutes every buffer after
it;
* the declared push-constant block fits the 128-byte PushConstantRange
pipeline.rs declares;
* for each kernel in REQUIRED_QUANT_KERNELS, the reflected block size
equals the byte length of the push_constants.rs helper that feeds it,
and the binding count equals the buffer count the dispatch site
passes.
Two documented exception lists keep the gates honest instead of loose,
and are asserted to be EXACT so a new offender fails:
* rms_norm_mul_rope_f32_{f32,f16} skip binding 2 by design
(rms_norm.comp: the rms_norm->rope handoff goes through shared
memory). Never dispatched from src/.
* add_rms_f32_f32_{f32,f16} (212 B) and rms_norm_mul_rope_f32_{f32,f16}
(228 B) declare push-constant blocks LARGER than the 128-byte range
pipeline.rs declares, i.e. larger than the Vulkan guaranteed-minimum
maxPushConstantsSize this fork assumes. Upstream ggml sizes its range
from the device limit. Also never dispatched.
One live divergence is recorded by the contract table:
paged_attn_decode_f16_sg's push block ends at `window_start` and has no
`ring_capacity` word, unlike its f32 twin, while both are fed by
sdpa_pc's 11 words -- so the f16-KV decode path silently ignores ring
addressing.
Fail-injection tests byte-corrupt real modules (binding hole,
out-of-range binding, bad magic, wrong SPIR-V version, truncation) and
feed deliberately wrong contracts through the same check the production
test runs, proving the gates can fail with a useful message.
REQUIRED_QUANT_KERNELS becomes pub(crate) so the new module reuses it
rather than duplicating the list.
…PIR-V
Adds two gates to the foundation job:
- lavapipe (Mesa's CPU Vulkan driver) is installed and selected for a second
test run, so the kernel conformance tests stop no-op'ing and actually
dispatch the compiled SPIR-V, checking each quantized matvec against a CPU
reference. The existing device-less run is kept: it is what proves the
graceful-skip path still works when no ICD is present.
- spirv-val checks every compiled module against the SPIR-V spec. This
complements the Rust reflection gates rather than duplicating them:
spirv-val validates the module in isolation, while reflection validates it
against the Rust dispatch contract (descriptor bindings, push-constant
block sizes), which spirv-val does not inspect.
The spirv-val --target-env must track scripts/compile_shaders.sh's --target-env;
the reflection tests assert the SPIR-V version from the other side, so a change
to one fails the other.
test_block_table_whole_batch_int64_conversion_is_faster_than_per_row failed on the macos-15 runner at 0.89x. The assertion is a bare wall-clock comparison (`assert whole_batch_elapsed < per_row_elapsed`) with no tolerance, so it demands one timing beat another on a shared, noisy runner where a 10% swing is ordinary scheduling jitter. It measures a real optimisation, but it is a benchmark rather than a correctness check, and as a hard gate it will flake indefinitely. Two sibling tests in the same file assert the same way and will flake the same way, so all three are marked together rather than one at a time: - test_cached_available_shaders_matches_uncached_reference_and_is_faster - test_paged_attn_decode_pc_resolved_spec_is_faster_than_relayout_lookup - test_block_table_whole_batch_int64_conversion_is_faster_than_per_row The `slow` marker is already declared in pyproject.toml as "opt-in performance or long-running tests", which is exactly what these are, and scripts/test.sh already runs `pytest -m "not slow"` -- so they drop out of CI while staying runnable on demand with `pytest -m slow`. The correctness half of each test (that the fast path matches the reference) is asserted before the timing comparison and is unaffected.
kdunn926
force-pushed
the
upstream/quant-matvec-shaders
branch
2 times, most recently
from
August 21, 2026 16:48
69516e5 to
8999f4c
Compare
kdunn926
changed the title
ericcurtin
approved these changes
Aug 22, 2026
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
3 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.
Second PR of the upstreaming stack. Adds the quantized matvec kernels that every model path dispatches, plus the two subgroup paged-attention variants:
scripts/compile_shaders.sh gains the full compile map plus a skip-when-absent guard, so one script drives any feature slice: it compiles the shaders present in the tree and cleanly skips the rest. build.rs derives the shader registry from the compiled output, so a skipped shader is simply absent from the registry and its dispatch path is never reached in that slice. This lets each following model PR add its own .comp files and have them compile with no script edit. Gate: cargo build --no-default-features --features multiple-pymethods, 117 SPIR-V shaders compiled.