CuTeDSL-Basics

Foreword

I had previously dabbled in CuTe but gave up——I didn’t stick with it (see CuTe Layout, that’s actually the only blog I’ve ever wrote about it…). Beyond the fact that it felt genuinely complex, the compilation times were painfully slow; I couldn’t quickly see the results of my experiments. The error messages were also nightmarish—a minor typo would flood the screen with errors, making it impossible to pinpoint the actual mistake.

I had actually come across CuTeDSL—an official Python version of CuTe—back then, but I hesitated because of the stereotype that Python code runs slowly. It wasn’t until I learned that FlashAttn4 was written entirely in CuTeDSL that I decided to switch over and give it a proper try (not that I had learned much CuTe to begin with).

It turns out that, Pythonic is the future!

Its short compilation times and native support for DLPack, in particular, makes prototyping experiments incredibly convenient. After examining the generated PTX, I found it comparable to code written directly in CUDA (practically identical, really), so I decided to dive deeper into it.

I plan to write a series of posts to document my learning journey. This article introduces the basics (not the fundamentals of CUDA programming itself; I assume the reader already has some background in CUDA).

Basics

Installation

First you need to install CuTeDSL. The official repo has given a detailed guide for installing. In short, yout can just install by:

pip install nvidia-cutlass

You can check installation by:

import cutlass
print(cutlass.__version__)

Basic Usages

First you need to import the library:

import cutlass
import cutlass.cute

Then, just like __device__ and __host__, the library provides two decorators to declare whether a function is a host-side or device-side function.

@cute.jit: host-side function
@cute.kernel: device-side function (Actually they are still slightly different from __device__ and __host__, but basically you can understand it this way for now)

You can launch a kernel inside the host-side functions by calling device_func().launch()

@cute.kernel
def device_func():
	print("Hello from kernel!")
	pass
	
@cute.jit
def host_func(x: cute.Tensor):
	print("Hello from host")
	device_func().launch(
		grid=(1, 1, 1),
		block=(16, 1, 1)
	)
	pass

Inside a kernel, like threadIdx and blockIdx, you can use cute.arch.thread_idx() and cute.arch.block_idx() to get thread and block index (similar for block and grid dimensions).

@cute.kernel
def device_func():
	tidx, _, _ = cute.arch.thread_idx()
	cute.printf(f"Hello from {tidx}!")

Here, what worth mentioning is that variables are devided into two types, static or dynamic, by whether they can be determined at compile time. And Python’s print can only print static variables correctly, and print ? for dynamic variables. To correctly print dynamic variables, you need to use cute.printf.

To run a kernel, you can either run directly to use JIT compilation or use cute.compile(func) to compile in advance. Both of the complie methods are really fast, compared to C++ CuTe…

host_func(x) # run with JIT compilation
 
compiled_func = cute.compile(host_func, x)
compiled_func(x) # run with compiled function

If you declare an argument of a jit function to be cutlass.Constexpr, then you shouldn’t include that arguments when you call the function after you compiled it, otherwise it will lead to a Segmentation fault.

@cute.jit
def test(a: cutlass.Constexpr):
	...
	
tmp = 1
func = cute.compile(test, tmp)
func() # correct
test(tmp) # correct
 
func(tmp) # wrong!

Tensor

One of the best things is that CuTeDSL is completly compatible with DLPack, which means you can directly create a tensor from torch or numpy.

from cutlass.cute.runtime import from_dlpack
 
a = torch.randn(8, 5, dtype=torch.float32)
a_ = from_dlpack(a) # cute.Tensor

Simple Example

People always use a elementwise add kernel as the CUDA version of “Hello World”. Here I use a even simpler naive add-one kernel as our first example.

import torch
import cutlass
import cutlass.cute as cute
 
@cute.kernel
def add_one_kernel(x: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    bdimx, _, _ = cute.arch.block_dim()
 
    thread_idx = bidx * bdimx + tidx # get the global thread idx
    _, n = x.shape
    row = thread_idx // n
    col = thread_idx % n
 
    x[row, col] += cute.BFloat16(1) # each thread handles only one element
    
@cute.jit
def add_one(x: cute.Tensor):
    num_threads_per_block = 256
    kernel = add_one_kernel(x)
    m, n = x.shape
    kernel.launch(
        grid=((m * n) // num_threads_per_block, 1, 1),
        block=(num_threads_per_block, 1, 1)
    )
	
M = 16384
N = 16384
 
x = torch.randn(M, N, dtype=torch.bfloat16)
x_plus1 = x + 1
x_ = from_dlpack(x)
add_one(x_)
 
torch.testing.assert_close(x, x_plus1) # check for correctness

Reference

Official examples, the notebook part is a great tutorial

Evan's blog

Explorer