GPU Programming
Made Simple

Wave Lang is a Python DSL that makes developing high-performance ML GPU kernels easier and more fun than ever before. Write elegant code that compiles to blazing-fast GPU kernels.

See Examples View on GitHub
vector_add.py 
import wave_lang.kernel.lang as tkl as tkl
import wave_lang.kernel.wave as tkw

@tkw.wave(constraints)
def vector_add(
    a: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16],
    b: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16],
    c: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16],
):
    lhs = tkw.read(a)
    rhs = tkw.read(b)
    res = lhs + rhs
    tkw.write(res, c)

Why Wave Lang?

⚡

Lightning Fast

Generate highly optimized GPU kernels with automatic memory management and hardware-aware optimizations.

🐍

Python Native

Write GPU kernels using familiar Python syntax with seamless PyTorch integration.

🎯

ML Focused

Purpose-built for machine learning workloads with optimizations for common ML patterns.

🔧

Easy to Use

Intuitive API that abstracts away the complexity of low-level GPU programming.

The Future of GPU Programming

Traditional GPU programming requires deep expertise in SIMT C, memory management, and hardware-specific optimizations. Wave Lang changes this by providing a high-level Python interface that automatically generates efficient GPU code.

Whether you're implementing attention mechanisms, custom activations, or novel optimization algorithms, Wave Lang lets you focus on the math, not the plumbing.

10x Faster Development
100% Python Compatible
0 SIMT C Knowledge Required

Ready to get started?

Install Wave Lang and start building high-performance GPU kernels today.

pip install wave-lang