243 citations · 438 across the 13 of their papers we have counts for
14 papers
Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU
Jevin Jiang, Ying Chen, Blake A. Hechtman +2
Large Language Model (LLM) deployment is increasingly shifting to cost-efficient accelerators like Google's Tensor Processing Units (TPUs), prioritizing both performance and total…
eXmY: A Data Type and Technique for Arbitrary Bit Precision Quantization
Aditya Agrawal, Matthew Hedlund, Blake Hechtman
eXmY is a novel data type for quantization of ML models. It supports both arbitrary bit widths and arbitrary integer and floating point formats. For example, it seamlessly supports…
TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s
Felix Chern, Blake Hechtman, Andy Davis +3
This paper presents a novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-the-art GPU algorithms with sim…
Unified Scaling Laws for Routed Language Models
Aidan Clark, Diego de las Casas, Aurelia Guy +23
The performance of a language model has been shown to be effectively modeled as a power-law in its parameter count. Here we study the scaling behaviors of Routing Networks: archite…
Scaling Language Models: Methods, Analysis & Insights from Training Gopher
Jack W. Rae, Sebastian Borgeaud, Trevor Cai +77
Language modelling provides a step towards intelligent communication systems by harnessing large repositories of written human knowledge to better predict and understand the world.…
GSPMD: General and Scalable Parallelization for ML Computation Graphs
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen +13
We present GSPMD, an automatic, compiler-based parallelization system for common machine learning computations. It allows users to write programs in the same way as for a single de…