TPU Kernel

Unlocking TPU Performance with Custom MatMul Kernels using JAX Pallas
Unlocking TPU Performance with Custom MatMul Kernels using JAX Pallas

Step-by-step optimization of TPU MatMul Kernels

Nov 27, 2024