From f6d50c78ac30e7236386a685a7a151a882cb3a50 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:18:26 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 4ac5a0c8b8..3d4a548a1f 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -37,10 +37,10 @@ def _cross_entropy_forward( logsumexp_ptr , labels_ptr , VOCAB_SIZE , - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING , + BLOCK_SIZE : tl.constexpr(tl.int32), + DO_SOFTCAPPING : tl.constexpr(tl.int1), SOFTCAP , - DO_LOGIT_SCALING , + DO_LOGIT_SCALING : tl.constexpr(tl.int1), LOGIT_SCALE , ): """ @@ -111,10 +111,10 @@ def _chunked_cross_entropy_forward( labels_ptr , VOCAB_SIZE , N_CHUNKS , - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING , + BLOCK_SIZE : tl.constexpr(tl.int32), + DO_SOFTCAPPING : tl.constexpr(tl.int1), SOFTCAP , - DO_LOGIT_SCALING , + DO_LOGIT_SCALING : tl.constexpr(tl.int1), LOGIT_SCALE , ): """ @@ -194,10 +194,10 @@ def _cross_entropy_backward( logsumexp_ptr , labels_ptr , VOCAB_SIZE , - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING , + BLOCK_SIZE : tl.constexpr(tl.int32), + DO_SOFTCAPPING : tl.constexpr(tl.int1), SOFTCAP , - DO_LOGIT_SCALING , + DO_LOGIT_SCALING : tl.constexpr(tl.int1), LOGIT_SCALE , ): """