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

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-08 00:16:05 +08:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
examples/python/CuTeDSL/jax/cutlass_call_basic.py
View File

@@ -30,7 +30,6 @@ from functools import partial
import jax
import jax.numpy as jnp
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
import cuda.bindings.driver as cuda
@@ -140,12 +139,12 @@ if __name__ == "__main__":
def run_cutlass_kernel(a, b, x, y):
call = cjax.cutlass_call(
launch_jax_wrapper,
# Jax requires output shapes/dtype information for each output
# 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
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
@@ -165,12 +164,11 @@ if __name__ == "__main__":
# 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),
# Jax requires output shapes/dtype information for each output
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
@@ -191,11 +189,12 @@ if __name__ == "__main__":
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# By default cutlass_call will treat all tensors as dynamic shape.
# 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 pass all tensors shapes and
# layouts known at compile time.
# 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,
@@ -209,19 +208,15 @@ if __name__ == "__main__":
@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),
),
# The modes of the layout for each tensor can be specified using the
# TensorSpec. By default modes will align with the physical layout
# but can be mapped to specific index position. If None is passed
# then the default mode is assumed for that tensor.
#
# Individual static/dynamic settings may also be applied. For example
# a specific tensor can be marked to have static shape.
input_spec=(
cjax.TensorSpec(mode=(1, 0, 2), static=True),
cjax.TensorSpec(mode=(3, 1, 2, 0)),
@@ -245,9 +240,8 @@ if __name__ == "__main__":
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, b.dtype),
),
# Can specify the input tensors that are aliasing outputs of this call.
# To avoid allocating separate output buffers. This is useful for kernels
# that update a tensor.
# 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,

195
examples/python/CuTeDSL/jax/cutlass_call_export.py
View File

@@ -26,45 +26,45 @@
# 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 pytest
from functools import partial
import argparse
import cuda.bindings.driver as cuda
import cutlass
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
from cutlass.jax.testing import create_tensor
"""
Examples of using jax.export APIs with functions using cutlass_call.
This example demonstrates the use of jax.export with CuTe DSL kernel. It assumes
familiarity with CuTe DSL concepts such as layouts and dynamic shapes as well as
Jax's exporting and serialization features:
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
# Run with defaults
python examples/jax/cutlass_call_export.py
python examples/jax/cutlass_call_export.py --M 512 --N 256
# Run with shape (1024, 512)
python examples/jax/cutlass_call_export.py --M 1024 --N 512
# Export with symbolic shapes.
python examples/jax/cutlass_call_export.py --export_symbolic
"""
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()
@@ -84,9 +84,6 @@ def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
@cute.jit
def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
print("mA: ", mA.layout)
print("mB: ", mB.layout)
print("mC: ", mC.layout)
num_threads_per_block = 256
m, n = mA.shape
kernel(mA, mB, mC).launch(
@@ -96,63 +93,100 @@ def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Ten
)
def run_example(M, N, export_symbolic_shapes):
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 ref_f(a, b):
return jax.nn.sigmoid(a + b)
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))
# Symbolic or partially shapes are supported by cutlass_call and cute.Tensor
# This allows export of functions calling Cut eDSL kernels w/o having to re-compile
# the kernel for each new shape.
if export_symbolic_shapes:
a, b = export.symbolic_shape("a, b")
export_shape_dtype = jax.ShapeDtypeStruct((a, b), jnp.float32)
else:
export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
print("Exporting with input signature: ")
print(f"({export_shape_dtype}, {export_shape_dtype})")
# jax.export can be used to export a jit function containing cutlass_call.
# The function get_export_disabled_safety_checks() returns a list of custom
# call targets that are used by cutlass_call not part of Jax's built-in
# list of stable custom calls.
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(export_shape_dtype, export_shape_dtype)
# Serialize the computation to a byte blob.
blob = traced.serialize()
print(f"Serialized computation is {len(blob)} bytes.")
# Deserialize and run
rehydrated = export.deserialize(blob)
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
a = create_tensor((M, N), dtype=jnp.float32, key=a_key)
b = create_tensor((M, N), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
# If the computation was exported with dynamic shapes then we can also
# call it with different shapes. The kernel will not be re-compiled
# even though the shapes are changing.
if export_symbolic_shapes:
a = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
a = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(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__":
@@ -161,8 +195,7 @@ if __name__ == "__main__":
)
parser.add_argument("--M", default=512, type=int)
parser.add_argument("--N", default=256, type=int)
parser.add_argument("--export_symbolic", action="store_true")
args = parser.parse_args()
run_example(args.M, args.N, args.export_symbolic)
run_example(args.M, args.N)
print("PASS")

2
examples/python/CuTeDSL/jax/cutlass_call_sharding.py
View File

@@ -27,14 +27,12 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from functools import partial
import argparse
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
import cutlass.cute as cute
import cutlass.jax as cjax
from cutlass.jax.testing import create_tensor

2
examples/python/CuTeDSL/jax/elementwise_apply_example.py
View File

@@ -30,7 +30,7 @@
import argparse
import operator
from functools import partial
from typing import List, Type
from typing import List
import cuda.bindings.driver as cuda
import cutlass