Update cross_entropy_loss.py

This commit is contained in:
Daniel Han 2024-11-03 21:42:13 -08:00
commit 5eddfd8cfe

View file

@ -36,7 +36,7 @@ def _cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING ,
SOFTCAP ,
@ -109,8 +109,8 @@ def _chunked_cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),
N_CHUNKS : tl.constexpr(tl.int32),
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING ,
SOFTCAP ,
@ -193,7 +193,7 @@ def _cross_entropy_backward(
dloss_row_stride ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING ,
SOFTCAP ,