fix more errors

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'triformer',
   packages = find_packages(),
-  version = '1.3.1',
+  version = '1.3.2',
   license='MIT',
   description = 'Transformer components in Triton',
   long_description=open('README.md').read(),

triformer/forward_mlp.py

@@ -96,6 +96,9 @@ def grid(meta):
 class TritonLinearFunction(Function):
     @staticmethod
     def forward(ctx, x, weight, bias, use_relu=True):
+        # Convert input to float16 if necessary
+        x = x.to(dtype=torch.float16)
+
         # Add batch dimension if necessary
         if x.ndim == 2:
             x = x.unsqueeze(0)
@@ -149,11 +152,12 @@ def __init__(self, in_features, out_features, use_relu=True):
         self.in_features = in_features
         self.out_features = out_features
         self.use_relu = use_relu
+        # Change weight and bias to float32
         self.weight = nn.Parameter(
-            torch.empty(out_features, in_features, device='cuda', dtype=torch.float16)
+            torch.empty(out_features, in_features, device='cuda', dtype=torch.float32)
         )
         self.bias = nn.Parameter(
-            torch.zeros(out_features, device='cuda', dtype=torch.float16)
+            torch.zeros(out_features, device='cuda', dtype=torch.float32)
         )
         self.reset_parameters()
 

triformer/utils.py

@@ -0,0 +1,14 @@
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def calc_num_warps(block_size):
+    num_warps = 4
+    if block_size >= 2048:
+        num_warps = 8
+    if block_size >= 4096:
+        num_warps = 16
+    return num_warps