Update cross_entropy_loss.py

unsloth/kernels/cross_entropy_loss.py

@@ -104,17 +104,18 @@ pass
 })
 @triton.jit
 def _chunked_cross_entropy_forward(
-    logits_ptr, logits_row_stride,
-    loss_ptr,
-    logsumexp_ptr,
-    labels_ptr,
-    VOCAB_SIZE      : tl.constexpr,
-    N_CHUNKS        : tl.constexpr,
-    BLOCK_SIZE      : tl.constexpr,
-    DO_SOFTCAPPING  : tl.constexpr,
-    SOFTCAP         : tl.constexpr,
-    DO_LOGIT_SCALING: tl.constexpr,
-    LOGIT_SCALE     : tl.constexpr,
+    logits_ptr        ,
+    logits_row_stride ,
+    loss_ptr          ,
+    logsumexp_ptr     ,
+    labels_ptr        ,
+    VOCAB_SIZE        ,
+    N_CHUNKS          ,
+    BLOCK_SIZE        : tl.constexpr,
+    DO_SOFTCAPPING    ,
+    SOFTCAP           ,
+    DO_LOGIT_SCALING  ,
+    LOGIT_SCALE       ,
 ):
     """
         256K vocab divided in 4 chunks
@@ -142,7 +143,7 @@ def _chunked_cross_entropy_forward(
     """
     row_idx   = tl.program_id(0)
     chunk_idx = tl.program_id(1)
-    logits_ptr    += row_idx * logits_row_stride.to(tl.int64)
+    logits_ptr    += row_idx * tl.cast(logits_row_stride, tl.int64)
     loss_ptr      += row_idx
     logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
     labels_ptr    += row_idx
@@ -151,14 +152,13 @@ def _chunked_cross_entropy_forward(
     mask = col_offsets < VOCAB_SIZE
 
     label_idx = tl.load(labels_ptr).to(tl.int32)
-    logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf"))
+    logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
 
     # Go logit scaling for Cohere: t * x
     if DO_LOGIT_SCALING: logits = LOGIT_SCALE * logits
     # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
     if DO_SOFTCAPPING:   logits = SOFTCAP * triton_tanh(logits / SOFTCAP)
 
-    logits = logits.to(tl.float32)
     c = tl.max(logits, 0)
     logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))