GPU Kernel Suite (gks)

High-performance GPU kernels written three ways — CUDA, Triton, and JAX/Pallas — with a unified Python harness for testing, benchmarking, autotuning, and roofline analysis.

Kernels

Kernel	CUDA	Triton	Pallas	Notes
Tiled SGEMM (FP32)	✓	✓	✓	shared-mem tiles, double-buffered cp.async, register tiling
WMMA SGEMM (FP16 → FP32)	✓	✓	—	Tensor Core mma.sync via nvcuda::wmma, FP32 accumulate
Online Softmax	✓	✓	✓	one-pass online algorithm with warp-shuffle reductions
RMSNorm	✓	✓	✓	one-pass mean-of-squares with warp-shuffle reduction, fused affine
Fused Softmax(RMSNorm(x))	✓	✓	✓	two reductions in one kernel, no round-trip to HBM
FlashAttention v1 forward	✓	✓	✓	online softmax over Q-blocks × K-blocks, FP16 inputs, FP32 accumulator, returns lse
Symmetric INT8 quantize/dequant	✓	✓	✓	per-row absmax scale, vectorized int4/int2 loads/stores

Each kernel uses some subset of these techniques:

• double-buffered async copy (cp.async + __pipeline_*) — SGEMM, FlashAttention
• Tensor Core WMMA — SGEMM (FP16), FlashAttention (matmul stages)
• coalesced global loads — every kernel; verified by inspecting transaction counts in Nsight
• bank-conflict-free shared layouts — SGEMM uses padded [BM][BK+1] tiles, FlashAttention uses XOR-swizzled K tiles
• warp-level reductions via __shfl_xor_sync — softmax, RMSNorm, FlashAttention
• vectorized affine loop tiling — every kernel reads/writes via float4/half2/int4
• symmetric INT8 quantization — quantize kernel + INT8-input GEMM variant

Layout

gpu-kernel-suite/
├── csrc/                # C++/CUDA sources, pybind11 module `_gks_cuda`
│   ├── include/         # headers shared across .cu files
│   ├── src/             # one .cu per kernel
│   └── bindings/        # pybind11 entry points
├── gks/                 # Python package
│   ├── triton/          # Triton ports of every kernel
│   ├── pallas/          # JAX/Pallas ports
│   ├── reference/       # PyTorch / NumPy reference impls
│   └── profile/         # autotune, roofline, Nsight wrappers
├── benchmarks/          # one bench script per kernel + run_all
├── tests/               # pytest suite, numerics + shapes
├── tools/               # IR dump, roofline plotter, ncu/nsys helpers
├── scripts/             # build + setup
└── examples/            # end-to-end usage demos

Build

Requires CUDA 12.x, a Hopper or Ampere GPU (SM ≥ 80 for cp.async, SM ≥ 90 unlocks the WGMMA path — not used here), Python 3.10+, and PyTorch 2.3+.

git clone <this repo> && cd gpu-kernel-suite
pip install -e .[dev]                  # builds the CUDA extension via setup.py + nvcc
pytest -q tests/                       # ~30s on a single A100
python benchmarks/run_all.py --json    # writes results/bench.json
python tools/roofline_plot.py results/bench.json

The Triton and Pallas backends are pure Python and are imported lazily; you can use the package with only Triton installed (e.g. on Colab) by skipping the C++ build:

GKS_SKIP_CUDA_BUILD=1 pip install -e .
python -c "from gks.triton import flash_attention_fwd"

Quick example

import torch
from gks import flash_attention_fwd                # routes to CUDA by default
from gks.triton import flash_attention_fwd as fa_triton
from gks.reference import flash_attention_fwd as fa_ref

q = torch.randn(2, 8, 1024, 64, device='cuda', dtype=torch.float16)
k = torch.randn_like(q); v = torch.randn_like(q)

o_cuda,   lse_cuda   = flash_attention_fwd(q, k, v, causal=True)
o_triton, lse_triton = fa_triton(q, k, v, causal=True)
o_ref,    lse_ref    = fa_ref(q, k, v, causal=True)

torch.testing.assert_close(o_cuda,   o_ref, atol=2e-3, rtol=2e-3)
torch.testing.assert_close(o_triton, o_ref, atol=2e-3, rtol=2e-3)

• tools/dump_triton_ir.py runs each Triton kernel with triton.compiler.compile and dumps ttir, ttgir, and ptx.