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

Authors: tianyi  zhang
jonah  yi
bowen  yao
zhaozhuo  xu
anshumali  shrivastava
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:49:27

Abstract: Large Language Model (LLM) inference on Central Processing Units (CPU) is challenging due to the vast quantities of  Multiply-Add (MAD) matrix operations in the attention computations. This paper highlights a rare gem in modern CPUs, Single-Instruction-Multiple-Data (SIMD) registers, which allows for ultra-low-latency lookups in a batch. We leverage this unique capability to propose NoMAD-Attention, an efficient attention algorithm that replaces MAD operations with in-register lookups. Through hardware-aware algorithmic designs, NoMAD-Attention achieves the computation of attention scores using repeated fast accesses to SIMD registers. NoMAD-Attention works with pre-trained attention-based LLMs without model finetuning. Extensive empirical evaluations demonstrate that NoMAD-Attention maintains the quality of the original LLMs well and speeds up the 4-bit quantized LLaMA-7B-based model by up to $2 \times$ at 16k context length.

Main paper: 
 introduction: 
 Auto-regressive transformer-based Large Language Models (LLM) demonstrate remarkable abilities across a wide range of natural language processing tasks without finetuning [35] and exhibit emergent abilities [49] for solving complex problems.The Need for Deploying LLM on CPUs. Despite the potential of LLMs, their deployment is costly [27] . Serving LLMs with billion-scale parameters requires specialized hardware such as NVIDIA A100 Graphics Processing Units (GPUs) [55] . However, mainstream personal devices are predominately equipped with only Central Processing Units (CPUs) [43] . As a result, making LLM-related services accessible to everyone remains a major challenge. Reducing the LLM inference latency on CPUs would significantly influence its accessibility and adoption.Expensive Multiply-add Operations for Attention in LLM Inference. LLM inference on CPUs is compute-bound, and the primary computational bottleneck is the calculation of attention scores [21] . Attention, a mechanism that models token interactions through all-pair dot products, heavily relies on the multiply-add (MAD) kernel on processors. The MAD operation involves computing the product of two numbers and adding that product to an accumulator [44] . MAD plays a crucial role in determining the attention score between tokens and subsequently blending their embeddings within the attention mechanism. The computational cost of attention grows quadratically with the sequence length due to the cumulative MAD operations. Since CPUs have limited parallel cores, they are inefficient for handling highly repetitive and parallel workloads. The extensive MAD operations required by the attention mechanism thus become the primary bottleneck during inference.Opportunities and Challenges from Modern CPUs: In-Register Lookups. The memory hierarchy of modern CPUs has undergone significant evolution, introducing a new type of registers optimized for Single-Instruction-Multiple-Data (SIMD) operations. The SIMD registers vary in size, ranging from 128 bits to 512 bits [41] , and support specialized SIMD instructions for high-throughput parallel processing [59] . SIMD registers have become a standard feature in commodity hardware, including laptops and mobile devices [13] . In this context, an in-register lookup (or shuffle) refers to the low-latency random access of information stored within SIMD registers. Storing information such as dot-product lookup tables (LUT) within SIMD registers, as opposed to cache memory, can accelerate LLM inference [4] . However, the limited size of SIMD registers poses challenges to fitting the computational paradigm of existing models.Algorithm 1 Attention Score Computation in LLM 1: Input: query q t , key k t , key cache K t-1 cache 2: let K t cache ← K t-1 cache k t ▷ Append the current key to key cache 3: return softmax(q t (K t cache ) ⊤ √ d )Then, the key and value embedding of the current token are appended to the key and value cache, respectively. The KV cache K t-1 cache , V t-1 cache of the step t -1 contains the key/value embeddings of all previous tokens, and after appending, the KV cache becomeK t cache = K t-1 cache k t =    k 1 k 2 . . . k t    , V t cache = V t-1 cache v t =    v 1 v 2 . . . v t   Finally, the attention output is computed asattention(e t ) = softmax q t (K t cache ) ⊤ √ d V t cachewhere d is the dimensionality of q t . We will refer to the result of softmax( qK ⊤ √ d ) as the attention scores since they dictate how much "attention" each token pays to other tokens. Computations in the prompt processing phase are similar to those in the decoding phase, except all the prompt tokens are computed in batch. LLMs use multi-head attention, which transforms the concatenation of the outputs of multiple single-head attentions to form an output embedding.MAD-based Attention. The attention mechanism models the interaction between tokens by performing all-pair dot products, where each dot product is computed via d Multiply-Add (MAD) operations. Since attention computes the interaction between all pairs of tokens, the amount of MAD operations scales quadratically with the sequence length, quickly overwhelming the computing capability of CPUs. CPUs are designed to handle complex workloads with granular control, while GPUs are optimized for processing simple and repetitive tasks in high throughput. Hence, the success of attention has largely been fueled by the development of highly parallel throughput-oriented processors, such as GPUs [12] .MAD-based Attention as Bottleneck of LLM Inference. The computation of attention scores becomes the bottleneck of LLM inference as the sequence length increases (see our analysis in Figure 4 ). At the t-th step of the decoding phase, the time complexity of computing attention score with MAD is O(t) due to t dot products, while all other components of LLMs such as MLP, skip connections, and normalization have a time complexity of O(1). We will focus on optimizing the efficiency of attention score computations in our proposed approach. Algorithm 1 presents the pseudocode for attention score computation, including key caching, for a single-head masked self-attention in LLM. This algorithm will serve as a point of comparison in our proposed approach. 
