post-train: M2b — batched KV-cache decode (G-way, token-identical)

The rollout long-pole fix deferred from M2a: decode the G samples of one prompt
in lockstep (one forward per step over the group → G× fewer kernel launches).

- rope_pos(x, positions[]): RoPE with a per-row absolute position (new forward-
  only kernel) — G rows share one decode position. Gate: == full rope for
  [0..n], == rope_at(P) per row for uniform P (bit-identical).
- generate_cached_batch: BatchKVCache [T, G·num_kv, hd] + batched decode_step.
  decode_attention is already batch-agnostic (bh = G·nh); repeat_kv(nh, batch=G)
  broadcasts per group. No finished-mask / ragged prompts yet (perf-only / next).
- Gate (tests/decode_batch.rs): all G greedy rows token-identical to the single-
  sequence decode (8 query / 2 kv heads → exercises repeat_kv batching).

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

crates/xtrain-model/src/decode.rs

@@ -265,3 +265,177 @@ fn argmax(row: &[f32]) -> usize {
         .unwrap()
         .0
 }
+
+// ===================================================================
+// M2b — batched KV-cache decode (G samples of one prompt, in lockstep)
+// ===================================================================
+
+/// Batched K/V cache: `G` sequences advancing together. Per layer, host-accumulates
+/// seq-major `[T, G·num_kv, head_dim]` (one step appends `G·num_kv·hd` f32), rebuilt
+/// to `[G·num_kv, T, hd]` per step. Same host-cache shape as M2a with a G dimension.
+struct BatchKVCache {
+    k: Vec<Vec<f32>>,
+    v: Vec<Vec<f32>>,
+}
+
+impl BatchKVCache {
+    fn new(n_layers: usize) -> Self {
+        Self {
+            k: vec![Vec::new(); n_layers],
+            v: vec![Vec::new(); n_layers],
+        }
+    }
+    fn append(&mut self, li: usize, k_tok: &[f32], v_tok: &[f32]) {
+        self.k[li].extend_from_slice(k_tok);
+        self.v[li].extend_from_slice(v_tok);
+    }
+}
+
+/// Batched KV-cache decode: roll out `n_samples` (G) completions of the SAME
+/// `prompt` in lockstep — all G share the prompt, so they advance at one common
+/// decode position each step (uniform RoPE via `rope_pos`). Returns G full token
+/// sequences (prompt + sampled continuation). The G-way batching amortises the
+/// per-step kernel launches across G (the rollout long-pole). Token-identical per
+/// row to G independent single-sequence decodes (gated by `tests/decode_batch.rs`).
+///
+/// `temperature == 0` ⇒ greedy (all G identical); `> 0` ⇒ independent samples
+/// (per-row draw from one shared `rng_state`). No finished-mask: all G generate
+/// `max_new` tokens; the caller cuts each at `<|endoftext|>` (a perf-only early
+/// stop is the M2b+ follow-up). Ragged (different-length prompts) is also deferred.
+pub fn generate_cached_batch(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    n_samples: usize,
+    max_new: usize,
+    temperature: f32,
+    rng_state: &mut u64,
+) -> Vec<Vec<i32>> {
+    assert!(!prompt.is_empty(), "prompt must be non-empty");
+    assert!(n_samples > 0, "n_samples must be > 0");
+    let cfg = model.config();
+    let cdt = model.compute_dtype();
+    let n_layers = cfg.n_layers;
+    let params: Vec<Tensor> = model.params().iter().map(|p| p.value()).collect();
+    let embed = &params[0];
+    let final_norm = &params[1 + n_layers * 11];
+    let lm_head = &params[1 + n_layers * 11 + 1];
+
+    let g = n_samples;
+    let mut cache = BatchKVCache::new(n_layers);
+    let mut seqs: Vec<Vec<i32>> = vec![prompt.to_vec(); g];
+
+    // Prefill: feed each prompt token (identical across G) at its position.
+    let mut logits = Vec::new(); // [G, vocab] flattened
+    for (pos, &tok) in prompt.iter().enumerate() {
+        let toks = vec![tok; g];
+        logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &toks, pos, embed, final_norm, lm_head);
+    }
+
+    let vocab = cfg.vocab;
+    for _ in 0..max_new {
+        let mut next = Vec::with_capacity(g);
+        for row in 0..g {
+            let lg = &logits[row * vocab..(row + 1) * vocab];
+            let t = if temperature <= 0.0 {
+                argmax(lg) as i32
+            } else {
+                sample_temperature(lg, temperature, rng_state) as i32
+            };
+            next.push(t);
+            seqs[row].push(t);
+        }
+        let pos = seqs[0].len() - 1; // all G are at the same position
+        logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &next, pos, embed, final_norm, lm_head);
+    }
+    seqs
+}
+
+/// One batched decode step: `toks` is one current token per sequence (`[G]`), all at
+/// absolute position `pos`. Appends each sequence's K/V and returns logits `[G·vocab]`.
+#[allow(clippy::too_many_arguments)]
+fn decode_step_batch(
+    params: &[Tensor],
+    cfg: &crate::Config,
+    cdt: DType,
+    device: Device,
+    cache: &mut BatchKVCache,
+    toks: &[i32],
+    pos: usize,
+    embed: &Tensor,
+    final_norm: &Tensor,
+    lm_head: &Tensor,
+) -> Vec<f32> {
+    let (nh, hd, num_kv) = (cfg.n_heads, cfg.head_dim, cfg.num_kv_heads);
+    let dim = cfg.dim;
+    let g = toks.len();
+    let g_kv = g * num_kv; // batch·num_kv heads in the cache
+    let scale = 1.0 / (hd as f32).sqrt();
+    let (theta, eps) = (cfg.rope_theta, cfg.eps);
+    let n_layers = cfg.n_layers;
+    // Uniform per-row position (all G at the same decode step).
+    let positions = Tensor::from_slice(&vec![pos as i32; g], &[g]).to_device(device);
+
+    let ids = Tensor::from_slice(toks, &[g]).to_device(device);
+    let mut h = embed.embedding(&ids); // [G, dim] f32
+    if cdt == DType::BF16 {
+        h = h.to_dtype(DType::BF16);
+    }
+
+    for li in 0..n_layers {
+        let base = 1 + li * 11;
+        let (attn_norm, wq, wk, wv) =
+            (&params[base], &params[base + 1], &params[base + 2], &params[base + 3]);
+        let (q_norm, k_norm, wo) = (&params[base + 4], &params[base + 5], &params[base + 6]);
+        let (ffn_norm, w_gate, w_up, w_down) =
+            (&params[base + 7], &params[base + 8], &params[base + 9], &params[base + 10]);
+
+        let normed = h.rms_norm(&gamma_t(cdt, attn_norm), eps).0; // [G, dim]
+
+        // Q: project → per-head QK-norm → RoPE at `pos` for every row.
+        let q = linear_t(cdt, &normed, wq).reshape(&[g, nh, hd]);
+        let q = q.reshape(&[g * nh, hd]).rms_norm(&gamma_t(cdt, q_norm), eps).0;
+        let q = q.reshape(&[g, nh, hd]).rope_pos(&positions, theta);
+        let q_bh = q.reshape(&[g * nh, 1, hd]); // bh = G·nh
+
+        let k = linear_t(cdt, &normed, wk).reshape(&[g, num_kv, hd]);
+        let k = k.reshape(&[g * num_kv, hd]).rms_norm(&gamma_t(cdt, k_norm), eps).0;
+        let k_tok = k.reshape(&[g, num_kv, hd]).rope_pos(&positions, theta);
+        let v_tok = linear_t(cdt, &normed, wv).reshape(&[g, num_kv, hd]);
+
+        let k_host = k_tok.to_dtype(DType::F32).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+        let v_host = v_tok.to_dtype(DType::F32).to_device(Device::Cpu).as_slice::<f32>().to_vec();
+        cache.append(li, &k_host, &v_host);
+
+        // Rebuild [T, G·num_kv, hd] → [G·num_kv, T, hd] → repeat_kv to [G·nh, T, hd].
+        let t_len = cache.k[li].len() / (g_kv * hd);
+        let build = |flat: &[f32]| -> Tensor {
+            let bh_kv = Tensor::from_slice(flat, &[t_len, g_kv, hd])
+                .to_device(device)
+                .transpose_3d01(); // [G·num_kv, T, hd], f32
+            let bh_kv = if cdt == DType::BF16 { bh_kv.to_dtype(DType::BF16) } else { bh_kv };
+            if num_kv == nh { bh_kv } else { bh_kv.repeat_kv(nh, g) } // [G·nh, T, hd]
+        };
+        let k_full = build(&cache.k[li]);
+        let v_full = build(&cache.v[li]);
+
+        let attn = q_bh.decode_attention(&k_full, &v_full, scale); // [G·nh, hd]
+        let attn = attn.reshape(&[g, dim]); // concat heads per sequence
+        let attn_out = linear_t(cdt, &attn, wo);
+        h = h.add(&attn_out);
+
+        let normed = h.rms_norm(&gamma_t(cdt, ffn_norm), eps).0;
+        let gate = linear_t(cdt, &normed, w_gate);
+        let up = linear_t(cdt, &normed, w_up);
+        let act = gate.silu().mul(&up);
+        let down = linear_t(cdt, &act, w_down);
+        h = h.add(&down);
+    }
+
+    let h = h.rms_norm(&gamma_t(cdt, final_norm), eps).0;
+    linear_t(cdt, &h, lm_head)
+        .to_dtype(DType::F32)
+        .to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .to_vec()
+}

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_cached, generate_greedy_cached};
+pub use decode::{generate_cached, generate_cached_batch, generate_greedy_cached};