TPU Acceleration & Distributed Topologies
We develop low-level JAX and Pallas kernels tailored for Google TPU v5e and Trillium (v6e) architectures. Our research optimizes high-bandwidth inter-chip interconnect (ICI) 3D torus networks, fusing sparse attention mechanisms directly into Vector Memory (VMEM) to minimize High-Bandwidth Memory (HBM) bottlenecks.
- Pallas Custom VMEM Sparse Aggregation Kernels
- 3D Torus Interconnect Optical Switch Routing
- Near-linear asymptotic efficiency on 512+ chip pods
- Low-precision FP8 & BF16 mixed-precision tensor pipelines