det-cross-entropy

det-cross-entropy is a bitwise-deterministic cross-entropy loss whose value and gradient are invariant to how the vocabulary is sharded, loadable through kernels. With det-train (the GEMM) and det-attn (the attention softmax), the loss is the third and last non-associative reduction in transformer training; pinning all three makes a loss curve a fixed function of the inputs. The reference baseline is torch.nn.functional.cross_entropy, matched to 1.9e-6 absolute with the gradient at 1.2e-10.

A tensor-parallel language model splits its vocabulary across devices and combines partial log-sum-exp states with a floating-point rescale, so the loss bits depend on the sharding, and the same model on a different world size trains along a different trajectory. This kernel computes the same loss exactly: shards combine by integer addition, so the loss and its gradient are the same bytes on one device or many.

$The vocabulary bar fractures into shards while the loss's hex digits stay identical and the float32 combine flips its last nibble$

One batch, the vocabulary split 1, 2, 4, 8 ways. Every token's det-cross-entropy loss is torch.equal at every split (one hex value); under the float32 logsumexp combine, 38 of 1,024 token losses change bits between splits.

Usage

import torch
from kernels import get_kernel

dce = get_kernel("phanerozoic/det-cross-entropy", version=1, trust_remote_code=True)

logits = model(input_ids)                       # [N, V] f32 or bf16
targets = labels.reshape(-1)                     # [N] int64
loss = dce.cross_entropy(logits, targets, reduction="mean")
loss.backward()                                  # deterministic dLogits

For a vocab-parallel LM head, the loss composes across shards exactly:

# each shard holds logits[:, v_lo:v_hi]; combine with no rescale
m = torch.stack([dce.rowmax(logits, lo, hi) for lo, hi in shards]).amax(0)
den = sum(dce.denom_digits(logits, m, lo, hi) for lo, hi in shards)
loss, denom = dce.finalize(den, logits, targets, m)      # bit-identical

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. reduction is "mean", "sum", or "none".

API

Symbol	Purpose
`cross_entropy(logits, targets, reduction)`	deterministic autograd cross-entropy loss
`forward(logits, targets)`	fused `(loss [N], rowmax [N], denom [N])`, the state reused by backward
`rowmax(logits, v_lo, v_hi)`	per-token max over a vocab range; combine shards by max
`denom_digits(logits, global_rowmax, v_lo, v_hi)`	exact denominator digits over a vocab range; combine by add
`finalize(den_digits, logits, targets, global_rowmax)`	`(loss, denom)` from composed digits

Method

The loss is a softmax reduction: logsumexp_v(x_iv) per token and the gradient softmax(x_i) - onehot(t_i). In floating point the vocabulary sum depends on order, and vocab-parallel partial sums combine by a non-associative rescale. This kernel removes the dependency with two exact reductions: a global row maximum m_i = max_v x_iv (order-invariant), then an exact denominator den_i = sum_v exp(x_iv - m_i) accumulated in a Kulisch long accumulator (signed int64 digits over the full fp32 exponent range). Because m_i is global, every exp(x_iv - m_i) is a fixed float and the sum is exact integer accumulation; shards compose by an int64 digit add and an elementwise max. The loss and gradient are fixed functions of the logits.

Correctness

Measured on L4, H200, and RTX PRO 6000 through get_kernel:

Matches torch: value agrees with torch.nn.functional.cross_entropy to a maximum absolute difference of 1.9e-6 (loss magnitude ~10), gradient with torch autograd to 1.2e-10.
Vocabulary-shard invariance: splitting into 2, 4, and 8 shards, composing row maxima (by max) and denominator digits (by int64 add), and finalizing gives a loss, denominator, and gradient torch.equal to the single-call result.
Deterministic: repeated forward and backward passes are bitwise identical.
f32 and bf16 logits.

Requirements and limits

NVIDIA GPU with compute capability 8.0+.
Exact accumulation on fp64/scalar arithmetic rather than a fused throughput kernel; built for reproducibility, evaluation, debugging, and audit. f32 or bf16 logits; targets int64.

References

The exact order-invariant softmax reduction of phanerozoic/det-attn and the Kulisch long-accumulator GEMM of phanerozoic/det-train; reproducible training under tensor and sequence parallelism.