Fused linear-CE log-prob from a PRE-COMPUTED post-finalNorm hidden
h [1, T, hidden] (the segmented-backward head: finalNorm is applied by the
caller). Mirrors responseOnlyLogpMean's response-span selection but routes
the head through the fused analytic CustomVjp, so the head term inside the
segmented backward is bounded to [chunk, V] instead of the full [M, V].
Composes with the segmented mlx_vjp (the CustomVjp nests inside it, like the
fused-GeGLU kernel). The CustomVjp + the response slice are pushed into sink,
disposed by the caller AFTER the head vjp is eval'd.
Fused linear-CE log-prob from a PRE-COMPUTED post-finalNorm hidden
h [1, T, hidden](the segmented-backward head: finalNorm is applied by the caller). MirrorsresponseOnlyLogpMean's response-span selection but routes the head through the fused analytic CustomVjp, so the head term inside the segmented backward is bounded to[chunk, V]instead of the full[M, V]. Composes with the segmentedmlx_vjp(the CustomVjp nests inside it, like the fused-GeGLU kernel). The CustomVjp + the response slice are pushed intosink, disposed by the caller AFTER the head vjp is eval'd.