Function cross_entropy_kernel

kernels/cross_entropy.py:21–67

Fused cross-entropy: log_softmax + nll_loss per row. One program per row (batch element).

(
    logits_ptr,
    targets_ptr,
    losses_ptr,
    n_cols,
    stride_logits_row,
    BLOCK_SIZE: tl.constexpr,
)

@triton.jit
def cross_entropy_kernel(
    logits_ptr,
    targets_ptr,
    losses_ptr,
    n_cols,
    stride_logits_row,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Fused cross-entropy: log_softmax + nll_loss per row.
    One program per row (batch element).
    """
    row_idx = tl.program_id(0)

    row_start = logits_ptr + row_idx * stride_logits_row
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # Load logits row in float32
    logits = tl.load(row_start + col_offsets, mask=mask, other=float("-inf")).to(tl.float32)

    # Numerically stable log-softmax
    # Step 1: find max
    row_max = tl.max(logits, axis=0)

    # Step 2: subtract max, exp, sum
    logits_shifted = logits - row_max
    exp_logits = tl.exp(logits_shifted)
    sum_exp = tl.sum(exp_logits, axis=0)
    log_sum_exp = tl.log(sum_exp)

    # log_softmax = logits_shifted - log_sum_exp
    # We only need the value at the target index

    # Load target for this row
    target = tl.load(targets_ptr + row_idx)

    # Get the logit at the target position
    # log_softmax[target] = logits[target] - max - log(sum(exp(logits - max)))
    target_logit = tl.load(row_start + target).to(tl.float32)
    log_softmax_target = (target_logit - row_max) - log_sum_exp

    # NLL loss = -log_softmax[target]
    loss = -log_softmax_target

    # Store per-sample loss
    tl.store(losses_ptr + row_idx, loss)


def kernel_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: