perf: keep bf16 logits (no persistent fp32 logits buffer)

At vocab 50257 the logits tensor [B*S, vocab] is ~1.6GB fp32 at batch
32 — held across the whole backward. Keep it bf16: cross_entropy
upcasts the bf16 logits to fp32 internally (transient) + caches fp32
probs, and its backward casts dx back to bf16 to chain into the
bf16 lm_head matmul backward. The sampler casts bf16 logits→f32 before
the host argmax/softmax. Halves the persistent logits activation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-model/src/model.rs

@@ -183,13 +183,10 @@ impl TinyTransformer {
         }
 
         let h = ops::rms_norm(&h, &self.norm_gamma(&self.final_norm), self.cfg.eps);
-        // lm_head matmul in compute dtype; cast logits back to fp32 for CE.
-        let logits = self.linear(&h, &self.lm_head); // [batch*seq, vocab]
-        if self.compute_dtype == DType::BF16 {
-            ops::cast(&logits, DType::F32)
-        } else {
-            logits
-        }
+        // lm_head matmul in compute dtype. Logits stay bf16 in bf16 mode — the
+        // cross_entropy op upcasts to fp32 internally (no persistent fp32 logits
+        // buffer, a real saving at vocab 50257), and its backward casts dx back.
+        self.linear(&h, &self.lm_head) // [batch*seq, vocab]
     }
 
     /// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast