diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index f256918746..f8b9b4245e 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -30,8 +30,8 @@ def _cross_entropy_forward(
     loss_ptr          ,
     logsumexp_ptr     ,
     labels_ptr        ,
-    VOCAB_SIZE        : tl.constexpr,
-    BLOCK_SIZE        : tl.constexpr,
+    VOCAB_SIZE        : tl.constexpr(tl.int32),
+    BLOCK_SIZE        : tl.constexpr(tl.int32),
     DO_SOFTCAPPING    : tl.constexpr(tl.int1),
     SOFTCAP           : tl.constexpr(tl.float32),
     DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
@@ -103,9 +103,9 @@ def _chunked_cross_entropy_forward(
     loss_ptr          ,
     logsumexp_ptr     ,
     labels_ptr        ,
-    VOCAB_SIZE        ,
-    N_CHUNKS          ,
-    BLOCK_SIZE        : tl.constexpr,
+    VOCAB_SIZE        : tl.constexpr(tl.int32),
+    N_CHUNKS          : tl.constexpr(tl.int32),
+    BLOCK_SIZE        : tl.constexpr(tl.int32),
     DO_SOFTCAPPING    : tl.constexpr(tl.int1),
     SOFTCAP           : tl.constexpr(tl.float32),
     DO_LOGIT_SCALING  : tl.constexpr(tl.int1),
@@ -187,8 +187,8 @@ def _cross_entropy_backward(
     dloss_row_stride  ,
     logsumexp_ptr     ,
     labels_ptr        ,
-    VOCAB_SIZE        ,
-    BLOCK_SIZE        : tl.constexpr,
+    VOCAB_SIZE        : tl.constexpr(tl.int32),
+    BLOCK_SIZE        : tl.constexpr(tl.int32),
     DO_SOFTCAPPING    : tl.constexpr(tl.int1),
     SOFTCAP           : tl.constexpr(tl.float32),
     DO_LOGIT_SCALING  : tl.constexpr(tl.int1),