GPU compute kernels for the Cube Memory FFN-replacement layer,
written in Rust and compiled to SPIR-V via
rust-gpu. Designed to be
loaded by llama.cpp's Vulkan backend (or any Vulkan host) at the
Phase 2 integration stage described in
../phase1/PLAN.md.
GLSL is the default in the llama.cpp Vulkan backend. We're using rust-gpu for new Cube Memory kernels because:
- Type safety on push-constant and buffer layouts catches mismatches at compile time instead of "wrong output, no error" at runtime.
- The bind/unbind/superpose primitives have clean Rust expressions; GLSL with manual complex math is harder to keep correct.
cargo testcan run the same kernels on the CPU via Rust's built-in execution path — useful for unit-testing FHRR algebra.
The host-side Vulkan plumbing (descriptor sets, pipeline creation,
dispatch) stays in C++ where llama.cpp lives. This crate produces
SPIR-V binaries that the C++ side loads via vkCreateShaderModule.
Requires a specific Rust nightly toolchain pinned in
rust-toolchain.toml. The pin tracks Rust-GPU/rust-gpu's HEAD
toolchain. Install once:
rustup install nightly-2026-04-11 \
--component rust-src --component rustc-dev --component llvm-toolsSystem dep (saves ~3 minutes on first build vs the bundled
use-compiled-tools feature):
sudo pacman -S spirv-tools # CachyOS / ArchThen:
cd ~/projects/cube-memory/shaders
cargo run -p cube-memory-shader-builder --releaseOn success the path to the produced SPIR-V file is printed. Actual
location:
target/spirv-builder/spirv-unknown-vulkan1.2/release/deps/cube_memory_shader.spv.
If the build fails with a panic inside the SPIR-V codegen backend,
check that the toolchain in rust-toolchain.toml matches the one
expected by the spirv-builder crate version pinned in Cargo.toml.
Mismatches are the most common cause of obscure failures.
The builder is configured with SpirvMetadata::None for ~5 KB output.
With SpirvMetadata::Full the produced binary embeds the full Rust
source and OpString debug info, ~50× the size (~330 KB). When you
need to map a SPIR-V validation error or driver crash back to a
source line, temporarily flip to Full in
cube-memory-shader-builder/src/main.rs and rebuild — spirv-dis
output will then include the source spans.
Element-wise complex multiplication. Two input phasor vectors of length N, one output. Workgroup size 64.
Dispatch:
groupCountX = ceil(N / 64)
groupCountY = 1
groupCountZ = 1
Buffers (descriptor_set=0):
binding 0: in_a — readonly [Vec2; N]
binding 1: in_b — readonly [Vec2; N]
binding 2: out — readwrite [Vec2; N]
Push constants:
struct { uint32_t n; }
In rough order of when each kernel is needed:
fhrr_bind— done (this skeleton)fhrr_unbind— same shape, conjugate the second operandfhrr_unitize— element-wise normalize to unit modulusfhrr_superpose— sum + unitize (the bundle operation)cube_memory_cleanup— argmax cosine match against a frozen codebook of M phasors, returning the snapped phasor + indexcube_memory_retrieve— top-k slot-key dot product + softmax-weighted slot-value gather. The full forward path.
The early kernels (1–4) validate the toolchain and let us write
unit tests against the Phase 0 CPU implementation
(../phase0/fhrr.py). Numerical parity between rust-gpu output
and the PyTorch reference is the gate before plugging into the
real ggml-vulkan integration.
(Phase 2 work, not done yet.) The expected path:
- Add a new ggml op
GGML_OP_CUBE_MEMORYinggml/include/ggml.h. - In
ggml/src/ggml-vulkan/ggml-vulkan.cpp, register a pipeline that loads the SPIR-V produced here. Useggml_vk_create_pipelinewith the bytes loaded fromcube-memory-shader.spvat startup, or vendored at build time by a small CMake step that runscargo run -p cube-memory-shader-builderbefore configure. - The op's
compute_forwardwalks src tensors and dispatches.
Two practical wrinkles:
- ggml's tensor types must round-trip the slot-store and codebook;
we'll likely need a new
GGML_TYPE_FHRR(complex64) or store phasors as (re, im) pairs inGGML_TYPE_F16and reinterpret. - The shader output dimensions need to match what the host computes for workgroup count; mismatches are silent and produce wrong output, not crashes.
MIT, matching the rest of the project.