feat: Add CUDA cross-entropy loss forward+backward kernel

TimDettmers · claude · TimDettmers · commit f0513d79bd0b · 2026-02-22T16:23:14.000-05:00
Cross-entropy loss with numerically stable logsumexp:
- Forward: per-sample loss via max-stabilized logsumexp
- Backward: softmax(logits) - one_hot(labels), scaled by grad_output
- One block per row, shared memory reduction for max and sum
- Supports ignore_index (-100 convention)
- Supports large vocabularies (tested up to 100K)
- fp16 and bf16 via C++ templates

Autograd wrapper computes mean loss with proper ignore_index
handling and per-sample gradient scaling.

10 tests pass covering forward correctness, gradient correctness,
ignore_index handling, large vocab, and both dtypes.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/bitsandbytes/_ops.py b/bitsandbytes/_ops.py
@@ -782,3 +782,37 @@ def _(
 @register_fake("bitsandbytes::rope_forward")
 def _(q: torch.Tensor, cos_cache: torch.Tensor, sin_cache: torch.Tensor, n_heads: int) -> None:
     pass
+
+
+# Cross-Entropy Loss forward: per-sample loss + logsumexp for backward
+torch.library.define(
+    "bitsandbytes::cross_entropy_forward",
+    "(Tensor logits, Tensor labels, int ignore_index) -> (Tensor, Tensor)",
+)
+
+
+@register_fake("bitsandbytes::cross_entropy_forward")
+def _(logits: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(logits.dim() == 2, lambda: "logits must be 2D [N, V]")
+    N = logits.shape[0]
+    losses = torch.empty(N, device=logits.device, dtype=torch.float32)
+    logsumexp = torch.empty(N, device=logits.device, dtype=torch.float32)
+    return losses, logsumexp
+
+
+# Cross-Entropy Loss backward: grad_logits from grad_output
+torch.library.define(
+    "bitsandbytes::cross_entropy_backward",
+    "(Tensor logits, Tensor labels, Tensor grad_output, Tensor logsumexp, int ignore_index) -> Tensor",
+)
+
+
+@register_fake("bitsandbytes::cross_entropy_backward")
+def _(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    grad_output: torch.Tensor,
+    logsumexp: torch.Tensor,
+    ignore_index: int,
+) -> torch.Tensor:
+    return torch.empty_like(logits)
diff --git a/bitsandbytes/autograd/training_kernels.py b/bitsandbytes/autograd/training_kernels.py
@@ -145,3 +145,76 @@ def rope(
         Rotated query tensor (same shape).
     """
     return RoPEFunction.apply(q, cos_cache, sin_cache, n_heads)
+
+
+class CrossEntropyFunction(torch.autograd.Function):
+    """Cross-entropy loss using CUDA kernel.
+
+    Forward: loss = -log_softmax(logits)[label] per row
+    Backward: grad_logits = (softmax(logits) - one_hot(label)) * grad_output
+
+    Stores logsumexp from forward for efficient backward (avoids recomputing).
+    """
+
+    @staticmethod
+    def forward(ctx, logits, labels, ignore_index=-100):
+        # Flatten to 2D for the CUDA kernel
+        orig_shape = logits.shape
+        logits_2d = logits.reshape(-1, logits.shape[-1]).contiguous()
+        labels_flat = labels.reshape(-1)
+
+        losses, logsumexp = torch.ops.bitsandbytes.cross_entropy_forward(
+            logits_2d, labels_flat, ignore_index,
+        )
+
+        ctx.save_for_backward(logits_2d, labels_flat, logsumexp)
+        ctx.ignore_index = ignore_index
+
+        # Compute mean loss (ignoring padding)
+        valid_mask = labels_flat != ignore_index
+        n_valid = valid_mask.sum()
+        if n_valid > 0:
+            mean_loss = losses[valid_mask].sum() / n_valid.float()
+        else:
+            mean_loss = losses.sum() * 0.0  # zero but with grad
+
+        return mean_loss
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        logits_2d, labels_flat, logsumexp = ctx.saved_tensors
+
+        # Expand scalar grad_output to per-sample
+        N = logits_2d.shape[0]
+        valid_mask = labels_flat != ctx.ignore_index
+        n_valid = valid_mask.sum()
+
+        grad_per_sample = torch.zeros(N, device=logits_2d.device, dtype=torch.float32)
+        if n_valid > 0:
+            grad_per_sample[valid_mask] = grad_output.float() / n_valid.float()
+
+        grad_logits = torch.ops.bitsandbytes.cross_entropy_backward(
+            logits_2d, labels_flat, grad_per_sample, logsumexp, ctx.ignore_index,
+        )
+
+        return grad_logits, None, None
+
+
+def cross_entropy(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Cross-entropy loss with CUDA kernel (autograd support).
+
+    Uses a numerically stable logsumexp-based implementation.
+
+    Args:
+        logits: Logit tensor (*, vocab_size), fp16 or bf16.
+        labels: Label tensor (*), int64.
+        ignore_index: Label value to ignore (default: -100).
+
+    Returns:
+        Scalar mean loss.
+    """
+    return CrossEntropyFunction.apply(logits, labels, ignore_index)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
@@ -1372,3 +1372,66 @@ def _(
             ct.c_int(n_heads),
             ct.c_int(head_dim),
         )
+
+
+@register_kernel("bitsandbytes::cross_entropy_forward", "cuda")
+def _(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    ignore_index: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(logits.dim() == 2, lambda: "logits must be 2D [N, V]")
+    torch._check(logits.is_contiguous(), lambda: "logits must be contiguous")
+    torch._check(
+        logits.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"cross_entropy supports float16/bfloat16, got {logits.dtype}",
+    )
+
+    N, V = logits.shape
+    losses = torch.empty(N, device=logits.device, dtype=torch.float32)
+    logsumexp = torch.empty(N, device=logits.device, dtype=torch.float32)
+    dtype_suffix = "fp16" if logits.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(logits):
+        fn = getattr(lib, f"ccross_entropy_forward_{dtype_suffix}_c")
+        fn(
+            get_ptr(logits),
+            get_ptr(labels),
+            get_ptr(losses),
+            get_ptr(logsumexp),
+            ct.c_int(N),
+            ct.c_int(V),
+            ct.c_int(ignore_index),
+        )
+
+    return losses, logsumexp
+
+
+@register_kernel("bitsandbytes::cross_entropy_backward", "cuda")
+def _(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    grad_output: torch.Tensor,
+    logsumexp: torch.Tensor,
+    ignore_index: int,
+) -> torch.Tensor:
+    torch._check(logits.is_contiguous(), lambda: "logits must be contiguous")
+
+    N, V = logits.shape
+    grad_logits = torch.empty_like(logits)
+    dtype_suffix = "fp16" if logits.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(logits):
+        fn = getattr(lib, f"ccross_entropy_backward_{dtype_suffix}_c")
+        fn(
+            get_ptr(logits),
+            get_ptr(labels),
+            get_ptr(grad_output),
+            get_ptr(logsumexp),
+            get_ptr(grad_logits),
+            ct.c_int(N),
+            ct.c_int(V),
+            ct.c_int(ignore_index),
+        )
+
+    return grad_logits
diff --git a/csrc/ops.cu b/csrc/ops.cu
@@ -3263,3 +3263,143 @@ void rope_forward(T* q, const T* cos_cache, const T* sin_cache,
 
 template void rope_forward<half>(half*, const half*, const half*, int, int, int);
 template void rope_forward<__nv_bfloat16>(__nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, int, int, int);
+
+// ---------- Cross-Entropy Loss forward+backward ----------
+// Forward: loss = -log_softmax(logits)[label] per row
+//        = -(logits[label] - logsumexp(logits))
+// Uses chunked logsumexp for numerical stability with large vocab.
+// Backward: grad_logits = (softmax(logits) - one_hot(label)) * grad_output
+// One thread block per row (sample in the batch).
+
+template <typename T, int BLOCK_SIZE>
+__global__ void kCrossEntropyForward(
+    const T* __restrict__ logits,   // [N, V]
+    const long* __restrict__ labels, // [N]
+    float* __restrict__ losses,      // [N]
+    float* __restrict__ logsumexp_out, // [N] stored for backward
+    int N, int V, int ignore_index
+) {
+    int row = blockIdx.x;
+    if (row >= N) return;
+
+    long label = labels[row];
+    if (label == ignore_index) {
+        losses[row] = 0.0f;
+        logsumexp_out[row] = 0.0f;
+        return;
+    }
+
+    const T* logits_row = logits + row * V;
+
+    // Phase 1: find max for numerical stability
+    __shared__ float shared[BLOCK_SIZE];
+    float thread_max = -1e30f;
+    for (int i = threadIdx.x; i < V; i += BLOCK_SIZE) {
+        float val = float(logits_row[i]);
+        thread_max = fmaxf(thread_max, val);
+    }
+    shared[threadIdx.x] = thread_max;
+    __syncthreads();
+
+    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (threadIdx.x < s) {
+            shared[threadIdx.x] = fmaxf(shared[threadIdx.x], shared[threadIdx.x + s]);
+        }
+        __syncthreads();
+    }
+    float row_max = shared[0];
+
+    // Phase 2: compute sum(exp(x - max))
+    float thread_sum = 0.0f;
+    for (int i = threadIdx.x; i < V; i += BLOCK_SIZE) {
+        thread_sum += expf(float(logits_row[i]) - row_max);
+    }
+    shared[threadIdx.x] = thread_sum;
+    __syncthreads();
+
+    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (threadIdx.x < s) {
+            shared[threadIdx.x] += shared[threadIdx.x + s];
+        }
+        __syncthreads();
+    }
+    float sum_exp = shared[0];
+
+    float lse = row_max + logf(sum_exp);
+    float logit_label = float(logits_row[label]);
+    float loss = -(logit_label - lse);
+
+    if (threadIdx.x == 0) {
+        losses[row] = loss;
+        logsumexp_out[row] = lse;
+    }
+}
+
+template <typename T, int BLOCK_SIZE>
+__global__ void kCrossEntropyBackward(
+    const T* __restrict__ logits,       // [N, V]
+    const long* __restrict__ labels,     // [N]
+    const float* __restrict__ grad_output, // [N] scalar per sample
+    const float* __restrict__ logsumexp,   // [N] from forward
+    T* __restrict__ grad_logits,         // [N, V]
+    int N, int V, int ignore_index
+) {
+    int row = blockIdx.x;
+    if (row >= N) return;
+
+    long label = labels[row];
+    const T* logits_row = logits + row * V;
+    T* grad_row = grad_logits + row * V;
+    float go = grad_output[row];
+
+    if (label == ignore_index) {
+        for (int i = threadIdx.x; i < V; i += BLOCK_SIZE) {
+            grad_row[i] = T(0.0f);
+        }
+        return;
+    }
+
+    float lse = logsumexp[row];
+
+    for (int i = threadIdx.x; i < V; i += BLOCK_SIZE) {
+        float softmax_i = expf(float(logits_row[i]) - lse);
+        float grad_i = softmax_i;
+        if (i == label) {
+            grad_i -= 1.0f;
+        }
+        grad_row[i] = T(grad_i * go);
+    }
+}
+
+// C wrapper functions
+template <typename T>
+void cross_entropy_forward(const T* logits, const long* labels, float* losses, float* logsumexp,
+                           int N, int V, int ignore_index) {
+    if (V <= 256) {
+        kCrossEntropyForward<T, 256><<<N, 256>>>(logits, labels, losses, logsumexp, N, V, ignore_index);
+    } else if (V <= 512) {
+        kCrossEntropyForward<T, 512><<<N, 512>>>(logits, labels, losses, logsumexp, N, V, ignore_index);
+    } else {
+        kCrossEntropyForward<T, 1024><<<N, 1024>>>(logits, labels, losses, logsumexp, N, V, ignore_index);
+    }
+    CUDA_CHECK_RETURN(cudaPeekAtLastError());
+}
+
+template <typename T>
+void cross_entropy_backward(const T* logits, const long* labels, const float* grad_output,
+                            const float* logsumexp, T* grad_logits,
+                            int N, int V, int ignore_index) {
+    if (V <= 256) {
+        kCrossEntropyBackward<T, 256><<<N, 256>>>(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+    } else if (V <= 512) {
+        kCrossEntropyBackward<T, 512><<<N, 512>>>(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+    } else {
+        kCrossEntropyBackward<T, 1024><<<N, 1024>>>(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+    }
+    CUDA_CHECK_RETURN(cudaPeekAtLastError());
+}
+
+template void cross_entropy_forward<half>(const half*, const long*, float*, float*, int, int, int);
+template void cross_entropy_forward<__nv_bfloat16>(const __nv_bfloat16*, const long*, float*, float*, int, int, int);
+template void cross_entropy_backward<half>(const half*, const long*, const float*, const float*, half*, int, int, int);
+template void cross_entropy_backward<__nv_bfloat16>(const __nv_bfloat16*, const long*, const float*, const float*, __nv_bfloat16*, int, int, int);
diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp
@@ -652,6 +652,29 @@ void crope_forward_bf16(__nv_bfloat16* q, const __nv_bfloat16* cos_cache, const
     rope_forward<__nv_bfloat16>(q, cos_cache, sin_cache, total_tokens, n_heads, head_dim);
 }
 
+// Cross-entropy loss forward declarations
+template <typename T> void cross_entropy_forward(const T*, const long*, float*, float*, int, int, int);
+template <typename T> void cross_entropy_backward(const T*, const long*, const float*, const float*, T*, int, int, int);
+
+void ccross_entropy_forward_fp16(const half* logits, const long* labels, float* losses,
+                                 float* logsumexp, int N, int V, int ignore_index) {
+    cross_entropy_forward<half>(logits, labels, losses, logsumexp, N, V, ignore_index);
+}
+void ccross_entropy_backward_fp16(const half* logits, const long* labels, const float* grad_output,
+                                  const float* logsumexp, half* grad_logits,
+                                  int N, int V, int ignore_index) {
+    cross_entropy_backward<half>(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+}
+void ccross_entropy_forward_bf16(const __nv_bfloat16* logits, const long* labels, float* losses,
+                                 float* logsumexp, int N, int V, int ignore_index) {
+    cross_entropy_forward<__nv_bfloat16>(logits, labels, losses, logsumexp, N, V, ignore_index);
+}
+void ccross_entropy_backward_bf16(const __nv_bfloat16* logits, const long* labels, const float* grad_output,
+                                  const float* logsumexp, __nv_bfloat16* grad_logits,
+                                  int N, int V, int ignore_index) {
+    cross_entropy_backward<__nv_bfloat16>(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+}
+
 #endif // BUILD_CUDA || BUILD_HIP (kbit unmangled)
 
 extern "C" {
@@ -1400,5 +1423,24 @@ void crope_forward_bf16_c(__nv_bfloat16* q, const __nv_bfloat16* cos_cache, cons
     crope_forward_bf16(q, cos_cache, sin_cache, total_tokens, n_heads, head_dim);
 }
 
+void ccross_entropy_forward_fp16_c(const half* logits, const long* labels, float* losses,
+                                   float* logsumexp, int N, int V, int ignore_index) {
+    ccross_entropy_forward_fp16(logits, labels, losses, logsumexp, N, V, ignore_index);
+}
+void ccross_entropy_backward_fp16_c(const half* logits, const long* labels, const float* grad_output,
+                                    const float* logsumexp, half* grad_logits,
+                                    int N, int V, int ignore_index) {
+    ccross_entropy_backward_fp16(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+}
+void ccross_entropy_forward_bf16_c(const __nv_bfloat16* logits, const long* labels, float* losses,
+                                   float* logsumexp, int N, int V, int ignore_index) {
+    ccross_entropy_forward_bf16(logits, labels, losses, logsumexp, N, V, ignore_index);
+}
+void ccross_entropy_backward_bf16_c(const __nv_bfloat16* logits, const long* labels, const float* grad_output,
+                                    const float* logsumexp, __nv_bfloat16* grad_logits,
+                                    int N, int V, int ignore_index) {
+    ccross_entropy_backward_bf16(logits, labels, grad_output, logsumexp, grad_logits, N, V, ignore_index);
+}
+
 #endif
 }
diff --git a/tests/test_training_kernels.py b/tests/test_training_kernels.py
