Add cute cache

stack-info: PR: #57, branch: drisspg/stack/15

Author: drisspg

test_cute_cache.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Test script to demonstrate CUTE kernel caching"""
+
+import torch
+from operator import add, mul
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+from transformer_nuggets.cute import cute_compile_and_cache, get_cache_stats, clear_cute_cache
+
+
+@cute.kernel
+def simple_elementwise_kernel(
+    op: cutlass.Constexpr,
+    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] = op(a_val, b_val)
+
+
+def cached_elementwise(
+    op: cutlass.Constexpr,
+    mA: cute.Tensor,
+    mB: cute.Tensor,
+    mC: cute.Tensor,
+):
+    # Define the kernel with @cute.jit
+    @cute.jit
+    def kernel(op, mA, mB, mC):
+        num_threads_per_block = 256
+        m, n = mA.shape
+
+        kernel_op = simple_elementwise_kernel(op, mA, mB, mC)
+        kernel_op.launch(
+            grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
+        )
+
+    # Use explicit caching
+    compiled_kernel = cute_compile_and_cache(kernel, op, mA, mB, mC)
+    return compiled_kernel(op, mA, mB, mC)
+
+
+def main():
+    print("CUTE Kernel Caching Demo")
+    print("=" * 50)
+
+    # Create test tensors
+    M, N = 1024, 1024
+    a = torch.randn(M, N, device="cuda", dtype=torch.float16)
+    b = torch.randn(M, N, device="cuda", dtype=torch.float16)
+    c = torch.zeros(M, N, device="cuda", dtype=torch.float16)
+
+    # Convert to CUTE tensors
+    a_ = from_dlpack(a, assumed_align=16)
+    b_ = from_dlpack(b, assumed_align=16)
+    c_ = from_dlpack(c, assumed_align=16)
+
+    print("\n1. First call with 'add' - should compile:")
+    cached_elementwise(add, a_, b_, c_)
+    torch.testing.assert_close(c, a + b)
+    print("✓ Addition result correct")
+
+    print("\n2. Second call with 'add' - should hit cache:")
+    c.zero_()
+    cached_elementwise(add, a_, b_, c_)
+
+    print("\n3. First call with 'mul' - should compile (different op):")
+    c.zero_()
+    cached_elementwise(mul, a_, b_, c_)
+    torch.testing.assert_close(c, a * b)
+    print("✓ Multiplication result correct")
+
+    print("\n4. Different tensor sizes - should compile:")
+    a2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
+    b2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
+    c2 = torch.zeros(512, 512, device="cuda", dtype=torch.float16)
+
+    a2_ = from_dlpack(a2, assumed_align=16)
+    b2_ = from_dlpack(b2, assumed_align=16)
+    c2_ = from_dlpack(c2, assumed_align=16)
+
+    cached_elementwise(add, a2_, b2_, c2_)
+
+    print("\n5. Same size as #4 - should hit cache:")
+    cached_elementwise(add, a2_, b2_, c2_)
+
+    print("\n" + "=" * 50)
+    print("Final Cache Statistics:")
+    stats = get_cache_stats()
+    print(f"  Total calls: {stats['total']}")
+    print(f"  Cache hits: {stats['hits']}")
+    print(f"  Cache misses: {stats['misses']}")
+    print(f"  Hit rate: {stats['hit_rate']:.2%}")
+    print(f"  Unique kernels cached: {stats['cache_size']}")
+
+    print("\nClearing cache...")
+    clear_cute_cache()
+    stats = get_cache_stats()
+    print(f"  Cache size after clear: {stats['cache_size']}")
+
+
+if __name__ == "__main__":
+    main()

transformer_nuggets/cute/__init__.py

@@ -0,0 +1,5 @@
+from transformer_nuggets.cute.cache import (
+    cute_compile_and_cache,
+    clear_cute_cache,
+    get_cache_stats,
+)

transformer_nuggets/cute/add.py

@@ -7,6 +7,7 @@
 from cutlass.cute.runtime import from_dlpack
 
 from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds
+from transformer_nuggets.cute.cache import cute_compile_and_cache, get_cache_stats
 
 
 @cute.kernel
@@ -72,25 +73,31 @@ def naive_elementwise_add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
     )
 
 
-@cute.jit
 def elementwise_op(
     op: cutlass.Constexpr,
     mA: cute.Tensor,
     mB: cute.Tensor,
     mC: cute.Tensor,
 ):
-    thr_layout = cute.make_layout((4, 32), stride=(32, 1))
-    val_layout = cute.make_layout((4, 8), stride=(8, 1))
-    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+    # Define the kernel with @cute.jit
+    @cute.jit
+    def kernel(op, mA, mB, mC):
+        thr_layout = cute.make_layout((4, 32), stride=(32, 1))
+        val_layout = cute.make_layout((4, 8), stride=(8, 1))
+        tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
 
-    gA = cute.zipped_divide(mA, tiler_mn)
-    gB = cute.zipped_divide(mB, tiler_mn)
-    gC = cute.zipped_divide(mC, tiler_mn)
+        gA = cute.zipped_divide(mA, tiler_mn)
+        gB = cute.zipped_divide(mB, tiler_mn)
+        gC = cute.zipped_divide(mC, tiler_mn)
 
-    elementwise_apply_kernel(op, gA, gB, gC, tv_layout).launch(
-        grid=[cute.size(gC, mode=[1]), 1, 1],
-        block=[cute.size(tv_layout, mode=[0]), 1, 1],
-    )
+        elementwise_apply_kernel(op, gA, gB, gC, tv_layout).launch(
+            grid=[cute.size(gC, mode=[1]), 1, 1],
+            block=[cute.size(tv_layout, mode=[0]), 1, 1],
+        )
+
+    # Use explicit caching
+    compiled_kernel = cute_compile_and_cache(kernel, op, mA, mB, mC)
+    return compiled_kernel(op, mA, mB, mC)
 
 
 def benchmark(callable, *, num_warmups=5, num_iterations=200):
@@ -123,7 +130,9 @@ def benchmark(callable, *, num_warmups=5, num_iterations=200):
     c.zero_()
 
     naive_elementwise_add_compiled = cute.compile(naive_elementwise_add, a_, b_, c_)
-    elementwise_add_compiled = cute.compile(elementwise_op, add, a_, b_, c_)
+
+    # Note: elementwise_op is now cached automatically!
+    # First call will compile, subsequent calls use cache
 
     print("\n" + "=" * 50)
     print("BENCHMARK RESULTS")
@@ -135,5 +144,12 @@ def benchmark(callable, *, num_warmups=5, num_iterations=200):
     print("2. PyTorch add (baseline):")
     benchmark(partial(torch.add, a, b, out=c))
 
-    print("3. Optimized elementwise add kernel:")
-    benchmark(partial(elementwise_add_compiled, a_, b_, c_))
+    print("3. Optimized elementwise add kernel (with caching):")
+    benchmark(partial(elementwise_op, add, a_, b_, c_))
+
+    print("\nCache Statistics:")
+    stats = get_cache_stats()
+    print(f"  Cache hits: {stats['hits']}")
+    print(f"  Cache misses: {stats['misses']}")
+    print(f"  Hit rate: {stats['hit_rate']:.2%}")
+    print(f"  Cache size: {stats['cache_size']}")

transformer_nuggets/cute/cache.py

@@ -0,0 +1,130 @@
+import hashlib
+import logging
+import threading
+from typing import Any
+from collections.abc import Callable
+
+import cutlass.cute as cute
+
+
+logger = logging.getLogger(__name__)
+
+
+class CuteKernelCache:
+    """Global cache for compiled CUTE kernels"""
+
+    def __init__(self):
+        self._cache: dict[str, Any] = {}
+        self._stats = {"hits": 0, "misses": 0}
+        self._lock = threading.Lock()
+
+    def get(self, key: str) -> Any | None:
+        with self._lock:
+            if key in self._cache:
+                self._stats["hits"] += 1
+                return self._cache[key]
+            self._stats["misses"] += 1
+            return None
+
+    def set(self, key: str, value: Any) -> None:
+        with self._lock:
+            self._cache[key] = value
+
+    def clear(self) -> None:
+        with self._lock:
+            self._cache.clear()
+            self._stats = {"hits": 0, "misses": 0}
+
+    def get_stats(self) -> dict[str, int | float]:
+        with self._lock:
+            total = self._stats["hits"] + self._stats["misses"]
+            hit_rate = self._stats["hits"] / total if total > 0 else 0
+            return {
+                **self._stats,
+                "total": total,
+                "hit_rate": hit_rate,
+                "cache_size": len(self._cache),
+            }
+
+
+# Global cache instance
+_kernel_cache = CuteKernelCache()
+
+
+def _generate_cache_key(*args, **kwargs) -> str:
+    """Generate a cache key from function arguments"""
+    key_parts = []
+
+    for arg in args:
+        if isinstance(arg, cute.Tensor):
+            # For cute tensors, use shape and dtype
+            key_parts.append(f"tensor_shape={arg.shape}_dtype={arg._dtype}")
+        elif hasattr(arg, "__name__"):
+            # For functions/operators like add, mul
+            key_parts.append(f"op={arg.__name__}")
+        else:
+            # For other types, use string representation
+            key_parts.append(str(arg))
+
+    # Add kwargs
+    for k, v in sorted(kwargs.items()):
+        key_parts.append(f"{k}={v}")
+
+    # Create hash of the key
+    key_str = "_".join(key_parts)
+    return hashlib.md5(key_str.encode()).hexdigest()
+
+
+def cute_compile_and_cache(func: Callable, *args, **kwargs):
+    """
+    Compile a @cute.jit decorated function and cache the result.
+
+    Args:
+        func: A function decorated with @cute.jit
+        *args: Arguments to pass to cute.compile and for cache key generation
+        **kwargs: Keyword arguments to pass to cute.compile and for cache key generation
+
+    Returns:
+        Compiled kernel that can be executed
+
+    Example:
+        @cute.jit
+        def my_kernel(a, b, c):
+            # kernel implementation
+            ...
+
+        # Cache the compilation
+        compiled_kernel = cute_compile_and_cache(my_kernel, tensor_a, tensor_b, tensor_c)
+        result = compiled_kernel(tensor_a, tensor_b, tensor_c)
+    """
+    # Generate cache key from function and arguments
+    cache_key = _generate_cache_key(*args, **kwargs)
+    cache_key = f"{func.__name__}_{cache_key}"
+
+    # Check cache
+    compiled_kernel = _kernel_cache.get(cache_key)
+
+    if compiled_kernel is not None:
+        logger.debug(f"Cache hit for {func.__name__} (key: {cache_key[:8]}...)")
+        return compiled_kernel
+
+    logger.debug(f"Cache miss for {func.__name__} (key: {cache_key[:8]}...) - Compiling...")
+
+    # Compile the kernel
+    compiled_kernel = cute.compile(func, *args, **kwargs)
+
+    # Cache the compiled kernel
+    _kernel_cache.set(cache_key, compiled_kernel)
+
+    return compiled_kernel
+
+
+# Utility functions
+def clear_cute_cache():
+    """Clear all cached kernels"""
+    _kernel_cache.clear()
+
+
+def get_cache_stats():
+    """Get cache statistics"""
+    return _kernel_cache.get_stats()