Implement the categorical cross-entropy loss with mean reduction in Triton. Given logits of shape (B,C) and integer targets of shape (B,):
loss=−B1i=0∑B−1(logits[i,ti]−logc=0∑C−1elogits[i,c])The log-sum-exp must be computed in the numerically stable form, subtracting the row maximum before the exponentials. This is a row-parallel reduction with a cross-program accumulation: one program per row computes the per-row loss and atomic-adds it into a single-element output buffer that the launcher pre-zeros and finalizes by dividing by B.
Input: logits = [[2.0, 1.0, 0.1]], target = [0]
Output: [0.4170]
Input: logits = [[1.0, 2.0], [2.0, 1.0]], target = [1, 0]
Output: [0.3133]
Each program is one batch row. Load the row of logits with other=-float(’inf’) on the column mask so padded lanes never affect the row max or the sum.
Compute row_max first, then lse=log∑exp(row−row_max)+row_max. The target logit is a scalar load at index ti.
Pre-zero the loss buffer on the host. The kernel adds (lse−target_logit) per row via one atomic per program; divide by B from the launcher after the kernel returns.
logits and loss_out are torch.float32; target is torch.int32 with each value in [0,C).BLOCK_SIZE as a tl.constexpr chosen as triton.next_power_of_2(C) in the launcher.atol = 1e-2 with rtol = 1e-2 because the final value sums a log-sum-exp across all rows and is sensitive to accumulation order.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement the categorical cross-entropy loss with mean reduction in Triton. Given logits of shape (B,C) and integer targets of shape (B,):
loss=−B1i=0∑B−1(logits[i,ti]−logc=0∑C−1elogits[i,c])The log-sum-exp must be computed in the numerically stable form, subtracting the row maximum before the exponentials. This is a row-parallel reduction with a cross-program accumulation: one program per row computes the per-row loss and atomic-adds it into a single-element output buffer that the launcher pre-zeros and finalizes by dividing by B.
Input: logits = [[2.0, 1.0, 0.1]], target = [0]
Output: [0.4170]
Input: logits = [[1.0, 2.0], [2.0, 1.0]], target = [1, 0]
Output: [0.3133]
Each program is one batch row. Load the row of logits with other=-float(’inf’) on the column mask so padded lanes never affect the row max or the sum.
Compute row_max first, then lse=log∑exp(row−row_max)+row_max. The target logit is a scalar load at index ti.
Pre-zero the loss buffer on the host. The kernel adds (lse−target_logit) per row via one atomic per program; divide by B from the launcher after the kernel returns.
logits and loss_out are torch.float32; target is torch.int32 with each value in [0,C).BLOCK_SIZE as a tl.constexpr chosen as triton.next_power_of_2(C) in the launcher.atol = 1e-2 with rtol = 1e-2 because the final value sums a log-sum-exp across all rows and is sensitive to accumulation order.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.