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

1260
examples/python/CuTeDSL/dsl_tutorials/jax/cute_dsl_jax.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

366
examples/python/CuTeDSL/dsl_tutorials/jax/cute_dsl_jax_kernels.py Normal file
View File

@@ -0,0 +1,366 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
import cuda.bindings.driver as cuda
"""
CuTe DSL kernels used by the ``cute_dsl_jax.ipynb`` notebook.
This module defines GPU kernels written in CuTe DSL (CUTLASS 4.x Python DSL)
that are called from JAX via ``cutlass.jax.cutlass_call``. ``cutlass_call`` is a
JAX primitive that triggers compilation of the kernel during lowering and embeds
it into the HLO computation, so XLA can launch it efficiently without callback
to Python.
Kernels provided:
- ``vector_add`` — element-wise c = a + b (3-D CuTe layout)
- ``saxpy`` — y = alpha * x + y
- ``relu`` — element-wise ReLU with flat indexing
- ``fused_bias_relu`` — fused bias addition + ReLU
- ``gemm`` — tiled matrix multiplication
- ``elementwise_add`` — 2-D element-wise add (flat indexing, ``jax.export``-compatible)
The notebook imports these kernels and wraps each one with ``cutlass_call``
inside ``@jax.jit`` functions. See ``cute_dsl_jax.ipynb`` for usage, validation,
and step-by-step explanations.
This module is imported by the notebook and by ``cute_dsl_jax.py``. It can also
be run directly to validate every kernel:
.. code-block:: bash
# Interactive notebook (recommended for learning)
jupyter lab cute_dsl_jax.ipynb
# Full demo as a standalone script
python cute_dsl_jax_kernels.py
"""
# ------------------------------------------------------------------ #
# Vector Add: c = a + b #
# ------------------------------------------------------------------ #
@cute.kernel
def vector_add_kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
"""Per-thread kernel: each thread adds one element."""
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
cute.autovec_copy(a[None, tidx, bidx], frgA)
cute.autovec_copy(b[None, tidx, bidx], frgB)
frgC.store(frgA.load() + frgB.load())
cute.autovec_copy(frgC, c[None, tidx, bidx])
@cute.jit
def launch_vector_add(
stream: cuda.CUstream,
a: cute.Tensor, b: cute.Tensor, c: cute.Tensor,
):
vector_add_kernel(a, b, c).launch(
grid=[a.shape[-1], 1, 1],
block=[a.shape[-2], 1, 1],
stream=stream,
)
# ------------------------------------------------------------------ #
# SAXPY: y = alpha * x + y #
# ------------------------------------------------------------------ #
@cute.kernel
def saxpy_kernel(x: cute.Tensor, y: cute.Tensor, out: cute.Tensor, alpha: float):
"""SAXPY: out[i] = alpha * x[i] + y[i]."""
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
frgX = cute.make_rmem_tensor(cute.size(x, mode=[0]), x.element_type)
frgY = cute.make_rmem_tensor(cute.size(y, mode=[0]), y.element_type)
frgO = cute.make_rmem_tensor(cute.size(out, mode=[0]), out.element_type)
cute.autovec_copy(x[None, tidx, bidx], frgX)
cute.autovec_copy(y[None, tidx, bidx], frgY)
frgO.store(alpha * frgX.load() + frgY.load())
cute.autovec_copy(frgO, out[None, tidx, bidx])
@cute.jit
def launch_saxpy(
stream: cuda.CUstream,
x: cute.Tensor, y: cute.Tensor, out: cute.Tensor,
*, alpha: float,
):
saxpy_kernel(x, y, out, alpha).launch(
grid=[x.shape[-1], 1, 1],
block=[x.shape[-2], 1, 1],
stream=stream,
)
# ------------------------------------------------------------------ #
# ReLU: out = max(0, x) #
# ------------------------------------------------------------------ #
@cute.kernel
def relu_kernel(x: cute.Tensor, out: cute.Tensor, N: int):
"""Per-thread kernel: each thread computes ReLU of one element."""
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdx, _, _ = cute.arch.block_dim()
idx = bidx * bdx + tidx
if idx < N:
val = x[idx]
out[idx] = cutlass.max(val, cutlass.Float32(0.0))
@cute.jit
def launch_relu(
stream: cuda.CUstream,
x: cute.Tensor, out: cute.Tensor,
*, N: int,
):
BLOCK_SIZE = 256
grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
relu_kernel(x, out, N).launch(
grid=[grid_size, 1, 1],
block=[BLOCK_SIZE, 1, 1],
stream=stream,
)
# ------------------------------------------------------------------ #
# Fused Bias + ReLU: out = max(0, x + bias[col]) #
# ------------------------------------------------------------------ #
@cute.kernel
def fused_bias_relu_kernel(
x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor, N: int, width: int,
):
"""Per-thread: out[i] = max(0, x[i] + bias[i % width])."""
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdx, _, _ = cute.arch.block_dim()
idx = bidx * bdx + tidx
if idx < N:
col = idx % width
val = x[idx] + bias[col]
out[idx] = cutlass.max(val, cutlass.Float32(0.0))
@cute.jit
def launch_fused_bias_relu(
stream: cuda.CUstream,
x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor,
*, N: int, width: int,
):
BLOCK_SIZE = 256
grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
fused_bias_relu_kernel(x, bias, out, N, width).launch(
grid=[grid_size, 1, 1],
block=[BLOCK_SIZE, 1, 1],
stream=stream,
)
# ------------------------------------------------------------------ #
# GEMM: D = A @ B #
# ------------------------------------------------------------------ #
@cute.kernel
def gemm_kernel(
A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
M: int, N: int, K: int, BLOCK_M: int, BLOCK_N: int,
):
"""Tiled GEMM: each thread accumulates output elements."""
tidx, _, _ = cute.arch.thread_idx()
bm, bn, _ = cute.arch.block_idx()
bdx, _, _ = cute.arch.block_dim()
for i in cutlass.range(tidx, BLOCK_M * BLOCK_N, bdx):
row = i // BLOCK_N
col = i % BLOCK_N
m_idx = bm * BLOCK_M + row
n_idx = bn * BLOCK_N + col
if m_idx < M and n_idx < N:
acc = cutlass.Float32(0.0)
for k in cutlass.range(K):
acc += A[m_idx * K + k] * B[k * N + n_idx]
D[m_idx * N + n_idx] = acc
@cute.jit
def launch_gemm(
stream: cuda.CUstream,
A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
*, M: int, N: int, K: int,
):
BLOCK_M, BLOCK_N = 64, 64
grid_m = (M + BLOCK_M - 1) // BLOCK_M
grid_n = (N + BLOCK_N - 1) // BLOCK_N
gemm_kernel(A, B, D, M, N, K, BLOCK_M, BLOCK_N).launch(
grid=[grid_m, grid_n, 1],
block=[256, 1, 1],
stream=stream,
)
# ------------------------------------------------------------------ #
# Element-wise Add (2-D, flat indexing) #
# ------------------------------------------------------------------ #
@cute.kernel
def elementwise_add_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
"""Per-thread kernel: 2-D element-wise add using flat indexing."""
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
thread_idx = bidx * bdim + tidx
m, n = gA.shape
ni = thread_idx % n
mi = thread_idx // n
a_val = gA[mi, ni]
b_val = gB[mi, ni]
gC[mi, ni] = a_val + b_val
@cute.jit
def launch_elementwise_add(
stream: cuda.CUstream,
mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor,
):
num_threads_per_block = 256
m, n = mA.shape
elementwise_add_kernel(mA, mB, mC).launch(
grid=((m * n) // num_threads_per_block, 1, 1),
block=(num_threads_per_block, 1, 1),
stream=stream,
)
# ------------------------------------------------------------------ #
# Self-tests #
# ------------------------------------------------------------------ #
if __name__ == '__main__':
import os
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
import jax
import jax.numpy as jnp
import numpy as np
BLOCK = 256
N_BLOCKS = 4
# ── Vector Add ────────────────────────────────────────────────────
# 3-D CuTe layout: (elems_per_thread, threads_per_block, num_blocks)
a = jax.random.normal(jax.random.PRNGKey(0), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
b = jax.random.normal(jax.random.PRNGKey(1), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
call = cjax.cutlass_call(
launch_vector_add,
output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
use_static_tensors=True,
)
c = jax.jit(call)(a, b)
np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
print('vector_add: PASSED')
# ── SAXPY ─────────────────────────────────────────────────────────
x = jax.random.normal(jax.random.PRNGKey(2), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
y = jax.random.normal(jax.random.PRNGKey(3), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
alpha = 2.5
call = cjax.cutlass_call(
launch_saxpy,
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
use_static_tensors=True,
alpha=alpha,
)
out = jax.jit(call)(x, y)
np.testing.assert_allclose(np.array(out), np.array(alpha * x + y), rtol=1e-5, atol=1e-5)
print('saxpy: PASSED')
# ── ReLU ──────────────────────────────────────────────────────────
N_ELEM = BLOCK * N_BLOCKS
x = jax.random.normal(jax.random.PRNGKey(4), (N_ELEM,), dtype=jnp.float32)
call = cjax.cutlass_call(
launch_relu,
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
N=N_ELEM,
)
out = jax.jit(call)(x)
np.testing.assert_allclose(np.array(out), np.array(jnp.maximum(x, 0)), rtol=1e-5, atol=1e-5)
print('relu: PASSED')
# ── Fused Bias + ReLU ─────────────────────────────────────────────
ROWS, COLS = 16, 64
x = jax.random.normal(jax.random.PRNGKey(5), (ROWS * COLS,), dtype=jnp.float32)
bias = jax.random.normal(jax.random.PRNGKey(6), (COLS,), dtype=jnp.float32)
call = cjax.cutlass_call(
launch_fused_bias_relu,
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
N=ROWS * COLS, width=COLS,
)
out = jax.jit(call)(x, bias)
ref = jnp.maximum(x.reshape(ROWS, COLS) + bias, 0).reshape(-1)
np.testing.assert_allclose(np.array(out), np.array(ref), rtol=1e-5, atol=1e-5)
print('fused_bias_relu: PASSED')
# ── GEMM ──────────────────────────────────────────────────────────
M, N, K = 128, 128, 64
A = jax.random.normal(jax.random.PRNGKey(7), (M * K,), dtype=jnp.float32)
B = jax.random.normal(jax.random.PRNGKey(8), (K * N,), dtype=jnp.float32)
call = cjax.cutlass_call(
launch_gemm,
output_shape_dtype=jax.ShapeDtypeStruct((M * N,), A.dtype),
M=M, N=N, K=K,
)
D = jax.jit(call)(A, B)
ref = A.reshape(M, K) @ B.reshape(K, N)
np.testing.assert_allclose(np.array(D.reshape(M, N)), np.array(ref), rtol=1e-2, atol=1e-2)
print('gemm: PASSED')
# ── Elementwise Add (2-D) ─────────────────────────────────────────
M, N = 16, 256
a = jax.random.normal(jax.random.PRNGKey(9), (M, N), dtype=jnp.float32)
b = jax.random.normal(jax.random.PRNGKey(10), (M, N), dtype=jnp.float32)
call = cjax.cutlass_call(
launch_elementwise_add,
output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
)
c = jax.jit(call)(a, b)
np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
print('elementwise_add: PASSED')
print('\nAll kernels passed.')

254
examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_basic.py Normal file
View File

@@ -0,0 +1,254 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from functools import partial
import jax
import jax.numpy as jnp
import cutlass.cute as cute
import cutlass.jax as cjax
import cuda.bindings.driver as cuda
"""
Examples of calling CuTe DSL from jax.jit function using cutlass_call.
cutlass_call is a Jax primitive the enables calling of CuTe DSL kernels within a
a jit-compiled Jax function. During the lowering process cutlass_call will
trigger compilation of the kernel and embed it into the HLO computation. It can
then be efficiently launched by XLA without callback to Python.
This example assumes familiarity with CuTe DSL concepts such as layouts and
dynamic shapes.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/cutlass_call_basic.py
"""
# This is a typical CuTe DSL kernel function that accepts both tensor and scalar values.
@cute.jit
def launch(
A: cute.Tensor,
B: cute.Tensor,
x: cute.Int32,
y: cute.Int32,
C: cute.Tensor,
D: cute.Tensor,
stream: cuda.CUstream,
):
# Print layouts
print("A layout: ", A.layout)
print("B layout: ", B.layout)
print("C layout: ", C.layout)
print("D layout: ", D.layout)
cute.printf("A layout: {}", A.layout)
cute.printf("B layout: {}", B.layout)
cute.printf("C layout: {}", C.layout)
cute.printf("D layout: {}", D.layout)
cute.printf("")
# Print non-tensor values
print("X is: ", x)
print("Y is: ", y)
cute.printf("X is: {}", x)
cute.printf("Y is: {}", y)
print()
# cutlass_call uses a fixed function signature to pass arguments between Jax and CuTeDSL kernel.
#
# Function Signature Requirement:
# stream, inputs, outputs, *, kwargs...
#
# The first argument must be the CUstream that the kernel is run. This stream is managed by the XLA runtime
# and is necessary to schedule and synchronize launches with the rest of your computation.
#
# The second set of arguments are the Jax arrays for inputs and outputs. Inputs must be passed before
# outputs.
#
# Lastly static arguments (i.e. static_argnums or static_argnames) values are passed as keyword only arguments
# by name.
#
# The the kernel does not match this signature a wrapper functions like the one shown below can be written
# or an inline lambda function can be used to rebind the arguments into the appropriate order.
@cute.jit
def launch_jax_wrapper(
stream: cuda.CUstream,
A: cute.Tensor,
B: cute.Tensor,
C: cute.Tensor,
D: cute.Tensor,
*,
x: cute.Int32,
y: cute.Int32,
):
launch(A, B, x, y, C, D, stream)
@cute.jit
def launch_aliased(
A: cute.Tensor, B: cute.Tensor, x: cute.Int32, y: cute.Int32, stream: cuda.CUstream
):
# Print layouts
print("A layout: ", A.layout)
print("B layout: ", B.layout)
cute.printf("A layout: {}", A.layout)
cute.printf("B layout: {}", B.layout)
cute.printf("")
# Print non-tensor values
print("X is: ", x)
print("Y is: ", y)
cute.printf("X is: {}", x)
cute.printf("Y is: {}", y)
print()
if __name__ == "__main__":
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel(a, b, x, y):
call = cjax.cutlass_call(
launch_jax_wrapper,
# Describe the shape and dtype of each output buffer.
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
# Returned value is a callable to invoke the kernel passing only jax arrays.
return call(a, b)
print("\nExample: example_basic_call_from_jit")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel(A, B, 0, 1)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_lambda(a, b, x, y):
call = cjax.cutlass_call(
# A lambda function may be used to wrap and bind arguments passed by jax
# to the kernel. Alternatively you can wrap using another separate cute.jit
# function.
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
# Returned value is a callable to invoke the kernel passing only jax arrays.
return call(a, b)
print("\nExample: run_cutlass_kernel_lambda")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_lambda(A, B, 1, 2)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_static_shapes(a, b, x, y):
call = cjax.cutlass_call(
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# By default cutlass_call treats all tensors as dynamic shape.
# Dynamic shapes are often expected for kernels so this default ensures
# the broadest support. If you know that a kernel can accept fully static
# tensors then you can enable this flag to compile all tensor shapes and
# layouts as constexpr values known at compile time.
# Individual tensors may opt out via .mark_layout_dynamic().
use_static_tensors=True,
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_static_shapes")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_static_shapes(A, B, 3, 4)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_with_modes(a, b, x, y):
# input_spec and output_spec accept TensorSpec values to attach layout
# metadata to tensors. mode remaps the logical dimension order seen by
# the kernel. static=True compiles that tensor's layout as constexpr.
call = cjax.cutlass_call(
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
input_spec=(
cjax.TensorSpec(mode=(1, 0, 2), static=True),
cjax.TensorSpec(mode=(3, 1, 2, 0)),
),
output_spec=(None, cjax.TensorSpec(mode=(0, 1, 3, 2))),
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_with_modes")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_with_modes(A, B, 5, 6)
@partial(jax.jit, static_argnums=[2, 3], donate_argnums=[0, 1])
def run_cutlass_kernel_aliased_outputs(a, b, x, y):
call = cjax.cutlass_call(
lambda stream, a, b, *, x, y: launch_aliased(a, b, x, y, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, b.dtype),
),
# Map input indices to output indices so XLA can reuse the input
# buffers for the outputs, avoiding extra allocations.
input_output_aliases={0: 0, 1: 1},
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_aliased_outputs")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
A, B = run_cutlass_kernel_aliased_outputs(A, B, 7, 8)

201
examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_export.py Normal file
View File

@@ -0,0 +1,201 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
Examples of using jax.export APIs with functions using cutlass_call.
This example demonstrates three export modes:
1. Concrete shapes -- shapes are fixed constants baked into the export.
2. Unconstrained symbolic shapes ("a, b")
3. Constrained symbolic shapes ("32*M, 16*N")
The JAX function being exported is the same in all three cases; only the
shape specification passed to jax.export differs.
It assumes familiarity with CuTe DSL concepts such as layouts and dynamic shapes
as well as JAX's exporting and serialization features:
https://docs.jax.dev/en/latest/export/index.html#export
To run this example:
.. code-block:: bash
python examples/jax/cutlass_call_export.py --M 512 --N 256
"""
import argparse
import cuda.bindings.driver as cuda
import cutlass.cute as cute
import jax
import jax.numpy as jnp
from jax import export
from cutlass.jax import cutlass_call, get_export_disabled_safety_checks, TensorSpec
from cutlass.jax.testing import create_tensor
# Simple element-wise addition kernel: gC[i,j] = gA[i,j] + gB[i,j]
@cute.kernel
def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
thread_idx = bidx * bdim + tidx
m, n = gA.shape
ni = thread_idx % n
mi = thread_idx // n
a_val = gA[mi, ni]
b_val = gB[mi, ni]
gC[mi, ni] = a_val + b_val
@cute.jit
def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
num_threads_per_block = 256
m, n = mA.shape
kernel(mA, mB, mC).launch(
grid=((m * n) // num_threads_per_block, 1, 1),
block=(num_threads_per_block, 1, 1),
stream=stream,
)
def _export_and_run(f, ref_f, input_shape_dtype, run_shapes):
"""Export f, serialize/deserialize, then run on each shape in run_shapes.
Both inputs (a, b) are assumed to share the same input_shape_dtype.
"""
print(f"Exporting with input signature: ({input_shape_dtype}, {input_shape_dtype})")
# jax.export can be used to export a jit function containing cutlass_call.
# CUTLASS custom call targets are not on JAX's built-in stable custom-call
# allowlist, so we pass them via disabled_checks to suppress that safety check.
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(input_shape_dtype, input_shape_dtype)
blob = traced.serialize()
print(f"Serialized computation is {len(blob)} bytes.")
rehydrated = export.deserialize(blob)
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
for shape in run_shapes:
a = create_tensor(shape, dtype=jnp.float32, key=a_key)
b = create_tensor(shape, dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b)), f"Mismatch at shape {shape}"
print(f" shape {shape}: OK")
def run_example(M, N):
@jax.jit
def ref_f(a, b):
return jax.nn.sigmoid(a + b)
# The same JAX function is used in all three examples below. The export
# mode is determined entirely by the shape spec passed to jax.export.
@jax.jit
def f(a, b):
call = cutlass_call(launch, output_shape_dtype=a)
return jax.nn.sigmoid(call(a, b))
# ── 1. Concrete shapes ────────────────────────────────────────────────────
# Shapes are fixed constants baked into the export. The deserialized
# computation only accepts exactly these dimensions at runtime.
print("\nConcrete shapes:")
input_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
_export_and_run(
f,
ref_f,
input_shape_dtype,
run_shapes=[(M, N)], # concrete exports reject any other shape
)
# ── 2. Unconstrained symbolic shapes ─────────────────────────────────────
# Both dimensions are fully dynamic. The exported computation accepts any
# (M, N) at runtime without recompilation.
print("\nUnconstrained symbolic shapes:")
a_sym, b_sym = export.symbolic_shape("a, b")
input_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)
_export_and_run(
f,
ref_f,
input_shape_dtype,
run_shapes=[(M, N), (M * 2, N * 4), (M * 4, N * 4)],
)
# ── 3. Constrained symbolic shapes (divisibility) ─────────────────────────
# Shapes are declared as multiples of a tile size via TensorSpec.divisibility.
# The symbolic expression "32*M, 16*N" tells jax.export that dim 0 is always
# a multiple of 32 and dim 1 is always a multiple of 16. This lets the
# compiler generate more efficient code (e.g. no remainder handling).
# Runtime shapes must satisfy these divisibility constraints.
print("\nConstrained symbolic shapes:")
@jax.jit
def f_divisible(a, b):
spec = TensorSpec(divisibility=(32, 16))
call = cutlass_call(
launch,
output_shape_dtype=a,
input_spec=(spec, spec),
output_spec=spec,
)
return jax.nn.sigmoid(call(a, b))
m_sym, n_sym = export.symbolic_shape("32*M, 16*N")
input_shape_dtype = jax.ShapeDtypeStruct((m_sym, n_sym), jnp.float32)
_export_and_run(
f_divisible,
ref_f,
input_shape_dtype,
run_shapes=[(M, N), (M * 2, N * 2), (M * 4, N * 4)],
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Demonstration of using jax.export with functions with cutlass_call"
)
parser.add_argument("--M", default=512, type=int)
parser.add_argument("--N", default=256, type=int)
args = parser.parse_args()
run_example(args.M, args.N)
print("PASS")

175
examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_sharding.py Normal file
View File

@@ -0,0 +1,175 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from functools import partial
import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P, AxisType
from jax.experimental.custom_partitioning import custom_partitioning
import cutlass.cute as cute
import cutlass.jax as cjax
from cutlass.jax.testing import create_tensor
import cuda.bindings.driver as cuda
"""
Examples of combining jax.jit, jax.shard_map and custom_partitioning for sharding
and executing kernels across multiple GPU devices.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/cutlass_call_sharding.py
"""
@cute.kernel
def kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
cute.autovec_copy(a[None, tidx, bidx], frgA)
cute.autovec_copy(b[None, tidx, bidx], frgB)
frgC.store(frgA.load() + frgB.load())
cute.autovec_copy(frgC, c[None, tidx, bidx])
@cute.jit
def launch(
stream: cuda.CUstream,
a: cute.Tensor,
b: cute.Tensor,
c: cute.Tensor,
):
cute.printf("a: {}", a.layout)
cute.printf("b: {}", b.layout)
cute.printf("c: {}", c.layout)
kernel(a, b, c).launch(
grid=[a.shape[-1], 1, 1], block=[a.shape[-2], 1, 1], stream=stream
)
def sharded_cutlass_call_impl(a_block, b_block):
"""The sharded implementation that operates on a single device."""
call = cjax.cutlass_call(
launch,
use_static_tensors=True,
output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype),
)
ref_result = a_block + b_block
return call(a_block, b_block), ref_result
@custom_partitioning
def custom_shared_call(a, b):
return sharded_cutlass_call_impl(a, b)
def custom_shared_call_partitioner(mesh, arg_shapes, result_shape):
arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
result_shardings = tuple([arg_shardings[0]] * len(result_shape))
def lower_fn(*args):
return sharded_cutlass_call_impl(*args)
return mesh, lower_fn, result_shardings, arg_shardings
custom_shared_call.def_partition(custom_shared_call_partitioner)
def run_example():
# Create a device mesh with one axis b
ngpu = jax.device_count()
mesh = jax.make_mesh((ngpu,), "b", axis_types=(AxisType.Explicit,))
if ngpu == 1:
print("Note: only 1 GPU was detected.")
# We will shard our 3D tensors over b
sharding = P("b", None, None)
named_sharding = NamedSharding(mesh, sharding)
print("Testing shard_map...")
@partial(
jax.jit, static_argnums=[0, 1], out_shardings=(named_sharding, named_sharding)
)
def allocate_sharded_tensors(shape, dtype):
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
a = create_tensor(shape, dtype, a_key)
b = create_tensor(shape, dtype, b_key)
return a, b
@jax.jit
def compute(a, b):
# This jax.shard_map partitions the cutlass_call over the mesh.
@partial(
jax.shard_map,
mesh=mesh,
in_specs=(sharding, sharding),
out_specs=(sharding, sharding),
)
def sharded_call(a_block, b_block):
return sharded_cutlass_call_impl(a_block, b_block)
return sharded_call(a, b)
# Allocate (32, 16, 64) on each GPU
shape = (32 * ngpu, 16, 64)
dtype = jnp.float32
a, b = allocate_sharded_tensors(shape, dtype)
c, c_ref = compute(a, b)
assert jnp.allclose(c, c_ref)
print("Testing custom_partitioning...")
# Test custom_partitioning implementation which should produce identical results
@jax.jit
def compute_cp(a, b):
return custom_shared_call(a, b)
c, c_ref = compute_cp(a, b)
assert jnp.allclose(c, c_ref)
if __name__ == "__main__":
run_example()
print("PASS")

329
examples/python/CuTeDSL/dsl_tutorials/jax/elementwise_apply_example.py Normal file
View File

@@ -0,0 +1,329 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import operator
from functools import partial
from typing import List
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
"""
An Elementwise Apply Example using CuTe DSL with cutlass.jax.cutlass_call
This example is similar to examples/ampere/elementwise_apply.py but demonstrates
how to run the code in a jax specific way using the cutlass_call primitive. It assumes
familiarity with basic CuTe DSL concepts as well as the cutlass_call primitive.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/elementwise_apply_example.py --M 1024 --N 512 --op add
# Run with multiplication operation
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op mul
# Run with subtraction operation
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op sub
"""
@cute.kernel
def elementwise_apply_kernel(
op: cutlass.Constexpr,
mInputs: List[cute.Tensor],
mC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout, # (tid, vid) -> logic coord
):
tidx, _, _ = cute.arch.thread_idx()
bidx, bidy, _ = cute.arch.block_idx()
###############################################################################
# Slice to local tile of thread block
###############################################################################
blk_crd = ((None, None), (bidx, bidy))
# Leverage the meta-programming capability of the DSL to slice the tensors for each input
# All for loops below on input tensors would be fully unrolled automatically at compile time
# logical coord -> memory address
gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN)
gC = mC[blk_crd] # (TileM, TileN)
gCrd = cC[blk_crd] # (TileM, TileN)
print("[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(gInputs)):
print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}")
print(f"[DSL INFO] gC = {gC.type}")
print(f"[DSL INFO] gCrd = {gCrd.type}")
###############################################################################
# Compose with thread block TV layout to map thread & value indices to memory address
###############################################################################
# (tid, vid) -> memory address
tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
tidfrgC = cute.composition(gC, tv_layout)
tidfrgCrd = cute.composition(gCrd, tv_layout)
# repeat None like vid to remove hierarchy of layout
thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
###############################################################################
# Slice to local tile of thread
###############################################################################
# vid -> address
thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_crd] # (V)
thrCrd = tidfrgCrd[thr_crd]
print("[DSL INFO] Sliced Tensors per thread:")
for i in cutlass.range_constexpr(len(thrInputs)):
print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}")
print(f"[DSL INFO] thrC = {thrC.type}")
print(f"[DSL INFO] thrCrd = {thrCrd.type}")
###############################################################################
# Compute predicate for out of boundary checks
###############################################################################
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] frgPred = {frgPred.type}")
for i in cutlass.range_constexpr(cute.size(frgPred)):
frgPred[i] = cute.elem_less(thrCrd[i], shape)
# if tidx == 0 and bidx == 0:
# cute.print_tensor(frgPred)
##########################################################
# Load data and compute result
##########################################################
# Load data before use. The compiler will optimize the copy and load
# operations to convert some memory ld/st into register uses.
result = op(*[thrInput.load() for thrInput in thrInputs])
thrC.store(result)
@cute.jit
def elementwise_apply(
op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
):
"""CUDA kernel applying binary operator on each element of two n-D input tensors in
CuTe Python and store to result tensor.
:param op: Binary operator or lambda function to apply element-wise
:type op: cutlass.Constexpr
:param a: First input tensor
:type a: cute.Tensor
:param b: Second input tensor
:type b: cute.Tensor
:param result: Output tensor to store the results of op(a, b)
:type result: cute.Tensor
:return: None
:rtype: None
"""
# Baseline: naive TV layout
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (512, 4) tile
# * tidx maps to mode-0 but input layout is contiguous on mode-1, performance will be bad
# tv_layout = cute.make_layout((128, (4, 4)), stride=(4, (512, 1)))
# cta_tiler = (512, 4)
# Opt-1: better TV layout with better 1D thread layout (SOL with 1D thread layout)
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (4, 512) tile
# * tidx maps to mode-1 which is leading mode of input tensor for coalesced load
# tv_layout = cute.make_layout((128, (4, 4)), stride=(16, (4, 1)))
# cta_tiler = (4, 512)
# Opt-2: 2D tile but worse
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (128, 16) logical tile
# * V layout is bad as contiguous mode is not on right-most
# * `cute.copy` only supports vectorize when stride-1 of v-layout on right-most )
# tv_layout = cute.make_layout(((32, 4), (4, 4)), stride=((4, 512), (1, 128)))
# cta_tiler = (128, 16)
# Opt-3: SOL with 2D thread tile
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (64, 256) logical tile
# * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
# Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
coalesced_ldst_bytes = 16
# Compile time validation: expect same element type for all input tensors
assert all(t.element_type == inputs[0].element_type for t in inputs)
dtype = inputs[0].element_type
thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
val_layout = cute.recast_layout(dtype.width, 8, val_layout)
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print("[DSL INFO] Input Tensors:")
for i, t in enumerate(inputs):
print(f"[DSL INFO] inputs{i} = {t}")
print(f"[DSL INFO] result = {result}")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
print("[DSL INFO] Tiled Tensors:")
mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
# ((TileM, TileN), (RestM, RestN))
mC = cute.zipped_divide(result, tiler_mn)
# (RestM, RestN) -> (RestN, RestM)
remap_block = cute.make_ordered_layout(
cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
)
for i, t in enumerate(mInputs):
print(f"[DSL INFO] gInputs{i} = {mInputs[i]}")
mInputs[i] = cute.composition(t, (None, remap_block))
print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}")
mC = cute.composition(mC, (None, remap_block))
print(f"[DSL INFO] gC = {mC}")
idC = cute.make_identity_tensor(result.shape)
cC = cute.zipped_divide(idC, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cC}")
# Launch the kernel asynchronously
# Group input tensors into a list as a single argument
elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
# Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
grid=cute.product_each(mC.shape[1]),
block=[cute.size(tv_layout, mode=[0]), 1, 1],
stream=stream,
)
@cutlass.dsl_user_op
def leaky_relu(x, alpha, *, loc=None, ip=None):
return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
def leaky_relu_ref(x, alpha):
import jax.numpy as jnp
return jnp.where(x > 0, x, alpha * x)
def run_and_verify(op, M, N, dtype, skip_ref_check=False):
import jax
import jax.numpy as jnp
import cutlass.jax as cjax
import cutlass.jax.testing as testing
if op == "leaky_relu":
op = partial(leaky_relu, alpha=0.01)
ref_op = partial(leaky_relu_ref, alpha=0.01)
num_inputs = 1
else:
op = getattr(operator, op)
ref_op = op
num_inputs = 2
# This jax function is transformed using jax.jit to compile its contents
# into an efficient HLO executable.
@partial(jax.jit, static_argnums=[1])
def jax_function(inputs, op):
call = cjax.cutlass_call(
# Bind jax arguments to kernel signature
lambda stream, inputs, output, *, op: elementwise_apply(
op, inputs, output, stream
),
# Specify output shape/dtype of result
output_shape_dtype=jax.ShapeDtypeStruct(inputs[0].shape, inputs[0].dtype),
# Pass static/constexpr values as kwargs
op=op,
)
# Call the kernel!
return call(inputs)
@partial(jax.jit, static_argnums=[1])
def jax_ref_function(inputs, op):
return op(*inputs)
print("\nRunning Elementwise Apply test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
jax_dtype = cjax.cutlass_to_jax_dtype(dtype)
keys = jax.random.split(jax.random.key(1435), num_inputs)
inputs = [testing.create_tensor((M, N), jax_dtype, key) for key in keys]
print("Input tensor shapes:")
for i in range(num_inputs):
print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
epsilon = 1.2
if op in (operator.truediv, operator.floordiv):
inputs[1] = jnp.where(inputs[1] == 0, epsilon, inputs[1])
# Call the jax.jit function which will compile the kernel
c = jax_function(inputs, op)
if not skip_ref_check:
print("Executing elementwise apply kernel...")
c = jax_function(inputs, op)
print("Verifying results...")
assert jnp.allclose(ref_op(*inputs), c)
print("Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3]}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Demonstration of calling a kernel with cutlass_call"
)
parser.add_argument("--M", default=4096, type=int)
parser.add_argument("--N", default=4096, type=int)
parser.add_argument("--op", default="add", type=str)
parser.add_argument("--skip_ref_check", action="store_true")
args = parser.parse_args()
run_and_verify(
args.op,
args.M,
args.N,
dtype=cutlass.Float32,
skip_ref_check=args.skip_ref_check,
)
print("\nPASS")