memory hierarchy of modern cpus: 
 CPU memory is organized into a pyramidal hierarchy, as shown in Figure 1 , with faster memory significantly smaller than slower memory. The memory unit with the fastest access speed is the registers. Each compute core accesses its dedicated registers in just 1-2 CPU cycles, but these registers are limited in size, usually not exceeding 64 bits. Modern processors include a new type of registers optimized for Single-Instruction-Multiple-Data (SIMD) operations. These SIMD registers range from 128 bits to 512 bits in size and support a limited set of SIMD instructions for throughput-oriented parallel processing. SIMD registers are common on commodity hardware, including laptops and mobile devices. Using SIMD operations can speed up deep learning models on CPUs by parallelizing matrix multiplications. However, due to the limited number of cores in a CPU, its efficiency in deep learning is still considerably worse than GPU. Prior works [42, 9] have resorted to sparsity and sampling-based approaches, but they require training models from scratch and may not apply to all architectures. Our approach exploits the SIMD registers to shift the computation paradigm from MAD to in-register lookups, which demonstrates significant speedup over MAD-based models. When describing our proposed algorithm, we assume SIMD registers are 128 bits wide. Some systems with wider SIMD registers support more parallelism, e.g., 256-bit registers in AVX-2 and 512-bit registers in AVX-512. However, the most universal form of SIMD registers uses 128 bits, which is supported by Arm NEON and AVX-compatible processors. 
