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

We sincerely thank the reviewer for their support of our paper and the thoughtful suggestions. We address the reviewer's concerns as follows.

[W1] Evaluation on additional and more recent open models.

• We conducted additional experiments on the LLaMA-3-8b model with NoMAD-Attention across a range of downstream tasks. Our results demonstrate that NoMAD-Attention with $d_{sub}=1$ effectively maintains model quality.

[W2, Q1] How well does the method generalize to other open models and to changes in the attention algorithm?

NoMAD-Attention can generalize to other pretrained models and attention variants such as grouped-query attention (GQA) [1] and Attention with Linear Biases (ALiBi) [2]. Given that GQA employs shared key heads across multiple query heads, we can adapt NoMAD-Attention to reuse the same key codes for performing in-register lookups. The ALiBi method adds a linear bias term to the query-key dot products, a process fully compatible with the NoMAD-Attention approach.

References

[1] Ainslie, Joshua, et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints." Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023.

[2] Press, Ofir, Noah Smith, and Mike Lewis. "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation." International Conference on Learning Representations.

We greatly appreciate the reviewer's careful review and valuable suggestions. We address their comments as follows.

[W1] Evaluations on harder tasks.

To further assess NoMAD-Attention, we conducted additional evaluations on the more challenging MMLU, GPQA, and MGSM (English) benchmarks. Our results demonstrate that NoMAD-Attention with $d_{sub}=1$ effectively maintains model quality across these diverse tasks.

We sincerely appreciate the reviewers' careful consideration of our paper and their valuable feedback. We have addressed each of their comments below.

[W1] 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.

• We thank the reviewer for their valuable suggestion. We will incorporate the latency of GPU decoding into the final version of the paper to provide a more comprehensive evaluation.

[W2 & W3] Values of $d_{sub}=1$.

• As demonstrated in Table 1, a $d_{sub}$ value of 2 still yields reasonable model performance, as measured by perplexity and accuracy.
• The $d_{sub}=1$ configuration already achieves significant speedups compared to the original model, with a 1.78x improvement observed at a context length of 16K (Figure 2).
• The efficiency improvement achieved when $d_{sub}=1$ is primarily due to two factors: 1. Leveraging Lower-Latency, Higher-Throughput Instructions. Unlike the vanilla multiply-add attention, which relies on batched multiplication and addition SIMD instructions (e.g., _mm256_mul_ps and _mm256_add_ps in AVX2), NoMAD-Attention utilizes the SIMD lookup instruction (_mm256_shuffle_epi8). This latter instruction operates on more elements at once (32 elements versus 8) and exhibits lower latency (1 cycle vs. 4 cycles on most architectures) [1], contributing significantly to the efficiency gains. 2. Minimized Data Movement. The product-quantized key cache employed in NoMAD-Attention effectively reduces the volume of data transferred between RAM and registers, hence speeding up computations.

[W4] How does the method perform compared to int8 or lower bit width quantization in terms of model accuracy and efficiency? --- Here is a comparison.

• We present additional experiments comparing NoMAD-Attention with INT8 and INT4 key cache quantization (q8_0 and q4_0 in llama.cpp). We report the accuracy on various benchmarks and the decoding latency at a context length of 16K tokens for LLaMA-7b.
• NoMAD-Attention ($d_{sub}$=1) demonstrates better model quality than INT4 quantized key cache and significantly higher speedup than INT4 and INT8 quantized key cache.

References

[1] https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html

We are grateful for the reviewers' careful review and insightful comments. We have addressed their feedback in detail below.

[W1] 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?

• We are sorry for the confusion. To clarify, $d_{sub}$ represents the dimensionality of each sub-quantizer, which is equivalent to the dimensionality of each cluster centroid or the number of dimensions encoded by each product-quantized key code. Notably, the special case where $d_{sub}=1$ can be considered a generalization of multi-dimensional product quantization.
• $d_{sub}=1$ results in a speedup primarily due to two reasons: 1. Leveraging Lower-Latency, Higher-Throughput Instructions. Unlike the vanilla multiply-add attention, which relies on batched multiplication and addition SIMD instructions (e.g., _mm256_mul_ps and _mm256_add_ps in AVX2), NoMAD-Attention utilizes the SIMD lookup instruction (_mm256_shuffle_epi8). This latter instruction operates on more elements at once (32 elements versus 8) and exhibits lower latency (1 cycle vs. 4 cycles on most architectures) [1], contributing significantly to the efficiency gains. 2. Minimized Data Movement. The product-quantized key cache employed in NoMAD-Attention effectively reduces the volume of data transferred between RAM and registers, hence speeding up computations.

[W2] Line 148 in this paper is overclaiming its contribution.

• We apologize for this oversight. We will revise line 148 in the final paper to the following: "Building upon previous work [5], we employ product quantization to approximate dot products within the attention mechanism."

[W3] Can you please clarify the additional memory overhead introduced by the codebook? --- Sure!

• The memory overhead introduced by the codebook can be computed as $l \times h \times d \times 16 \times 4$ bytes, where $l$ represents the number of layers, $h$ the number of attention key heads, $d$ the dimensionality in an attention head, $16$ the number of centroids, and $4$ the number of bytes for storing each centroid parameter.
• Therefore, the codebook memory overhead for LLaMA-7b/LLaMA-2-7b with $d_{sub} \in \\{1,2,4\\}$ is 8.4MB, and the memory overhead for LLaMA-13b/LLaMA-2-13b with $d_{sub} \in \\{1,2,4\\}$ is 13.1MB, which is a relatively modest memory footprint compared to the overall model size.
• We thank the reviewer for highlighting this point, and we will incorporate a detailed discussion of codebook memory overhead into the final paper.

[W4] "double quantization" of PQ has been introduced in prior work and is not cited.

• We thank the reviewer for bringing PQA to our attention. We will incorporate a citation to this relevant work in the final version of our paper.

References

[1] https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html