diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 17168b230a..15a148b00e 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -36,12 +36,12 @@ def _cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -109,13 +109,13 @@ def _chunked_cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, - N_CHUNKS : tl.constexpr, + VOCAB_SIZE , + N_CHUNKS , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ 256K vocab divided in 4 chunks @@ -193,12 +193,12 @@ def _cross_entropy_backward( dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x)