Skip to content

perf(ce): single-shot dx_y gradient correction #1268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vaibhavjindal merged 4 commits into from
Jul 28, 2026
Merged

perf(ce): single-shot dx_y gradient correction #1268

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ec0ad32
perf(ce): single post-loop dx_y correction (drop per-element tl.where)
justinhh4 Jun 25, 2026
f517468
Merge branch 'main' into ce-opt-3-dxy
vaibhavjindal Jul 13, 2026
60c1d22
Merge branch 'main' into ce-opt-3-dxy
vaibhavjindal Jul 28, 2026
347e38e
address nit and removed comments
justinhh4 Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions src/liger_kernel/ops/cross_entropy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,8 +216,6 @@ def liger_cross_entropy_kernel(
X_block += 2 * lse_square_scale * lse * X_block
# smoothing term
X_block += -eps
# special handle dx_y
X_block = tl.where(X_offsets != y, X_block, X_block - (1 - label_smoothing))
# reduction scale
if reduction == "mean":
X_block = X_block / n_non_ignore
Expand All @@ -226,8 +224,6 @@ def liger_cross_entropy_kernel(
softmax_X = tl.exp2((X_block - m) * LOG2_E) / d
# derivative of original_loss
dloss_ori = (1 - label_smoothing) * softmax_X
# specially handle dx_y
dloss_ori = tl.where(X_offsets != y, dloss_ori, dloss_ori - (1 - label_smoothing))
dloss_ori = dloss_ori * weight_y
# derivative of smooth_loss
dloss_smooth = eps * (-weight_block + softmax_X * weight_sum)
Expand All @@ -249,6 +245,22 @@ def liger_cross_entropy_kernel(

tl.store(X_ptr + X_offsets, X_block, mask=X_offsets < n_cols)

# dx_y correction: apply the -(1 - label_smoothing) term once at index y
# (replaces the per-element tl.where removed above). Barrier first so the loop's
# in-place store to X[y] is visible before we read it back.
tl.debug_barrier()
dxy = -(1 - label_smoothing)
if HAS_WEIGHT:
dxy = dxy * weight_y
if reduction == "mean":
dxy = dxy / sum_non_ignore_weight
elif reduction == "mean":
dxy = dxy / n_non_ignore
if HAS_SOFTCAPPING:
t_y = ori_X_y / softcap
dxy = dxy * (1 - t_y * t_y)
tl.store(X_ptr + y, tl.load(X_ptr + y) + dxy)

# We need tl.debug_barrier() to ensure the new result of X_ptr is written as mentioned in
# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/ops/cross_entropy.py#L34
tl.debug_barrier()
Expand Down
Loading