methodology: 
 This section describes our proposed approach, NoMAD-Attention, which replaces MAD operations with in-register lookups to enable faster attention computations on CPUs. NoMAD-Attention utilizes three techniques to enable lookup-based attention: 1. transforming dot product computations to memory lookups through product quantization, 2. compressing lookup tables into SIMD registers for low-latency access, 3. reorganizing the memory layout of key cache for batch parallel dot product lookups.Previous works have shown that inexact attention scores in transformers work well for sequence modeling [53] . NoMAD leverages Product Quantization (PQ) [23] to compute high-quality estimations of dot products through in-register lookups. PQ, originally designed for compressing high-dimensional vectors, quantizes a floating-point vector into discrete codes. It makes use of sub-quantizers; for a d-dimensional vector space, a vector is divided evenly in dimension into S sub-vectors, where each sub-vector has dimension d sub = d S , and each sub-vector space is quantized independently. We use π s (e), where s ∈ {1 . . . S}, to denote the function that maps a d-dimensional vector e to its d sub -dimensional sub-vector of the s-th sub-quantizer. Codebooks are used to quantize sub-vectors to codes, which are collections of cluster centroids learned from a set of calibration vectors. We use b s,c to denote the c-th centroid in the codebook of the s-th sub-quantizer. For a given vector e, the product-quantized codes of e, denoted c 1 , . . . , c S , are the indexes of the nearest centroid of each sub-quantizer, i. Once base vectors have been product-quantized to codes, PQ leverages asymmetric distance computation to keep estimation errors low. In the computed distances, the original query vector is used while the quantized base vectors are used, hence the asymmetry. For a given query q, the distances to the centroids of each sub-quantizer are computed and stored in a lookup table (LUT). Then, the corresponding distances in the LUT are looked up based on the codes of base vectors and accumulated to produce the final distance estimation. More concretely, denoting the distance between query q and the c-th centroid for the s-th sub-quantizer using LUT s [c] = dist π s (q), b s,c , then the estimated distance between query q and a product-quantized base vector e, where PQ(e) = [c 1 . . . c S ], isdist(q, e) = S s=1 LUT s [c s ]Building upon previous work [7] , we employ product quantization to approximate dot products within the attention mechanism. We product-quantize the key vectors in attention to produce key codes, which will be stored in place of the key cache of LLM attention. The codebooks are learned by performing clustering on a set of key vectors from a calibration set. The key vectors are quantized to the nearest centroid with respect to L2 distance. For a given query, the query-dependent LUT is computed to hold dot products with respect to centroids. Dot products of sub-vectors are retrieved from the LUT based on key codes and accumulated to produce the final dot product estimates. This procedure allows us to compute attention scores through lookups.Estimating dot products through PQ mostly eliminates the use of MAD kernels in the computation of attention scores. However, this approach yields limited speedup over dot-product attention since a high proportion of the CPU cycles are wasted due to cache/memory access stalling. L1-cache-resident LUT is not enough to offer high-performance PQ [3] . The full potential of lookup-based attention can only be unlocked by having the LUT stored in registers, which take only 1-2 CPU cycles to access. However, the highly limited size of registers poses a challenge to fitting the LUT. In PQ, each sub-quantizer commonly uses 256 centroids, which translates to 8-bit codes. Combined with 32-bit floating-point (FP32) dot products, the LUT for each sub-quantizer consumes 8192 bits of memory, while the SIMD registers are only 128 bits wide. We leverage hardware-aware techniques proposed in [4] to enable low-latency retrieval from register-resident LUT.8-bit Quantized Dot Products in LUT Due to the 128-bit width of SIMD registers, the FP32 representation of dot product is too costly to store. Adopting FP32 dot products in LUT implies that each codebook can only contain up to 4 centroids, which will lead to significant quantization errors. Therefore, we adopt the 8-bit dynamically quantized representation of dot products. Compressing beyond 8-bit is infeasible since most SIMD instruction sets do not support parallel lookups below 8 bits. The quantization is done dynamically for each query to minimize quantization errors. For a given query and sub-quantizer, dot products to centroids are first computed in full FP32 precision. Then, the quantization range is determined by the minimum and maximum dot products to the centroids. Finally, the range is evenly divided into 2 8 buckets, and dot products are quantized to the bucket they fall into. More formally, suppose dp min = min c (π s (q) • b s,c ) and dp max = max c (π s (q) • b s,c ) are the minimum and maximum dot products of the query q to the centroids of the s-th sub-quantizer, then the LUT stores the quantized dot products to centroid c asEQUATIONThe quantization and de-quantization process can be done efficiently without much computational overhead, and the quantization error is kept low thanks to dynamic query-dependent quantization (we analyze its effects in Table 3 ).Constrained Codebook Size By adopting 8-bit quantized dot products in LUT, we can fit 16 dot products on 128-bit SIMD registers. This implies that the codebook size of each sub-quantizer is constrained to 16 centroids, and evidence suggests this limited size may work well with attention: it has been shown that the output of attention loses rank extremely quickly [15] , implying that the intermediate embeddings of transformers may exhibit clear clustering structures.Quantized dot products and constrained codebooks enable LUT to be stored in SIMD registers, but the layout format of the key cache needs to be reorganized for SIMD instructions. The original key cache in LLM attention stores each key vector contiguously in a row to optimize single vector reads. NoMAD-Attention uses the key-code cache in place of the key cache, which stores the quantized codes of keys. To allow fast lookups of LUT entries based on key codes, we store the key codes in a transposed blocked format. A comparison between the LLM key cache and the NoMAD key-code cache is given in Figure 1 .The storage format of the NoMAD key-code cache is transposed: stored in column-major order instead of row-major, and blocked: with 32 keys as a block. The SIMD instruction shuffle, leveraged for performing low-latency batch lookups, takes a batch of byte-size integers as input and retrieves the values held in the registers corresponding to the integer indices. The original storage format of the key cache stores all dimensions of a key contiguously, which precludes efficient use of shuffle. To maximize the usage of the LUT held in registers, we store key codes belonging to the same sub-quantizer contiguously in rows of 32 codes. Since shuffle performs lookups in a batch size of 16, the keys within the same block are stored in alternating order. Each quantized code occupies half a byte as there are 16 centroids in a codebook, while the shuffle instruction uses each byte as an input argument. By performing SIMD bit-shifting and bit-masking on a block of alternating keys, we obtain the key codes in the original order, ready for use with shuffle. Section A in the Appendix describes additional details on how shuffle is performed on each block of key-code cache, and provides pseudocode.By combining these three techniques, NoMAD-Attention achieves fast MAD-free attention score computations through SIMD in-register lookups. For a given query, first, LUTs with 8-bit quantized dot products are computed for each sub-quantizer. Then, an LUT is loaded into registers, followed by SIMD shuffle instructions to retrieve dot products in the LUT in batch based on key codes. The loading and lookup are repeated for all sub-quantizers, and the retrieved dot products are accumulated in batch through SIMD add. Finally, the quantized dot products accumulated over all sub-quantizers are de-quantized, scaled, and fed through softmax to produce the attention scores. The pseudocode for NoMAD-Attention score computations is given in Algorithm 2.Compressing each segment of key activation embeddings into 4-bit codes, without degrading model quality, is challenging. The straightforward approach is to first cache key activation embeddings on a calibration dataset, and then learn the centroids through k-means clustering [29] on the embeddings. However, this first-cut approach significantly degrades model quality, especially for d sub > 1 (we analyze its effects in Section 4.2). The first-cut approach performs sub-optimally since it is uninformed; clustering is performed with the aim to minimize the reconstruction error with all tokens. However, as shown previously [60, 28] , the key cache of certain tokens are more pivotal for preserving model quality. Therefore, we leverage the Fisher Information Matrix (FIM) to bias the centroids towards important activations. Specifically, we approximate the Hessian using diagonals of the FIM, which is the element-wise square of the gradient, and use it to weigh the reconstruction errors. Using the FIM to minimize reconstruction error was first proposed in [26] . Our optimization objective for EQUATION, whereEQUATIONwhere k 1 . . . k n are cached key activation embeddings on the calibration dataset. We use weighted k-means++ [5] to optimize the objective.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: After reviewing the NeurIPS Code of Ethics, we can confirm that the paper conforms with the code in every respect.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
