Write and implement GPU kernels using NVIDIA CuTe DSL (CUTLASS 4.x Python API) — NOT for Triton, CUDA C++, or conceptual explanations. Trigger only when the user wants to write or implement a kernel, not when asking questions about CuTe DSL concepts or layouts. CuTe DSL uses cute.jit/cute.kernel decorators and cutlass.cute imports. Covers element-wise kernels, GEMM patterns, reductions, memory hierarchy (global/shared/register/TMA), MMA tensor core operations, software pipelining, and framework integration.
Write and implement GPU kernels using NVIDIA CuTe DSL (CUTLASS 4.x Python API) — NOT for Triton, CUDA C++, or conceptual explanations. Trigger only when the user wants to write or implement a kernel, not when asking questions about CuTe DSL concepts or layouts. CuTe DSL uses cute.jit/cute.kernel decorators and cutlass.cute imports. Covers element-wise kernels, GEMM patterns, reductions, memory hierarchy (global/shared/register/TMA), MMA tensor core operations, software pipelining, and framework integration.
license
Apache-2.0
metadata
{"author":"NVIDIA Corporation"}
CuTe DSL
CuTe DSL is a Python-based domain-specific language for GPU kernel development,
part of CUTLASS 4.x. It provides Python abstractions over CUTLASS C++ templates
with JIT compilation to optimized CUDA kernels via MLIR and ptxas.
When to Use
Triggers:
Writing CUDA kernels in Python (element-wise, GEMM, custom ops)
For any non-trivial kernel (GEMM, attention, pipelined, fused ops), start by
finding the most similar existing example to use as a starting point — study
its structure, then rework it for your use case. Do not copy examples verbatim;
they target specific dtypes, architectures, and problem shapes that likely differ.
Pick the closest example from the index below.
Prefer examples matching the target GPU architecture (check with
torch.cuda.get_device_capability()) when the operation is similar.
Fetch via web_fetch with base URL
https://raw.githubusercontent.com/NVIDIA/cutlass/main/examples/python/CuTeDSL
Note which dtype/arch it targets (many examples are fp16/bf16-specific)
Check if it uses APIs tied to a specific arch (TMA → SM90+, tcgen05 → SM100)
Rework for the user's workload (do not copy-paste):
Change shapes, data types, tile sizes to match requirements
Replace compute logic (epilogue, activation fusion) as needed
If dtype differs (e.g., example is fp16, need fp32), expect vectorization
and layout changes — the scalar-loop patterns in references/ may be a
better starting point than adapting a vectorized example
Runtime wrapper must be lightweight: kernel_fn() should only call
from_dlpack() + the compiled kernel. Never allocate intermediate tensors,
copy data, or re-compile per call — these belong in one-time setup
Apply optimizations from this skill's reference docs
⛔ Blackwell/Hopper GEMM + extra tensors — STOP:
If the target GPU is SM90+ (Hopper/Blackwell) and the GEMM requires
extra tensors beyond A, B, C in the epilogue (e.g., bias vector, activation
inputs), do not attempt it. These examples use TMA descriptors for all
data movement — adding tensors requires modifying TMA descriptor setup,
which is prohibitively complex. Instead, tell the user this limitation and
suggest a two-kernel approach: run the GEMM kernel as-is, then apply
bias + activation in a separate element-wise kernel (Workflow 1).
Plain GEMM (just A×B→C with scalar alpha/beta) on Hopper/Blackwell is fine.
The kernel file must export kernel_fn, reference_fn, and get_inputs().
When to skip examples: Pure element-wise operations (Workflow 1) have
complete patterns in references/patterns-elementwise.md — no need to fetch
external examples.
Reduction kernels (softmax, layernorm, RMSNorm): Use
references/patterns-reduction.md which provides complete, proven patterns
for float32 reductions using scalar loops + butterfly shuffle + shared memory.
Workflow 1: Element-wise Kernel
For unary/binary/in-place operations that map inputs to outputs 1:1.
Select pattern from references/patterns-elementwise.md (Variations A–E)
Write kernel applying all four invariant principles:
P1: from_dlpack(tensor, assumed_align=16) for vector loads
P2: Derive vec_size from element_type.width
P3: cute.zipped_divide(mA, tiler) for coalesced access
P4: cutlass.dynamic_expr(thread_idx < total) for bounds
Critical rules: No early return, no a * 2 (use a + a), no cute.math.sigmoid
Pre-compile with cute.compile(): Always pre-compile the kernel once
using cute.compile() so that kernel_fn calls the compiled object, not
@cute.jit directly. Without pre-compilation, every call recompiles
(~20-50ms overhead). Use .mark_layout_dynamic() so a single compiled
kernel handles arbitrary input shapes without recompilation:
# Compile once with dynamic layouts — works for any shape
fake_x = from_dlpack(torch.empty(1, 1, dtype=torch.float16, device="cuda"),
assumed_align=16).mark_layout_dynamic()
fake_out = from_dlpack(torch.empty(1, 1, dtype=torch.float16, device="cuda"),
assumed_align=16).mark_layout_dynamic()
compiled_kernel = cute.compile(host_fn, fake_x, fake_out)
defkernel_fn(x):
out = torch.empty_like(x)
compiled_kernel(from_dlpack(x, assumed_align=16).mark_layout_dynamic(),
from_dlpack(out, assumed_align=16).mark_layout_dynamic())
return out
Pipeline: Use PipelineTmaAsync (Hopper) or PipelineTmaUmma (Blackwell).
⚠️ TMA-based pipelines manage data movement via TMA descriptors — adding
extra tensors (bias, activation inputs) to the epilogue requires modifying
descriptor setup, which is prohibitively complex. See the stop condition in
Workflow 0 step 4.
Epilogue: Predicated store with alpha/beta scaling
Pre-compile with cute.compile(): Always pre-compile the GEMM kernel
so kernel_fn calls the compiled object, not @cute.jit directly.
Without pre-compilation, every call recompiles (~20-50ms overhead).
Autotune: Search over tile sizes, cluster shapes, pipeline depths
Workflow 3: Framework Integration
For wrapping CuTe DSL kernels as PyTorch/JAX custom operators.
Write kernel using Workflow 1 or 2
Create wrapper: Accept torch.Tensor, convert via from_dlpack, call host fn
For production: Compile with TVM FFI for zero-overhead tensor passing:
compiled = cute.compile(host_fn, *fake_tensors, options="--enable-tvm-ffi")
compiled(torch_a, torch_b) # Direct torch.Tensor, no from_dlpack
For deployment: Use AOT compilation → export to .o → load at runtime
Workflow 4: Debugging & Profiling
Set environment: CUTE_DSL_PRINT_IR=1, CUTE_DSL_KEEP_PTX=1
Use cute.printf() for runtime values (not Python print)
get_inputs() — returns a list of CUDA tensors for testing
# Example kernel.py contractimport torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
defkernel_fn(x):
out = torch.empty_like(x)
# ... call compiled cute kernel ...return out
defreference_fn(x):
return torch.nn.functional.gelu(x)
defget_inputs():
return [torch.randn(1024, 512, dtype=torch.float16, device="cuda")]
Examples
Example: 2D Unary Element-wise (ReLU)
import torch, cutlass, cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.kerneldefrelu_kernel(gA: cute.Tensor, gC: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
idx = bidx * bdim + tidx
m, n = gA.shape[1]
total = m * n
if cutlass.dynamic_expr(idx < total):
a = gA[(None, (idx // n, idx % n))].load()
gC[(None, (idx // n, idx % n))] = cute.where(a > 0, a, 0)
@cute.jitdefrelu_host(mA: cute.Tensor, mC: cute.Tensor):
vec = 16 // (mA.element_type.width // 8)
gA = cute.zipped_divide(mA, (1, vec))
gC = cute.zipped_divide(mC, (1, vec))
T = 256
N = cute.size(gA.shape[1])
relu_kernel(gA, gC).launch(grid=((N+T-1)//T,1,1), block=(T,1,1))
x = torch.randn(1024, 512, dtype=torch.float16, device="cuda")
out = torch.empty_like(x)
relu_host(from_dlpack(x, assumed_align=16), from_dlpack(out, assumed_align=16))
Error Handling
Error
Cause
Fix
MLIR function requires a Context
Called @kernel from Python
Launch via @cute.jit host function
DSLAstPreprocessorError on return
Early return in @kernel
Use if cutlass.dynamic_expr(cond):
Type mismatch on store
a * 2 promotes FP16→FP32
Use a + a or .to(cutlass.Float16)
could not get source code
Kernel in exec() context
Write to file and import
Scalar loads in Nsight
Missing alignment hint
Add assumed_align=16 to from_dlpack
Missing required argument
Not all @jit params passed
Pass ALL declared parameters
AttributeError: sigmoid
No cute.math.sigmoid
Use 1.0/(1.0+cute.math.exp(-x))
See references/troubleshooting.md for the full error table and limitations.
Debugging rule: Never delete kernel.py during debugging. Use backup_file
to save a checkpoint, then edit_file to iterate. If stuck, revert_file to
restore the backup. A partially-working kernel is always better than no kernel.
Finding More Information
Tier 1: This File (SKILL.md)
Workflows above cover element-wise kernels, GEMM, framework integration, and
debugging. Search this file first for procedural questions.
Tier 2: references/ Directory
Grep for keywords across references/. Headers are grep-friendly.
