59 KiB
59 KiB
In [ ]:
!nvidia-smiIn [ ]:
import subprocess
def get_compute_capability():
"""Query the compute capability of the first visible GPU."""
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], text=True
)
major, minor = out.strip().split("\n")[0].split(".")
return int(major), int(minor)
SM_MAJOR, SM_MINOR = get_compute_capability()
print(f"Detected compute capability: SM {SM_MAJOR}.{SM_MINOR}")
if SM_MAJOR < 8:
print("WARNING: CuTe DSL requires SM 8.0+ (Ampere or newer).")
print("Some examples may not run on this GPU.")
else:
print("GPU is compatible with CuTe DSL.")In [ ]:
%pip install "nvidia-cutlass-dsl[cu13]==4.4.0.dev1" --quietIn [ ]:
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" # suppress TF/XLA info & warnings
os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/local/cuda"
import cutlass
from importlib.metadata import version as _pkg_version
print(f"CUTLASS version: {_pkg_version('nvidia-cutlass-dsl')}")
import jax
import jax.numpy as jnp
import numpy as np
print(f"JAX version: {jax.__version__}")
print(f"JAX devices: {jax.devices()}")In [ ]:
import cutlass.jax as cjax
from cute_dsl_jax_kernels import (
launch_vector_add, launch_saxpy, launch_gemm,
launch_relu, launch_fused_bias_relu,
launch_elementwise_add,
)
print("Imported: launch_vector_add, launch_saxpy, launch_gemm, launch_relu, launch_fused_bias_relu, launch_elementwise_add")In [ ]:
BLOCK = 256 # threads per block for vector add: 256 is a practical default:
# large enough to expose parallelism, small enough to scale
# well across different GPUs, and aligned with the hardware’s
# 32-thread warp execution model.
@jax.jit
def jax_vector_add(a, b):
"""JAX-compatible vector add using CUTLASS kernel."""
N = a.shape[0]
padded = ((N + BLOCK - 1) // BLOCK) * BLOCK
a_pad = jnp.pad(a, (0, padded - N))
b_pad = jnp.pad(b, (0, padded - N))
# Reshape to (1, BLOCK, num_blocks) for the CuTe kernel
a_3d = a_pad.reshape(1, BLOCK, padded // BLOCK)
b_3d = b_pad.reshape(1, BLOCK, padded // BLOCK)
call = cjax.cutlass_call(
launch_vector_add,
output_shape_dtype=jax.ShapeDtypeStruct(a_3d.shape, a_3d.dtype),
use_static_tensors=True,
)
c_3d = call(a_3d, b_3d)
return c_3d.reshape(-1)[:N]
print("jax_vector_add defined.")In [ ]:
# Test vector add
N = 1024
key = jax.random.PRNGKey(0)
a = jax.random.normal(key, (N,), dtype=jnp.float32)
b = jax.random.normal(jax.random.PRNGKey(1), (N,), dtype=jnp.float32)
c = jax_vector_add(a, b)
c_ref = a + b
np.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)
print(f"Vector Add PASSED (N={N})")
print(f" Max error: {float(jnp.max(jnp.abs(c - c_ref))):.2e}")In [ ]:
from functools import partial
@partial(jax.jit, static_argnums=(2,))
def jax_saxpy(x, y, alpha=2.0):
"""JAX-compatible SAXPY using CUTLASS kernel."""
N = x.shape[0]
padded = ((N + BLOCK - 1) // BLOCK) * BLOCK
x_pad = jnp.pad(x, (0, padded - N))
y_pad = jnp.pad(y, (0, padded - N))
x_3d = x_pad.reshape(1, BLOCK, padded // BLOCK)
y_3d = y_pad.reshape(1, BLOCK, padded // BLOCK)
call = cjax.cutlass_call(
launch_saxpy,
output_shape_dtype=jax.ShapeDtypeStruct(x_3d.shape, x_3d.dtype),
use_static_tensors=True,
alpha=alpha,
)
out_3d = call(x_3d, y_3d)
return out_3d.reshape(-1)[:N]
print("jax_saxpy defined.")In [ ]:
# Test SAXPY
N = 2048
ALPHA = 2.5
key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (N,), dtype=jnp.float32)
y = jax.random.normal(jax.random.PRNGKey(43), (N,), dtype=jnp.float32)
result = jax_saxpy(x, y, alpha=ALPHA)
ref = ALPHA * x + y
np.testing.assert_allclose(np.array(result), np.array(ref), rtol=1e-5)
print(f"SAXPY PASSED (N={N}, alpha={ALPHA})")
print(f" Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}")In [ ]:
@jax.jit
def jax_relu(x):
"""JAX-compatible ReLU using CUTLASS kernel."""
N = x.size
x_flat = x.reshape(-1)
call = cjax.cutlass_call(
launch_relu,
output_shape_dtype=jax.ShapeDtypeStruct(x_flat.shape, x_flat.dtype),
N=N,
)
out_flat = call(x_flat)
return out_flat.reshape(x.shape)
print("jax_relu defined.")In [ ]:
# Test ReLU
N = 2048
key = jax.random.PRNGKey(7)
x = jax.random.normal(key, (N,), dtype=jnp.float32)
result = jax_relu(x)
ref = jax.nn.relu(x)
np.testing.assert_allclose(np.array(result), np.array(ref), rtol=1e-5)
print(f"ReLU PASSED (N={N})")
print(f" Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}")
print(f" Sample: x[:6] = {x[:6]}")
print(f" out[:6] = {result[:6]}")In [ ]:
from functools import partial
@partial(jax.jit, static_argnums=(2,))
def jax_fused_bias_relu(x, bias, width):
"""JAX-compatible fused Bias+ReLU using CUTLASS kernel.
Args:
x: Input matrix of shape (batch, width), flattened to 1-D for the kernel.
bias: Bias vector of shape (width,).
width: Number of columns (static, passed as constexpr to the kernel).
"""
N = x.size
x_flat = x.reshape(-1)
call = cjax.cutlass_call(
launch_fused_bias_relu,
output_shape_dtype=jax.ShapeDtypeStruct(x_flat.shape, x_flat.dtype),
N=N, width=width,
)
out_flat = call(x_flat, bias)
return out_flat.reshape(x.shape)
print("jax_fused_bias_relu defined.")In [ ]:
# Test Fused Bias+ReLU
BATCH, WIDTH = 64, 512
key = jax.random.PRNGKey(99)
x = jax.random.normal(key, (BATCH, WIDTH), dtype=jnp.float32)
bias = jax.random.normal(jax.random.PRNGKey(100), (WIDTH,), dtype=jnp.float32)
result = jax_fused_bias_relu(x, bias, WIDTH)
ref = jnp.maximum(0, x + bias[None, :])
np.testing.assert_allclose(np.array(result), np.array(ref), rtol=1e-5)
print(f"Fused Bias+ReLU PASSED (batch={BATCH}, width={WIDTH})")
print(f" Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}")In [ ]:
@jax.jit
def jax_cutlass_gemm(a, b):
"""JAX wrapper for the CUTLASS GEMM kernel."""
M, K = a.shape
_, N = b.shape
a_flat = a.reshape(-1)
b_flat = b.reshape(-1)
call = cjax.cutlass_call(
launch_gemm,
output_shape_dtype=jax.ShapeDtypeStruct((M * N,), a.dtype),
M=M, N=N, K=K,
)
d_flat = call(a_flat, b_flat)
return d_flat.reshape(M, N)
print("jax_cutlass_gemm defined.")In [ ]:
# Test GEMM
M, N, K = 256, 256, 128
key = jax.random.PRNGKey(0)
A = jax.random.normal(key, (M, K), dtype=jnp.float32)
B = jax.random.normal(jax.random.PRNGKey(1), (K, N), dtype=jnp.float32)
D = jax_cutlass_gemm(A, B)
D_ref = jnp.matmul(A, B)
np.testing.assert_allclose(np.array(D), np.array(D_ref), rtol=1e-2, atol=1e-2)
print(f"GEMM PASSED (M={M}, N={N}, K={K})")
print(f" Max error: {float(jnp.max(jnp.abs(D - D_ref))):.2e}")In [ ]:
import time
M, N, K = 512, 512, 512
A = jax.random.normal(jax.random.PRNGKey(0), (M, K), dtype=jnp.float32)
B = jax.random.normal(jax.random.PRNGKey(1), (K, N), dtype=jnp.float32)
# Warmup
_ = jax_cutlass_gemm(A, B).block_until_ready()
_ = jnp.matmul(A, B).block_until_ready()
NUM_RUNS = 20
# Time CUTLASS GEMM
start = time.perf_counter()
for _ in range(NUM_RUNS):
_ = jax_cutlass_gemm(A, B).block_until_ready()
cutlass_time = (time.perf_counter() - start) / NUM_RUNS
# Time JAX matmul
start = time.perf_counter()
for _ in range(NUM_RUNS):
_ = jnp.matmul(A, B).block_until_ready()
jax_time = (time.perf_counter() - start) / NUM_RUNS
print(f"Matrix size: {M}x{N}x{K}")
print(f"CUTLASS GEMM: {cutlass_time*1000:.3f} ms")
print(f"JAX jnp.matmul: {jax_time*1000:.3f} ms")
print(f"Ratio (CUTLASS / JAX): {cutlass_time / jax_time:.2f}x")
print()
print("Note: Our simple tiled kernel is not expected to beat cuBLAS.")
print("CuTe DSL's value is in specialized kernels cuBLAS doesn't provide.")In [ ]:
import warnings
from functools import partial
from jax.sharding import PartitionSpec as P
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
from jax.experimental.shard_map import shard_map
num_devices = len(jax.devices())
print(f"Number of devices: {num_devices}")
if num_devices > 1:
mesh = jax.make_mesh((num_devices,), "x")
# Kernel expects 3-D tensors: (elems_per_thread, threads, blocks)
# Shard along the blocks axis (last dim)
sharding = P(None, None, "x")
@jax.jit
def sharded_vector_add(a, b):
@partial(
shard_map,
mesh=mesh,
in_specs=(sharding, sharding),
out_specs=sharding,
)
def _add(a_shard, b_shard):
call = cjax.cutlass_call(
launch_vector_add,
output_shape_dtype=jax.ShapeDtypeStruct(
a_shard.shape, a_shard.dtype
),
use_static_tensors=True,
)
return call(a_shard, b_shard)
return _add(a, b)
# Create 3-D tensors: (1, 256, total_blocks) with total_blocks divisible by device count
blocks_per_device = 16
total_blocks = blocks_per_device * num_devices
shape = (1, BLOCK, total_blocks)
a_m = jax.random.normal(jax.random.PRNGKey(10), shape, dtype=jnp.float32)
b_m = jax.random.normal(jax.random.PRNGKey(11), shape, dtype=jnp.float32)
c_m = sharded_vector_add(a_m, b_m)
np.testing.assert_allclose(np.array(c_m), np.array(a_m + b_m), rtol=1e-5)
N_total = int(np.prod(shape))
print(f"Sharded Vector Add PASSED across {num_devices} devices (N={N_total})")
else:
print("Only 1 device detected. Skipping multi-GPU example.")
print("On a multi-GPU system, shard_map distributes CUTLASS kernels across devices.")In [ ]:
from jax import export
from cutlass.jax import get_export_disabled_safety_checks
# Define a function that uses a CUTLASS kernel + JAX ops.
# We use launch_elementwise_add which accepts 2-D tensors directly
# with flat indexing — compatible with jax.export's tracing.
@jax.jit
def f(a, b):
call = cjax.cutlass_call(launch_elementwise_add, output_shape_dtype=a)
return jax.nn.sigmoid(call(a, b))
# Reference implementation (pure JAX)
@jax.jit
def ref_f(a, b):
return jax.nn.sigmoid(a + b)
# --- Export with concrete shapes ---
M, N = 512, 256
export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
print(f"Exporting with input signature: ({export_shape_dtype}, {export_shape_dtype})")
# Export the function — get_export_disabled_safety_checks() tells JAX
# that CUTLASS custom call targets are safe to include
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(export_shape_dtype, export_shape_dtype)
# Serialize to a byte blob
blob = traced.serialize()
print(f"Serialized computation: {len(blob):,} bytes")
# Deserialize and run — this works independently of the original function
rehydrated = export.deserialize(blob)
key = jax.random.PRNGKey(1123)
a = jax.random.normal(key, (M, N), dtype=jnp.float32)
b = jax.random.normal(jax.random.PRNGKey(456), (M, N), dtype=jnp.float32)
c = rehydrated.call(a, b)
c_ref = ref_f(a, b)
np.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)
print(f"Export + Deserialize PASSED (M={M}, N={N})")
print(f" Max error: {float(jnp.max(jnp.abs(c - c_ref))):.2e}")In [ ]:
# --- Export with symbolic shapes ---
a_sym, b_sym = export.symbolic_shape("a, b")
symbolic_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)
print(f"Exporting with symbolic signature: ({symbolic_shape_dtype}, {symbolic_shape_dtype})")
exported_sym = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced_sym = exported_sym(symbolic_shape_dtype, symbolic_shape_dtype)
blob_sym = traced_sym.serialize()
print(f"Serialized computation: {len(blob_sym):,} bytes")
rehydrated_sym = export.deserialize(blob_sym)
# Call with different shapes — no recompilation needed.
# The same serialized blob works for any (M, N) where M*N is a
# multiple of the kernel's block size (256).
for shape in [(512, 256), (1024, 512), (2048, 1024)]:
a = jax.random.normal(jax.random.PRNGKey(42), shape, dtype=jnp.float32)
b = jax.random.normal(jax.random.PRNGKey(43), shape, dtype=jnp.float32)
c = rehydrated_sym.call(a, b)
c_ref = ref_f(a, b)
np.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)
print(f" Symbolic export PASSED for shape {shape}")
print("All symbolic shape tests passed.")