@triton.jit
def triton_kl_backward(target_ptr, target_stride, grad_output_ptr,
    grad_output_stride, num_classes, BLOCK_SIZE: tl.constexpr, log_target:
    tl.constexpr=False):
    row_id = tl.program_id(0).to(tl.int64)
    target_ptr += row_id * target_stride
    grad_output_ptr += row_id * grad_output_stride
    base_offsets = tl.arange(0, BLOCK_SIZE)
    mask = base_offsets < num_classes
    for i in range(0, num_classes, BLOCK_SIZE):
        offsets = i + base_offsets
        mask = offsets < num_classes
        target_val = tl.load(target_ptr + offsets, mask=mask, other=0.0)
        if not log_target:
            grad = target_val * -1
        else:
            grad = -tl.exp(target_val)
        tl.store(grad_output_ptr + offsets, grad, mask=mask)
