Hello, world!

Recently, I’ve been working on optimizing CUDA kernels for my own LLM implementation. Most of them are now performing quite well, especially considering that I’m running everything on an RTX 3050 Laptop GPU with only 4 GB of VRAM.

There is still one kernel that keeps frustrating me: GQA attention. Its GPU utilization is much lower than the rest, and I’m not sure how to optimize it further.

Here are my current benchmark results:

RMSNorm

• Input: [512, 4096]
• Data type: FP16
• Execution time: 0.049 ms
• GPU throughput: 89%

SwiGLU

• Input: [512, 4096]
• Data type: FP16
• Execution time: 0.074 ms
• GPU throughput: 90%

Linear layer

• Input: [512, 4096]
• Weights: [4096, 4096]
• Output: [512, 4096]
• Data type: FP16
• Execution time: 1.04 ms
• GPU throughput: 60%

Note: The reported throughput is relatively low because this benchmark runs GEMM on a comparatively small input matrix. At this problem size, the GPU cannot reach the same utilization as it would with larger batches or longer sequences.

Rotary Position Embedding — RoPE

• Query: [1, 512, 32, 128]
• Key: [1, 512, 8, 128]
• Cosine and sine cache: [512, 64]
• Data type: FP16
• Execution time: 0.06 ms
• GPU throughput: 87%

Grouped Query Attention — GQA

• Query and output: [1, 512, 32, 128]
• Key and value: [1, 512, 8, 128]
• Data type: FP16
• Execution time: 3.32 ms
• GPU throughput: 10%

This is currently the main bottleneck. I have already applied tiling, online softmax, and shared-memory optimizations, but the utilization is still very low.

Embedding lookup

• Number of input tokens: 512
• Vocabulary size: 32,000
• Hidden size: 4,096
• Output: [512, 4096]
• Data type: FP16
• Execution time: 0.05 ms
• GPU throughput: 88%

Summary

Overall, I’m satisfied with the results, but the attention kernel still needs a lot of work. (I will probably use cuDNN). The hardest thing before me - backprop. But now I can play with kernels and model layers.