Beating FP16 with 4-bit Weights: A Portable W4A16 GEMM in Triton

Code: github.com/bassrehab/triton-kernels · Kernel: huggingface.co/kernels/bassrehab/w4a16

Last time I fused an entire Mixture-of-Experts forward pass into Triton and it beat Megablocks. That kernel went after the routing. This one goes after the other half of the memory-bound inference problem: the weights.

LLM inference is bottlenecked on weight loading. A 7B model in FP16 is 14GB that you drag across the memory bus on every forward pass. Quantize those weights to 4 bits and you move a quarter of the bytes. That is the entire pitch behind GPTQ and AWQ, and it is why 4-bit weight-only quantization is how most people actually deploy LLMs in 2026.

There is a catch. The fast 4-bit kernels, Marlin and AWQ and exllama, are all CUDA. If you run on AMD, or you just do not want to bet your inference stack on a single vendor, you are stuck.

So I wrote a W4A16 GEMM in pure Triton. FP16 activations, 4-bit weights, dequantized inside the kernel. It runs on NVIDIA and AMD with no CUDA and no vendor intrinsics.

Here is the honest version of the result. My first correct kernel was three times slower than a plain FP16 matmul. Portable, numerically perfect, and completely useless. Getting it to actually beat FP16 took two rewrites and taught me more about memory bandwidth than the MoE kernel did. The final version beats cuBLAS FP16 by 1.1 to 1.3x in the decode regime, which is exactly the regime that matters when you generate tokens one at a time.

Why 4-bit weight-only is the right lever

When you decode text, you produce one token at a time. Batch size is tiny, often one. At that size the GPU has almost nothing to compute, so the whole forward pass is a race to read weights out of HBM. The math is nearly free. The memory is everything.

This is the memory-bound regime, and it is where quantization pays off. Four-bit weights are a quarter the size of FP16, so you move a quarter of the bytes. If the kernel were perfectly bandwidth-bound, that would be a clean 4x.

One number kept me honest the entire time. At batch size 1, cuBLAS FP16 already hits about 71% of the A100’s peak memory bandwidth. That is not a lazy baseline. To beat it you have to move less data and keep the pipes just as full. Moving less data is trivial. Keeping the pipes full while unpacking 4-bit weights on the fly is the hard part, and it is where all the difficulty lives.

The design

The weights are quantized to 4 bits with one scale and one zero-point per group of 128 rows, GPTQ and AWQ style. Eight 4-bit values pack into a single int32 along the contraction dimension. Dequantization happens inside the GEMM loop, so the weights cross the memory bus at 4 bits and only expand to FP16 once they are already on chip, in registers.

The inner loop is the whole kernel:

# Load one packed int32 (8 weights), unpack all 8 nibbles in registers,
# dequantize with the group's scale, then matmul. No 4-bit weight is read twice.
pk = tl.load(b_ptrs)                                    # (BLOCK_K // 8, BLOCK_N) int32
w  = (pk[:, None, :] >> shifts[None, :, None]) & 0xF    # (BLOCK_K // 8, 8, BLOCK_N)
w  = tl.reshape(w, (BLOCK_K, BLOCK_N)).to(tl.float16)   # row 8*i + j maps to K = 8*i + j
w  = (w - zero) * scale                                 # dequant, still in registers
acc += tl.dot(x, w, out_dtype=tl.float32)               # FP16 tensor cores, FP32 accumulate

That “no 4-bit weight is read twice” comment is doing a lot of work. It is the single change that took the kernel from useless to competitive, and I will come back to why in the war-stories section.

The other half of the design is a split-K path for decode. A skinny matmul, one or a few rows tall, launches very few thread blocks, so most of the GPU sits idle and you never saturate HBM. Splitting the contraction dimension across many blocks, each computing a partial sum and accumulating into the output, brings the parallelism back. The kernel dispatches to split-K for small batches and to the plain tiled kernel once the batch is large enough to fill the machine on its own.

Results

Measured on an A100-SXM4-80GB against torch FP16, which is cuBLAS underneath. Group size 128, on current-generation models. Speedup is W4A16 over FP16 at decode batch sizes.

The kernel beats FP16 across the decode regime on both the dense FFN and the MoE-expert projections, and it does it while using a quarter of the weight memory. The speedups track each other closely across models from 2024 to 2026, because a weight-only GEMM is set by its matrix dimensions and those have barely moved between generations. Once the batch grows past a handful of tokens, into prefill, the unpack overhead stops hiding behind the weight traffic and cuBLAS pulls back ahead. That crossover is fine. Prefill runs once per request; decode runs once per generated token. Decode is where the bill is.