experiments: 
 This section evaluates the effectiveness of our proposed NoMAD-Attention in preserving model quality and speeding up LLM inference on CPUs. We first introduce the software, hardware, models, benchmarks, and baselines, then provide detailed results and discussions, and finally perform an ablation study to validate each component of our proposal.Software and Hardware Our implementation of NoMAD-Attention is built in C and C++, based on the open-source projects llama.cpp [20] and FAISS [16] . We also built a GPU implementation of NoMAD-Attention for quick prototyping and key-compression learning, which is based on PyTorch [33] and HuggingFace Transformers [51] . Experiments for latency and throughput are performed on a Linux server equipped with an Intel Xeon E5-2695 V3 14-core CPU, which supports AVX2 SIMD instructions, and 512GB of DDR4 RAM. Experiments for accuracy and perplexity are performed on two NVIDIA A100-40GB GPUs.We evaluate the quality of NoMAD-Attention with 4 popular LLMs:1. LLaMA-7b 2. LLaMA-13b [45] 3. LLaMA-2-7b 4. LLaMA-2-13b [46] . We measure the model quality with perplexity on WikiText-2 [30] and C4 [14] at the context length of 2048, and zero-shot accuracy (using the default configurations of lm-evaluation-harness [19] ) on SciQ [50] , Arc Easy (Arc-E), Arc Challenge (Arc-C) [11] , Hellaswag [54] , WinoGrande [38] , and PIQA [6] . The centroids for key compression of NoMAD-Attention are learned on a calibration set of 16 sequences from WikiText-2, each with 2048 tokens. To test the model efficiency, we benchmark the latency and throughput of CodeLlama-7b [37] (with 16-bit weights and 4-bit q4_0 quantized weights), which has a longer context length of 16,384 than the LLaMA family of models. We compare the efficiency of NoMAD-Attention-based models (with d sub ∈ {1, 2, 4}) against Attention-based models with a llama.cpp-based implementation. Model Quality Table 1 presents the perplexity and accuracy of NoMAD-Attention on different benchmarks with comparisons to Attention. NoMAD-Attention (d sub = 1) incurs negligible perplexity and accuracy loss, d sub = 2 incurs minimal degradation, and d sub = 4 incurs some degradation due to the high key compression factor.The latency and throughput for decoding 16,384 tokens for NoMAD-Attentionbased and Attention-based CodeLlama-7b are presented in Figure 2 . For 4-bit CodeLlama-7b, NoMAD-Attention achieves 1.78×, 1.95×, and 2.07× higher throughput (tokens per second) than Attention for d sub = 1, 2, 4, respectively. For 16-bit CodeLlama, NoMAD-Attention achieves 1.46×, 1.54×, and 1.56× higher throughput (tokens per second) than Attention for d sub = 1, 2, 4, respectively.Prompt Processing Efficiency We investigate the efficiency of NoMAD-Attention for prompt processing. We generate prompts of length 4,000, 8,000, and 16,000 tokens and record the time until the first token is decoded. Figure 3 presents prompt processing time for CodeLlama-7b. At the prompt length of 16k tokens, NoMAD-Attention speeds up prompt processing time by 1.63 -1.78× for the 4-bit model, and 1.64 -1.79× for the 16-bit model. Prompt processing times are similar for 16-bit and 4-bit model, since weight quantization is ineffective for speeding up batch processing.We perform a set of ablative experiments to study the latency breakdown, and the effects of FIMinformed clustering and 8-bit LUT quantization on model quality.We investigate the makeup of latency in each decoding step. Figure 4 presents the latency breakdown of 4-bit CodeLlama-7b for decoding 16,384 tokens, and Figure 7 in MLP and Attention stays constant as context length increases, and the latency of key caching is insignificant compared to other components. The latency of dot-products between queries and keys, as well as the multiplications between attention scores and values, grows linearly with increasing context length. In the original Attention, computing attention scores quickly becomes the latency bottleneck as context length increases. NoMAD-Attention mitigates this bottleneck by speeding up query-key dot-product computation by 5.24×, 9.75×, and 14.77× for d sub = 1, 2, 4, respectively.We study the effects of FIM-informed centroid learning for preserving model quality. Table 2 presents the perplexity of NoMAD-Attention-based LLaMA-7b on WikiText-2 with uninformed centroids and FIM-informed centroids. FIM-informed centroids consistently achieve better model quality, shown by the lower perplexity.LUT Quantization We examine the effects of 8-bit quantized dot products in LUT on model quality. Table 3 presents the perplexity of NoMAD-Attention-based LLaMA-7b on WikiText-2 with 8-bit quantized LUT and 32-bit unquantized LUT. 8-bit quantization of LUT incurs negligible loss in perplexity.  
