perf(model): batched causal-mask native_train_step (~10.75x), gradient-checked

[InauguralPhysicist]

The transformer trainer native_train_step — the model-extension kernel iLambdaAi's ~24-hour training runs on — re-ran a full forward+backward over the growing prefix for every target position in each window: ~31 nested passes over prefixes of length 1…31 per 32-token window.

The attention is causal (model_infer.c masks j > i), so each position's activations are independent of prefix length — all 31 passes recompute identical work. Standard transformer training does all positions in one causal-masked pass.

Fix

One causal-masked forward over the whole window + a cheap per-position loss/seed loop + one backward. Same forward/backward kernels, called once instead of seq times. Gradient-identical, because every position already trains against the same start-of-step weights, so the masked batch computes the identical summed gradient.

Measured

EIGS_TRAIN_PROFILE, d=32/L=4 model:

	per-position	batched
ms/window	179	16.65
speedup	—	10.75×
~24h run		~2.2h

The profiler showed 97% of the old cost was the redundant forward+backward compute (3% alloc/update), so the win is realized, not eaten by overhead.

Verified against its own oracle

The per-position path is the gradient-check oracle. From a common checkpoint, both paths produce float-identical, non-diverging loss sequences — and that's the decisive test, because a real gradient error compounds exponentially and would visibly diverge within a few steps. It doesn't.

• EIGS_TRAIN_PERPOS=1 restores the oracle path.
• New test_native_train_gradcheck.sh (suite [47c], model build): tiny model, six steps each path, asserts max loss diff < 1e-3 — pins the two paths together so the equivalence can't silently regress.

Gates

• Full http+model suite 3164/3164 (includes [47c])
• Default (non-model) build unaffected — model_train.c isn't compiled there, so the standard suite and ASan legs are untouched
• Gated profiler (EIGS_TRAIN_PROFILE) kept as the A/B instrument, zero-cost off

Note: this makes the search for a working iLambdaAi model ~10× cheaper per experiment — the point is A/B iteration speed, not a faster working model (there isn't one yet).

🤖 Generated with Claude Code