Blockwise FP8 on Ada GPUs, part 2: my kernel was correct and 4x too slow

In my last post I showed that block-quantized FP8 models on Ada GPUs (L4, L40S, RTX 4090) don't use the FP8 tensor cores in vLLM. vLLM falls back to Marlin, which does its math in 16-bit. I ended that post with "I'll write up the kernel itself once it has real numbers."

It has numbers now. I opened the kernel as vLLM PR #61224. It's still waiting for review, so this is work in progress, not something that's in vLLM yet.

The short version: the first version was correct on every test and about 4x slower than the Triton kernel vLLM already had. This post is about why, and what fixed it.

The starting point

CUTLASS ships an Ada example for blockwise FP8 (example 94). It uses one activation scale per 128×128 block. vLLM's checkpoints use one activation scale per token per 128 values along K (1×128 groups), with 128×128 blocks for the weights. So I adapted the example to per-token activation scales and wired it into vLLM's CUTLASS path for SM89.

The core of a blockwise kernel is the scaling inside the main loop:

for each K block kb (128 wide):
    partial = Σ A_fp8 · B_fp8              # FP8 tensor cores, fp32 accumulate
    acc    += partial · As[row, kb] · Bs[col/128, kb]

On Oct 6 I ran it on a rented L4. Everything passed. A standalone test against a float reference passed with a relative error around 2e-3, which is normal for FP8. vLLM's own CUTLASS block-FP8 tests, which are skipped on Ada today, passed too.

Then I looked at the speed.

Correct, and slow

For o_proj of Qwen3-8B (N = K = 4096) on the L4:

About 4x slower than Triton, and slower than the Marlin fallback I was trying to replace. At M = 1 it took about 175 µs, where the weights alone (16 MB) should stream in a fraction of that.

The profiler showed one kernel per call and nothing else. No extra copies, no hidden transposes. So the problem was inside the kernel.

Reading the kernel name

The mangled kernel name tells you the exact config that got compiled. Mine was:

• thread block tile 64×128×128, warp tile 64×64×128
• that's 2 warps per thread block
• 3 pipeline stages, about 72 KB of shared memory, so 1 thread block per SM

Two warps per SM is thin. My first guess was "too few warps", but vLLM's existing per-tensor FP8 kernel on SM89 uses the same 64×128×128 tile with 2 warps for large M, and it's fast. So the tile alone wasn't it.

The real clue came from cuobjdump -res-usage:

REG:255 STACK:600

255 registers is the hardware cap per thread. A 600-byte stack frame means the compiler ran out of registers and spilled to local memory, inside the hottest loop of the kernel.

Where the registers went

The warp tile is 64×64 in fp32, which is already 128 accumulator registers per thread. That's normal.

The problem was how the scaling was done. The FP8 MMA instruction on Ada works on 32-wide slices of K, so one 128-wide scale block is 4 slices. My adapted main loop did this for every 32-wide slice:

1. MMA into a fresh fp32 fragment
2. multiply that whole fragment by the scales
3. add it to the main accumulator

So two or three full fp32 fragments were alive at the same time, which is where the spills came from. And the scaling work on the CUDA cores was 4x what it needed to be, because the scales don't change inside a 128-wide block.

The fix

Since the scales are constant across a 128-wide K block, there's no reason to scale each 32-wide slice:

for each K block kb (128 wide):
    block_acc = 0
    for each 32-wide slice in kb:
        block_acc += A_fp8 · B_fp8          # MMA only, no scaling
    acc += block_acc · As[row, kb] · Bs[col/128, kb]   # scale once per block

One temporary fragment per block instead of one per slice, and a quarter of the scaling work. With that, the kernel compiles without spills.

Before spending money on GPU time, I built a small tuning kit: 16 tile configs × 3 main-loop variants, all compiled without a GPU (in Docker) to check register use and spills first, so I knew which ones spilled before paying for the L4.

On Oct 8, all 48 combinations passed the correctness check against a CPU reference. The best config ran the same o_proj at M = 4096 in 1374 µs, down from 6536 µs.

I also found and fixed a latent bug on the way: the original code picked one weight scale per thread block, which is only correct when the N tile is exactly 128 wide. It now looks the scale up per warp.

One tile doesn't fit all M

Decode (M = 1 to 64) and prefill (M in the thousands) want very different tiles. Small M needs short, narrow tiles so enough thread blocks exist to keep the GPU busy. Large M wants big tiles for data reuse. So the kernel picks a tile by M:

The swizzle mattered more than I expected. It changes the order thread blocks are launched in, so neighbouring blocks share data in L2. For down_proj (K = 12288) at M = 4096 it cut the time from 7144 µs to 5614 µs.

Results on the L4

Microbenchmark inside vLLM, Qwen3-8B shapes, µs (median). Triton and this kernel include quantizing the activations to FP8; Marlin doesn't need that step:

End to end, Qwen/Qwen3-8B-FP8 on the same L4, --max-model-len 4096. "main" is precompiled vLLM at the same base commit, which logs "Selected MarlinFP8ScaledMMLinearKernel". With the PR it logs "Selected CutlassFp8BlockScaledMMKernel":

The gsm8k difference is inside the ±0.02 error bar, so accuracy is the same.

Tests: vLLM's test_block_fp8.py gives 226 passed, 255 skipped, 0 failed on the L4. The 5 CUTLASS block-FP8 tests that skip on Ada today now run and pass.

What's not good yet

Two gaps, both stated in the PR description:

• Batch-1 decode is ~3.6% slower end to end. It comes from gate_up_proj and o_proj at M = 1. gate_up_proj is the largest layer, and at M = 1 it's purely limited by memory bandwidth, where Marlin is very well tuned.
• down_proj at large M is ~10% behind Triton (6159 vs 5584 µs at M = 4096).

Both look like per-shape tuning problems, and I'd rather keep the PR to the kernel and its wiring and handle tuning in a follow-up.

What I learned

• Check registers before tuning tiles. I spent a while thinking about warp counts. cuobjdump -res-usage answered it in one line.
• Compile-check on a free machine first. Building all 48 variants without a GPU told me which ones spilled before I paid for any GPU time.
• Correct is the easy part. The first version passed everything. It took two more L4 sessions to make it worth using.

Links

• Kernel PR: vLLM #61224 (open, waiting for review)
• Issue: vLLM #58241
• SM89 block-FP8 test enablement: vLLM #60171
• Part 1: Blockwise FP8 on Ada GPUs: why your L4 falls back to Marlin, measured