related works: 
 Approximate Attention and Efficient Transformer Since the introduction of attention in transformers [47] , there has been a body of work on approximating the attention mechanism for efficient training and inference. Dynamically sparse attention was achieved using LSH [24] , Nyström method [52] , and random sampling [53] . Low-rank attention has also been extensively explored [48, 10, 8] and shown to have compute-and memory-efficiency advantages. Hardware-aware attention mechanisms such as FlashAttention [12] propose to mitigate the IO bottleneck in GPUs. In large language models, multiple approaches [60, 28, 58] have been proposed to reduce the high memory overhead of the KV cache. For CPU-only environments, [39] proposes to speed up LLM inference through weight quantization. Weight quantization [18, 56] is an effective approach for accelerating LLM inference and fine-tuning [57] by mitigating the IO bottleneck.Approximate matrix multiplication is applicable in a wide range of computational problems, and its optimization has been a topic of interest for years [32] . Modern researchers have begun optimizing matrix multiplication around the specific limitations of computers, including mitigating the IO bottleneck between the CPU and main memory [25] . Compression techniques were developed to accelerate large-scale matrix multiplications and scale up the size of multiplication [31, 7, 25, 1] . However, many of these algorithms still had limitations that make them too inaccurate or costly for approximating matrix multiplications in LLMs [7, 25] . 
