Skip to content

Commit 93cb8ff

Browse files
drisspgdrisspg
committed
Add cute cache
stack-info: PR: #57, branch: drisspg/stack/15
1 parent 3bd431f commit 93cb8ff

4 files changed

Lines changed: 336 additions & 29 deletions

File tree

‎test/test_cute.py‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
#!/usr/bin/env python3
2+
"""Test script to demonstrate CUTE kernel caching"""
3+
4+
import torch
5+
from operator import add, mul
6+
import cutlass.cute as cute
7+
from cutlass.cute.runtime import from_dlpack
8+
import pytest
9+
10+
from transformer_nuggets.cute import cute_compile_and_cache, get_cache_stats, clear_cute_cache
11+
12+
13+
@cute.kernel
14+
def simple_add_kernel(
15+
gA: cute.Tensor,
16+
gB: cute.Tensor,
17+
gC: cute.Tensor,
18+
):
19+
tidx, _, _ = cute.arch.thread_idx()
20+
bidx, _, _ = cute.arch.block_idx()
21+
bdim, _, _ = cute.arch.block_dim()
22+
23+
thread_idx = bidx * bdim + tidx
24+
25+
m, n = gA.shape
26+
ni = thread_idx % n
27+
mi = thread_idx // n
28+
29+
a_val = gA[mi, ni]
30+
b_val = gB[mi, ni]
31+
32+
gC[mi, ni] = a_val + b_val
33+
34+
35+
@cute.kernel
36+
def simple_mul_kernel(
37+
gA: cute.Tensor,
38+
gB: cute.Tensor,
39+
gC: cute.Tensor,
40+
):
41+
tidx, _, _ = cute.arch.thread_idx()
42+
bidx, _, _ = cute.arch.block_idx()
43+
bdim, _, _ = cute.arch.block_dim()
44+
45+
thread_idx = bidx * bdim + tidx
46+
47+
m, n = gA.shape
48+
ni = thread_idx % n
49+
mi = thread_idx // n
50+
51+
a_val = gA[mi, ni]
52+
b_val = gB[mi, ni]
53+
54+
gC[mi, ni] = a_val * b_val
55+
56+
57+
def cached_elementwise(
58+
op,
59+
mA: cute.Tensor,
60+
mB: cute.Tensor,
61+
mC: cute.Tensor,
62+
):
63+
# Define the kernel with @cute.jit
64+
@cute.jit
65+
def add_kernel(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
66+
num_threads_per_block = 256
67+
m, n = mA.shape
68+
69+
simple_add_kernel(mA, mB, mC).launch(
70+
grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
71+
)
72+
73+
@cute.jit
74+
def mul_kernel(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
75+
num_threads_per_block = 256
76+
m, n = mA.shape
77+
78+
simple_mul_kernel(mA, mB, mC).launch(
79+
grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
80+
)
81+
82+
# Use explicit caching based on operation
83+
if op == add:
84+
compiled_kernel = cute_compile_and_cache(add_kernel, mA, mB, mC)
85+
return compiled_kernel(mA, mB, mC)
86+
elif op == mul:
87+
compiled_kernel = cute_compile_and_cache(mul_kernel, mA, mB, mC)
88+
return compiled_kernel(mA, mB, mC)
89+
else:
90+
raise ValueError(f"Unsupported operation: {op}")
91+
92+
93+
def test():
94+
print("CUTE Kernel Caching Demo")
95+
print("=" * 50)
96+
97+
# Create test tensors
98+
M, N = 1024, 1024
99+
a = torch.randn(M, N, device="cuda", dtype=torch.float16)
100+
b = torch.randn(M, N, device="cuda", dtype=torch.float16)
101+
c = torch.zeros(M, N, device="cuda", dtype=torch.float16)
102+
103+
# Convert to CUTE tensors
104+
a_ = from_dlpack(a, assumed_align=16)
105+
b_ = from_dlpack(b, assumed_align=16)
106+
c_ = from_dlpack(c, assumed_align=16)
107+
108+
print("\n1. First call with 'add' - should compile:")
109+
cached_elementwise(add, a_, b_, c_)
110+
torch.testing.assert_close(c, a + b)
111+
print("✓ Addition result correct")
112+
113+
print("\n2. Second call with 'add' - should hit cache:")
114+
c.zero_()
115+
cached_elementwise(add, a_, b_, c_)
116+
117+
print("\n3. First call with 'mul' - should compile (different op):")
118+
c.zero_()
119+
cached_elementwise(mul, a_, b_, c_)
120+
torch.testing.assert_close(c, a * b)
121+
print("✓ Multiplication result correct")
122+
123+
print("\n4. Different tensor sizes - should compile:")
124+
a2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
125+
b2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
126+
c2 = torch.zeros(512, 512, device="cuda", dtype=torch.float16)
127+
128+
a2_ = from_dlpack(a2, assumed_align=16)
129+
b2_ = from_dlpack(b2, assumed_align=16)
130+
c2_ = from_dlpack(c2, assumed_align=16)
131+
132+
cached_elementwise(add, a2_, b2_, c2_)
133+
134+
print("\n5. Same size as #4 - should hit cache:")
135+
cached_elementwise(add, a2_, b2_, c2_)
136+
137+
print("\n" + "=" * 50)
138+
print("Final Cache Statistics:")
139+
stats = get_cache_stats()
140+
print(f" Total calls: {stats['total']}")
141+
print(f" Cache hits: {stats['hits']}")
142+
print(f" Cache misses: {stats['misses']}")
143+
print(f" Hit rate: {stats['hit_rate']:.2%}")
144+
print(f" Unique kernels cached: {stats['cache_size']}")
145+
146+
print("\nClearing cache...")
147+
clear_cute_cache()
148+
stats = get_cache_stats()
149+
print(f" Cache size after clear: {stats['cache_size']}")
150+
151+
152+
if __name__ == "__main__":
153+
pytest.main([__file__])

‎transformer_nuggets/cute/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from transformer_nuggets.cute.cache import (
2+
cute_compile_and_cache,
3+
clear_cute_cache,
4+
get_cache_stats,
5+
)

‎transformer_nuggets/cute/add.py‎

Lines changed: 66 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@
77
from cutlass.cute.runtime import from_dlpack
88

99
from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds
10+
from transformer_nuggets.cute.cache import cute_compile_and_cache, get_cache_stats
11+
12+
13+
# init_logging()
1014

1115

1216
@cute.kernel
@@ -31,6 +35,17 @@ def naive_elementwise_add_kernel(
3135
gC[mi, ni] = a_val + b_val
3236

3337

38+
@cute.jit
39+
def naive_elementwise_add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
40+
num_threads_per_block = 256
41+
m, n = mA.shape
42+
43+
kernel = naive_elementwise_add_kernel(mA, mB, mC)
44+
kernel.launch(
45+
grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
46+
)
47+
48+
3449
@cute.kernel
3550
def elementwise_apply_kernel(
3651
op: cutlass.Constexpr,
@@ -62,23 +77,7 @@ def elementwise_apply_kernel(
6277

6378

6479
@cute.jit
65-
def naive_elementwise_add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
66-
num_threads_per_block = 256
67-
m, n = mA.shape
68-
69-
kernel = naive_elementwise_add_kernel(mA, mB, mC)
70-
kernel.launch(
71-
grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
72-
)
73-
74-
75-
@cute.jit
76-
def elementwise_op(
77-
op: cutlass.Constexpr,
78-
mA: cute.Tensor,
79-
mB: cute.Tensor,
80-
mC: cute.Tensor,
81-
):
80+
def elem_kernel(op: cutlass.Constexpr, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
8281
thr_layout = cute.make_layout((4, 32), stride=(32, 1))
8382
val_layout = cute.make_layout((4, 8), stride=(8, 1))
8483
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
@@ -93,13 +92,41 @@ def elementwise_op(
9392
)
9493

9594

96-
def benchmark(callable, *, num_warmups=5, num_iterations=200):
95+
def elementwise_op(
96+
op: cutlass.Constexpr,
97+
a: torch.Tensor,
98+
b: torch.Tensor,
99+
c: torch.Tensor,
100+
*,
101+
assumed_align: int = 16,
102+
):
103+
"""
104+
Apply an elementwise operation using CUTE kernels.
105+
106+
Args:
107+
op: Cutlass constexpr operation (e.g., add, mul)
108+
a: Input tensor A (PyTorch tensor)
109+
b: Input tensor B (PyTorch tensor)
110+
c: Output tensor C (PyTorch tensor, will be modified in-place)
111+
assumed_align: Memory alignment assumption for dlpack conversion
112+
"""
113+
# Convert PyTorch tensors to CUTE tensors
114+
mA = from_dlpack(a, assumed_align=assumed_align)
115+
mB = from_dlpack(b, assumed_align=assumed_align)
116+
mC = from_dlpack(c, assumed_align=assumed_align)
117+
118+
# Compile and execute the kernel
119+
compiled_kernel = cute_compile_and_cache(elem_kernel, op, mA, mB, mC)
120+
return compiled_kernel(mA, mB, mC)
121+
122+
123+
def benchmark(callable, tensor_a, *, num_warmups=5, num_iterations=200):
97124
time = benchmark_cuda_function_in_microseconds(callable)
98125
avg_time = time / 1e3
99126

100127
print(f"Average execution time: {avg_time:.4f} ms")
101128

102-
total_bytes = 3 * a.numel() * a.element_size()
129+
total_bytes = 3 * tensor_a.numel() * tensor_a.element_size()
103130
throughput_gb_s = total_bytes / (avg_time / 1000) / 1e9
104131
print(f"Throughput: {throughput_gb_s:.2f} GB/s")
105132
print()
@@ -112,28 +139,38 @@ def benchmark(callable, *, num_warmups=5, num_iterations=200):
112139
b = torch.randn(M, N, device="cuda", dtype=torch.float16)
113140
c = torch.zeros(M, N, device="cuda", dtype=torch.float16)
114141

115-
a_ = from_dlpack(a, assumed_align=16)
116-
b_ = from_dlpack(b, assumed_align=16)
117-
c_ = from_dlpack(c, assumed_align=16)
118-
119-
elementwise_op(mul, a_, b_, c_)
142+
# Test the new API that takes regular PyTorch tensors
143+
elementwise_op(mul, a, b, c)
120144
torch.testing.assert_close(c, mul(a, b))
121145
print("✓ Multiplication test passed")
122146

123147
c.zero_()
124148

149+
# For naive kernel, we still need CUTE tensors
150+
a_ = from_dlpack(a, assumed_align=16)
151+
b_ = from_dlpack(b, assumed_align=16)
152+
c_ = from_dlpack(c, assumed_align=16)
125153
naive_elementwise_add_compiled = cute.compile(naive_elementwise_add, a_, b_, c_)
126-
elementwise_add_compiled = cute.compile(elementwise_op, add, a_, b_, c_)
154+
155+
# Note: elementwise_op is now cached automatically!
156+
# First call will compile, subsequent calls use cache
127157

128158
print("\n" + "=" * 50)
129159
print("BENCHMARK RESULTS")
130160
print("=" * 50 + "\n")
131161

132162
print("1. Naive elementwise add kernel:")
133-
benchmark(partial(naive_elementwise_add_compiled, a_, b_, c_))
163+
benchmark(partial(naive_elementwise_add_compiled, a_, b_, c_), a)
134164

135165
print("2. PyTorch add (baseline):")
136-
benchmark(partial(torch.add, a, b, out=c))
166+
benchmark(partial(torch.add, a, b, out=c), a)
167+
168+
print("3. Optimized elementwise add kernel (with caching):")
169+
benchmark(partial(elementwise_op, add, a, b, c, assumed_align=128), a)
137170

138-
print("3. Optimized elementwise add kernel:")
139-
benchmark(partial(elementwise_add_compiled, a_, b_, c_))
171+
print("\nCache Statistics:")
172+
stats = get_cache_stats()
173+
print(f" Cache hits: {stats['hits']}")
174+
print(f" Cache misses: {stats['misses']}")
175+
print(f" Hit rate: {stats['hit_rate']:.2%}")
176+
print(f" Cache size: {stats['cache_size']}")

0 commit comments

Comments
 (0)