diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index fd34b9a4df..20b57b30b7 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -31,7 +31,7 @@ def _cross_entropy_forward(
     logsumexp_ptr     ,
     labels_ptr        ,
     VOCAB_SIZE        ,
-    BLOCK_SIZE        ,
+    BLOCK_SIZE        : tl.constexpr,
     DO_SOFTCAPPING    ,
     SOFTCAP           ,
     DO_LOGIT_SCALING  ,
@@ -105,7 +105,7 @@ def _chunked_cross_entropy_forward(
     labels_ptr        ,
     VOCAB_SIZE        ,
     N_CHUNKS          ,
-    BLOCK_SIZE        ,
+    BLOCK_SIZE        : tl.constexpr,
     DO_SOFTCAPPING    ,
     SOFTCAP           ,
     DO_LOGIT_SCALING  ,
@@ -188,7 +188,7 @@ def _cross_entropy_backward(
     logsumexp_ptr     ,
     labels_ptr        ,
     VOCAB_SIZE        ,
-    BLOCK_SIZE        ,
+    BLOCK_SIZE        : tl.constexpr,
     DO_SOFTCAPPING    ,
     SOFTCAP           ,
     DO_LOGIT_SCALING  ,