Overview
Characterizing TPU v5e/v6e arithmetic and reproducing its matrix-multiplication behavior in a deterministic CPU simulator.
Problem
Accelerator arithmetic can differ from idealized numerical descriptions in ways that affect reproducibility, porting, and model behavior.
Approach
The project measures TPU Matrix Multiply Unit behavior at the bit level and uses those observations to implement a CPU-side simulator of the numerical path.
Results & current status
- TPU v5e and v6e numerical behavior has been characterized across the investigated path.
- A CPU simulator reproduces real Llama-3.1-8B matrix multiplications bit-for-bit on a native pass.
- The simulator enables deterministic study of TPU arithmetic without treating the accelerator as a black box.
Technical details
The work spans floating-point behavior, matrix-multiplication traces, accelerator validation, JAX execution, and CPU reference simulation.
Reproducibility
Methods and code will be linked when the characterization is ready for public technical review.