@triton.jit
def triton_kl_forward(y_pred_ptr, y_pred_stride, y_true_ptr, y_true_stride,
    output_loss_ptr, output_loss_stride, num_classes, epsilon, BLOCK_SIZE:
    tl.constexpr, log_target: tl.constexpr=False, reduction_mode: tl.
    constexpr=REDUCE_BATCH_MEAN):
    row_id = tl.program_id(0).to(tl.int64)
    y_pred_ptr += row_id * y_pred_stride
    y_true_ptr += row_id * y_true_stride
    output_loss_ptr += row_id * output_loss_stride
    base_offsets = tl.arange(0, BLOCK_SIZE)
    loss_sum = 0.0
    for i in range(0, num_classes, BLOCK_SIZE):
        offsets = i + base_offsets
        mask = offsets < num_classes
        y_pred = tl.load(y_pred_ptr + offsets, mask=mask, other=0.0)
        y_true = tl.load(y_true_ptr + offsets, mask=mask, other=0.0)
        if not log_target:
            loss = y_true * (tl.log(tl.maximum(y_true, epsilon)) - y_pred)
        else:
            loss = tl.exp(y_true) * (y_true - y_pred)
        if reduction_mode == REDUCE_NONE:
            tl.store(output_loss_ptr + offsets, loss, mask=mask)
        else:
            loss_sum += tl.sum(loss, axis=0)
    if reduction_mode != REDUCE_NONE:
        tl.store(output_loss_ptr, loss_sum)
