Skip to content

cuda: fmt=6 (E8/IQ3) decode kernels, so E8 containers can reach VRAM (#452) - #627

Merged
JustVugg merged 1 commit into
JustVugg:devfrom
ZacharyZcR:feat/fmt6-cuda-kernel
Jul 26, 2026
Merged

cuda: fmt=6 (E8/IQ3) decode kernels, so E8 containers can reach VRAM (#452)#627
JustVugg merged 1 commit into
JustVugg:devfrom
ZacharyZcR:feat/fmt6-cuda-kernel

Conversation

@ZacharyZcR

Copy link
Copy Markdown
Contributor

qt_cuda_upload refused fmt=6 outright (colibri.c:300, "no CUDA kernel yet"), so every E8 tensor stayed on the CPU. That inverts the format's purpose: E8 buys βˆ’25% expert bytes to fit more experts in VRAM, but as shipped it pushed all of them onto the CPU tier β€” which is exactly the ~57 ms/token wall the #431 A/B identified as the real bound at hit β‰ˆ100%.

This is the "CUDA sibling kernels" step of the #452 ladder, riding the GroupDesc pipeline #451 opened.

Decode

A super-block is 98 bytes per 256 weights: 64 codebook indices, 8 words of (4Γ—7 sign bits + a 4-bit sub-scale), one fp16 super-scale. Decode is inherently per 32-weight sub-block, so threads stride over sub-blocks rather than elements β€” expanding once per 32 weights instead of redoing the word/parity work for every element.

Two things that would silently produce wrong numbers if done the obvious way:

  • Nothing inside a 98-byte block is aligned. Block b starts at base + b*98, so the 4-byte sign/scale word and the 2-byte fp16 land on arbitrary offsets. Both are assembled byte-wise; dereferencing a uint32_t* there is UB and misreads on the wrong stride.
  • fp16 conversion mirrors e8_fp16_to_f32 rather than calling __half2float, so the CPU and GPU decoders cannot drift on subnormals.

The rotation split

fmt=6 stores W@Q, so activations need Q^T. Where that happens matters a lot:

who rotates cost on GLM dims
gate/up input the caller, once per layer (all routed experts share the row) ~1.4 ms
down input the device, per expert (it is the silu product β€” not shareable) β€”

Doing the gate/up rotation per expert instead would cost ~11 ms against ~1.4 ms (the placement rule in quant.h's comment). So the GPU path mirrors the CPU split at moe() exactly: the caller hands over pre-rotated activations, and coli_cuda_expert_mlp rotates only the silu product before the down matmul.

e8_rot_rows_kernel stages each power-of-two block in shared memory (capping n at 4096 floats, which covers 6144 β†’ 2048+4096 and 1536 β†’ 512+1024) and regenerates the sign stream in-kernel from the same xorshift64* seeded 417+n that e8_signs uses β€” so no rotation data is stored, uploaded, or able to drift.

What moe() was missing

moe() fed x straight into the GPU at three packing sites β€” the overlap issue path, the sync group path, and the group-failure fallback β€” all of which continue past the rotation that the per-expert CPU path does. All three now go through one E8_XE helper that materialises the rotated copy at most once per call. The fallback's CPU recompute was also missing the down-input rotation its sibling path does two hundred lines up.

Two symmetry fixes

Both were unreachable before, because fmt=6 could never be uploaded:

  • coli_cuda_tensor_free charged for a scale buffer fmt=6 never allocates. The over-subtraction tripped the >= guard, so the branch was skipped entirely and the tensor's bytes stayed on the device counter forever β€” this is what the new leak assertion caught.
  • coli_cuda_tensor_update would cudaMemcpy O floats out of a NULL scales pointer on a repin (the scale_count ? : O fallback defeats a zero count).

Portability

cudaMemcpyToSymbol had no HIP mapping β€” added, or the HIP build breaks. The Windows loader resolves the new entry point optionally: an older coli_cuda.dll leaves fmt=6 on the CPU rather than taking the entire GPU backend down over one missing symbol.

Tests

test_backend_cuda.cu gains fmt=6 coverage for both the matmul and the full expert MLP including the device-side rotation, checked against a decoder + FWHT written from the format description and deliberately not shared with quant.h, so a common-mode error cannot hide in both. It uploads a synthetic codebook, which keeps the test self-contained while still exercising the upload path the real table travels. Placed after the existing exact tensor-count assertions, with its own leak check.

Also verified in isolation first, before wiring anything: the decode kernel against a double-precision decode of the same container, since CPU sequential accumulation and GPU tree reduction are both float32 and neither is the oracle. Across 8 shapes the GPU is never materially worse and on long rows is better (I=6144, S=4: 2.97e-06 vs CPU 1.14e-05) β€” tree reduction's O(log n) error against sequential's O(n). The FWHT matches to 1.68e-07, exact at dim=4096.

make check 211 passed. make cuda-test green on 1 and 2 devices (RTX 5090, sm_120, CUDA 13.3).

What is NOT verified

No end-to-end run against a real fmt=6 container. There isn't one to run: the GLM fmt=6 quality gate is still blocked on a genuine fp16 source (#452), and glm_tiny can't be converted since E8 needs I % 256 == 0. So this is verified at the kernel and API level β€” decode semantics, the rotation split, upload/free/update accounting β€” but not yet on a converted model. A GLM or OLMoE e8 container would be the last mile, and I'd rather land the kernels reviewed than sit on them until one exists.

fmt=6 had no GPU kernel, so qt_cuda_upload refused it outright and every
E8 tensor stayed CPU-side. That inverts the format's whole purpose: E8
buys -25% expert bytes to fit MORE experts in VRAM (JustVugg#452), but as shipped
it pushed all of them onto the CPU tier -- the ~57 ms/token wall measured
in JustVugg#431.

Decode is per 32-weight sub-block (a 98-byte super-block holds 64 codebook
indices, 8 words of 4x7 signs + 4-bit sub-scale, and one fp16 scale), so
threads stride over sub-blocks rather than elements and expand once per 32
weights. Nothing inside a 98-byte block is guaranteed aligned, hence the
byte-wise word assembly; the fp16 conversion mirrors e8_fp16_to_f32 rather
than calling __half2float so the two decoders cannot drift on subnormals.

The activation rotation splits exactly as it does on the CPU path: all
routed experts of a layer share the gate/up input, so the CALLER rotates
it once per layer (~1.4 ms vs ~11 ms per-expert on GLM dims), while the
down projection's input is the per-expert silu product and must be rotated
on the device -- e8_rot_rows_kernel, staging each power-of-two block in
shared memory and regenerating the sign stream in-kernel from the same
xorshift64* as quant.h, so no rotation data is stored or uploaded.

moe() fed x straight to the GPU at three packing sites (the overlap issue
path, the sync group path, and the group-failure fallback); all three now
take the rotated copy through one E8_XE helper. The fallback's CPU
recompute was also missing the down-input rotation the main CPU path does.

Two symmetry fixes the format needs, both previously unreachable because
fmt=6 could not be uploaded at all: the tensor accounting in
coli_cuda_tensor_free charged for a scale buffer fmt=6 never allocates,
which tripped the >= guard and leaked the byte count permanently; and
coli_cuda_tensor_update would memcpy O floats out of a NULL scales pointer
on a repin.

The codebook is published from quant.h to each device rather than copied
into the backend, so there is one source of truth. The Windows loader
resolves the entry point optionally -- an older DLL degrades fmt=6 to the
CPU instead of disabling the whole GPU backend.

Tests: fmt=6 matmul and the full expert MLP (including the device-side
rotation) against a decoder written from the format description, not
shared with quant.h, so a common-mode error cannot hide. The synthetic
codebook also exercises the upload path the real table travels.
make check 211; cuda-test green on 1 and 2 devices (RTX 5090, sm_120).

Refs JustVugg#452, JustVugg#431
@JustVugg
JustVugg merged commit 8245766 into JustVugg:dev Jul 26, 2026
10 checks passed
steve-m added a commit to steve-m/colibri that referenced this pull request Jul 27, 2026
…each VRAM

CUDA gained an fmt=6 kernel in JustVugg#627; on Vulkan the format was still refused at
upload, so every E8 tensor stayed CPU-side. That inverts what the container is
for: E8 costs 3.0625 bpw against int3-g64's real 3.5 (3.00 weights + 0.5 scale),
so it exists to fit ~12% MORE experts in the same VRAM β€” but as shipped it put
all of them on the CPU tier instead.

Decode strides 32-weight SUB-blocks rather than elements: that is the
granularity at which the sign/scale word changes, and the two 4-dim grid rows a
lane touches stay in L1 across its four 8-weight lanes. Every field inside a
98-byte super-block is 4-aligned relative to the block START, but 98 % 4 == 2,
so blocks alternate 4- and 2-aligned within the row β€” hence the byte-wise word
assembly, with the straddle shift only ever 0 or 16. The fp16 super-scale is
read 16 bits wide on purpose: its 32-bit window can reach past the last row of
the slab, its 2 bytes never do. The four lanes are expanded and consumed as two
vec4 chains for the same reason dot16_* splits its accumulator.

The rotation splits exactly as on the CPU and CUDA paths. All routed experts of
a layer share the gate/up input, so the CALLER rotates it once per layer; the
down projection's input is the per-expert silu product, produced and consumed on
the GPU, so e8_rot.comp rotates it in place between the two phases of the same
expert-group submit. It runs over the whole hidden buffer at once β€” the
transform is per row, so every expert's slice rides one set of dispatches
instead of paying the per-expert cost. Blocks tile block-diagonally as
e8_rot_rows does, each staged in shared memory.

Only the ROUTED expert projections are ever fmt=6: --ebits/--io-bits are integer
bit widths and only --xbits/--{up,gate,down}-bits take "e8", so attention and
dense tensors cannot reach the format and VK_FMT_OK is left alone. The
converter also refuses any input dim that is not a whole number of 256-weight
blocks, which upload_tensor re-checks β€” so the shader has no ragged tail.

moe() fed x straight to the tier at three sites inside the Vulkan block; all
three now take the rotated copy. The registry serves an expert with NO RAM slot,
so its format is only knowable from the resident tensor β€” hence
coli_vk_tensor_fmt. The block's two CPU fallbacks (the CPU share and the
device-lost recompute) were also missing the down-input rotation the main CPU
path does; unreachable before, since fmt=6 could not be uploaded at all.

shaders/e8_grid.glsl carries a second copy of quant.h's codebook because GLSL
cannot include the C header β€” the same arrangement tools/iq3xxs_grid.json has
for the Python codec. Without shaders/e8_rot.spv, coli_vk_e8_ok() reports 0 and
fmt=6 stays on the CPU rather than decoding against an unrotated input.

Tests: make vk-test checks the shader decode against matmul_e8 and the full
expert MLP against the CPU path, with int4 on one side to isolate the fused
gate_up branch from the rotation. Reported as relative L2, which is the
meaningful metric for these sign-cancelling dot products; fmt=6 lands at
3e-7..5e-7, at or below the int4 baseline's 1e-6. Green on RADV GFX1201
(RX 9070) and llvmpipe β€” both subgroup widths. make check 220, and the existing
VK_TEST harness is unchanged for fmt 1/2/4/5 on both.

Refs JustVugg#452, JustVugg#627
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants