Skip to content
Sunday, 4 October 2026
Tenesys AI News
Subscribe

All News

TENESYS AI NEWS tracks the most important AI news and explains what changed, why it matters and whether a technology is worth testing.

#JAX — 2 articles ✕

Google Reproduces Ai2's Olmo 3 7B Training in MaxText on TPUs

Google's TPU engineering team reproduced the Allen Institute for AI's fully open Olmo 3 7B model from scratch using MaxText, a JAX/XLA training framework, matching Ai2's PyTorch/GPU reference run on held-out metrics rather than just the training loss curve. The reproduction covered stage-1 pre-training (~5.93T tokens, 1.41M steps) and stage-2 mid-training annealing, including porting Olmo 3's non-standard architecture (reordered-norm block, QK-norm, 3:1 sliding/global attention) to JAX. The team also caught and fixed a subtle data-loader bug that caused training loss to falsely appear better due to memorization, verified by held-out evaluation.

Google Developers

Google Shows How to Speed Up Video Diffusion Attention on TPUs by 2.4x

Google engineers describe how they implemented sparse spatio-temporal attention for video diffusion models on TPU v6e chips, converting the theoretical sparsity of the Sparse VideoGen (SVG) approach into real hardware speedups. Through a progression of kernel optimizations (full/boundary tile specialization, tile-size tuning, and mask-tile alignment), they reduced attention kernel latency from 96.37ms (naive sparse) to 32.76ms, a 2.40x speedup over dense Splash Attention on a single TPU v6e chip, using 75.6K tokens, 10 heads, and head dimension 128.

Google Developers