Updates for dynamic

stack-info: PR: #58, branch: drisspg/stack/16

Author: drisspg

transformer_nuggets/__init__.py

@@ -1,15 +1,16 @@
 from transformer_nuggets import quant as quant, utils as utils
+import logging
 
 
-def init_logging():
+def init_logging(level=logging.INFO):
     """
     Configure logging for transformer_nuggets library at INFO level.
     Adds a StreamHandler if none exists.
     """
     import logging
 
     logger = logging.getLogger("transformer_nuggets")
-    logger.setLevel(logging.INFO)
+    logger.setLevel(level)
 
     if not logger.handlers:
         handler = logging.StreamHandler()

transformer_nuggets/cute/cache.py

@@ -68,7 +68,19 @@ def _generate_cache_key(*args, **kwargs) -> str:
 
     for arg in args:
         if isinstance(arg, cute.Tensor):
-            key_parts.append(f"tensor_shape={arg.shape}_dtype={arg._dtype}")
+            # Get string representation and extract the shape:stride pattern
+            tensor_str = str(arg)
+            # Format is: Tensor<address@mem o (shape):(stride)>
+            # We want just the (shape):(stride) part
+
+            if " o " in tensor_str and ")>" in tensor_str:
+                # Extract everything after ' o ' and before '>'
+                inner_part = tensor_str.split(" o ")[1].rstrip(">")
+                # inner_part should be like "(?,?):(?,1)"
+                key_parts.append(f"tensor_{inner_part}_dtype={arg._dtype}")
+            else:
+                # Fallback if format is different
+                key_parts.append(f"tensor_shape={arg.shape}_dtype={arg._dtype}")
         elif hasattr(arg, "__name__"):
             key_parts.append(f"op={arg.__name__}")
         else:
@@ -78,6 +90,7 @@ def _generate_cache_key(*args, **kwargs) -> str:
         key_parts.append(f"{k}={v}")
 
     key_str = "_".join(key_parts)
+    logger.debug(f"Generated cache key: {key_str}")
     return hashlib.sha256(key_str.encode()).hexdigest()[:16]
 
 

transformer_nuggets/cute/dynamic_args.py

@@ -0,0 +1,94 @@
+import torch
+from operator import add
+
+import cutlass
+import cutlass.cute as cute
+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
+from transformer_nuggets.cute.element_wise import elementwise_apply_kernel
+from rich import print
+from transformer_nuggets import init_logging
+import logging
+
+init_logging(logging.INFO)
+
+
+@cute.jit
+def elem_kernel_parameterized(
+    op: cutlass.Constexpr,
+    mA: cute.Tensor,
+    mB: cute.Tensor,
+    mC: cute.Tensor,
+    thr_m: cutlass.Constexpr,
+    thr_n: cutlass.Constexpr,
+    val_m: cutlass.Constexpr,
+    val_n: cutlass.Constexpr,
+):
+    """Parameterized kernel that accepts layout dimensions as constexpr"""
+    thr_layout = cute.make_layout((thr_m, thr_n), stride=(thr_n, 1))
+    val_layout = cute.make_layout((val_m, val_n), stride=(val_n, 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)
+
+    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],
+    )
+
+
+def elementwise_op_dynamic(
+    op: cutlass.Constexpr,
+    a: torch.Tensor,
+    b: torch.Tensor,
+) -> torch.Tensor:
+    M, N = a.shape
+    c = torch.empty(M, N, device="cuda", dtype=torch.float16)
+
+    # Choose parameters based on size
+    total_elements = M * N
+    if total_elements < 1024 * 1024:
+        thr_m, thr_n, val_m, val_n = 8, 32, 2, 8
+    elif total_elements < 16 * 1024 * 1024:
+        thr_m, thr_n, val_m, val_n = 4, 64, 4, 8
+    else:
+        thr_m, thr_n, val_m, val_n = 2, 128, 8, 8
+
+    # Create tensors with optimization hints
+    mA = from_dlpack(a, assumed_align=16).mark_layout_dynamic(1)
+    mB = from_dlpack(b, assumed_align=16).mark_layout_dynamic(1)
+    mC = from_dlpack(c, assumed_align=16).mark_layout_dynamic(1)
+
+    # Convert to compile-time constants - this creates separate kernels for each configuration
+    compiled_kernel = cute_compile_and_cache(
+        elem_kernel_parameterized, op, mA, mB, mC, thr_m, thr_n, val_m, val_n
+    )
+    compiled_kernel(mA, mB, mC)
+    return c
+
+
+if __name__ == "__main__":
+    shapes = [(2**i, 2**i) for i in range(8, 14)]
+    shapes.extend([(1000, 1000), (1234, 5678), (3333, 7777), (999, 1001)])
+    for M, N in shapes:
+        a = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        b = torch.randn(M, N, device="cuda", dtype=torch.float16)
+
+        # Test the new API that takes regular PyTorch tensors
+        out = elementwise_op_dynamic(add, a, b)
+        torch.testing.assert_close(out, add(a, b))
+
+        time_torch = benchmark_cuda_function_in_microseconds(lambda: add(a, b))
+        time_cute = benchmark_cuda_function_in_microseconds(
+            lambda: elementwise_op_dynamic(add, a, b)
+        )
+        print(f"M = {M}, N = {N}")
+        print(f"torch GB/s = {M * N * 3 * out.element_size() / time_torch * 1e-3}")
+        print(f"cute GB/s = {M * N * 3 * out.element_size() / time_cute * 1e-3}")
+
+    stats = get_cache_stats()
+    print(f"Cache stats: {stats}")

transformer_nuggets/cute/add.py transformer_nuggets/cute/element_wise.pytransformer_nuggets/cute/add.py renamed to transformer_nuggets/cute/element_wise.py

@@ -120,7 +120,7 @@ def elementwise_op(
     return compiled_kernel(mA, mB, mC)
 
 
-def benchmark(callable, tensor_a, *, num_warmups=5, num_iterations=200):
+def benchmark(callable, tensor_a):
     time = benchmark_cuda_function_in_microseconds(callable)
     avg_time = time / 1e3
 
@@ -166,7 +166,7 @@ def benchmark(callable, tensor_a, *, num_warmups=5, num_iterations=200):
     benchmark(partial(torch.add, a, b, out=c), a)
 
     print("3. Optimized elementwise add kernel (with caching):")
-    benchmark(partial(elementwise_op, add, a, b, c, assumed_align=128), a)
+    benchmark(partial(elementwise_op, add, a, b, c, assumed_align=16), a)
 
     print("\nCache Statistics:")
     stats = get_cache_stats()