A custom attention kernel for Gemma 4 on the RTX PRO 6000

vLLM runs attention with FlashAttention wherever it can, because on most models and GPUs it is the fastest kernel available. Gemma 4 on NVIDIA’s RTX PRO 6000 is an exception. FlashAttention has no kernel for the model’s attention layers with the largest heads on this GPU, so vLLM runs them on a general kernel that leaves most of the GPU’s matrix throughput unused. We wrote a kernel for those layers, and it cut their attention time to about a quarter of vLLM’s default.

The deployment

We serve Gemma 4 31B instruct with FP8 weights and a draft model that proposes 4 tokens per step, on one RTX PRO 6000. We optimize its time to first token. For a long prompt, most of that time is the prefill, the pass that reads the prompt. On vLLM’s default, a 20k-token prompt took 4.5 s to the first token.

Fifty of the model’s 60 layers attend only to the last 1,024 tokens, with heads of 256 dimensions, so their attention cost grows in proportion to the prompt. The other 10, which we call the global layers, attend to the whole prompt with heads of 512 dimensions, so their cost per token grows with the context before it.

Finding the layer type to rewrite

We profiled the prefill of one 20k-token prompt on vLLM’s default and split its GPU time by kind of work (Figure 1). Attention in the global layers took 38% of the prefill’s GPU time, more than all the FP8 matrix multiplications together. It ran at 74 TFLOP/s, 17% of the 446 TFLOP/s we measured for BF16 matrix multiplication on this GPU’s tensor cores, while the FP8 matrix multiplications ran at about 765 TFLOP/s.

The same check works on any model. Divide each layer type’s attention FLOPs by its time and compare the rate with the GPU’s measured matrix rate. Rewrite the layer type that runs at a small fraction of that rate and takes a large share of the time.

Existing kernels first

vLLM 0.31 offers no FlashAttention kernel for heads of 512 dimensions on this GPU. FlashAttention 2 supports heads of up to 256 dimensions, and vLLM uses FlashAttention 4 only on B200-class GPUs. vLLM picks one attention backend for the whole model, so all 60 layers run on its general Triton kernel. The startup log says so: “FA4 not available, forcing TRITON_ATTN backend.”

Upstream tunes the Triton kernel’s prefill settings only for the B200 and heads of 256 dimensions, so on this GPU it ran the global layers with 16-row query blocks. Blocks of 32 rows, 4 warps and 2 pipeline stages gave bit-identical output and took 41% less time per call. FlashInfer 0.7’s paged prefill matched the tuned Triton kernel. cuDNN frontend 1.31’s new kernel for this GPU ran at 323–358 TFLOP/s, but its error against an fp32 reference was 1.5× Triton’s and its first call at each new prompt length took 0.8–2.8 s, so we did not use it.

The kernel in three steps

The RTX PRO 6000 is a Blackwell GPU without the B200’s tensor-core instructions, so a kernel for it uses Ampere-style mma.sync instructions and at most 99 KB of shared memory per thread block. Our kernel fits those limits with the layout llama.cpp, FlashMLA and cuDNN use for large heads (Figure 2).

1. Each thread block takes 64 query rows, 8 tokens times the 8 query heads that share a key-value head, so every key and value tile it loads into shared memory serves 64 rows. Triton’s kernel serves 32.
2. The 512 dimensions of Q and the output do not fit in one warp’s registers, so each pair of warps splits them. The pair adds its two halves of the scores through shared memory, and each warp then runs the softmax and multiplies by its half of V.
3. Two warp groups run in ping-pong, so one group’s matrix multiplications overlap the other’s softmax.

The kernel runs at 283–292 TFLOP/s per call at 0, 8k and 16k tokens of context, 2.5–2.6× the tuned Triton kernel and 64–65% of the measured tensor-core rate.

What it changed end to end

We measured time to first token in one server lifetime, one uncached 19,985-token prompt at a time, with greedy sampling. The kernel took 18% off it on top of the tuned Triton kernel. The floor is the time the 131 TFLOP of causal attention in the global layers would take at the measured BF16 rate.

The two smaller changes are a shorter first pass of the draft model and exact parallel tokenization.

The time left inside the kernel

Timers inside the first version of the kernel, before the ping-pong, show where each 16-key tile goes (Figure 3). The matrix multiplications take two thirds of the tile’s time and run at about 92% of the tensor-core rate. The other third goes to waiting for the next key and value tile, the score exchange and the softmax, which all 8 warps reach at the same moment. Q and the output take 192 of 255 registers per thread, so only one block fits per SM, and 99 KB of shared memory holds only two stages of keys and values. Serving every key and value from L2 gained only 1–2%, so memory bandwidth does not set the limit. TMA loads, multicast across a cluster and a four-way split of the dimensions all ran slower.

Which version we kept, and why

We kept the BF16 ping-pong version because it is as accurate as Triton. Against an fp64 reference both have a mean absolute error of 5.05 × 10⁻⁴ and a maximum of 0.0637. Outputs can still differ in the last bit because the summation order differs. A faster variant that rebuilt part of K from V was 1.6–1.8× farther from the reference, so we dropped it. An 8-bit version later took this attention to 286 ms per request and cut time to first token by another 6.8%. It changes numerics, so we accepted it only after its output drift to the BF16 model matched the FP8 model’s (KL 0.0357 against 0.0362). With these and other changes, the configuration we ship takes 2,309 ms against 4,480 ms for vLLM 0.31 in the same container.

The kernel is sized to 8 query heads per key-value head, 512 dimensions, 255 registers and 99 KB of shared memory, and a different model or GPU changes all four. It runs only for prefill steps of 64 or more tokens.