The roofline also shows what I did not manage. Peak achieved bandwidth on the weight stream tops out around 24% of the A100’s HBM. I am collecting maybe a third of the 4x the format promises; most of it is still on the table. More on that in what’s next.

Things I got wrong along the way

The 8x redundant load. My first kernel indexed the packed weights by offs_k // 8, which meant eight consecutive rows of the contraction dimension all loaded the same int32. Read the same word eight times and you have moved eight times the packed data. That is more traffic than FP16, not less. The kernel was correct and roughly three times slower than a plain matmul, achieving about 4% of peak bandwidth. The fix is the loop above: load each int32 once, unpack all eight nibbles in registers. Obvious in hindsight, and it took me an afternoon and a bandwidth counter to see it.

Parity is not a win. After the load-once fix the kernel reached parity with cuBLAS on the large matmuls and stayed slower on the small ones. Correct, portable, and exactly as fast as the thing it was supposed to beat. The bandwidth counter said 19% of peak, which meant the kernel was no longer memory-bound at all. It was overhead-bound, spending its time on the unpack and the register shuffle rather than on memory. Beating FP16 needed a different lever entirely, and that lever was parallelism, not bandwidth.

Split-K only helps when the matmul is skinny. Adding split-K got me past FP16 at batch 1. It also made the kernel slower at batch 32, because the extra blocks each write a partial sum through an FP32 atomic add, and once there is enough work to fill the GPU that atomic traffic costs more than the parallelism buys. So split-K is a dispatch decision, not a default.

The best split factor is not monotonic in the shape. I assumed a bigger contraction dimension wanted more splits. Wrong. The attention projection wanted eight splits, the FFN wanted four, and they have the same contraction dimension. The optimum depends on the full shape in a way I could not capture with a clean rule, so I stopped trying to be clever and handed it to triton.autotune. The subtlety: autotuning a kernel that accumulates through atomic add corrupts its own output during the timing trials, so you have to pass reset_to_zero to re-zero the accumulator before each trial. Miss that and every autotuned result is garbage.

Naming the file after the package. This one is embarrassing. When I packaged the kernel for the Hugging Face Kernel Hub I put the code in w4a16/w4a16.py inside a package also named w4a16. The build’s import check failed with a circular import, because from .w4a16 import ... resolved to the half-initialized package instead of the module. Renamed the file to gemm.py and it built. Twenty minutes of staring at a stack trace that was entirely my fault.

What’s next

The 4x is still out there. Collecting it means a Marlin-style layout: pre-permute the weights offline so they feed the tensor cores in exactly the order the MMA wants them, without the runtime reshape, and pipeline the loads with async copies so the unpack hides under memory latency instead of competing with it. That is a real project, not a tweak, and a Triton port probably lands somewhere short of Marlin’s hand-tuned CUDA. But it should turn the decode win from 1.3x into something a lot bigger.

Also on the list: FP8 activations for a W4A8 variant, and correctness validation on an AMD MI300X to make the cross-platform claim concrete rather than structural. The kernel is pure Triton with no vendor code, so it runs on AMD by construction. I just want to prove it on the hardware.

Code

Everything is in the repo:

• triton_kernels/w4a16.py: the kernel, the split-K path, and the weight-packing utilities.
• reference/w4a16_reference.py: the plain-PyTorch ground truth.
• tests/test_w4a16.py: correctness across shapes, group sizes, and symmetric and asymmetric quantization.
• benchmarks/roofline/w4a16_roofline.py: the roofline plot.
• docs/w4a16.md: the full writeup.

Or load it straight from the Hub, no install:

from kernels import get_kernel

w4a16 = get_kernel("bassrehab/w4a16", version=1, trust_remote_code=True)
packed, scales, zeros = w4a16.quantize_weight_int4_grouped(weight, group_size=128)
y = w4a16.w4a16_gemm(x, packed, scales, zeros, group_size=128)

Takeaways

1. The naive 4-bit kernel is a trap. It is easy to write a correct W4A16 kernel that moves more memory than FP16, not less. Watch the bytes you actually move, not just the output. If your weight traffic is not below FP16’s, you have bought nothing, you are just paying for the unpack.

2. Decode is a parallelism problem, not only a bandwidth problem. At batch 1 the bottleneck is that the GPU is empty, not that the bus is slow. Split-K fills it. That was the counterintuitive part for me: the fix for a memory-bound kernel was more parallelism.

3. Stop hand-tuning shapes. When the optimum is non-monotonic in the shape, a heuristic will fight you forever. Autotune it, and remember reset_to_zero if you accumulate through atomics.

4. Portable and competitive are compatible, but the last 4x is CUDA’s home turf. Pure Triton got me a real decode win over cuBLAS on any vendor. Matching Marlin is a separate, harder fight, and an honest kernel says so.