xtrain/crates/xtrain-model/tests/ragged_batch.rs

// M2d gate: does forward_batched on RIGHT-PADDED ragged sequences reproduce the
// per-sequence single-seq forward on the real (non-pad) rows? The batched GRPO
// training-side forwards depend on this "right-pad is free under causal attention"
// property — a real completion row is at an earlier position than the trailing pad,
// and causal masking forbids attending forward, so its logits should be unchanged.
//
// Tested in fp32 (exact) over both SDPA cores (composed + fused flash), since the
// bench uses flash and a kernel could in principle leak the pad keys into the online
// softmax.
#![cfg(not(no_cuda))]

use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor};
use xtrain_tensor::{DType, Device, Tensor};

fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
    let mut state = seed.wrapping_mul(2862933555777941757).wrapping_add(3037000493);
    (0..n)
        .map(|_| {
            state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
        })
        .collect()
}

fn build(cfg: Config, device: Device, dtype: DType, flash: bool) -> TinyTransformer {
    let mut seed = 1u64;
    let m = TinyTransformer::new(cfg, device, |shape| {
        seed = seed.wrapping_add(1);
        let n: usize = shape.iter().product();
        if shape.len() == 1 {
            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
        } else {
            fill(n, seed, 0.08)
        }
    });
    m.with_compute_dtype(dtype).with_flash(flash)
}

fn host(t: &Tensor) -> Vec<f32> {
    t.to_dtype(DType::F32).to_device(Device::Cpu).as_slice::<f32>().to_vec()
}

#[test]
fn forward_batched_ragged_matches_looped() {
    if device::device_count().unwrap_or(0) == 0 {
        eprintln!("no CUDA device; skipping");
        return;
    }
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    let mut cfg = Config::tiny();
    cfg.vocab = 32;
    cfg.n_layers = 2;
    let vocab = cfg.vocab;

    // Ragged lengths incl. one crossing the flash tile (>32) and short ones.
    let lens = [6usize, 40, 9, 4];
    let lmax = *lens.iter().max().unwrap();
    let n = lens.len();
    let seqs: Vec<Vec<i32>> = lens
        .iter()
        .enumerate()
        .map(|(b, &l)| (0..l).map(|i| ((b * 7 + i * 3 + 1) % vocab) as i32).collect())
        .collect();

    for (dtype, tol) in [(DType::F32, 2e-3f32), (DType::BF16, 3e-1f32)] {
      for flash in [false, true] {
        let m = build(cfg, device, dtype, flash);
        // Looped: each sequence on its own (the ground truth).
        let looped: Vec<Vec<f32>> = seqs.iter().map(|s| host(&m.forward(&ids_tensor(s, device)).value())).collect();

        // Batched: right-pad each to lmax (pad id 0), one forward_batched(batch = n).
        let mut flat = vec![0i32; n * lmax];
        for (i, s) in seqs.iter().enumerate() {
            flat[i * lmax..i * lmax + s.len()].copy_from_slice(s);
        }
        let ids = Tensor::from_slice(&flat, &[n * lmax]).to_device(device);
        let batched = host(&m.forward_batched(&ids, n).value()); // [n*lmax, vocab]

        let mut dmax = 0f32;
        for (i, s) in seqs.iter().enumerate() {
            for r in 0..s.len() {
                for c in 0..vocab {
                    let a = looped[i][r * vocab + c];
                    let b = batched[(i * lmax + r) * vocab + c];
                    dmax = dmax.max((a - b).abs());
                }
            }
        }
        println!("dtype={dtype:?} flash={flash}: ragged right-pad vs looped, max|Δlogit| (real rows) = {dmax:.3e}");
        assert!(dmax < tol, "dtype={dtype:?} flash={flash}: right-pad NOT free under causal — max|Δ| = {dmax}");
      }
    }
    println!("forward_batched_ragged_matches_looped OK: right-pad is free under causal (fp32+bf16, composed + flash)");
}