post-train: M4 — GRPO actor-learner loop + cached temperature rollout

train_grpo: the online, critic-free RL loop — per step sample B prompts, roll
out G completions each, score with the rule-based checker (reward 0/1), compute
group-relative advantage A=(r−mean)/(std+ε), then K inner clipped_pg_loss
epochs with a KL leash to the frozen reference. Reward = pure 0/1 correctness
(KL is the format protector, the M3 collapse lesson). Tracks mean rollout reward
(the falsifiable "it learns" signal). Periodic checkpoint save.

decode: generate_cached adds temperature sampling to the KV-cache engine (M2) —
single-row [1,vocab] logits per step vs the naive sampler's [seq,vocab], far
lighter on the caching allocator (the naive sampler fragments it over a long
rollout). generate_greedy_cached now routes through it (temp 0); decode_kv
token-identical gate still passes.

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

crates/xtrain-model/src/decode.rs

@@ -83,6 +83,24 @@ pub fn generate_greedy_cached(
     device: Device,
     prompt: &[i32],
     max_new: usize,
+) -> Vec<i32> {
+    let mut rng = 0u64;
+    generate_cached(model, device, prompt, max_new, 0.0, &mut rng)
+}
+
+/// KV-cache decode with temperature sampling (`temperature == 0` → greedy argmax,
+/// matching [`generate_greedy_cached`]; otherwise sample from `softmax(logits/T)`).
+/// The KV-cache rollout the GRPO loop uses: each step allocates only a single-row
+/// `[1, vocab]` logits buffer (vs the naive sampler's `[seq, vocab]`), so it is far
+/// lighter on memory + the allocator — the naive sampler fragments the caching
+/// allocator over a long rollout, which is the M4 "rollout is the long pole" wall.
+pub fn generate_cached(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    max_new: usize,
+    temperature: f32,
+    rng_state: &mut u64,
 ) -> Vec<i32> {
     assert!(!prompt.is_empty(), "prompt must be non-empty");
     let cfg = model.config();
@@ -116,7 +134,11 @@ pub fn generate_greedy_cached(
     }
 
     for _ in 0..max_new {
-        let next = argmax(&logits) as i32;
+        let next = if temperature <= 0.0 {
+            argmax(&logits) as i32
+        } else {
+            sample_temperature(&logits, temperature, rng_state) as i32
+        };
         tokens.push(next);
         let pos = tokens.len() - 1; // absolute position of the token just appended
         logits = decode_step(&params, cfg, cdt, device, &mut cache, next, pos, embed, final_norm, lm_head);
@@ -124,6 +146,26 @@ pub fn generate_greedy_cached(
     tokens
 }
 
+/// Sample a token from `softmax(logits / temperature)` (numerically stable). Same
+/// LCG + inverse-CDF scheme as the naive `sample::sample_temperature`.
+fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
+    let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exps: Vec<f32> = row.iter().map(|&x| ((x - max) / temperature).exp()).collect();
+    let sum: f32 = exps.iter().sum();
+    *rng_state = rng_state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    let r = ((*rng_state >> 32) as f32 / u32::MAX as f32) * sum;
+    let mut acc = 0.0;
+    for (i, &e) in exps.iter().enumerate() {
+        acc += e;
+        if acc >= r {
+            return i;
+        }
+    }
+    exps.len() - 1
+}
+
 /// One incremental decode step for token `tok` at absolute position `pos`: append
 /// its K/V to the cache and return the next-token logits as host f32 `[vocab]`.
 #[allow(clippy::too_many_arguments)]

crates/xtrain-model/src/lib.rs

@@ -29,4 +29,4 @@ pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};
 #[cfg(not(no_cuda))]
 pub mod decode;
 #[cfg(not(no_cuda))]
-pub use decode::generate_greedy_cached;
+pub use decode::{generate_cached, generate_greedy_cached};