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

29 Jul, 2026

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 60 to 125 tokens, a full paragraph about a thousand. Mamba2's chunk_size is 256. So a normal Zonos prompt is a single partial chunk. The whole reason the chunked scan exists, splitting a long sequence so chunks run in parallel and only the cross-chunk state stays sequential, never engages at Zonos's real prompt lengths.

That closed the library framing's hook into Zonos and reframed the phase. If the kernel barely runs in the model I care about, the point is not a speedup. The point is writing a four-stage tensor-core algorithm from scratch in a tile DSL, validating it hard, and characterizing it honestly. So I benchmark on synthetic long sequences to exercise the path, and hold to a plain bar: working, validated, characterized, written up. Not "beat Tri Dao." Keep that in mind, because the kernels lose, and under this bar that is a finding.

Read the reference, then rewrite it

Mamba2's prefill scan, SSD for state space duality, ships in mamba_ssm as four Triton kernels. I read Tri Dao's source before writing anything, the same discipline as the roofline above, and turned it into notes (docs/ssd_things.md). The four stages:

• bmm: a batched matmul building the within-chunk, attention-like score matrix. A pure tiled GEMM.
• chunk_state: each chunk's end-state, a decay-weighted matmul of that chunk's inputs.
• state_passing: the one sequential stage, a scalar-decay recurrence threading state across chunks. No matmul.
• chunk_scan: the finale, two matmuls into one accumulator (the cross-chunk state path and the intra-chunk attention path), plus a 2D decay, a causal mask, and the D skip and z gate.

I wrote them in order of increasing pain. bmm is the Rosetta Stone, one TiledMMA and nothing else. chunk_state is bmm plus one line, a per-position decay on the operand before the dot. chunk_scan is those two stacked plus the genuinely new work, a decay that depends on both the query and key position and a causal triangle. state_passing has no matmul, and I did it last because after chunk_scan it was a rest.

The muscle that mattered most was not the matmul, it was the edge masking. Real Zonos prefill is a single partial chunk, so the kernels have to handle 100 valid rows out of 256 without reading the 156 rows of garbage past the end. Every kernel got a clean test and a "ragged" partial-chunk test, and the ragged one earned its keep: the first bmm passed every clean-divide shape and then silently emitted NaNs on the unpadded partial chunk, which is exactly the shape the model runs. Green on the easy case is necessary and not sufficient. The partial chunk is the case that ships.

I chained the four into one op, combined_scan, that mirrors mamba_chunk_scan_combined, reusing Triton's cumsum for the preprocessing since a cumsum is not one of the four and not the point. Then I checked it three ways. Against the reference: rel L2 of 3.3e-3 on the output, and on the final SSM state that seeds the decode cache my check reports zero relative error (0.0e+00), at the real single-partial-chunk shape, not just the clean ones. In the model: I monkeypatched it into Zonos (scripts/zonos_prefill_eartest.py), it ran in all 34 Mamba2 prefill layers with zero fallbacks, and produced clean, coherent speech. And the token check, for the second time in this post: a 98% patched-versus-stock flip rate that, again, means nothing. My kernel only touches prefill, decode uses the unpatched step path, the seeded state carries no measurable error, so the only difference is a 3e-3 perturbation in the first frame's logits, and that flips the first sampled token and the rollout compounds from there. Same brittleness as the GEMV's 4.6%, 20x larger only because the whole rollout hangs off one perturbed frame. The audio is the judge, and the audio is fine.

Where Triton is hard to beat

Now the title. Benchmarked against the Triton it reimplements, on synthetic sequences long enough to exercise the chunked path (kernels/bench_combined.py):

Two things in that table are the finding. First, I lose everywhere, 1.6x to 2.7x. Second, the gap widens with length, and that is the mechanism, not noise. Triton is nearly flat, 0.72 to 0.89 ms as the work grows 32x, while mine climbs from 1.15 to 2.43. Triton barely pays for the extra chunks; I pay full price for each one.

That flat-versus-climbing curve is where Triton is hard to beat, and I read it as software pipelining, with the same caution I gave the base gap. Triton's kernels are staged to overlap the next chunk's HBM loads with the current chunk's compute; mine are single-buffered, loads then compute in sequence. Exposed per-chunk load latency that pipelining hides is exactly the cost that stays flat for Triton and grows linearly for me, so that is the mechanism I believe. But a curve is consistent with more than one story: more grid waves and a longer sequential state-passing pass also scale with chunk count, and I did not isolate it with NCU. Leading suspect, strongly implied by the shape, not a proven cause. What is certain is the input, Triton is pipelined and mine is not.

The base gap, the 0.63x you see even at a single chunk before pipelining is in play, is the other half, and I will name the suspect rather than overclaim it. My decay, scale, mask, and gate are per-output-element loops that call exp one element at a time, the scalar glue between the matmuls; Triton fuses those into vectorized passes. That per-element transcendental work is the likely constant-factor cost. I did not run the full NCU decomposition to split it precisely, exp recompute versus scheduling versus my single fixed MMA config with no tile autotuning, and that is the honest edge of what I measured. The pipelining half, the curve proves; the rest is a named suspect, not a number.

And end to end it costs nothing. Per call at the real prefill shape I am 2.2x slower, 1.65 ms against 0.73, but wired into a 200-token generation that is +1.5% wall (2365 ms to 2400 ms), because prefill is about 1% of a generation. The only real cost is a 1.7 second CuTeDSL compile on the first call, a one-request TTFA hit that ahead-of-time compilation removes and that has nothing to do with the kernels. Prefill was never a Zonos bottleneck, the same way the decode scan was not. The kernel-writing was the deliverable, not a Zonos win.

Where a hand kernel is worth writing

Put the two kernels side by side and they answer the same question from opposite ends. The GEMV beat cuBLAS by 1.07 to 1.36x because cuBLAS was solving a more general problem than mine, padding an M=2 matmul out to a GEMM tile and leaving bandwidth on the floor at exactly the shape I cared about. The scan lost to Triton by 2x because Triton is not solving a more general problem, it is Tri Dao's kernel solving this exact one, already autotuned and software-pipelined. The hand kernel wins where the library is forced to generalize and loses where the library is already specialized. Knowing which situation you are in, before you spend three weeks, is the whole skill.

The series has a shape now: I found the obvious decode kernel was 7% and already at the roofline and did not write it, found the shape cuBLAS handles worst and wrote a GEMV that beat it, then took on the hard four-stage scan, got it correct against the reference and coherent in the model, and measured exactly where it loses and why. None of it required beating Tri Dao, and the value was never going to be that I did. A team that writes serving kernels wants someone who knows when not to write one, can write the one that pays, and can write the hard one and then tell you without flinching that it is 2x off the reference because it lacks software pipelining, not because it is wrong. And underneath both kernels, the same recurring lesson: a deep autoregressive model is pinned to its reference's exact rounding, so a drop-in accelerator, faster and even more accurate, still moves tokens, and you gate on the ear, not on matching bit for bit.

Tile by tile. One faster, one slower, both sure of why.

The code, the plan, and every profiling doc referenced above live in the repo: github.com/Vishal-Padia/Sonata.

As always, happy to chat if anything here is unclear or wrong. Just ping me on Twitter.