From 1e98cc3eb2414228ce5fd7285ccf9c025956f176 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 3 Nov 2024 21:25:32 -0800
Subject: [PATCH] typing

---
 unsloth/kernels/cross_entropy_loss.py |  2 +-
 unsloth/kernels/rms_layernorm.py      | 14 ++++++++++----
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 09eb0854e8..d396538e69 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -190,7 +190,7 @@ def _cross_entropy_backward(
     logits_ptr        ,
     logits_row_stride : tl.constexpr(tl.int64),
     dloss_ptr         ,
-    dloss_row_stride  : tl.constexpr(tl.int32),
+    dloss_row_stride  ,
     logsumexp_ptr     ,
     labels_ptr        ,
     VOCAB_SIZE        : tl.constexpr(tl.int32),
diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py
index 13faf08d6a..c0fb222b8a 100644
--- a/unsloth/kernels/rms_layernorm.py
+++ b/unsloth/kernels/rms_layernorm.py
@@ -53,7 +53,7 @@ def _rms_layernorm_forward(
 pass
 
 
-@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],})
+@triton.heuristics({"GEMMA": lambda args: bool(args["GEMMA"]),})
 @triton.jit
 def _rms_layernorm_backward(
     dY, dY_row_stride,
@@ -130,11 +130,15 @@ pass
 
 class Fast_RMS_Layernorm(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, X, W, eps, gemma = False):
+    def forward(ctx, X, W, eps :float, gemma : bool = False):
         shape = X.shape
-        dim = shape[-1]
+        dim : int = shape[-1]
         X = X.view(-1, dim)
+        n_rows : int
+        n_cols : int
         n_rows, n_cols = X.shape
+        BLOCK_SIZE : int
+        num_warps  : int
         BLOCK_SIZE, num_warps = calculate_settings(n_cols)
 
         Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
@@ -161,9 +165,11 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
     @staticmethod
     def backward(ctx, dY):
         shape = dY.shape
-        dim = shape[-1]
+        dim : int = shape[-1]
         dY = dY.view(-1, dim)
         X, W, r = ctx.saved_tensors
+        n_rows : int
+        n_cols : int
         n_rows, n_cols = dY.shape
         dW = X