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

# pip install -U kernels
from kernels import get_kernel

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

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

Note on loading. kernels resolves kernels against the Hub's dedicated kernel repo type. Creating one is currently access-gated for this namespace, so this repo is published as a model repo carrying the standard kernel layout. Until it is recreated under the kernel repo type the call above will not resolve; clone the repo and load it locally instead:

from pathlib import Path
from kernels import get_local_kernel

flce = get_local_kernel(Path("/path/to/clone/of/fused-linear-ce"))

The layout, metadata.json and variant resolution are already exactly what a kernel repo needs, so recreating it under the kernel repo type is a straight copy of these files.

Full signature:

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, {"mean_token_accuracy": ..., "entropy": ...})

hidden, weight and bias all receive gradients.

Metrics for free

Trainers that switch to a fused loss normally lose their mean_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:

loss, metrics = fused_linear_cross_entropy(hidden, weight, targets, return_metrics=True)
metrics["mean_token_accuracy"]  # scalar, over non-ignored positions
metrics["entropy"]              # scalar, mean softmax entropy in nats

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 (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.