🍎 One Kernel a Day, Keeps High Latency Away.
A hands-on CUDA learning path from novice to expert.
Welcome to your daily dose of GPU performance fuel. Vitamin-CUDA is a curated collection of 170+ hand-tuned CUDA kernels, ranging from basic grid-stride loops for beginners to cutting-edge asynchronous TMA pipelines for performance enthusiasts(SGEMM, HGEMM Flash Attention/Decoding). Every kernel is seamlessly wrapped as a PyTorch C++ extension, delivering raw hardware performance with zero-overhead Python integration.
No over-engineered abstractions. No black boxes. Just pure hardware saturation.
- [2026.05.21] topk_topp_sampling -Released template-specialized (split-K) sampling kernels optimized for small-K/BS regimes using register-level insertion sort and hierarchical merge. Performance insights from this implementation helped drive a 2~4x speedup optimization in FlashInfer(see #3389).
- [2026.04.26] flash_decode_sm120 — Released a custom sm120 decode kernel leveraging
TMA + double Ks buffers + single Vs bufferto maximize memory bandwidth utilization. - [2026.04.16] flash_attn_sm120 — Released a custom sm120 kernel that outperforms torch.sdpa(FA2) across all prefill sequence lengths(1.2x in best case). Leverages TMA to drastically reduce register pressure and maximize L2 cache locality.
- [2026.04.04] hgemm_120 — FP16/BF16 Tensor-Core kernel pushing the absolute compute limit of the GPU. Outperforms cuBLAS via
TMA copy + 3-stage pipeline + double register buffer + swizzle + ldmatrix + mma. - [2026.03.18] hgemm — FP16/BF16 Tensor-Core kernel outperforming cuBLAS via
cp.async + double SMEM buffer + swizzle + ldmatrix + mma. - [2026.03.10] sgemm_tf32 — TF32 Tensor-Core kernel outperforming cuBLAS via
cp.async + double SMEM buffer + swizzle + ldmatrix + mma. - [2026.02.27] sgemm — SIMT kernel outperforming cuBLAS via
SMEM + swizzle + double buffer + coalesced read/write.
- NVIDIA GPU (Compute Capability 6.0+)
- CUDA Toolkit 11.0+
- C++ Compiler (GCC / Clang / MSVC)
- CMake 3.18+ (optional but recommended)
- PyTorch (for Python bindings and performance comparison)
Quick Start — We recommend using the NVIDIA PyTorch NGC Docker image for a ready-to-go environment.
All kernels were tested on an RTX 5060 Laptop GPU (unless otherwise noted) and benchmarked against PyTorch 2.9+.
- elementwise: element-wise add
-
elementwise_add— FP32 / FP16 -
elementwise_add_fp16x2— FP16 vectorized (×2) -
elementwise_add_fp16x8— FP16 vectorized (×8) -
elementwise_add_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- sigmoid
-
sigmoid— FP32 / FP16 -
sigmoid_fp16x2— FP16 vectorized (×2) -
sigmoid_fp16x8— FP16 vectorized (×8) -
sigmoid_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- swish
-
swish— FP32 / FP16 -
swish_fp16x2— FP16 vectorized (×2) -
swish_fp16x8— FP16 vectorized (×8) -
swish_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- relu
-
relu— FP32 / FP16 -
relu_fp16x2— FP16 vectorized (×2) -
relu_fp16x8— FP16 vectorized (×8) -
relu_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- relu6
-
relu6— FP32 / FP16 -
relu6_fp16x2— FP16 vectorized (×2) -
relu6_fp16x8— FP16 vectorized (×8) -
relu6_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- elu
-
elu— FP32 / FP16 -
elu_fp16x2— FP16 vectorized (×2) -
elu_fp16x8— FP16 vectorized (×8) -
elu_fp16x8— FP16 vectorized (×8, packed r/w; ~2× speedup viahalf2intrinsics) - PyTorch op binding & correctness check
-
- gelu
-
gelu— FP32 / FP16 -
gelu_fp16x2— FP16 vectorized (×2) -
gelu_fp16x8— FP16 vectorized (×8) -
gelu_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- hardswish
-
hardswish— FP32 / FP16 -
hardswish_fp16x2— FP16 vectorized (×2) -
hardswish_fp16x8— FP16 vectorized (×8) -
hardswish_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- embedding
-
embedding— FP32 / FP16 -
embedding_fp32x4— FP32 vectorized (×4) -
embedding_fp32x4— FP32 vectorized (×4, packed r/w) -
embedding_fp16x2— FP16 vectorized (×2) -
embedding_fp16x8— FP16 vectorized (×8) -
embedding_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- rope — Rotary Position Embedding
- PyTorch naive RoPE
- PyTorch RoPE with cos/sin look-up table
-
rope_fp32— ~10× faster than the PyTorch naive implementation -
rope_fp32x4— FP32 vectorized; tens of times faster at larger sizes - PyTorch op binding & correctness check
- reduce — Warp-shuffle-based reductions
-
reduce_sum— FP32 / FP16 -
reduce_sum_fp16x2— FP16 vectorized (×2) -
reduce_sum_fp16x8_packed— FP16 vectorized (×8, packed r/w) -
reduce_sum— INT8 -
reduce_sum_i8x16_packed— INT8 vectorized (×16, packed r/w) -
reduce_sum_i8x16_packed— INT8 vectorized (×16, packed r/w,dp4a; tens of times faster than naive PyTorch) -
reduce_sum_i8x64_packed— INT8 vectorized (×64, packed r/w,dp4a) - PyTorch op binding & correctness check
-
- dot_product
-
dot_product— FP32 / FP16 -
dot_product_fp32x4— FP32 vectorized (×4) -
dot_product_fp16x2— FP16 vectorized (×2) -
dot_product_fp16x8— FP16 vectorized (×8, packed r/w) - PyTorch op binding & correctness check
-
- cumsum - Warp-shuffle-based scan
- naive Torch cumsum
-
cumsum_fp32 -
cumsum_fp32x4— FP32 vectorized (×4) -
cumsum_bf16 -
cumsum_bf16x8_packed— BF16 vectorized (×8, packed r/w) -
cumsum_fp32x4_split_k— split-k two-pass -
cumsum_fp32x4_multi_cta_scan— multi-CTA single kernel (decoupled look-back) - pytorch op bindings && diff check
- normalization - (x - mean) / std
- naive Torch norm
-
norm_fp32 -
norm_fp32x4— FP32 vectorized (×4) -
norm_fp16 -
norm_fp16x8_packed— FP16 vectorized (×8, packed r/w) -
norm_fp32x4_split_k— split-k two-pass - pytorch op bindings && diff check
- softmax — Safe online softmax
- One-pass
- Safe online softmax — FP32 / FP16
- Safe online softmax — FP32 vectorized (×4)
- Safe online softmax — FP16 vectorized (×8, pure register, packed r/w)
- Safe online softmax medium — FP16 vectorized (moderate register + SMEM, packed r/w)
- Safe online softmax extreme — FP16 vectorized (max register + SMEM, packed r/w)
- Two-pass (large rows beyond single-block capacity)
- Safe online softmax arbitrary — FP16 vectorized (max register + SMEM, packed r/w)
- Safe online softmax split-k — FP16 vectorized (max register + SMEM, packed r/w)
- PyTorch op binding & correctness check
- One-pass
- rmsnorm
- PyTorch naive RMSNorm (baseline)
-
rmsnorm— FP32 / FP16 -
rmsnorm_fp32x4— FP32 vectorized (×4) -
rmsnorm_fp32x4_smem— with shared memory -
rmsnorm_fp16x8— FP16 vectorized (×8, packed r/w) -
rmsnorm_fp16x8_smem— FP16 vectorized (×8, packed r/w, shared memory) - PyTorch op binding & correctness check
- transpose — Matrix transpose
-
transpose_coalesced_read— coalesced global loads (input-centric) -
transpose_coalesced_write— coalesced global stores (output-centric) -
transpose_smem— shared-memory tiled transpose -
transpose_smem_bcf— SMEM bank-conflict-free -
transpose_smem_packed_bcf— SMEM bank-conflict-free,float4vectorized r/w -
transpose_smem_swizzled_packed— SMEM swizzled,float4vectorized r/w - PyTorch op binding & correctness check
-
- sgemv — Single-precision GEMV
-
gemv_fp32— baseline -
gemv_fp32x4— vectorized loads (×4) -
gemv_fp32x4_split_k— split-k two-pass - PyTorch op binding & correctness check
- PyTorch op binding & correctness check
-
- sgemm — Single-precision GEMM (SIMT)
-
sgemm_cublas— cuBLAS FP32 baseline -
sgemm_tiling— vectorized r/w + block tiling with SMEM -
sgemm_at_tiling— A-transpose into SMEM (4-way bank conflict, inner-loopfloat4loads) -
sgemm_at_bcf_swizzling— A-transpose + swizzle, bank-conflict-free -
sgemm_at_bcf_swizzling_rw— + coalesced C write-back -
sgemm_at_bcf_swizzling_dbf_rw— + double buffer pipeline, outperforms cuBLAS - PyTorch op binding & correctness check
-
- sgemm_tf32 — TF32 Tensor-Core GEMM
-
sgemm_cublas_tf32— cuBLAS TF32 baseline -
sgemm_tf32_bt— vectorized loads for A/B, B-transpose into SMEM,ldmatrix+mma -
sgemm_tf32_bt_swizzle— + A-SMEM zero bank conflicts -
sgemm_tf32_bt_swizzle_dbf— + grid swizzling, 97–102% of cuBLAS -
sgemm_tf32_swizzle_bcf—cp.async, warp-shuffle B-register transpose, A/B SMEM conflict-free, grid swizzling -
sgemm_tf32_swizzle_bcf_dbf— + double buffer, outperforms cuBLAS - PyTorch op binding & correctness check
-
- hgemm
-
hgemm_cublas— cuBLAS BF16/FP16 baseline -
hgemm_naive— BF16/FP16,ldmatrix+mma -
hgemm_bcf— BF16/FP16,ldmatrix+mma, A/B SMEM swizzle conflict-free, 95–99% of cuBLAS -
hgemm_bcf_dbf— + double buffer, outperforms cuBLAS -
hgemm_bcf_dbf_rw— + coalesced GMEM read/write, outperforms cuBLAS - PyTorch op binding & correctness check
-
- hgemm_sm120
- hgemm_cublas baseline
- hgemm_bcf_dbf_rw baseline: cp.async + ldmatrix + mma + double buffer + coalesced gmem r/w
- hgemm_k_stages: cp.async + ldmatrix + mma + 2/3 stages buffer + coalesced gmem
- hgemm_tma_r_k_stages_64: TMA read + ldmatrix + mma: one block, 3 stages, 128x128x64, double buffer register
- hgemm_tma_r_k_stages_32: TMA read + ldmatrix + mma: two blocks, 3 stages, 128x128x32, double buffer register
- pytorch op bindings && diff check
- flash_attantion_sm120
- fmha: sigle buffer (BMxBN = 64x64, better for small seq_len size)
- pytorch op bindings && diff check
- flash_decoding_sm120
- flash_decode_tma_128
- pytorch op bindings && diff check
- topk_topp_sampling
- sampling_topk_topp_batched
- sampling_topk_topp_split_k
- pytorch op bindings && diff check
- vector_add:
- tilelang/triton kernels
- pytorch/cuda op compare && diff check
- softmax
- tilelang online softmax (naive two pass)
- pytorch/cuda op compare && diff check