leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.5 tag update (#3202)

Browse Source 
* Python DSL examples reorganization.

* v4.5 tag update.
This commit is contained in:
Junkai-Wu
2026-05-06 08:55:27 +08:00
committed by GitHub
parent f74fea9ce3
commit cb37157db5
351 changed files with 36688 additions and 8117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
media/docs/pythonDSL/cute_dsl.rst
View File

@@ -11,6 +11,7 @@ CuTe DSL
Control Flow <cute_dsl_general/dsl_control_flow.rst>
JIT Argument Generation <cute_dsl_general/dsl_jit_arg_generation.rst>
JIT Argument: Layouts <cute_dsl_general/dsl_dynamic_layout.rst>
Struct-like JIT Arguments <cute_dsl_general/dsl_struct_types.rst>
JIT Caching <cute_dsl_general/dsl_jit_caching.rst>
JIT Compilation Options <cute_dsl_general/dsl_jit_compilation_options.rst>
JIT Types <cute_dsl_general/types.rst>

20
media/docs/pythonDSL/cute_dsl_api/changelog.rst
View File

@@ -2,6 +2,24 @@
Changelog for CuTe DSL API changes
======================================
`4.4.0 <https://github.com/NVIDIA/cutlass/releases/tree/main>`_ (2026-03-24)
==============================================================================
* Added native support for ``typing.NamedTuple`` as a JIT function argument.
- A NamedTuple whose fields are DSL scalar types (``Int32``, ``Float32``, …)
can be passed directly to ``@cute.jit`` / ``cute.compile`` without any
protocol implementation.
- Fields are flattened field-by-field through the existing pytree system and
reconstructed via the NamedTuple constructor on entry to the kernel body.
Field attribute access (``tup.a``, ``tup.b``, …) works as in native Python.
- NamedTuple fields are **immutable** (tuple subclass). To replace a field,
construct a new NamedTuple inside the kernel. Use ``@native_struct`` when
mutable fields are required.
- See :doc:`../cute_dsl_general/dsl_struct_types` for a guide to NamedTuple,
``@native_struct``, and other struct-like JIT argument types.
`4.3.0 <https://github.com/NVIDIA/cutlass/releases/tree/main>`_ (2025-10-20)
==============================================================================
@@ -73,7 +91,7 @@ Changelog for CuTe DSL API changes
- Introduce S2T CopyOps in `tcgen05/copy.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py>`_.
- Introduce BlockScaled layout utilities in `blockscaled_layout.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/utils/blockscaled_layout.py>`_ for creating the required scale factor layouts in global memory, shared memory and tensor memory.
* ``cutlass.cute.compile`` now supports compilation options. Refer to `JIT compilation options <https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html>`_ for more details.
* ``cutlass.cute.compile`` now supports compilation options. Refer to `JIT compilation options <https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html>`_ for more details.
* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-assertions`` as compilation option to enable.
* ``cutlass.cute.make_tiled_copy`` is now deprecated. Please use ``cutlass.cute.make_tiled_copy_tv`` instead.
* Shared memory capacity query

76
media/docs/pythonDSL/cute_dsl_general/debugging.rst
View File

@@ -39,10 +39,11 @@ CuTe DSL provides environment variables to control logging level:
# Enable console logging (default: False)
export CUTE_DSL_LOG_TO_CONSOLE=1
# Log to file instead of console (default: False)
export CUTE_DSL_LOG_TO_FILE=my_log.txt
# Log to file instead of console (default: False).
# Set to 1/True to enable; the log file path is chosen automatically by the DSL.
export CUTE_DSL_LOG_TO_FILE=1
# Control log verbosity (0, 10, 20, 30, 40, 50, default: 10)
# Control log verbosity (0=disabled, 1=all messages (debug and above), 10=debug, 20=info, 30=warning, 40=error, 50=critical; default: 1)
export CUTE_DSL_LOG_LEVEL=20
@@ -68,40 +69,53 @@ Similar to standard Python logging, different log levels provide varying degrees
+--------+-------------+
Dump the generated IR
~~~~~~~~~~~~~~~~~~~~~
Save generated artifacts to files
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For users familiar with MLIR and compilers, CuTe DSL supports dumping the Intermediate Representation (IR).
This helps you verify whether the IR is generated as expected.
CuTe DSL can save generated artifacts (IR, PTX, CUBIN, …) to files for offline inspection.
Use ``CUTE_DSL_KEEP`` with a comma-separated list of artifact tokens:
.. code:: bash
# Dump Generated CuTe IR (default: False)
# Save clean IR (after canonicalize+cse, human-readable) to a .mlir file
export CUTE_DSL_KEEP=ir
# Save raw IR (before any passes) to a .mlir file
export CUTE_DSL_KEEP=ir-debug
# Save PTX assembly to a .ptx file
export CUTE_DSL_KEEP=ptx
# Save CUBIN binary to a .cubin file
export CUTE_DSL_KEEP=cubin
# Save LLVM IR to a file
export CUTE_DSL_KEEP=llvm
# Save multiple artifacts at once
export CUTE_DSL_KEEP=ir,ptx,cubin
# Save all supported artifacts
export CUTE_DSL_KEEP=all
Files are written to the current working directory by default. Use ``CUTE_DSL_DUMP_DIR``
to redirect them (see `Change the dump directory`_ below).
.. note::
The ``sass`` token requires ``nvdisasm`` (or ``nvdisasm_internal``) to be available
in your ``PATH``. It is usually installed with the CUDA toolkit.
Print the generated IR to the console
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
To print the IR directly to the console (without writing a file):
.. code:: bash
# Print generated IR to stdout (default: False)
export CUTE_DSL_PRINT_IR=1
# Keep Generated CuTe IR in a file (default: False)
export CUTE_DSL_KEEP_IR=1
Dump the generated PTX & CUBIN
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For users familiar with PTX and SASS, CuTe DSL supports dumping the generated PTX and CUBIN.
.. code:: bash
# Dump generated PTX in a .ptx file (default: False)
export CUTE_DSL_KEEP_PTX=1
# Dump generated cubin in a .cubin file (default: False)
export CUTE_DSL_KEEP_CUBIN=1
To further get SASS from cubin, users can use ``nvdisasm`` (usually installed with CUDA toolkit) to disassemble the cubin.
.. code:: bash
nvdisasm your_dsl_code.cubin > your_dsl_code.sass
Access the dumped contents programmatically
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

24
media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst
View File

@@ -101,18 +101,11 @@ The result:
|DSL| bridges Python and GPU hardware through a three-stage pipeline.
.. _fig-dsl-modes:
.. figure:: dsl_modes.png
:width: 400
.. figure:: dsl_compilation.png
:width: 600
:align: center
*Left*: tracing mode records only the path that executed.
*Right*: preprocessor mode emits structured |IR| for every branch and loop
before tracing the arithmetic.
The default |DSL| compilation pipeline (mode 2): Python source flows through AST preprocessing
The |DSL| compilation pipeline: Python source flows through AST preprocessing
and interpreter-driven tracing to produce |IR|, which is then lowered and
compiled to device code.
@@ -258,8 +251,8 @@ Practical Implications
4. |DSL| Code-Generation Modes
------------------------------
CuTe's Python front-end combines the techniques above into **two mutually
exclusive modes** (see :ref:`fig-dsl-modes`), selectable with the ``preprocessor`` flag of the
CuTe’s Python front-end combines the techniques above into **two mutually
exclusive modes**, selectable with the ``preprocessor`` flag of the
``@jit`` decorator:
1. Tracing mode ``@jit(preprocess=False)`` – tracing only.
@@ -272,3 +265,10 @@ optimisation problems of pure tracing; tracing then fills in the arithmetic.
This hybrid “preprocessor” pipeline is unique to |DSL| and was designed
specifically to overcome the disadvantages identified above.
.. figure:: dsl_modes.png
:width: 400
:align: center
*Left*: tracing mode records only the path that executed.
*Right*: preprocessor mode emits structured |IR| for every branch and loop
before tracing the arithmetic.

6
media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst
View File

@@ -117,6 +117,12 @@ Defines GPU kernel functions, compiled as specialized GPU symbols through |DC|.
- ``False`` (default) — Standard kernel launch.
- ``True`` — Cooperative kernel launch.
- ``smem_merge_branch_allocs``
Enables mutually exclusive control flow branches (sequentially executed if-else) to reuse the same shared memory.
- ``False`` (default) — Shared memory is allocated additively across all branches (default CUDA C++ behavior).
- ``True`` — Merge shared-memory allocations across branches (experimental feature, recommended for mega-kernels).
Calling Conventions
-------------------

134
media/docs/pythonDSL/cute_dsl_general/dsl_struct_types.rst Normal file
View File

@@ -0,0 +1,134 @@
.. _dsl_struct_types:
Struct-like JIT Arguments
=========================
|DSL| supports several struct-like Python types as JIT function arguments.
Each provides a different trade-off between mutability, syntax convenience,
and low-level control.
.. |DSL| replace:: CuTe DSL
.. contents:: On this page
:local:
:depth: 2
Overview
--------
.. list-table::
:header-rows: 1
:widths: 25 15 60
* - Type
- Mutable fields?
- Notes
* - ``typing.NamedTuple``
- **No**
- Tuple subclass — fields fixed at construction.
Flattened field-by-field through the pytree system.
* - ``@dataclass(frozen=True)``
- **No**
- Frozen dataclass — treated as a read-only pytree container,
similar to ``NamedTuple``.
NamedTuple
----------
A ``typing.NamedTuple`` whose fields are DSL scalar types (``Int32``,
``Float32``, etc.) can be passed directly to ``@cute.jit`` /
``cute.compile`` without any boilerplate or protocol implementation.
**How it works.** NamedTuples are registered as pytree containers in the DSL
tree system. Each field is flattened individually through the existing DSL
type paths and reconstructed by calling the NamedTuple constructor on the way
into the kernel body. Field attribute access (``tup.a``, ``tup.b``, …)
works exactly as in native Python.
Basic usage
^^^^^^^^^^^
.. code-block:: python
from typing import NamedTuple
import cutlass
import cutlass.cute as cute
class Vec3(NamedTuple):
x: cutlass.Int32
y: cutlass.Int32
z: cutlass.Int32
@cute.jit
def print_vec(v: Vec3):
cute.printf("x=%d y=%d z=%d\n", v.x, v.y, v.z)
v = Vec3(x=cutlass.Int32(1), y=cutlass.Int32(2), z=cutlass.Int32(3))
cute.compile(print_vec, v)(v)
Control flow on fields
^^^^^^^^^^^^^^^^^^^^^^
Fields are DSL values inside the kernel, so they work in ``if``/``else``
branches and ``for`` loops:
.. code-block:: python
@cute.jit
def clamp_positive(v: Vec3, out: cute.Tensor):
"""Write max(field, 0) for each component."""
out[0] = cutlass.Int32(0) if v.x < cutlass.Int32(0) else v.x
out[1] = cutlass.Int32(0) if v.y < cutlass.Int32(0) else v.y
out[2] = cutlass.Int32(0) if v.z < cutlass.Int32(0) else v.z
@cute.jit
def triangular_sum(v: Vec3, out: cute.Tensor):
"""Sum 0..v.x-1 into out[0], and so on."""
s = cutlass.Int32(0)
for i in range(v.x):
s = s + i
out[0] = s
Creating a new NamedTuple value inside the kernel
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NamedTuple fields are **immutable** — the same constraint as native Python
tuples. Assigning ``tup.x = ...`` inside a kernel raises ``AttributeError``.
To "update" a field, construct a replacement NamedTuple:
.. code-block:: python
@cute.jit
def scale(v: Vec3, factor: cutlass.Int32, out: cute.Tensor):
# Construct a new Vec3 with all fields scaled
scaled = Vec3(x=v.x * factor, y=v.y * factor, z=v.z * factor)
out[0] = scaled.x
out[1] = scaled.y
out[2] = scaled.z
Choosing the right type
-----------------------
.. list-table::
:header-rows: 1
:widths: 35 65
* - Use case
- Recommended type
* - Read-only config / parameters passed into a kernel
- ``NamedTuple`` or ``@dataclass(frozen=True)``
* - Accumulator or running state updated inside a kernel
- ``@native_struct``
* - Want Python-native immutable semantics (hashable, unpackable)
- ``NamedTuple``
* - Need fine-grained LLVM struct control (packing, zero-init)
- ``@native_struct``
See also
--------
* :doc:`dsl_jit_arg_generation` — overview of JIT function argument protocols
* :doc:`dsl_dynamic_layout` — passing ``Layout`` objects as JIT arguments

6
media/docs/pythonDSL/cute_dsl_general/framework_integration.rst
View File

@@ -79,6 +79,12 @@ by reducing register usage and the number of address calculation instructions. W
to True, a runtime check is performed to ensure that the layout does not overflow. Please note that this parameter
only has an effect when the tensor's layout is marked as dynamic.
For packed subbyte torch dtypes such as ``torch.float4_e2m1fn_x2``, ``from_dlpack`` exposes the
logical element layout expected by CuTe instead of the packed storage layout. For example, a torch
tensor with shape ``(128, 128)`` and dtype ``torch.float4_e2m1fn_x2`` is exposed as a logical FP4
tensor with shape ``(128, 256)``. The same logical reinterpretation also applies when the leading
dimension is not the last mode.
Code Example
~~~~~~~~~~~~

4
media/docs/pythonDSL/cute_dsl_general/resources.rst
View File

@@ -15,7 +15,7 @@ Conference Talks
An introduction to the |DSL| architecture, covering the hybrid AST-rewrite and
tracing approach, MLIR code generation, and integration with CUTLASS.
* `LLVM Video <https://www.youtube.com/watch?v=5NXd6MbKYNQ>`_
* `Video <https://www.youtube.com/watch?v=5NXd6MbKYNQ>`__
* `Slides (PDF) <https://llvm.org/devmtg/2025-10/slides/technical_talks/ozen.pdf>`_
----
@@ -25,4 +25,4 @@ tracing approach, MLIR code generation, and integration with CUTLASS.
Learn how to leverage Tensor Cores directly from Python using CUTLASS 4.0's
new DSL front-end, enabling rapid kernel development without writing CUDA C++.
* `GTC Video <https://www.nvidia.com/en-us/on-demand/session/gtc25-s74639/>`_
* `Video <https://www.nvidia.com/en-us/on-demand/session/gtc25-s74639/>`__

6
media/docs/pythonDSL/functionality.rst
View File

@@ -3,11 +3,7 @@
Functionality
====================
The CUTLASS DSL 4.0 release supports **Python 3.12** only. It shares the same driver requirements
as the `CUDA Toolkit 12.9 <https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html>`__.
Specifically, the driver version must be 575.51.03 or later.
Currently, only Linux x86_64 is supported. Additional platform support will be added in future releases.
For dependency version requirements, refer to the :doc:`quick_start` section.
Supported MMA Operations
---------------------------------

64
media/docs/pythonDSL/limitations.rst
View File

@@ -217,18 +217,68 @@ Programming Model
**CuTe Layout algebra in native Python**
Entirety of CuTe Layout algebra operations and APIs require JIT compilation. These
functionalities are exclusively available within JIT-compiled functions and cannot be
Entirety of CuTe Layout algebra operations and APIs require JIT compilation. These
functionalities are exclusively available within JIT-compiled functions and cannot be
accessed in standard Python execution environments.
Additionally, there exists a restricted set of data types that can be passed as arguments
to JIT-compiled functions, which further constrains their usage in native Python contexts.
Only following CuTe algebra types are supported as JIT function arguments: ``Tensor``, ``Pointer``,
Additionally, there exists a restricted set of data types that can be passed as arguments
to JIT-compiled functions, which further constrains their usage in native Python contexts.
Only following CuTe algebra types are supported as JIT function arguments: ``Tensor``, ``Pointer``,
``Shape``, ``Stride``, ``Coord`` and ``IntTuple``. For ``Stride``, we don't support ``ScacledBasis``
from native Python Context. Unfortunately, in the first release, we don't support
from native Python Context. Unfortunately, in the first release, we don't support
passing ``Layout`` under native Python Context.
**Block-level Utilities (block_copy)**
The block-level utility ``block_copy`` provides a high-level abstraction
for common copy patterns, but has the following limitations:
**block_copy limitations:**
- **Limited copy op support**: Currently only ``TmaCopyOp``-based tiled copies
(TMA loads/stores) and S2T copies (SMEM to TMEM, e.g., ``tcgen05.Cp*Op``) are
supported. Other ``TiledCopy`` ops will raise ``NotImplementedError``. Support
for additional copy ops may be added in future releases.
**Global variables**
CuTe DSL does not support global variables.
It is not allowed to use ``global`` in the DSL.
The following example illustrates functionality in Python that is not supported in the DSL:
.. code:: python
@cute.jit
def foo():
global x
x = 1
foo()
The example above fails to compile because ``global x`` is not supported in the DSL.
**Nonlocal variables**
The use of the ``nonlocal`` keyword is restricted in CuTe DSL. CuTe DSL does not support capturing variables
from an outer (enclosing) scope that is outside of the JIT-compiled function. If you try to use ``nonlocal``
to refer to a variable defined in Python code that is not tracked by current JIT context, a runtime error will be raised.
.. code:: python
def outer():
x = 1
@cute.jit
def inner():
nonlocal x # Not supported
x = 2
inner()
The above code will fail with a runtime error because ``x`` is defined in a scope not managed
by the CuTe DSL's JIT compilation. Nonlocal variables must be managed within the same JIT context;
otherwise, a runtime error will be raised.
Suggestions
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

28
media/docs/pythonDSL/quick_start.rst
View File

@@ -3,19 +3,19 @@
Quick Start Guide
=======================
The CUTLASS DSL 4.4 release currently supports **Linux** and **Python 3.10 - 3.14** only. To install CUTLASS DSLs (limited to CuTe DSL for now), use the following command
Compatibility Requirements
---------------------------------
The CUTLASS DSL 4.4 release currently supports **Linux** and **Python 3.10 - 3.14** only.
Only Linux x86_64 and aarch64 are supported. Additional platform support will be added in future releases.
CUTLASS DSL supports the same NVIDIA driver version as the corresponding `CUDA Toolkit <https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html>`_
(CUDA Toolkit 12.9 or CUDA Toolkit 13.1). Specifically, for 12.9, the driver version must be 575.51.03 or later.
Installation
-----------------------
Before installing the latest version, you need to uninstall any previous CUTLASS DSL Installation.
.. code-block:: bash
pip uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 -y
To ensure compatibility with the examples and code on `GitHub <https://github.com/NVIDIA/cutlass/tree/main>`_,
use the `setup.sh <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/setup.sh>`_ file from the corresponding commit in the repository.
@@ -38,12 +38,10 @@ If you just want to try out the last known stable release of the CUTLASS DSL (ma
pip install nvidia-cutlass-dsl
# For CUDA Toolkit 13.1:
pip install nvidia-cutlass-dsl[cu13]
pip install "nvidia-cutlass-dsl[cu13]"
The ``nvidia-cutlass-dsl`` wheel includes everything needed to generate GPU kernels. It requires
the same NVIDIA driver version as the corresponding `CUDA Toolkit <https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html>`_
(CUDA Toolkit 12.9 or CUDA Toolkit 13.1).
The ``nvidia-cutlass-dsl`` wheel includes everything needed to generate GPU kernels.
Recommended Dependencies
---------------------------------
@@ -52,9 +50,7 @@ To run examples and begin development, we recommend installing:
.. code-block:: bash
pip install torch jupyter
We recommend installing JAX with CUDA support at version 0.8.1 to run JAX examples.
pip install torch jupyter mypy==1.19.1
Recommended Python environment variables for jupyter notebooks
--------------------------------------------------------------