The kernels that were there: where a hand kernel beats the library, and where it can't

Sequel to "The kernel that wasn't there". Last time I profiled my way out of writing a kernel and shipped streaming instead. This time I wrote them: a skinny bf16 GEMV that beats cuBLAS by up to 1.36x, and a four-stage SSD scan that loses to Triton by about 2x. The two outcomes together are the point.

Last post ended on a promise. Profiling a Mamba2 decode step told me the scan everyone points at is 7% of the layer and already at the memory roofline, so I did not write the fused decode kernel I set out to write. The 85% that actually costs is the two projection matmuls, memory-bound at batch 1 to 2 on reading the weight bytes, where cuBLAS runs a GEMM kernel that pads the tiny M dimension out to a tensor-core tile and wastes most of the multiply. And I flagged prefill as the other place kernels might matter: the tensor-core-bound chunked scan, a genuinely different regime.

So this post is two kernels in those two regimes. The decode projections, where a hand-written GEMV wins because cuBLAS is solving a more general problem than mine. And the prefill scan, where a hand-written CuTeDSL port loses because the Triton it copies is already specialized and pipelined for exactly this. Every number below is on an A10G (sm86), same as before.

The GEMV: measure the roofline before you believe the headroom

The first thing I did was almost talk myself out of it. I had been quoting headroom against the A10G's 600 GB/s spec. But the number that matters is the bandwidth you can actually hit, and since this kernel streams the weight and writes almost nothing, the right ceiling is peak read bandwidth. Measured, that is 520 GB/s. Against it, cuBLAS is already at ~95% on the big-N projections. There is nothing to take there. The real headroom is out_proj (cuBLAS well below the read roof) and the cfg-doubled M=2 path. So the win was smaller and more specific than the spec math implied, and worth exactly the shape cuBLAS handles worst.

I wrote it in CuTeDSL. The design follows from "this is load-bound, so do not use tensor cores": parallelize over the output columns, one CTA per column tile, each warp strides the K dimension so adjacent threads read adjacent weights (coalesced), accumulate in fp32, store bf16. Two specializations, because M is only ever 1 (pure streaming) or 2 (cfg): the M=2 kernel loads each weight once and feeds both rows' accumulators, so classifier-free guidance does not double the weight traffic.

The first version was correct but sat at 40% of roof. NCU said why in one line: memory-bound as designed, occupancy fine at 82%, compute idle, but DRAM only at 66%. The scalar two-byte loads were not driving the memory system hard enough. So I vectorized the loads to 128-bit, 8 bf16 per thread:

Now at the memory limit. Graph-timed against cuBLAS on the M=2 projection shapes:

A hand-written GEMV beating cuBLAS is not magic. cuBLAS is optimized for real GEMMs, not for M=2, and it leaves bandwidth on the floor exactly where I measured it would. The full write-up with all shapes and NCU is in docs/gemv_kernel.md.

The bench that lied, and the integration it forced

Before any of those numbers were real, the bench told me the kernel ran in 2.5 microseconds. That is 27 TB/s, roughly 50x the memory the chip has. Physically impossible, which is the useful kind of wrong.

The decode loop runs under a CUDA graph, so I was timing under torch.cuda.graph. The graph captured, replayed, and reported 2.5 us because it had recorded nothing. Correctness still "passed" only because the warmup calls before capture had left the right answer in the output buffer. CuTeDSL, by default, launches on stream 0 and then synchronizes; torch was capturing a different stream, so the kernel launched onto a stream that was never recorded, and the graph replayed empty.

The fix was to thread the launch stream through the launcher and pass torch's current stream, which is the capture stream inside the graph. Two lines. But it is not a benchmarking detail, it is the exact prerequisite for putting the kernel into Zonos's decode graph at all. If the launch cannot be captured, it cannot ship. The empty graph was the integration test arriving early.

Integration is a monkeypatch that swaps the decode-shaped projection Linears for the kernel and falls back to cuBLAS for prefill (where M is large and cuBLAS is right). Stock versus fully patched, on Zonos's own CUDA-graph decode path:

Minus 12.9% on the decode step, captured correctly inside Zonos's graph (I check the output has hundreds of distinct codes, not the handful an empty-graph replay would produce), on the step my profiling had called mostly irreducible. That is the result I set out to get. Then two things I did not set out to find.

Two surprises: the ear, and TTFA

I expected the swap to be numerically invisible. It is the same math as cuBLAS at bf16, so my teacher-forced probe should have shown near-zero token flips. It showed 4.6% of the coarse-codebook tokens flipping, as many as int8 quantization flipped in the last post.

The reflex is "the kernel is buggy." It is not. A direct per-op check says my kernel is at least as accurate as cuBLAS against fp32, and on out_proj slightly more accurate; the cute-versus-cuBLAS difference is 0.01 to 0.28%. That tiny rounding difference, run through a 34-layer autoregressive stack, amplifies into 4.6% different token choices. The model is effectively pinned to cuBLAS's exact rounding, and any deviation, even a more accurate one, sends the rollout down a different valid path. So the flip-rate gate that was right for quantization (a real degradation) gives a misleading RED here. The honest quality gate is not token-match, it is the ear: I generated audio and listened. Clean, no artifacts, just a slightly different reading of the same sentence. Hold onto that lesson, because it comes back at 20x the magnitude in part two.

The second surprise: I re-ran the streaming sweep with the kernel, and steady-state per-frame latency dropped about 13%, which pushes sustained streaming across RTF 1.0. A clean throughput win. And time-to-first-audio got worse by about 40 ms, consistently, at every chunk size. TTFA is the headline metric for a voice agent, so this matters. The cause is graph capture: Zonos captures the decode graph once per generate call, and the patched capture records 116 of my kernel launches instead of 46 cuBLAS launches. That heavier capture is paid by the first audio inside the call. I tried the obvious fix, pre-binding buffers to kill a per-call host op, and it did nothing and slightly hurt steady-state. The cost is not host-side glue I can reach from Python, it is CuTeDSL's own per-launch capture marshaling. The real fix is to stop re-capturing the graph every request, the way a persistent server captures once and replays across utterances. That is a bigger change than a wrapper tweak, and I am calling it honestly rather than hiding a 40 ms regression behind a 13% headline.

That is the decode kernel: a real win where cuBLAS was weakest, integrated into the live graph, with two lessons a microbenchmark would never have surfaced. Now the other regime.

The scan: the plan said library, the measurement said learning

Prefill is the tensor-core-bound chunked scan, and my plan for it was the ambitious one: port Mamba2's SSD scan to CuTeDSL, package it as a reusable, benchmarked kernel suite, and show it matching or beating Tri Dao's Triton across shapes. Same as every phase in this project, the first measurement rearranged the ambition. Where the plan and the numbers disagree, read it as chronology.

Zonos prefill is short. The scan runs over the conditioning plus the text prompt, and I measured what that is: a sentence is…