Upload custom kernels

build/torch-universal/triton_llama_mlp/__init__.py

@@ -0,0 +1,2 @@
+from . import layers
+__all__ = ["layers"]

build/torch-universal/triton_llama_mlp/layers.py

@@ -0,0 +1,3 @@
+from .mlp import TritonLlamaMLP
+
+__all__ = ["TritonLlamaMLP"]

build/torch-universal/triton_llama_mlp/mlp.py

@@ -0,0 +1,210 @@
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+from typing import Callable, Optional
+
+
+@triton.jit
+def matmul_kernel(
+        # Pointers to matrices
+        a_ptr, b_ptr, c_ptr,
+        # Matrix dimensions
+        M, N, K,
+        # The stride variables represent how much to increase the ptr by when moving by 1
+        # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
+        # by to get the element one row down (A has M rows).
+        stride_am, stride_ak,  #
+        stride_bk, stride_bn,  #
+        stride_cm, stride_cn,
+        # Meta-parameters
+        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
+        GROUP_SIZE_M: tl.constexpr,  #
+        ACTIVATION: tl.constexpr = None #
+):
+    """Kernel for computing the matmul C = A x B.
+    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    # See above `L2 Cache Optimizations` section for details.
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+    # See above `Pointer Arithmetic` section for details
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+
+    # -----------------------------------------------------------
+    # Iterate to compute a block of the C matrix.
+    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+    # of fp32 values for higher accuracy.
+    # `accumulator` will be converted back to fp16 after the loop.
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Load the next block of A and B, generate a mask by checking the K dimension.
+        # If it is out of bounds, set it to 0.
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+        # We accumulate along the K dimension.
+        accumulator += tl.dot(a, b)
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+    # You can fuse arbitrary activation functions here
+    # while the accumulator is still in FP32!
+
+    c = accumulator.to(tl.float32)
+
+    # -----------------------------------------------------------
+    # Write back the block of the output matrix C with masks.
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`.
+@triton.jit
+def silu(x):
+    return x * tl.sigmoid(x)
+
+def matmul(a, b, activation=""):
+    # Check constraints.
+    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
+    assert a.is_contiguous(), "Matrix A must be contiguous"
+    M, K = a.shape
+    K, N = b.shape
+    BLOCK_SIZE_M = 32
+    BLOCK_SIZE_N = 32
+    BLOCK_SIZE_K = 32
+    GROUP_SIZE_M = 8
+    # Allocates output.
+    c = torch.empty((M, N), device=a.device, dtype=torch.float)
+    # 1D launch kernel where each block gets its own program.
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
+    matmul_kernel[grid](
+        a, b, c,  #
+        M, N, K,  #
+        a.stride(0), a.stride(1),  #
+        b.stride(0), b.stride(1),  #
+        c.stride(0), c.stride(1),  #
+        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,  #
+        GROUP_SIZE_M,  #
+        ACTIVATION=activation  #
+    )
+    return c
+
+
+class TritonLlamaMLP(nn.Module):
+    """LlamaMLP implementation using Triton kernels for matrix multiplication"""
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+    act_fn: Callable
+    def forward(self, x):
+        # Replace nn.Linear with matmul using triton kernel
+        # Save original shape for reshaping back later
+        original_shape = x.shape
+        # Reshape input to 2D for matmul: (*, hidden_size) -> (batch_size*seq_len, hidden_size)
+        x_2d = x.reshape(-1, x.size(-1))
+        
+        # Gate projection
+        gate_output = matmul(x_2d, self.gate_proj.weight.t())
+        if self.gate_proj.bias is not None:
+            gate_output += self.gate_proj.bias
+            
+        # Up projection
+        up_output = matmul(x_2d, self.up_proj.weight.t())
+        if self.up_proj.bias is not None:
+            up_output += self.up_proj.bias
+            
+        # Apply activation function and element-wise multiplication
+        intermediate_output = self.act_fn(gate_output) * up_output
+        
+        # Final projection
+        down_output = matmul(intermediate_output, self.down_proj.weight.t())
+        if self.down_proj.bias is not None:
+            down_output += self.down_proj.bias
+            
+        # Reshape back to original dimensions: (batch_size*seq_len, hidden_size) -> (*, hidden_size)
+        return down_output.reshape(original_shape)
+
+
+# def test_triton_llama_mlp():
+#     """Test that TritonLlamaMLP produces the same output as LlamaMLP from transformers."""
+#     import torch
+#     import torch.nn.functional as F
+    
+#     # Skip test if CUDA is not available
+#     if not torch.cuda.is_available():
+#         print("CUDA not available, skipping test")
+#         return True
+    
+#     # Define test parameters
+#     batch_size = 2
+#     seq_len = 4
+#     hidden_size = 128
+#     intermediate_size = 256
+    
+#     # Create input tensor
+#     x = torch.randn(batch_size, seq_len, hidden_size, device="cuda", dtype=torch.float)
+    
+#     # Create a standard PyTorch implementation for comparison
+#     class StandardLlamaMLP(nn.Module):
+#         def __init__(self, hidden_size, intermediate_size):
+#             super().__init__()
+#             self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+#             self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+#             self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+#             self.act_fn = F.silu
+            
+#         def forward(self, x):
+#             return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+    
+#     # Initialize models
+#     standard_mlp = StandardLlamaMLP(hidden_size, intermediate_size).to("cuda").to(torch.float)
+    
+#     # Create our Triton implementation
+#     triton_mlp = TritonLlamaMLP()
+#     triton_mlp.gate_proj = standard_mlp.gate_proj
+#     triton_mlp.up_proj = standard_mlp.up_proj
+#     triton_mlp.down_proj = standard_mlp.down_proj
+#     triton_mlp.act_fn = standard_mlp.act_fn
+    
+#     # Run both implementations
+#     with torch.no_grad():
+#         standard_output = standard_mlp(x)
+#         triton_output = triton_mlp(x)
+    
+#     # Compare outputs
+#     max_diff = torch.max(torch.abs(standard_output - triton_output))
+#     print(f"Maximum difference between standard and Triton implementation: {max_diff}")
+    
+#     # Check if outputs are close enough (allowing for some floating point differences)
+#     is_close = torch.allclose(standard_output, triton_output, rtol=1e-2, atol=1e-2)
+#     print(f"Outputs match within tolerance: {is_close}")
+    
+#     return is_close
+
+# if __name__ == "__main__":
+#     test_triton_llama_mlp()
+

torch-ext/triton_llama_mlp/mlp.py

@@ -71,7 +71,7 @@ def matmul_kernel(
     # You can fuse arbitrary activation functions here
     # while the accumulator is still in FP32!
 
-    c = accumulator.to(tl.float16)
+    c = accumulator.to(tl.float32)
 
     # -----------------------------------------------------------
     # Write back the block of the output matrix C with masks.
@@ -98,7 +98,7 @@ def matmul(a, b, activation=""):
     BLOCK_SIZE_K = 32
     GROUP_SIZE_M = 8
     # Allocates output.
-    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
+    c = torch.empty((M, N), device=a.device, dtype=torch.float)
     # 1D launch kernel where each block gets its own program.
     grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
     matmul_kernel[grid](
@@ -166,7 +166,7 @@ class TritonLlamaMLP(nn.Module):
 #     intermediate_size = 256
     
 #     # Create input tensor
-#     x = torch.randn(batch_size, seq_len, hidden_size, device="cuda", dtype=torch.float16)
+#     x = torch.randn(batch_size, seq_len, hidden_size, device="cuda", dtype=torch.float)
     
 #     # Create a standard PyTorch implementation for comparison
 #     class StandardLlamaMLP(nn.Module):
@@ -181,7 +181,7 @@ class TritonLlamaMLP(nn.Module):
 #             return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
     
 #     # Initialize models
-#     standard_mlp = StandardLlamaMLP(hidden_size, intermediate_size).to("cuda").to(torch.float16)
+#     standard_mlp = StandardLlamaMLP(hidden_size, intermediate_size).to("cuda").to(torch.float)
     
 #     # Create our Triton implementation
 #     triton_mlp = TritonLlamaMLP()