From b6192e52fc0c8116c477e6608ca2ed2cfabd22c3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 30 Oct 2024 15:33:37 -0700
Subject: [PATCH] Update cross_entropy_loss.py

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

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index e23f818a11..fd34b9a4df 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -25,17 +25,17 @@ from packaging.version import Version
 # })
 @triton.jit
 def _cross_entropy_forward(
-    logits_ptr        : tl.pointer_type,
-    logits_row_stride : tl.constexpr(tl.int64),
-    loss_ptr          : tl.pointer_type(tl.float32),
-    logsumexp_ptr     : tl.pointer_type(tl.float32),
-    labels_ptr        : tl.const_pointer_type(tl.int32),
-    VOCAB_SIZE        : tl.constexpr,
-    BLOCK_SIZE        : tl.constexpr,
-    DO_SOFTCAPPING    : tl.constexpr(tl.int1),
-    SOFTCAP           : tl.constexpr(tl.float32),
-    DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
-    LOGIT_SCALE       : tl.constexpr(tl.float32),
+    logits_ptr        ,
+    logits_row_stride ,
+    loss_ptr          ,
+    logsumexp_ptr     ,
+    labels_ptr        ,
+    VOCAB_SIZE        ,
+    BLOCK_SIZE        ,
+    DO_SOFTCAPPING    ,
+    SOFTCAP           ,
+    DO_LOGIT_SCALING  ,
+    LOGIT_SCALE       ,
 ):
     """
         Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
@@ -98,18 +98,18 @@ pass
 # })
 @triton.jit
 def _chunked_cross_entropy_forward(
-    logits_ptr        : tl.const_pointer_type,
-    logits_row_stride : tl.constexpr(tl.int64),
-    loss_ptr          : tl.pointer_type(tl.float32),
-    logsumexp_ptr     : tl.pointer_type(tl.float32),
-    labels_ptr        : tl.const_pointer_type(tl.int32),
-    VOCAB_SIZE        : tl.constexpr,
-    N_CHUNKS          : tl.constexpr,
-    BLOCK_SIZE        : tl.constexpr,
-    DO_SOFTCAPPING    : tl.constexpr(tl.int1),
-    SOFTCAP           : tl.constexpr,
-    DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
-    LOGIT_SCALE       : tl.constexpr,
+    logits_ptr        ,
+    logits_row_stride ,
+    loss_ptr          ,
+    logsumexp_ptr     ,
+    labels_ptr        ,
+    VOCAB_SIZE        ,
+    N_CHUNKS          ,
+    BLOCK_SIZE        ,
+    DO_SOFTCAPPING    ,
+    SOFTCAP           ,
+    DO_LOGIT_SCALING  ,
+    LOGIT_SCALE       ,
 ):
     """
         256K vocab divided in 4 chunks
@@ -181,18 +181,18 @@ pass
 # })
 @triton.jit
 def _cross_entropy_backward(
-    logits_ptr        : tl.pointer_type,
-    logits_row_stride : tl.constexpr(tl.int64),
-    dloss_ptr         : tl.const_pointer_type(tl.float32),
-    dloss_row_stride  : tl.constexpr,
-    logsumexp_ptr     : tl.const_pointer_type(tl.float32),
-    labels_ptr        : tl.const_pointer_type(tl.int32),
-    VOCAB_SIZE        : tl.constexpr,
-    BLOCK_SIZE        : tl.constexpr,
-    DO_SOFTCAPPING    : tl.constexpr(tl.int1),
-    SOFTCAP           : tl.constexpr,
-    DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
-    LOGIT_SCALE       : tl.constexpr,
+    logits_ptr        ,
+    logits_row_stride ,
+    dloss_ptr         ,
+    dloss_row_stride  ,
+    logsumexp_ptr     ,
+    labels_ptr        ,
+    VOCAB_SIZE        ,
+    BLOCK_SIZE        ,
+    DO_SOFTCAPPING    ,
+    SOFTCAP           ,
+    DO_LOGIT_SCALING  ,
+    LOGIT_SCALE       ,
 ):
     """
         CE_i = -y log(P) = y * (log[sum(exp(x))] - x)