NoMAD-Attention: Efficient LLM Inference on CPUs Through Multiply-add-free Attention

Summary

The paper presents NoMAD-Attention, using SIMD registers in CPU, to speedup LLM inference in CPU.

Strengths

(1) Important problem with interesting solution. (2) The idea is straightforward, and the figure is very clear. (3) Good system speedup and maintaining the original performance of attention.

Weaknesses

The ML benchmarks seem a bit weak (e.g. perplexity, and PIQA-like easy benchmark). Can you evaluate on harder tasks like MT-Bench?

Questions

Please see weakness.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors include the limitation that this only works on CPUs with SIMD, but the reviewer thinks this is not a significant limitation.

Summary

This paper proposes using SIMD instructions on CPUs to speed up Transformers by removing multiply-add instructions. The paper replaces the attention operation a lookup-table based alternative. The paper motivates the application well for Transformer inference on CPUs.

Strengths

Strengths: * Important problem (efficiency) on an understudied platform (CPUs) * Shows speedup compared to standard attention * The quality studies are strong, with evaluation on upstream and downstream benchmarks. There is some degradation in quality, but it can be adjusted based on hyperparameters. I have also reviewed this paper in a previous conference (ICML) and I see that the paper is much improved from that version. I recommend accept based on the strength of the motivation and quality of the evaluation.

Weaknesses

Evaluation on additional and more recent open models (LLaMa-2, LLaMa-3, Gemma, QWEN) would make the paper more convincing. Two open questions from the current evaluation are how well the method generalizes to other pretrained models, and whether changes in the attention algorithm from more recent models change it. These are not necessary for the paper, but are natural follow-ups for the work and would improve the paper's impact.

Questions

How well does the method generalize to other open models and to changes in the attention algorithm?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Yes

Summary

This paper utilizes product quantization (PQ) to replace dot product operations in the matrix multiplications involved in the attention mechanism of transformers with memory lookup operations, showcased on language models. To my knowledge, this technique has been first introduced by Blalok et al (reference [5] in the paper). This paper goes further into the systems aspect of deploying PQ on GPUs by specifically utilizing SIMD registers to store the codebook used in PQ. This promises to significantly improve the memory access overhead even compared to alternative implementations that keep the PQ codebook in the L1 cache. However, this poses very stringent codebook size constraints that are addressed in the paper. PQ is done completely post-training unlike prior work, but uses a calibration dataset.

Strengths

- This paper is addressing an important topic: LLM optimization on commodity CPUs. - The use of SIMD registers is innovative and results in significant speedup. - PQ has been demonstrated a number of times in the past couple of years, but never with strong results on LLMs.

Weaknesses

- d_sub is very confusing to me. What is the length or your codeword? if d_sub = 1, is the size of the codeword = 1? how is that product quantization? How does it result in a speedup? - the proposed product quantization has been proposed to replace MAD in matrix multiplications in reference [5]. Line 148 in this paper is overclaiming its contribution suggesting that this submitted paper has proposed this method. - Can you please clarify the additional memory overhead introduced by the codebook? - "double quantization" of PQ has been introduced in prior work and is not cited: Abouelhamayed et al: PQA: Exploring the Potential of Product Quantization in DNN Hardware Acceleration.

Questions

see above

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

Summary

This paper proposes an algorithm to compute vector inner product efficiently on CPU by exploiting the fast access speed of in-register memory for fast self-attention computation on CPU for model inference. Instead of use multiply and add to compute the inner product between query and key vectors, the authors propose to break down the inner product between two vectors to the summation of multiple inner product between sub-vectors. For each sub-vector product, given a query vector, the algorithm enumerates all possible dot-product results with all possible key sub-vectors in a lookup table during preprocessing. The key sub-vectors are quantized so that the the lookup table is small enough to fit into the in-register memory. The authors further optimize the efficiency via optimizing key cache memory layout.

Strengths

1. The algorithm uses lookup operation to avoid multiplication and allows calculating multiple dimensions at once when the dimension of sub-vectors is larger than 1. 2. The algorithm is able to utilize an extremely small memory (128 bits) to accelerate the calculation.

Weaknesses

1. It would be better to also give some data about the latency of GPU decoding so that we know how far CPU is behind the GPU in Figure 2. 2. When the $d_{sub} > 1$, as shown in Table 1, the performance of these models got a severe hit, so it seems that the algorithm only works well on enumerating all possible scalar products for $d_{sub} = 1$. 3. Even when $d_{sub} = 1$, there are still latency improvement on Figure 2, but I think the number of operations should be the same as the vanilla multiply-add attention, except that each multiply is replaced with a lookup. I am wondering where is the efficiency improvement. 4. For $d_{sub} = 1$, the algorithm is quite relevant to quantized matmul for self-attention computation, I am wondering how the method performs compared to int8 or lower bit width quantization in terms of model accuracy and efficiency.

Questions

see weakness section.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors discussed limitations and societal impact.

Thank you for the response

The clarifications proposed by the authors should improve the paper writeup. I will maintain my positive-leaning score for this paper..

Response

Thank you for the rebuttal - I will be keeping my high score. Good work!