vkML 0.1.0

The Vulkan backend

How a graph node becomes a dispatch, and the device limits that shape every decision along the way.

Three required features, and what each buys

A device missing any of these cannot run vkML at all, and vulkan_unavailable_reason() names the first one found:

FeatureWhat depends on it
bufferDeviceAddress Every buffer in every shader. vkML binds no descriptor sets at all — all 24 shaders take their operands as uint64_t addresses inside the push-constant block.
scalarBlockLayout Lets the push block use C-like scalar packing, so the GLSL struct and the C++ struct agree without padding rules diverging.
timelineSemaphore Host-side completion. One monotonically increasing counter per stream — waiting for value N means "everything up to submit N has finished", which replaces fences entirely: no fence pool, no reset, no per-submit object.
ⓘ Note

Zero descriptor sets is unusual enough to be worth stating plainly, and it explains why the push-constant budget is the recurring constraint in this backend: every operand costs 8 bytes of the 128 Vulkan guarantees before any shape metadata is packed alongside it.

The push-constant budget

maxPushConstantsSize is guaranteed to be only 128 bytes. The development GPU reports 256, which is exactly why this went wrong once: three blocks exceeded the guarantee, producing 19 failing tests on an AMD Windows driver reporting 128 and nothing at all locally.

Every block now fits, and the constraint is enforced from both sides:

Two of those assertions were missing until recently, on the two smallest blocks — which is exactly where a missing guard hides, since they are the least likely to grow.

The largest blocks today are binary at 124 B, softmax at 120 B and where at 116 B. Fitting them took per-operator work rather than one general fix: where and softmax store shared extents once instead of per operand, and cat derives its operands' extents from the output's, which brought it from 144 down to 112.

Workgroup width adapts to the device

No shader declares a literal workgroup size. common.glsl declares layout(local_size_x_id = 0) in; once and all 24 inherit it, so the width is a specialisation constant resolved at pipeline creation.

That is the mechanism behind the clamp: the general width is min(256, maxComputeWorkGroupInvocations). Vulkan guarantees only 128, and this asked for 256 unconditionally — so a conformant minimum-spec device could create almost no pipeline at all. Devices allowing 256 keep it and are unaffected.

Varying it is safe because the width reaches the shader as a constant and the grid is derived from it, so every kernel on that path adapts, and shared-memory requests scale with it — a narrower workgroup asks for less, never more.

⚠ Warning

Clamping alone would not have been enough. Both blocked GEMM kernels hardcode 256 invocations — deliberately, so the three variants stay comparable — so a floor device would have had every element-wise, reduction and movement operator and still no matmul: the appearance of support without the ability to train. matmul therefore falls back to the naive kernel, which takes the clamped width and fits by construction.

The dispatch grid, and the 65535 ceiling

maxComputeWorkGroupCount[x] is guaranteed to be only 65535. The development GPU reports 2³²−1, so every elementwise operation above 64 MiB failed on a driver reporting the floor and nothing showed locally.

choose_dispatch_grid folds an oversized grid into the second dimension. It takes the limits as parameters rather than reading them from a device, which is what makes the guaranteed floor testable on any machine — the test for it compiles unconditionally, including in the three CPU-only CI jobs.

The reduction path hits that ceiling far sooner than the element-wise one, because it dispatches one workgroup per output row: 32 images of 3×64×64 average-pooled is 98,304 rows against a guaranteed 65,535, an ordinary batch — where the element-wise path needs 16.7 million elements to get there.

Synchronisation: two primitives, both the simplest correct choice

A global memory barrier between dependent dispatches — a single vkCmdPipelineBarrier with shaderRead|shaderWrite on both sides, not a per-buffer barrier. Tracking per-buffer hazards costs more CPU time than the barrier costs GPU time, for a graph where almost every node depends on its predecessor anyway. It is conservative, ordering more than strictly necessary, which is the right default when correctness comes first.

The barrier is not optional: without it a dispatch may read a buffer another is still writing, because the GPU does not serialise dispatches on its own. Making it selective requires a planner that knows which nodes alias, and that is the one place skipping it will be justified.

A timeline semaphore for host completion, as described above.

The recorder makes no decisions

Recorder records and submits. It does not choose an order, does not allocate, and does not decide which kernel runs — those belong to the executor above it. Keeping the split means a future lowered execution graph can drive the same recorder unchanged, simply calling dispatch in a different order.

Pipelines and specialisation

KernelConfig exists so the runtime never hardcodes a workgroup or subgroup decision. llama.cpp demonstrates why that matters on this exact GPU: its RDNA1 table pins wave64 for softmax, argmax and matrix-vector work and wave32 everywhere else, because the right width differs per kernel and the device's default — 64 here — is not the right answer for most of them.

A subgroup size of 0 means "let the driver choose"; anything else creates the pipeline with VkPipelineShaderStageRequiredSubgroupSizeCreateInfo, which needs subgroupSizeControl. Shared-memory requests are validated against maxComputeSharedMemorySize before creation rather than failing at dispatch.

Autotuning is explicitly not implemented — but when it is, it becomes a search over these fields rather than a redesign, because nothing above reads a hardcoded constant.

Testing against the floor

VKML_MIN_SPEC=1 makes any device report the Vulkan 1.3 Required Limits. It only ever reports limits smaller than the hardware has, so it can make vkML more conservative and never less — which is what makes it a testing facility rather than a tuning knob.

It exists because the alternative was hand-editing vk_device.cpp, which CI cannot run and a contributor cannot reproduce. Run it before claiming a limit is satisfied: it is the cheapest way to find the next instance of this project's most common bug.

vkML — Vulkan-first machine learning in C++20. Apache-2.0. Signatures on this page are generated from the installed module.