Overview
A TPU implementation of the DSA attention path, developed as a foundation for a GLM-5.2 port.
Problem
Long-context model ports require attention kernels that preserve reference behavior while remaining operational at substantially longer sequence lengths.
Approach
The implementation is validated against the reference at the bit level before being exercised across progressively longer contexts.
Results & current status
- The current implementation is bit-exact against the reference through 8K tokens.
- It runs through 256K tokens.
- The kernel is the foundation of ongoing GLM-5.2 TPU porting work.
Technical details
The project combines JAX/Pallas kernel work, attention validation, long-context execution, and TPU numerical debugging.
Reproducibility
The implementation remains under active development. A reproducible release will follow once interfaces and validation coverage are stable.