wip: T10 batched forward (validation)

crates/xtrain-model/tests/parity.py

@@ -55,10 +55,13 @@ NH = int(cfg["n_heads"])
 HD = int(cfg["head_dim"])
 EPS = float(cfg["eps"])
 THETA = float(cfg["rope_theta"])
+# Batched: B sequences of length SEQ, flattened sequence-major to [B*SEQ] ids.
+B = int(cfg.get("batch", "1"))
+SEQ = int(cfg["seq"])
 
 ids = read_ids("ids.txt")
 targets = read_ids("targets.txt")
-SEQ = len(ids)
+assert len(ids) == B * SEQ, f"ids {len(ids)} != B*SEQ {B*SEQ}"
 
 # Load params as leaf tensors requiring grad (float64 for a clean reference).
 P = {}
@@ -76,15 +79,16 @@ def rms_norm(x, gamma):
     return x * torch.rsqrt(ms + EPS) * gamma
 
 
-def rope(x):  # x: [seq, nh, hd], position = token index, matching the kernel
+def rope(x):  # x: [B*SEQ, nh, hd], position = (row % SEQ) — resets per sequence
     half = HD // 2
     out = torch.empty_like(x)
     i = torch.arange(half, dtype=torch.float64)
     freq = THETA ** (-(2.0 * i) / HD)  # [half]
-    pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1)  # [seq,1]
-    ang = pos * freq  # [seq, half]
-    c = torch.cos(ang).reshape(SEQ, 1, half)
-    s = torch.sin(ang).reshape(SEQ, 1, half)
+    # Position within each sequence: rows 0..SEQ for seq 0, 0..SEQ for seq 1, ...
+    pos = (torch.arange(B * SEQ, dtype=torch.float64) % SEQ).reshape(B * SEQ, 1)
+    ang = pos * freq  # [B*SEQ, half]
+    c = torch.cos(ang).reshape(B * SEQ, 1, half)
+    s = torch.sin(ang).reshape(B * SEQ, 1, half)
     x0 = x[..., :half]
     x1 = x[..., half:]
     out[..., :half] = x0 * c - x1 * s
@@ -102,26 +106,30 @@ for l in range(NL):
                     "ffn_norm", "w_gate", "w_up", "w_down"]})
 
 idx = torch.tensor(ids, dtype=torch.long)
+# Per-sequence causal mask (broadcast over the batch); NO cross-sequence attention.
 mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)
 
-h = emb[idx]  # [seq, dim]
+h = emb[idx]  # [B*SEQ, dim]  (everything stays flattened, matching the Rust path)
 for L in layers:
     # Attention
     x = rms_norm(h, L["attn_norm"])
-    q = (x @ L["wq"]).reshape(SEQ, NH, HD)
-    k = (x @ L["wk"]).reshape(SEQ, NH, HD)
-    v = (x @ L["wv"]).reshape(SEQ, NH, HD)
+    q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
+    k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
+    v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
     # Per-head QK-norm (Qwen3-style), before RoPE.
     q = rms_norm(q, L["q_norm"])
     k = rms_norm(k, L["k_norm"])
-    q = rope(q).transpose(0, 1)  # [nh, seq, hd]
-    k = rope(k).transpose(0, 1)
-    v = v.transpose(0, 1)
+    q = rope(q)  # [B*SEQ, nh, hd]
+    k = rope(k)
+    # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
+    q = q.reshape(B, SEQ, NH, HD).transpose(1, 2)  # [B, nh, seq, hd]
+    k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
+    v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
     scale = 1.0 / math.sqrt(HD)
-    scores = (q @ k.transpose(-1, -2)) * scale + mask  # [nh, seq, seq]
+    scores = (q @ k.transpose(-1, -2)) * scale + mask  # [B, nh, seq, seq]
     probs = torch.softmax(scores, dim=-1)
-    out = probs @ v  # [nh, seq, hd]
-    out = out.transpose(0, 1).reshape(SEQ, DIM)  # [seq, dim]
+    out = probs @ v  # [B, nh, seq, hd]
+    out = out.transpose(1, 2).reshape(B * SEQ, DIM)  # [B*SEQ, dim]
     attn = out @ L["wo"]
     h = h + attn
     # MLP
@@ -133,7 +141,7 @@ for L in layers:
     h = h + mlp
 
 h = rms_norm(h, final_norm)
-logits = h @ lm_head  # [seq, vocab]
+logits = h @ lm_head  # [B*SEQ, vocab]
 
 loss = torch.nn.functional.cross_entropy(
     logits, torch.tensor(targets, dtype=torch.long), reduction="mean")