From 8942b2fc76c59069af9b79acff7b1336bccad1f5 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 31 Oct 2024 01:23:41 -0700
Subject: [PATCH] Update cross_entropy_loss.py

---
 unsloth/kernels/cross_entropy_loss.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 8cbdbf2a2c..a13475e227 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -92,10 +92,10 @@ def _cross_entropy_forward(
 pass
 
 
-# @triton.heuristics({
-#     "DO_SOFTCAPPING":   lambda args: args["DO_SOFTCAPPING"  ],
-#     "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
-# })
+@triton.heuristics({
+    "DO_SOFTCAPPING":   lambda args: args["DO_SOFTCAPPING"  ],
+    "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
+})
 @triton.jit
 def _chunked_cross_entropy_forward(
     logits_ptr        ,
@@ -106,9 +106,9 @@ def _chunked_cross_entropy_forward(
     VOCAB_SIZE        ,
     N_CHUNKS          ,
     BLOCK_SIZE        : tl.constexpr,
-    DO_SOFTCAPPING    ,
+    DO_SOFTCAPPING    : tl.constexpr(tl.int1),
     SOFTCAP           ,
-    DO_LOGIT_SCALING  ,
+    DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
     LOGIT_SCALE       ,
 ):
     """
@@ -189,9 +189,9 @@ def _cross_entropy_backward(
     labels_ptr        ,
     VOCAB_SIZE        ,
     BLOCK_SIZE        : tl.constexpr,
-    DO_SOFTCAPPING    ,
+    DO_SOFTCAPPING    : tl.constexpr(tl.int1),
     SOFTCAP           ,
-    DO_LOGIT_SCALING  ,
+    DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
     LOGIT_SCALE       ,
 ):
     """
@@ -347,7 +347,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
             SOFTCAP          = ctx.logit_softcapping,
             DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING,
             LOGIT_SCALE      = ctx.logit_scaling,
-            num_warps      = 8,
+            num_warps        = 8,
         )
         return logits, None, None, None,
     pass