From 0ab731d40b3e87af323c9fdb6f3f266ccced1e7e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 5 Nov 2024 00:05:56 -0800
Subject: [PATCH] typing

---
 unsloth/kernels/rms_layernorm.py  | 17 +++++++++--------
 unsloth/kernels/rope_embedding.py | 18 ++++++++++++++++--
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py
index 0846a09de6..4b22f8c3e5 100644
--- a/unsloth/kernels/rms_layernorm.py
+++ b/unsloth/kernels/rms_layernorm.py
@@ -60,7 +60,7 @@ def _rms_layernorm_backward(
     X,   X_row_stride,
     W,   W_row_stride,
     r,   r_row_stride,
-    dW, dW_row_stride,
+    # dW, dW_row_stride,
     n_cols, eps,
     GEMMA      : tl.constexpr,
     BLOCK_SIZE : tl.constexpr,
@@ -130,7 +130,7 @@ pass
 
 class Fast_RMS_Layernorm(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, X, W, eps :float, gemma : bool = False):
+    def forward(ctx, X : torch.Tensor, W : torch.Tensor, eps : float, gemma : bool = False):
         shape = X.shape
         dim : int = shape[-1]
         X = X.view(-1, dim)
@@ -163,7 +163,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
     pass
 
     @staticmethod
-    def backward(ctx, dY):
+    def backward(ctx, dY : torch.Tensor):
         shape = dY.shape
         dim : int = shape[-1]
         dY = dY.view(-1, dim)
@@ -171,14 +171,14 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
         n_rows : int
         n_cols : int
         n_rows, n_cols = dY.shape
-        dW = X
+        # dW = X
 
         _rms_layernorm_backward[(n_rows,)](
             dY, dY.stride(0),
             X,  X .stride(0),
             W,  W .stride(0),
             r,  r .stride(0),
-            dW, dW.stride(0),
+            # dW, dW.stride(0),
             n_cols, ctx.eps,
             GEMMA      = ctx.GEMMA,
             BLOCK_SIZE = ctx.BLOCK_SIZE,
@@ -190,10 +190,11 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
 pass
 
 
+# [TODO] Unsure why RMS Layernorm is not torch.compiling properly
 @torch.compiler.disable
-def fast_rms_layernorm(layernorm, X, gemma = False):
-    W   = layernorm.weight
-    eps = layernorm.variance_epsilon if \
+def fast_rms_layernorm(layernorm, X : torch.Tensor, gemma : bool = False):
+    W : torch.Tensor = layernorm.weight
+    eps : float = layernorm.variance_epsilon if \
         hasattr(layernorm, "variance_epsilon") \
         else layernorm.eps
     out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)
diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py
index 2934ac41c9..44a7cda12f 100644
--- a/unsloth/kernels/rope_embedding.py
+++ b/unsloth/kernels/rope_embedding.py
@@ -18,7 +18,7 @@ import torch
 from .utils import calculate_settings
 ROPE_GROUP_SIZE = 4
 
-@triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],})
+@triton.heuristics({"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),})
 @triton.jit
 def _rope_embedding(
     Q,     Q_row_stride,
@@ -75,8 +75,14 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
     @staticmethod
     def forward(ctx, Q, cos, sin):
         cos, sin = cos.squeeze(), sin.squeeze()
+        batch    : int
+        seq_len  : int
+        n_heads  : int
+        head_dim : int
         batch, seq_len, n_heads, head_dim = Q.shape
         Q = Q.view(batch*seq_len, n_heads*head_dim)
+        n_rows : int
+        n_cols : int
         n_rows, n_cols = Q.shape
         assert(seq_len <= cos.shape[0])
 
@@ -85,8 +91,10 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
         BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2)
         
         # group_size = 4 # 4 or 8, too large group_size can hurt performance.
+        div : int
+        mod : int
         div, mod = divmod(n_heads, ROPE_GROUP_SIZE)
-        n_groups = div + (mod != 0)
+        n_groups : int = div + (mod != 0)
 
         _rope_embedding[(n_rows, n_groups, )](
               Q,   Q.stride(0),
@@ -108,9 +116,15 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
 
     @staticmethod
     def backward(ctx, dY):
+        batch    : int
+        seq_len  : int
+        n_heads  : int
+        head_dim : int
         batch, seq_len, n_heads, head_dim = dY.shape
         dY = dY.reshape(batch*seq_len, n_heads*head_dim)
         # Must be reshape not view
+        n_rows : int
+        n_cols : int
         n_rows, n_cols = dY.shape
 
         cos = ctx.cos