Python DSL Overview

​

Python Packages

​

CUTLASS Python Interface (cutlass_cppgen)

​

Key Features

• High-level interfaces requiring only a few parameters
• Automatic selection of sensible default configurations
• Enumeration of known working configurations
• Descriptive Python exceptions instead of C++ compile errors
• Easy export to framework extensions (PyTorch CUDA extensions)

​

Quick Example

import cutlass
import numpy as np

plan = cutlass.op.Gemm(element=np.float16, layout=cutlass.LayoutType.RowMajor)
A, B, C, D = [np.ones((1024, 1024), dtype=np.float16) for i in range(4)]
plan.run(A, B, C, D)

​

Supported Operations

• GEMMs - General matrix multiplication
• Fused Epilogues - GEMMs with elementwise operations (e.g., ReLU) for pre-SM90
• Stream K - Stream K swizzling for pre-SM90 kernels
• Grouped GEMM - Multiple GEMMs in a single kernel for pre-SM90

​

Design Philosophy

​

Goals

​

Non-Goals

1. Select optimal kernel configurations - Default selections may not achieve highest performance. Users should:
   • Profile different parameter combinations, or
   • Use optimized libraries like cuBLAS
2. Act as a fast container - Does not minimize Python overhead. For deployment:
   • Use the emitted C++ code directly, or
   • Use framework extension emitters
3. Be a JIT compilation engine - Enables CUTLASS in Python but doesn’t aim to be a Python-to-CUDA JIT framework

​

Comparison to PyCUTLASS

​

CuTe DSL

​

Key Capabilities

• Write CUDA kernels in Python using @cute.kernel decorator
• Express complex tensor layouts with layout algebra
• Utilize hardware features (Tensor Cores, TMA, async pipelines)
• JIT compilation to optimized PTX/SASS
• Integration with PyTorch, JAX, and NumPy

​

Quick Example

import cutlass
import cutlass.cute as cute

@cute.kernel
def elementwise_add(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    
    # Partition work across threads
    tiled_copy = cute.make_tiled_copy_tv(...)
    thr_copy = tiled_copy.get_slice(tidx)
    
    # Load, compute, store
    frgA = cute.make_fragment_like(thr_copy.partition_S(gA[bidx]))
    frgB = cute.make_fragment_like(thr_copy.partition_S(gB[bidx]))
    cute.copy(thr_copy, gA[bidx], frgA)
    cute.copy(thr_copy, gB[bidx], frgB)
    
    result = frgA.load() + frgB.load()
    frgC.store(result)
    cute.copy(thr_copy, frgC, gC[bidx])

​

Installation

pip install nvidia-cutlass

pip install nvidia-cutlass-dsl

cd /path/to/cutlass/python/CuTeDSL
pip install .

# Clone CUTLASS repository
git clone https://github.com/NVIDIA/cutlass.git
cd cutlass

# Install cutlass_cppgen
pip install .

# Or install in editable mode for development
pip install -e .

# Use NGC PyTorch container (recommended)
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.08-py3

# Install CUTLASS packages inside container
pip install nvidia-cutlass nvidia-cutlass-dsl

​

Requirements

​

Next Steps

​

Resources

• GitHub Repository
• Python Examples
• API Reference