[CLEANUP]

examples/cython_tests/mqa.pyx

@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+cimport cython
+
+cdef class MultiQueryAttention:
+    cdef int embed_dim
+    cdef int num_heads
+    cdef int head_dim
+    cdef object query_proj  # Treat nn.Linear as a Python object
+    cdef object key_proj    # Treat nn.Linear as a Python object
+    cdef object value_proj  # Treat nn.Linear as a Python object
+    cdef object out_proj    # Treat nn.Linear as a Python object
+
+    def __cinit__(self, int embed_dim, int num_heads):
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        # Initialize nn.Linear layers as regular Python objects
+        self.query_proj = nn.Linear(embed_dim, embed_dim)
+        self.key_proj = nn.Linear(embed_dim, self.head_dim)
+        self.value_proj = nn.Linear(embed_dim, self.head_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def forward(self, query, key, value):
+        cdef int batch_size, seq_len, _
+
+        # Assuming the input tensors are torch.Tensor objects
+        batch_size, seq_len, _ = query.size()
+
+        # Linear projections
+        queries = self.query_proj(query)
+        keys = self.key_proj(key)
+        values = self.value_proj(value)
+
+        # Reshape for multi-head attention
+        queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        keys = keys.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+        values = values.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+
+        # Scaled dot-product attention
+        scores = torch.matmul(queries, keys.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        attn_weights = torch.nn.functional.softmax(scores, dim=-1)
+        attn_output = torch.matmul(attn_weights, values)
+
+        # Concatenate and project the output
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
+        output = self.out_proj(attn_output)
+
+        return output

examples/cython_tests/mqa_test.py

@@ -0,0 +1,54 @@
+import timeit
+import torch
+from zeta import MultiQueryAttention as PyTorchMQA
+from mqa import MultiQueryAttention as CythonMQA
+
+# Input parameters
+batch_size = 32
+seq_len = 128
+embed_dim = 512
+num_heads = 8
+
+# Create sample input tensors
+query = torch.randn(batch_size, seq_len, embed_dim)
+key = torch.randn(batch_size, seq_len, embed_dim)
+value = torch.randn(batch_size, seq_len, embed_dim)
+
+# Initialize the PyTorch Multi-Query Attention layer (from zeta package)
+pytorch_mqa = PyTorchMQA(dim=embed_dim, heads=num_heads)
+
+# Initialize the Cython Multi-Query Attention layer (from mqa module)
+cython_mqa = CythonMQA(embed_dim, num_heads)
+
+
+# Define functions for benchmarking
+def run_pytorch_mqa():
+    output, _, _ = pytorch_mqa(query)
+    return output
+
+
+def run_cython_mqa():
+    output = cython_mqa.forward(query, key, value)
+    return output
+
+
+# Warm-up runs (important to avoid cold start issues)
+for _ in range(20):
+    run_pytorch_mqa()
+    run_cython_mqa()
+
+# Benchmark PyTorch implementation
+pytorch_time = timeit.timeit(
+    "run_pytorch_mqa()", globals=globals(), number=1000
+)
+
+# Benchmark Cython implementation
+cython_time = timeit.timeit("run_cython_mqa()", globals=globals(), number=1000)
+
+# Print the results
+print(f"PyTorch MQA execution time: {pytorch_time:.6f} seconds")
+print(f"Cython MQA execution time: {cython_time:.6f} seconds")
+if cython_time < pytorch_time:
+    print(f"Cython is faster by: {pytorch_time / cython_time:.2f}x")
+else:
+    print(f"PyTorch is faster by: {cython_time / pytorch_time:.2f}x")

examples/cython_tests/new_c_example.py

@@ -0,0 +1,10 @@
+import torch
+import torch_extension  # Import the compiled Cython module
+
+# Create a PyTorch tensor
+input_tensor = torch.tensor([0.0, 1.0, 2.0, 3.0])
+
+# Use the Cython function to apply the sin function
+output_tensor = torch_extension.apply_sin(input_tensor)
+
+print(output_tensor)

examples/cython_tests/setup.py

@@ -0,0 +1,15 @@
+from setuptools import setup, Extension
+from torch.utils.cpp_extension import BuildExtension
+from Cython.Build import cythonize
+
+setup(
+    name="mqa",
+    ext_modules=cythonize(
+        Extension(
+            "mqa",
+            sources=["mqa.pyx"],
+            language="c++",
+        )
+    ),
+    cmdclass={"build_ext": BuildExtension},
+)

examples/cython_tests/torch_extension.pyx

@@ -0,0 +1,20 @@
+import torch  # Use standard Python import for PyTorch
+cimport cython
+import numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def apply_sin(input_tensor):
+    cdef int i
+    cdef int size = input_tensor.numel()
+
+    # Convert the PyTorch tensor to a NumPy array
+    np_array = input_tensor.numpy()
+
+    # Apply sin element-wise using NumPy
+    np_output = np.sin(np_array)
+
+    # Convert the NumPy array back to a PyTorch tensor
+    output_tensor = torch.from_numpy(np_output)
+
+    return output_tensor