---
library_name: kernels
license: apache-2.0
tags:
- kernel
- triton
- cross-entropy
---

# fused-linear-ce

A Triton **fused linear cross-entropy** kernel: it computes the cross-entropy of
`hidden @ weight.T (+ bias)` against `targets` **without ever materialising the `(N, V)` logit
matrix**. At a large vocabulary this turns the single most expensive step of LM training from a
tens-of-gigabytes operation into a flat, small one.

## Usage

```python
# pip install -U kernels
from kernels import get_kernel

flce = get_kernel("qgallouedec/fused-linear-ce", version=1, trust_remote_code=True)

loss = flce.fused_linear_cross_entropy(hidden, lm_head.weight, labels)
```

Full signature:

```python
fused_linear_cross_entropy(
    hidden,             # (N, H)   float32 / bfloat16 / float16
    weight,             # (V, H)   same dtype as hidden
    targets,            # (N,)     int64
    bias=None,          # (V,) or None
    ignore_index=-100,
    logit_scale=1.0,
    softcap=None,
    reduction="mean",   # "mean" | "sum" | "none"
    return_metrics=False,
) -> loss  |  (loss, {"num_correct_tokens": ..., "entropy_sum": ...})
```

`hidden`, `weight` and `bias` all receive gradients.

### Metrics for free

Trainers that switch to a fused loss normally lose their token-accuracy and entropy logging,
because there are no logits left to compute them from. This kernel carries both through the same
vocabulary sweep it already performs — the argmax rides along with the running row max, and the
entropy term `sum_j p_j * logit_j` is accumulated with the same online rescaling as the
log-sum-exp:

```python
loss, metrics = fused_linear_cross_entropy(hidden, weight, targets, return_metrics=True)
metrics["num_correct_tokens"]  # int64 scalar, count of argmax == target
metrics["entropy_sum"]         # float32 scalar, summed softmax entropy in nats
```

**Both are sums over non-ignored positions, not means.** Sums are the more general contract: a
per-rank mean cannot be recombined across distributed ranks holding different token counts, so a
distributed trainer must gather sums before dividing, and the accuracy count stays an exact
integer instead of a float that callers have to multiply back up. A single-process caller
recovers the means by dividing by a count it already has:

```python
n = (targets != -100).sum()
mean_token_accuracy = metrics["num_correct_tokens"] / n
mean_entropy = metrics["entropy_sum"] / n
```

Both are forward-only. `return_metrics=True` leaves the loss and the gradients **bit-identical**
to `return_metrics=False`.

## Algorithm

Implements the method described in
[*Cut Your Losses in Large-Vocabulary Language Models*](https://huggingface.co/papers/2411.09009)
(Wijmans et al., 2024).

The forward pass tiles over the vocabulary. For each tile it computes the logit tile with `tl.dot`
into an fp32 accumulator, applies bias / `logit_scale` / softcap, and merges the tile into a
running log-sum-exp using the online-softmax rescaling that flash-attention uses. Nothing of size
`(N, V)` is ever written. The backward pass **recomputes** each logit tile rather than storing it,
forms `p - onehot`, and accumulates `grad_hidden = (p - onehot) @ weight` and
`grad_weight = (p - onehot).T @ hidden`. Recomputation is cheaper than the memory traffic it
avoids.

## Clean-room statement

This is a **clean-room implementation**, written from the published algorithm and from the
standard online-softmax formulation. It contains **no code derived from `apple/ml-cross-entropy`**
or from any other implementation of the method. The upstream package was used only as a
black-box numerical and performance oracle. Licensed **Apache-2.0**.

## Backends

`build/torch-cuda` only. The sources are pure Triton and carry no CUDA-specific intrinsics, but
ROCm and XPU builds are not shipped because they were not tested.