conclusion: 
 This study addresses the challenges of large language model inference on Central Processing Units (CPUs), particularly the difficulties associated with the expensive Multiply-Add (MAD) matrix operations in attention mechanisms. The investigation highlighted the untapped potential of Single-Instruction-Multiple-Data (SIMD) registers and their fast in-register lookup capabilities within CPUs. The proposed NoMAD-Attention algorithm serves as an efficient alternative to traditional MADbased approaches, leveraging in-register lookups and optimizing memory access to SIMD registers. The implementation of NoMAD-Attention resulted in a significant acceleration of LLaMA-7B-based model inference, achieving up to a 2× speedup on CPUs. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We include a Limitations and Broader Impacts section to detail the potential societal impacts of our work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: With no new dataset or model, the paper poses no such risk.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The assets used in the paper were properly credited and their licenses and terms were listed and respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: New assets introduced in the paper are documented wherever applicable.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: The paper does not involve crowdsourcing nor reserach with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The SIMD shuffle presented in Algorithm 2 is a simplification of the actual hardware implementation. We give full details of lines 5 to 11 in Algorithm 2 in Algorithm 3. Keys are stored in blocks of 32, in which keys are stored in an alternating order (see Figure 1 for an illustration). After a LUT is loaded into registers, the row of key codes in the block corresponding to the sub-quantizer is used to perform shuffle. First, each byte in the row is bit-shifted to the right by 4 bits via a SIMD instruction, which produces the codes of the first 16 keys in the block. The codes are fed to shuffle to retrieve the quantized dot products of the first 16 keys from the LUT. Then, the first 4 bits of each byte in the row are masked out via a SIMD instruction, which produces the code of the last 16 Figure 5 illustrates the process of mapping and compressing key vector k t to construct K t cache . For an input key vector k t , functions π s , where s ∈ 1 . . . S, split the vector into sub-vectors k t = (π 1 (k t ), π 2 (k t ), . . . , π S (k t )). Subsequently, each sub-quantizer π s (k t ) is mapped to its nearest centroid c t s by referencing the codebook b s , where i ∈ 1 . . . S, among 16 centroids in the codebook. The resulting values are then stored in the key cache K t cache . Similarly, Figure 6 illustrates the process of mapping and compressing query vector q t to construct the Look-up Tables (LUT). Given a query vector q t , functions π s , where s ∈ 1 . . . S, first split the query into sub-queries q = (π 1 (q t ), π 2 (q t ), . . . , π S (q t )). Subsequently, the distances between each sub-query π s (q t ) and the 16 centroids from the codebook b s are computed and then quantized to values within the range of 0-255. Lastly, the quantized vectors are converted into 8-bit codes and stored in LUT t s . • • • • • • • • • • • • • • • • • • •Figure 5 : Illustration demonstrating the mapping of an input key k t to its s-th sub-quantizer using π s (k t ), where s ∈ 1 . . . S. Subsequently, each sub-quantizer maps to its closest centroid c t i , where i ∈ 1 . . . S, and the results are stored in the key cache K t cache .Compute distance from the subquery to the 16 centroids of the corresponding codebook• • • • • • • • • • • • • • • • • • • • • • • • • • Quantize distances to numbers between 0-255 • • • • • • • • • • • • • • • • • • • • • • • • • • Convert quantized distance to 8 bit codes • • • • • • • • • • • • • • • • • • • •Figure 6 : Illustration depicting the mapping of a query vector q to its s-th sub-quantizer using π s (q), where s ∈ 1 . . . S. Subsequently, the distance between π s (q) and 16 centroids is computed. This distance is quantized to a value within the range of 0-255, and the resulting quantized distance is further converted into 8-bit codes, which are stored in LU T s .The learned centroids of NoMAD-Attention introduce a small amount of memory overhead. The storage overhead of centroids can be calculated as l × h × d × 16 × 4 bytes, where l represents the number of layers, h the number of attention key heads, d the dimensionality of each 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 = {1, 2, 4} is 8.4MB, and the memory overhead for LLaMA-13b/LLaMA-2-13b with d sub = {1, 2, 4} is 13.1MB, which is a modest memory footprint compared to the overall model size. We empircally compare NoMAD-Attention and Attention with integer quantized key cache. We test the accuracy of LLaMA-7b, using NoMAD-Attention (d sub = 1) and 8-bit and 4-bit integer quantized key cache (q8_0 and q4_0 in llama.cpp). Furthermore, we test decoding latency at a context length of 16K tokens, and measure the speedup compared to the full model. As shown in Table 5 , 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. F Results on Llama-3We evaluate the effectiveness of NoMAD-Attention on the Llama-3-8b [17] model, which uses grouped query attention [2] . Table 6 presents the accuracy of Llama-3-8b with Attention and NoMAD-Attention across a range of downstream tasks. The results demonstrate that NoMAD-Attention with d sub = 1 effectively maintains model quality for Llama 3. We evaluate NoMAD-Attention on the more challenging MMLU [22] , GPQA [36] , and MGSM (English) [40] benchmarks. We use lm-evaluation-harness [19] for accuracy evaluation, and the task names are mmlu_stem, mmlu_social_sciences, mmlu_humanities, mmlu_other, gpqa_main_zeroshot, mgsm_direct_en. As shown in Table 7 , NoMAD-Attention with d sub = 1 effectively maintains model quality across these diverse challenging tasks. NoMAD-Attention achieves speedup relative to full Attention when d sub = 1 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.NoMAD-Attention can generalize to other pretrained transformer models and attention variants such as grouped-query attention (GQA) [2] and Attention with Linear Biases (ALiBi) [34] . Given that GQA employs a shared key head 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.NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html 


Decision: This paper presents an effective method for accelerating LLMS on commodity CPUs, and the reviews clearly indicate acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: see weakness section.
limitations: limitations: The authors discussed limitations and societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: see above
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: How well does the method generalize to other open models and to changes in the attention algorithm?
limitations: limitations: Yes
paper_summary: paper_summary: The paper presents NoMAD-Attention, using  SIMD registers in CPU, to speedup LLM inference in CPU.
strength_weakness: strength_weakness: 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: questions: Please see weakness.
limitations: limitations: The authors include the limitation that this only works on CPUs with SIMD, but the reviewer thinks this is not a significant limitation.

Comments: title: Response
Thank you for the response
comment: Thank you for the rebuttal - I will be keeping my high score. Good work!
Thank you! This is good. I raised the score to 6. Please consider accept this paper.
The clarifications proposed by the authors should improve the paper writeup.
I will maintain my positive-leaning score for this paper..