Transformer tricks: Precomputing the first layer

TL;DR

A trick to speed up inference of transformers with RoPE (such as LLaMA, Mistral, PaLM, and Gemma) by optimizing only one layer, which results in slightly lower latency and lower cost-per-token.

Abstract

This micro-paper describes a trick to speed up inference of transformers with RoPE (such as LLaMA, Mistral, PaLM, and Gemma). For these models, a large portion of the first transformer layer can be precomputed, which results in slightly lower latency and lower cost-per-token. Because this trick optimizes only one layer, the relative savings depend on the total number of layers. For example, the maximum savings for a model with only 4 layers (such as Whisper tiny) is limited to 25%, while a 32-layer model is limited to 3% savings. See https://github.com/OpenMachine-ai/transformer-tricks for code and more transformer tricks.

Figures (2)

• Figure 1: First layer of parallel transformer (a) without precompute; and (b) with precompute of FFN and linear layers Q, K, and V.
• Figure 2: First transformer layer. (a) Vanilla with pre-normalization and vanilla PE; (b) Vanilla with RoPE; (c) Precomputing linear layers Q, K, V.