From 943cc1c3df944eff31be40641fe64fab2553bd5b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 4 Nov 2024 23:16:01 -0800
Subject: [PATCH] Update utils.py

---
 unsloth/kernels/utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py
index a8c20c75a4..b394d122fd 100644
--- a/unsloth/kernels/utils.py
+++ b/unsloth/kernels/utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import triton
-MAX_FUSED_SIZE = 65536
+MAX_FUSED_SIZE : int = 65536
 next_power_of_2 = triton.next_power_of_2
 
 # torch.cuda.amp.custom_fwd is deprecated >= 2.4
@@ -40,12 +40,12 @@ else:
 pass
 
 
-def calculate_settings(n):
-    BLOCK_SIZE = next_power_of_2(n)
+def calculate_settings(n : int) -> (int, int,):
+    BLOCK_SIZE : int = next_power_of_2(n)
     if BLOCK_SIZE > MAX_FUSED_SIZE:
         raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "\
                            f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
-    num_warps = 4
+    num_warps : int = 4
     if   BLOCK_SIZE >= 32768: num_warps = 32
     elif BLOCK_SIZE >=  8192: num_warps = 16
     elif BLOCK_SIZE >=  2048: num_warps = 8