OptionalflashWhen set (with fused), route the head through these flash-CCE Metal kernels
(makeFlashCceHeadVjp), composed once for the run: in-kernel quantized logits
dh accumulation (bwd), so neither [M,V] logits
nor a dequantized head touch HBM — the lowest-memory head, and the fastest on
large vocab (e4b/CPM). The Apple-CCE coeff filter is on by default in the
backward.OptionalfusedWhen set, use the FUSED linear-CE head (fusedLogpMeanB1: one CustomVjp with
an analytic softmax−onehot backward, no autograd through the head and no
retained [M,V] logits) instead of the per-chunk Checkpoint head. Same
token chunkSize; the fused path bounds large-vocab memory in both
directions and is the Liger/CCE structure (the moat). Off → the Checkpoint
token-chunked head (unchanged).
Disposables whose lifetime must extend past the enclosing value_and_grad
backward: the per-chunk Checkpoints / the fused CustomVjp (they recompute
in the backward) AND the hidden h / chunk slices they read (these
deliberately drop their input activations, so freeing them early corrupts the
recompute — and even the forward eval). The caller (orpoLoop) disposes these
only AFTER the step's grads are eval'd (the afterMicroEval hook).
Token-chunking context for the ORPO loss head. When
chunkSize > 0, the B=1 response-only head is computed in token-chunks, each wrapped in a Checkpoint so its[chunk, vocab]logits are RECOMPUTED in the backward instead of all[M, vocab]being retained — bounding the dominant large-vocab term to[chunkSize, vocab]. The checkpoints recompute during the enclosing value_and_grad backward, so they must outliveorpoLoss: they are pushed intosink, which the caller (orpoLoop) disposes AFTERvag.applyreturns (i.e. after the backward). Exact — rematerialization is numerically identical, and MLX computes the (quantized) head backward itself.