Sparse paged GQA decode attention with FP8 KV cache on SM89: CUDA vs Triton vs CuTe DSL vs Mojo vs eager PyTorch
I wrote the same sparse paged GQA decode kernel in six ways. I compared them on an RTX 4060 Laptop GPU (SM89, Ada). All code, charts, and Nsight Compute reports are in the repo
Workload
- 256 requests, 32 query heads, 8 KV heads, head dimension 128
- Top-K 512, BF16 query/gate/output, FP8 E4M3FN K/V cache
- Page size 16, 32 pages per request (256 MiB for K and V together)
- Operation: sparse scaled dot-product attention over selected tokens, then a sigmoid gate
- Nominal FLOPs: 4 × B × TopK × Hq × D = 2.15 GFLOP
Implementations
- Raw CUDA C++: one 128-thread CTA for each request and KV head. One warp for each query head. Each lane owns 4 dimensions. The QK reduction uses 5 butterfly shuffles. FP8 to FP16 conversion uses inline PTX (cvt.rn.f16x2.e4m3x2).
- Optimized CuTe DSL: the same mapping as raw CUDA.
- Baseline CuTe DSL: one CTA for each query head (8,192 CTAs). Four warps share one dot product through shared memory and barriers.
- Triton: one program for each request and query head. 32-token tiles with online softmax.
- Mojo (MAX): the same mapping as raw CUDA.
- Eager PyTorch: gather K/V once for each KV head, then batched matmul, softmax, and gate.
Optimized CuTe and CUDA fold k_scale / sqrt(D) × log2(e) into one score scale and use base-2 exponentials. They apply v_scale once at the end.
Results (CUDA events, one run)
|Implementation|Time|Derived rate|
|:-|:-|:-|
||
|Raw CUDA C++|2.1641 ms|0.992 TFLOP/s|
|Triton|2.4420 ms|0.879 TFLOP/s|
|Optimized CuTe DSL|4.4411 ms|0.484 TFLOP/s|
|Baseline CuTe DSL|15.4536 ms|0.139 TFLOP/s|
|PyTorch eager|31.0188 ms|0.069 TFLOP/s|
|Mojo (MAX)|12.985 ms (enqueue + sync, host timer)|not ranked|
Nsight Compute (basic set)
|Kernel|Occupancy|Registers/thread|SM throughput|DRAM throughput|
|:-|:-|:-|:-|:-|
||
|Raw CUDA|94.98%|35|74.92%|40.16%|
|Triton|57.81%|72|86.31%|41.35%|
|Optimized CuTe|95.15%|38|75.19%|27.23%|
|Baseline CuTe|98.86%|40|70.13%|15.46%|
Observations
1. Every fused kernel is much faster than eager PyTorch. The eager time also includes Python dispatch gaps.
2. The grouped mapping makes optimized CuTe 3.5x faster than the baseline CuTe kernel. It has 2,048 CTAs, not 8,192. It also avoids shared memory and barriers for QK.
3. Triton is only 13% slower than raw CUDA. It has the lowest occupancy but high SM throughput.
4. Occupancy alone does not predict latency.
Open problem: The grouped kernels reduce the CTA count. They do not share K/V loads. Each warp still loads its own K/V rows. The logical K/V request is about 1 GiB for the fused kernels, compared with about 256 MiB for the eager gather. I found no low-risk change in the source that removes this duplicate load. How would you share K/V data across the 4 query heads in a CTA?
Limits
- These are single-run results on a laptop GPU. Close rankings can change between runs.
- The CUDA and Mojo checks use constant fixtures. The other versions use random input and a PyTorch reference.
- The Mojo timer includes host enqueue cost. It is not comparable to CUDA events.
- NCU durations are profiler-instrumented. They are not normal latencies.
- The results are for synthetic data and fixed Top-K. They do not show end-to-end model serving.
I welcome corrections to the benchmark method.
6 points · 1 comment
Mentions
Sparse paged GQA decode attention with FP8 KV cache on SM89: CUDA vs Triton vs CuTe DSL vs Mojo vs eager PyTorch I wrote the same sparse paged GQA decode kernel in six ways. I compared them on an RTX 4060 Laptop GPU (SM89, Ada). All code, charts, and Nsight Compute reports are in the repo Workload - 256 requests, 32 query heads, 8 KV heads, head dimension 128 - Top-K 512, BF16 query/gate/output, FP8 E4M3FN K/V cache - Page size 16, 32 pages per request (256 MiB for K and V together) - Operation: sparse scaled dot-product attention over selected tokens, then a sigmoid gate - Nominal FLOPs: 4 × B × TopK × Hq × D = 2.15 GFLOP Implementations - Raw CUDA C++: one 128-thread CTA for each request and KV head. One warp for each query head. Each lane owns 4 dimensions. The QK reduction uses 5 butterfly shuffles. FP8 to FP16 conversion uses inline PTX (cvt.rn.f16x2.e4m3x2). - Optimized CuTe DSL: the same mapping as raw CUDA. - Baseline CuTe DSL: one CTA for each query head (8,192 CTAs). Four warps share one dot product through shared memory and barriers. - Triton: one program for each request and query head. 32-token tiles with online softmax. - Mojo (MAX): the same mapping as raw CUDA. - Eager PyTorch: gather K/V once for each KV head, then batched matmul, softmax, and gate. Optimized CuTe and CUDA fold k_scale / sqrt(D) × log2(e) into one score scale and use base-2 exponentials. They apply v_scale once at the end. Results (CUDA events, one run) |Implementation|Time|Derived rate| |:-|:-|:-| || |Raw CUDA C++|2.1641 ms|0.992 TFLOP/s| |Triton|2.4420 ms|0.879 TFLOP/s| |Optimized CuTe DSL|4.4411 ms|0.484 TFLOP/s| |Baseline CuTe DSL|15.4536 ms|0.139 TFLOP/s| |PyTorch eager|31.0188 ms|0.069 TFLOP/s| |Mojo (MAX)|12.985 ms (enqueue + sync, host timer)|not ranked| Nsight Compute (basic set) |Kernel|Occupancy|Registers/thread|SM throughput|DRAM throughput| |:-|:-|:-|:-|:-| || |Raw CUDA|94.98%|35|74.92%|40.16%| |Triton|57.81%|72|86.31%|41.35%| |Optimized CuTe|95.15%|38|75.19%|27.23%| |Baseline CuTe|98.86%|40|70.13%|15.46%| Observations 1. Every fused kernel is much faster than eager PyTorch. The eager time also includes Python dispatch gaps. 2. The grouped mapping makes optimized CuTe 3.5x faster than the baseline CuTe kernel. It has 2,048 CTAs, not 8,192. It also avoids shared memory and barriers for QK. 3. Triton is only 13% slower than raw CUDA. It has the lowest occupancy but high SM throughput. 4. Occupancy alone does not predict latency. Open problem: The grouped kernels reduce the CTA count. They do not share K/V loads. Each warp still loads its own K/V rows. The logical K/V request is about 1 GiB for the fused kernels, compared with about 256 MiB for the eager gather. I found no low-risk change in the source that removes this duplicate load. How would you share K/V data across the 4 query heads in a CTA? Limits - These are single-run results on a laptop GPU. Close rankings can change between runs. - The CUDA and Mojo checks use constant fixtures. The other versions use random input and a PyTorch reference. - The Mojo timer includes host enqueue cost. It is not comparable to CUDA events. - NCU durations are profiler-instrumented. They are not normal latencies. - The results are for synthetic data and fixed Top-K. They do not show end-to-end model serving. I welcome corrections to the benchmark method.