Head-level KV cache compression for contextual question answering

Head-level KV cache compression in LLMs optimizes memory allocation by prioritizing important heads based on retrieval and reasoning abilities, enhancing computational efficiency and reducing latency.

WO2026152363A1PCT designated stage Publication Date: 2026-07-23MICROSOFT TECHNOLOGY LICENSING LLC
View PDF 0 Cites 0 Cited by

Patent Information

Authority / Receiving Office
WO · WO
Patent Type
Applications
Current Assignee / Owner
MICROSOFT TECHNOLOGY LICENSING LLC
Filing Date
2025-01-17
Publication Date
2026-07-23

AI Technical Summary

Technical Problem

As input lengths increase, memory usage and latency grow significantly in Large Language Models (LLMs) due to self-attention in transformers, conventional KV cache compression techniques optimize memory retention per layer, failing to fully optimize KV cache allocation for downstream tasks that require both retrieval and reasoning abilities.

Method used

Head-level KV cache compression allocates memory based on importance scores, estimating retrieval and reasoning attributes of individual heads in multi-head attention models, dynamically distributing cache budgets to prioritize more critical heads, and retaining relevant cache entries.

Benefits of technology

This approach significantly reduces memory usage and decoding latency while preserving the model's retrieval and reasoning capabilities, outperforming conventional layer-level allocation strategies.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN2025072958_23072026_PF_FP_ABST
    Figure CN2025072958_23072026_PF_FP_ABST
Patent Text Reader

Abstract

Head-level key-value (KV) cache compression may be performed by allocating memory to individual heads of a multi-head attention model based on importance scores assigned to the individual heads of a multi-head attention model. The importance scores may be calculated based on both retrieval and reasoning attributes of the individual heads of the multi-head attention model. The retrieval attributes and reasoning attributes of the individual heads of the multi-head attention model may be determined based on an updated Needle-in-a-Haystack test that incorporates a Retrieval-Reasoning example. It should be appreciated that the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model.
Need to check novelty before this filing date? Find Prior Art

Description

HEAD-LEVEL KV CACHE COMPRESSION FOR CONTEXTUAL QUESTION ANSWERINGBACKGROUND

[0001] Modern Large Language Models (LLMs) increasingly support long inputs. For instance, ChatGPT 4.0 (GPT-4) supports up to 128K tokens and Claude 3.5 supports up to 1 million tokens. These extended capacities improve performance on tasks like dialogue generation, question answering, and summarization. As input lengths increase, memory usage and latency grow significantly due to self-attention in transformers. To improve inference speed and efficiency, most LLM inference include two phases: prefilling for input processing and decoding for token generation.SUMMARY

[0002] This Summary is provided to introduce a selection of concepts in a simplified form that are further described below in the Detailed Description. This Summary is not intended to identify key features or essential features of the claimed subject matter, nor is it intended to be used as an aid in determining the scope of the claimed subject matter.

[0003] Some examples provide a method for operating a large language model (LLM) . The method includes receiving an input prompt, converting the input prompt into a sequence of input tokens, processing the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, and outputting the response to the input prompt. Processing the sequence of input tokens using the multi-head attention model includes allocating key-value cache memory to individual heads of the multi-head attention model.

[0004] Other examples provide a computer for supporting a large language model (LLM) that includes a processor and computer-readable medium storing programming instructions that, when executed by the processor, cause the processor to perform a series of steps. The steps include receiving an input prompt, converting the input prompt into a sequence of input tokens, processing the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, and outputting the response to the input prompt. Processing the sequence of input tokens using the multi-head attention model includes allocating key-value cache memory to individual heads of the multi-head attention model.

[0005] Yet other examples provide a computer program product, for supporting a large language model (LLM) , that stores programming instructions that, when executed by a processor, cause the processor to perform a series of steps. The steps include receiving an input prompt, converting the input prompt into a sequence of input tokens, processing the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, and outputting the response to the input prompt. Processing the sequence of input tokens using the multi-head attention model includes allocating key-value cache memory to individual heads of the multi-head attention model.BRIEF DESCRIPTION OF THE DRAWINGS

[0006] FIG. 1 is a diagram of an LLM system;

[0007] FIG. 2 is a diagram of an LLM decoder;

[0008] FIG. 3 is a diagram of key-value cache in a multi-head attention model;

[0009] FIG. 4 is a diagram of a head-level KV cache compression technique;

[0010] FIG. 5 is a flowchart of a method for LLM processing;

[0011] FIG. 6 is a diagram of retrieval and reasoning examples for head identification;

[0012] FIG. 7 is a table of experimental results for comparing LongBench and LooGLE benchmarks for Llama-3-8B-Instruct and Mistral-7B-Instruct models;

[0013] FIG. 8 is a graph of experimental results for the Llama-3-8B-Instruct model using different KV cache allocation techniques;

[0014] FIG. 9 is a graph of experimental results for the Mistral-7B-Instruct model using different KV cache allocation techniques;

[0015] FIGS. 10A-10C are graphs of head visualization for Llama-3-8B-Instruct results;

[0016] FIG. 11 is a table of ablation study results for LongBench benchmarks;

[0017] FIGS. 12A-12F are graphs of Needle-in-a-Haystack test results on the Llama-3-8B-Instruct model;

[0018] FIG. 13 is a table of test results from Reasoning-in-a-Haystack test results;

[0019] FIGS. 14A-14B are graphs of decoding and latency test results for different KV cache compression techniques;

[0020] FIG. 15 is a table of example tasks used for Reasoning-in-a-Haystack tests;

[0021] FIG. 16 is a functional block diagram of a computing apparatus.

[0022] Corresponding reference characters indicate corresponding parts throughout the drawings. Any of the figures may be combined into a single example or embodiment.DETAILED DESCRIPTION

[0023] A more detailed understanding can be obtained from the following description, presented by way of example, in conjunction with the accompanying drawings. The entities, connections, arrangements, and the like that are depicted in, and in connection with the various figures, are presented by way of example and not by way of limitation. As such, any and all statements or other indications as to what a particular figure depicts, what a particular element or entity in a particular figure is or has, and any and all similar statements, that can in isolation and out of context be read as absolute and therefore limiting, can only properly be read as being constructively preceded by a clause such as “In at least some examples, ... ” For brevity and clarity of presentation, this implied leading clause is not repeated ad nauseum.

[0024] As mentioned above, LLM inference typically includes LLM decoding following tokenization. One technique for reducing the amount of computation cycles that are performed during decoding is key value caching (KV caching) . In particular, decoding typically relies on autoregressive attention models, meaning that past values (e.g., past tokens) are used to compute future value (e.g., future tokens) . KV caching generally refers to the storage (caching) of key-values associated with past tokens such that those values do not need to be recalculated when computing future tokens. As input length increases, KV cache memory grows rapidly, posing significant challenges for storage and efficiency. In general, key-values are stored in a head-layer matrix, with multiple heads being used to compute values over multiple layers. Conventional techniques for reducing memory requirements compress KV caches at the layer-level based on attention scores. These technique are aimed at optimizing retention per layer during prefilling.

[0025] Aspects of this disclosure compress KV caches at the head-level based on importance scores, which are different than attention scores. In particular, importance scores are calculated based on both retrieval and reasoning, with heads that are deemed more important receiving larger KV cache budget allocations. In addition to allocating KV cache across attention heads (rather than layers) , aspects of this disclosure distribute cache budgets based on head importance scores. In some examples, importance scores are estimated for each individual head of a multi-head attention model based on both retrieval attributes and reasoning attributes of the individual heads. The retrieval and reasoning attributes of the individual heads of the multi-head attention model may be determined based on an updated Needle-in-a-Haystack test that incorporates Retrieval-Reasoning example (s) (aka Retrieval-Reasoning evaluation (s) ) . The updated Needle-in-a-Haystack test may be used to estimate importance scores for the individual heads by evaluating both the retrieval attributes and the reasoning attributes of the individual heads of the multi-head attention model. It should be appreciated that the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model. It should be further appreciated that key-value cache memory is allocated to the individual heads of the multi-head attention model at a head-level of the multi-head attention model, rather than a layer-level of the multi-head attention model.

[0026] FIG. 1 is a diagram of an LLM system 100. As shown, the LLM system 100 includes a user 110 and an LLM 170. The user 110 submits an input prompt 120 to the LLM 170. The input prompt 120 may include a question that requires the LLM 170 to perform both reasoning and retrieval functions to process input prompt 120. As an example, the input prompt 120 may state “What is the favorite thing of the younger one between John and Mary? ” as discussed in some of the examples below. The LLM 170 may perform tokenization 180 on the input prompt 120 to generate a sequence of tokens 130, which may then be processed by the LLM decoder 190 to form a response 140 to the input prompt 120. The response 140 may be output by the LLM 170 responsive to the input prompt 120.

[0027] FIG. 2 is a diagram of an LLM decoder 290, which may correspond to the LLM decoder 190. As shown, the LLM decoder 290 includes an embedding layer 292, a positional encoding layer 293, decoding layers 295, a linear layer 296, a softmax layer 297, and LLM head (s) 298. The embedding layer 292, the positional encoding layer 293, the decoding layers 295, the linear layer 296, and the softmax layer 297 may collectively apply weights and transforms to a sequence of input tokens 230 to transform the sequence of input tokens 230 into logits 240. The logits 240 are then processed by the LLM head (s) 298 to obtain a response to a corresponding input response. In some applications, the LLM head (s) 298 may include a multi-head attention model that relies on key-value caching to accomplish an LLM inference function.

[0028] FIG. 3 is a diagram of key-value (KV) cache 300 in a multi-head attention model. As shown, the KV cache 300 is made up of heads 311, 312, 313, 314, 315 and layers 321, 322, 323, 324, 325. Conventional KV cache compression techniques allocate memory at the layer-level such that the memory is allocated amongst different layers 321, 322, 323, 324, 325, but generally does not vary between different heads 311, 312, 313, 314, 315. In contrast, head-level KV cache compression techniques provided by this disclosure allocate memory at the head-level such that the memory is allocated amongst different heads 311, 312, 313, 314, 315.

[0029] Aspects of this disclosure integrate head importance measures with head-level KV cache compression. Allocating larger KV cache budgets solely to retrieval heads may be insufficient for tasks that require both retrieval and reasoning to handle long input contexts, such as contextual Question Answering (QA) . To address this, aspects of this disclosure employ importance-score estimation that jointly evaluates each head's retrieval and reasoning capabilities for KV cache allocation. Questions may be constructed that require both the retrieval and reasoning abilities of an LLM. For instance, the question “What is the favorite thing of the younger one between John and Mary? ” may require the model to first identify “who is younger between John and Mary” (John) by referencing the context of their ages prior to retrieving “John's favorite thing. ” In such an instance, the importance score of each head may be estimated based on the attention scores generated by the model. Using the estimated importance scores, the KV cache budget may then be allocated for each individual head such that heads demonstrating greater importance in retrieval and reasoning retain a larger portion of the KV cache. Within each head, the most relevant KV cache entries may be retained.

[0030] FIG. 4 is a diagram of a head-level KV cache compression technique. As shown, head-level KV cache compression may include two steps: (1) Head-Level Importance Score Estimation which identifies important heads that contribute to the contextual reasoning ability using an updated Needle-in-a-Haystack test; and (2) Head-Level KV Cache Allocation which allocates KV cache budgets for each head based on the importance score distribution.

[0031] FIG. 5 is a flowchart of a method 500 for LLM processing. As shown, the method 500 begins at step 510, where the LLM receives an input prompt from a user. At step 520, the LLM converts the input prompt into a sequence of input tokens. At step 530, the LLM processes, or otherwise decodes, the sequence of input tokens using a multi-head attention model to obtain a response to the input prompt. This may generally include allocating key-value cache memory to individual heads of the multi-head attention model. At step 540, the LLM outputs the response to the input prompt to a user.

[0032] In some examples, processing the sequence of input tokens using the multi-head attention model also includes assigning importance scores to the individual heads of the multi-head attention model. In such examples, the key-value cache memory is allocated to the individual heads of the multi-head attention model based on the importance scores assigned to the individual heads of the multi-head attention model. Notably, the importance scores may be assigned to the individual heads of the multi-head attention model based on both retrieval attributes and reasoning attributes of the individual heads of the multi-head attention model. The retrieval and reasoning attributes of the individual heads of the multi-head attention model may be determined based on an updated Needle-in-a-Haystack test that includes, while the reasoning attributes of the individual heads of the multi-head attention model may be determined based on an updated Reasoning-in-a-Haystack test. It should be appreciated that the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model. It should further be appreciated that key-value cache memory is allocated to the individual heads of the multi-head attention model at a head-level of the multi-head attention model, rather than a layer-level of the multi-head attention model.

[0033] FIG. 6 is a diagram of retrieval and reasoning examples for head identification. Testing has been conducted based on various benchmarks requiring both retrieval and reasoning abilities, including QA tasks, using back-bone models. Results demonstrate that the head-level KV cache compression techniques outperform conventional approaches. Furthermore, KV cache allocation based on estimated importance scores, reflecting both retrieval and reasoning abilities, is shown to outperform conventional allocation strategies based on attention (i.e., retrieval ability) alone. In challenging scenarios, like Needle-in-a-Haystack and Reasoning-in-a-Haystack tests, embodiment methods effectively preserve the model's retrieval and reasoning capabilities. Also, experiments on memory and latency reveal show that embodiment approaches significantly reduce memory usage and decoding latency.

[0034] Aspects of this disclosure provide KV cache compression for multi-head attention models. Such approaches may considerably improve computational efficiency for LLMs, particularly for extremely long inputs. This may address one distinct challenge with KV caching, which is the tendency for memory usage to increase linearly with input length, leading to significant memory management advantages. Embodiments of this disclosure may significantly improve performance and efficiency over conventional approaches, which largely depend on rigid layer-level allocation strategies that may not fully optimize KV cache allocation for downstream tasks. Aspects of this disclosure may allocate KV cache primarily based on head-level importance scores, independent of layer constraints, resulting in more effective compression and outperforming baselines without increasing decoding latency or memory usage.

[0035] Embodiment head-level KV cache compression techniques may include three key components, including: (1) identifying important heads and calculating head-level importance score distributions; (2) using these distributions to efficiently allocate KV cache budgets across heads; and (3) determining which Key and Value vectors to retain within each head.

[0036] Accurate budget allocation at the head level requires identifying which heads are most and least important for the given task. By leveraging this importance distribution, it is possible to assign larger KV cache budgets to more critical heads and smaller budgets to less significant ones. To achieve this, aspects of this disclosure estimate the importance of each attention head. The Needle-in-a-Haystack test (e.g., test 2) , as well as custom retrieval examples, may be used to estimate the importance score for each retrieval head. In these examples, a question q that cannot be answered using the model's parametric knowledge is paired with an answer (k) , e.g., the “Needle” , and inserted into a haystack (c) at different positions (pi) . The model may then be required to retrieve the exact answer k from the combined input (k, c) . During each decoding step (t) , a head (h) earns a fraction of an importance score if (1) the token with the highest attention score arg max (ah) matches the generated token, and (2) the token is part of the inserted answer k. The final importance score for each head (h) may be calculated according to the following function:

[0037] where N is the length of the inserted answer (k) , and ah is the attention score on the combined input from head (h) . By using various settings of the inserted position (pi) , it is possible to obtain the head-level Retrieval Heads distribution.

[0038] Directly using this distribution may face two challenges: (1) it focuses on the retrieval-and-paste mechanism, lacking consideration for the contextual and reasoning skills needed for complex questions; and (2) the distribution is too sparse for effective head-level KV cache allocation, with nearly 70%of heads receiving an importance score of zero due to the strict exact-match requirement.

[0039] For retrieval-reasoning (R2) heads, aspects of this disclosure provide an embodiment importance score estimation method that accounts for both the retrieval and reasoning abilities of the heads, enabling a more accurate assessment of their significance. First, retrieval-reasoning examples may be constructed by adding explicit contextual reasoning steps to the retrieval examples (as shown in the Retrieval-Reasoning Example part of FIG. 6) . Additionally, the inserted needle may be modified based on the following equation: k = (r, c1, c2) , where r is the reasoning step, and c1 and c2 are different answers to the refined question q. The model may reason with (r) to retrieve and generate the correct answer (c2) , avoiding the wrong answer (c1) . Second, the estimation method may be refined by focusing on the entire correct answer (c2) , since all tokens are relevant to the question (q) . This approach aligns with contextual QA, which requires both retrieval and reasoning abilities. By considering the full correct answer, the importance score for each head (h) does not depend solely on the token (s) with the highest attention score. The importance score for head h is calculated as follows:

[0040] where ai ∈ ah represents the i-th highest attention score from head h, and top-i (ah) is the token with the i-th highest score. The importance score may be computed by considering the entire correct answer and increasing the number of tokens evaluated per head. Heads with higher attention on the correct answer (k) may receive higher importance scores. The importance scores may be further refined using attention weights in order to obtain a more accurate distribution.

[0041] Importance score estimation for each head may allow the model to identify the key heads and allocate the KV cache budget accordingly. These distributions may be incorporated into head-level KV cache allocation. In multi-head attention models, for each head (h) in layer (1) , the embedded input may be represented as follows: X = {x1, x2, ..., xn} ∈ Rn×d. The embedded input may be mapped into different subspaces using the query key and value matrices:

[0042] To optimize memory and enhance efficiency, KV cache compression methods may be employed to discard less important KV cache entries while preserving performance. For each head (h) , the compressed KV cache may be reduced to and where s << n, resulting in a significant improvement to computational efficiency.

[0043] Head-level allocation may differ from conventional KV cache compression during the prefill phase, which is limited to layer-level allocation treating all heads within a layer equally. In particular, head-level allocation may offer a comprehensive KV cache allocation strategy, where each head (h) is initially assigned a fixed KV cache size (b) with an associated importance score (Sh) . To allow dynamic allocation, a shared budget pool (B) is provided by extracting a portion of the budget from each head, leaving the remainder as the basic budget. The budget pool (B) is then distributed among the heads in proportion to their importance scores (Sh) . The final head-level KV cache allocation is as follows:

[0044] where b is the initial fix budget for each head, β is a hyper-parameter to control the size of the dynamic budget pool, L is the number of layers of the current LLM, and H is the number of heads of the current LLM. The last α instruct tokens are preserved before forming the dynamic budget pool B to guide the selection process. The retained KV cache for each head includes: (1) the basic budget (b -b ) , (2) the dynamic budget Sh × B, proportional to its importance score, and (3) the last α instruct tokens.

[0045] After determining the number of KV cache entries to retain, an attention-based selection strategy may be employed to select the most relevant entries for retainment. Specifically, the last α instruction tokens (local windows) guide KV cache selection for each head. Attention scores from these local windows to the remaining tokens are aggregated through a pooling layer, with higher-scoring tokens considered more important and retained in the cache.

[0046] Experiments have been conducted to determine the KV cache baselines and implementation details, as well as to: (1) emphasize the importance of enhancing contextual reasoning in importance score estimation; (2) use the Needle-in-a-Haystack and Reasoning-in-a-Haystack tests to demonstrate how head-level KV cache compression improves long-context retrieval and reasoning; and (3) provide a comprehensive comparison with conventional layer-level allocation methods.

[0047] Head-level KV cache compression may be used in conjunction with Llama-3-8B-Instruct and Mistral-7B-Instruct models. FIG. 7 is a table of experimental results for comparing LongBench and LooGLE benchmarks for Llama-3-8B-Instruct and Mistral-7B-Instruct models. For LongBench, datasets from the Single-Doc QA and Multi-Doc QA categories are used to assess contextual reasoning. For LooGLE, the Long Dependency QA task is evaluated, which includes four QA-related tasks.

[0048] To determine baselines and settings, three conventional KV cache compression techniques are evaluated, including: (1) SnapKV which uses the last et tokens as local windows to guide KV cache selection. Attention scores from these windows to the remaining tokens are pooled to cluster information and guide the selection process; (2) PyramidKV which follows a pyramid attention pattern, allocating more KV cache to lower layers to retain key information, while reducing the budget for higher layers where information is already aggregated; and (3) Ada-KV which dynamically allocates budgets to heads within each layer based on their concentration degrees. Ada-KV can be combined with SnapKV or PyramidKV. Ada-SnapKV is used as the baseline due to its superior performance over Ada-PyramidKV.

[0049] FIG. 8 is a graph of experimental results for the Llama-3-8B-Instruct model using different KV cache sizes and KV cache allocation techniques. FIG. 9 is a graph of experimental results for the Mistral-7B-Instruct model using different KV cache sizes and KV cache allocation techniques. Results are shown for different KV cache sizes (64, 128, 256, 512, 1024) , showing average accuracy across six datasets from the LongBench benchmark with an average input length of 8, 683 tokens. Notably, a KV cache size of 64 retains just 0.7%of the total tokens.

[0050] Head-level KV cache compression techniques may be implemented with a strategy to guide KV cache selection after allocating the budget. One such strategy may be to set the size of the local windows (e.g., α = 8) for both the baselines and head-level KV cache compression. The hyper-parameter (β) , which controls the size of the shared budget pool, may be chosen from the set {1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10} . Head-level KV cache compression consistently outperforms the baselines, including with 64 and 128 KV cache configurations. Layer-level methods allocate a fixed cache size to all heads within a layer, making it difficult to retain essential information. Ada-SnapKV improves this by allowing dynamic allocation within layers, but still relies on fixed layer-level budgets. In contrast, head-level KV cache compression dynamically allocates across individual heads, retaining critical information by adjusting the budget based on the importance score associated with each head. Global KV cache allocation was performed using both retrieval heads distribution (HeadKV-R) and retrieval-reasoning heads distribution (HeadKV-R2) . Integrating HeadKV-R2 significantly improves results over HeadKV-R. Indeed, HeadKV-R2 surpasses the Full-KV cache average (32.90) , achieving 32.95 with a 1024 KV cache. Overall, both the head-level allocation strategy and HeadKV-R2 provide significant performance gains.

[0051] FIGS. 10A-10C are graphs of head visualization for Llama-3-8B-Instruct results. As shown, the retrieval heads distribution is relatively sparse, while Retrieval-Reasoning Heads has denser distribution which allows for more effective differentiation between heads.

[0052] FIG. 11 is a table of ablation study results for LongBench benchmarks. As shown, HeadKV-R leverages the standard Retrieval Heads distribution, while HeadKV-ER uses retrieval examples with importance score estimation. HeadKV-R2 leverages both the proposed importance score estimation and retrieval-reasoning examples. Alongside the standard HeadKV-R and Retrieval-Reasoning Heads distributions, Enhanced-Retrieval Heads distribution (HeadKV-ER) is introduced using retrieval examples with importance score estimation. Comparing HeadKV-R and HeadKV-ER reveals that focusing on the entire needle, rather than specific tokens, improves performance. As shown, HeadKV-R is sparser than HeadKV-ER. Due to strict constraints, HeadKV-R results in most heads receiving a score of zero, leading to a sparser results than when incorporating HeadKV-R2. While the HeadKV-ER improves performance slightly, it relies heavily on retrieval examples and lacks full consideration of contextual reasoning. In contrast, HeadKV-R2 reflects both retrieval and reasoning abilities, and consistently outperforms HeadKV-R and HeadKV-ER. The updated Needle-in-a-Haystack test (which incorporates Retrieval-Reasoning example (s) ) may be conducted to assess the long-context retrieval ability of different KV cache compression methods. FIGS. 12A-12F are graphs of Needle-in-a-Haystack test results on the Llama-3-8B-Instruct model. The KV size has been set to 128 for experimental results discussed herein. Results from the Llama-3-8B-Instruct demonstrate that head-level KV cache compression effectively retains important information compared with baselines. FIG. 13 is a table of test results from Reasoning-in-a-Haystack test results.

[0053] Final results are averaged across QA1-QA5 tasks for each length. Unlike the Needle-in-a-Haystack test, this test inserts multiple needles into the haystack, requiring the model to reason through them to extract the correct answer. This long-context contextual reasoning ability is crucial for many tasks, such as question answering (QA) and summarization. Therefore, it is helpful to conduct a Reasoning-in-a-Haystack test to evaluate the long-context reasoning ability of each KV cache compression method across different scenarios.

[0054] Reasoning-in-a-Haystack testing was setup based on the format set forth in the non-patent literature (NPL) publication entitled “Testing the limits of llms with long context reasoning-in-a-haystack” (hereinafter “Kuratov 2024” ) , which is incorporated herein by reference as if reproduced in its entirety (<https:  / / arxiv. org / abs / 2406.10149>) . As described therein, reasoning needles are inserted into the haystack and the model retrieves and reasons through the needles to generate a correct answer. Datasets from Kuratov et al. (2024) are used to average results across QA1-QA5 tasks for evaluation. FIG. 15 is a table of example tasks used for Reasoning-in-a-Haystack tests.

[0055] Head-level KV cache compression significantly outperforms baselines. By incorporating retrieval-reasoning examples, head-level KV cache compression achieves better accuracy, particularly with the Llama-3-8B-Instruct model. Notably, combining head-level KV cache allocation with the HeadKV-R yields improved results over baselines. This is due to overlap between HeadKV-R and HeadKV-ER (indicating heads may serve multiple roles) as well as the bAbI benchmark containing an answer within the inserted needle (emphasizing retrieval alone helps head-level KV cache compression locate the needle) .

[0056] Head-level KV cache compression significantly improves memory and latency performance. The computational efficiency of head-level KV cache compression is evaluated using Mistral-7B-Instruct with a maximum sequence length of 32K and FlashAttention, which is discussed in the NPL publication entitled “Faster attention with better parallelism and work partitioning, ” which is incorporated herein by reference as if reproduced in its entirety. To assess the decoding latency of each method, the 32K-length data input from Kuratov 2024 with various generation lengths (1, 512, 1024, 2048, 4096) is used. FIGS. 14A-14B are graphs of decoding and latency test results for different KV cache compression techniques. As shown, head-level KV cache compression achieves near identical decoding latency as other KV cache compression techniques (i.e., HeadKV is essentially not distinguishable from SnapKV, PyramidKV, and Ada-SnapKV in FIG. 14A) , while maintaining performance close to Full KV cache (as shown in FIG. 14B) . Decoding time includes both the pre-filling time and the decoding time. The pre-filling time for head-level KV cache compression is comparable to PyramidKV, Ada-SnapKV (i.e., the increase in prefilling time is negligible) .

[0057] Notably, head-level KV cache compression also provides comparable memory usage during the decoding phase when compared to conventional KV cache compression techniques.

[0058] Exemplary Operating Environment

[0059] The present disclosure is operable with a computing apparatus according to an embodiment as a functional block diagram 1600 in FIG. 16. In an example, components of a computing apparatus 1618 are implemented as a part of an electronic device according to one or more embodiments described in this specification.

[0060] The computing apparatus 1618 comprises one or more processors 1619 which can be microprocessors, controllers, or any other suitable type of processors for processing computer executable instructions to control the operation of the electronic device. Alternatively, or in addition, the processor 1619 is any technology capable of executing logic or instructions, such as a hardcoded machine. In some examples, platform software comprising an operating system 1620 or any other suitable platform software is provided on the apparatus 1618 to enable application software 1621 to be executed on the device. It should be appreciated that the computing apparatus 1618 may be a computer for supporting a large language model (LLM) .

[0061] In some examples, computer executable instructions are provided using any computer-readable medium or media accessible by the computing apparatus 1618. Computer-readable media include, for example, computer storage media such as a memory 1622 and communications media. Computer storage media, such as a memory 1622, include volatile and non-volatile, removable, and non-removable media implemented in any method or technology for storage of information such as computer readable instructions, data structures, program modules or the like. Computer storage media include, but are not limited to, Random Access Memory (RAM) , Read-Only Memory (ROM) , Erasable Programmable Read-Only Memory (EPROM) , Electrically Erasable Programmable Read-Only Memory (EEPROM) , persistent memory, phase change memory, flash memory or other memory technology, Compact Disk Read-Only Memory (CD-ROM) , digital versatile disks (DVD) or other optical storage, magnetic cassettes, magnetic tape, magnetic disk storage, shingled disk storage or other magnetic storage devices, or any other non-transmission medium that can be used to store information for access by a computing apparatus. In contrast, communication media may embody computer readable instructions, data structures, program modules, or the like in a modulated data signal, such as a carder wave, or other transport mechanism. As defined herein, computer storage media do not include communication media. Therefore, a computer storage medium does not include a propagating signal. Propagated signals per se are not examples of computer storage media. Although the computer storage medium (the memory 1622) is shown within the computing apparatus 1618, it will be appreciated by a person skilled in the art, that, in some examples, the storage is distributed or located remotely and accessed via a network or other communication link (e.g., using a communication interface 1623) .

[0062] Further, in some examples, the computing apparatus 1618 comprises an input / output controller 1624 configured to output information to one or more output devices 1625, for example a display or a speaker, which are separate from or integral to the electronic device. Additionally, or alternatively, the input / output controller 1624 is configured to receive and process an input from one or more input devices 1626, for example, a keyboard, a microphone, or a touchpad. In one example, the output device 1625 also acts as the input device. An example of such a device is a touch sensitive display. The input / output controller 1624 in other examples outputs data to devices other than the output device, e.g., a locally connected printing device. In some examples, a user provides input to the input device (s) 1626 and / or receives output from the output devices 1625.

[0063] The functionality described herein can be performed, at least in part, by one or more hardware logic components. The computing apparatus 1618 is configured by the program code when executed by the processor 1619 to execute the embodiments of the operations and functionality described. Alternatively, or in addition, the functionality described herein can be performed, at least in part, by one or more hardware logic components. For example, and without limitation, illustrative types of hardware logic components that can be used include Field-programmable Gate Arrays (FPGAs) , Application-specific Integrated Circuits (ASICs) , Program-specific Standard Products (ASSPs) , System-on-a-chip systems (SOCs) , Complex Programmable Logic Devices (CPLDs) , Graphics Processing Units (GPUs) .

[0064] At least a portion of the functionality of the various elements in the figures may be performed by other elements in the figures, or an entity (e.g., processor, web service, server, application program, computing device, etc. ) not shown in the figures.

[0065] Although described in connection with an exemplary computing system environment, examples of the disclosure are capable of implementation with numerous other general purpose or special purpose computing system environments, configurations, or devices.

[0066] Examples of well-known computing systems, environments, and / or configurations that are suitable for use with aspects of the disclosure include, but are not limited to, mobile or portable computing devices (e.g., smartphones) , personal computers, server computers, hand-held (e.g., tablet) or laptop devices, multiprocessor systems, gaming consoles or controllers, microprocessor-based systems, set top boxes, programmable consumer electronics, mobile telephones, mobile computing and / or communication devices in wearable or accessory form factors (e.g., watches, glasses, headsets, or earphones) , network PCs, minicomputers, mainframe computers, distributed computing environments that include any of the above systems or devices, and the like. In general, the disclosure is operable with any device with processing capability such that it can execute instructions such as those described herein. Such systems or devices accept input from the user in any way, including from input devices such as a keyboard or pointing device, via gesture input, proximity input (such as by hovering) , and / or via voice input.

[0067] Examples of the disclosure may be described in the general context of computer-executable instructions, such as program modules, executed by one or more computers or other devices in software, firmware, hardware, or a combination thereof. The computer-executable instructions may be organized into one or more computer-executable components or modules. Generally, program modules include, but are not limited to, routines, programs, objects, components, and data structures that perform particular tasks or implement particular abstract data types. Aspects of the disclosure may be implemented with any number and organization of such components or modules. For example, aspects of the disclosure are not limited to the specific computer-executable instructions, or the specific components or modules illustrated in the figures and described herein. Other examples of the disclosure include different computer-executable instructions or components having more or less functionality than illustrated and described herein.

[0068] In examples involving a general-purpose computer, aspects of the disclosure transform the general-purpose computer into a special-purpose computing device when configured to execute the instructions described herein.

[0069] Any range or device value given herein may be extended or altered without losing the effect sought, as will be apparent to the skilled person.

[0070] While no personally identifiable information is tracked by aspects of the disclosure, examples have been described with reference to data monitored and / or collected from the users. In some examples, notice may be provided to the users of the collection of the data (e.g., via a dialog box or preference setting) and users are given the opportunity to give or deny consent for the monitoring and / or collection. The consent can take the form of opt-in consent or opt-out consent.

[0071] Although the subject matter has been described in language specific to structural features and / or methodological acts, it is to be understood that the subject matter defined in the appended claims is not necessarily limited to the specific features or acts described above. Rather, the specific features and acts described above are disclosed as example forms of implementing the claims.

[0072] It will be understood that the benefits and advantages described above may relate to one embodiment or may relate to several embodiments. The embodiments are not limited to those that solve any or all of the stated problems or those that have any or all of the stated benefits and advantages. It will further be understood that reference to ′an′ item refers to one or more of those items.

[0073] The embodiments illustrated and described herein as well as embodiments not specifically described herein but within the scope of aspects of the claims constitute exemplary means for receiving a first search request, the first search request including one or more search terms; identifying one or more product categories as output from a machine learning classification model in response to inputting of the one or more search terms; identifying a first plurality of products that are assigned to the one or more product categories, each product of the first plurality of products including a plurality of product titles and a plurality of product short descriptions in a natural language; applying the plurality of product titles and the plurality of product short descriptions as input to a second machine learning model that is configured to generate a plurality of recommended searches, each recommended search of the plurality of recommended searches including at least one search term; scoring each recommended search of the plurality of recommended searches; selecting one or more recommended searches of the plurality of recommended searches based on the scoring; and causing the one or more recommended searches to be displayed as user-interactable components on a graphical user interface, each user-interactable component being configured to execute a second search request upon user interaction with the user-interactable component.

[0074] At least a portion of the functionality of the various elements in FIG. 1 to FIG. 9 can be performed by other elements.

[0075] In some examples, the operations described herein can be implemented as software instructions encoded on a computer-readable medium, in hardware programmed or designed to perform the operations, or both. For example, aspects of the disclosure can be implemented as a system on a chip or other circuitry including a plurality of interconnected, electrically conductive elements.

[0076] While the aspects of the disclosure have been described in terms of various examples with their associated operations, a person skilled in the art would appreciate that a combination of operations from any number of different examples is also within scope of the aspects of the disclosure.

[0077] The term “Wi-Fi” as used herein refers, in some examples, to a wireless local area network using high frequency radio signals for the transmission of data. The term  as used herein refers, in some examples, to a wireless technology standard for exchanging data over short distances using short wavelength radio transmission. The term “NFC” as used herein refers, in some examples, to a short-range high frequency wireless communication technology for the exchange of data over short distances.

[0078] The term “comprising” is used in this specification to mean including the feature (s) or act (s) followed thereafter, without excluding the presence of one or more additional features or acts.

[0079] In some examples, the operations illustrated in the figures are implemented as software instructions encoded on a computer readable medium, in hardware programmed or designed to perform the operations, or both. For example, aspects of the disclosure are implemented as a system on a chip or other circuitry including a plurality of interconnected, electrically conductive elements.

[0080] The order of execution or performance of the operations in examples of the disclosure illustrated and described herein is not essential, unless otherwise specified. That is, the operations may be performed in any order, unless otherwise specified, and examples of the disclosure may include additional or fewer operations than those disclosed herein. For example, it is contemplated that executing or performing a particular operation before, contemporaneously with, or after another operation is within the scope of aspects of the disclosure.

[0081] When introducing elements of aspects of the disclosure or the examples thereof, the articles "a, " "an, " "the, " and "said" are intended to mean that there are one or more of the elements. The terms "comprising, " "including, " and "having" are intended to be inclusive and mean that there may be additional elements other than the listed elements. The term “exemplary” is intended to mean “an example of. ” The phrase “one or more of the following: A, B, and C” means “at least one of A and / or at least one of B and / or at least one of C. "

[0082] Within the scope of this application, it is expressly intended that the various aspects, embodiments, examples, and alternatives set out in the preceding paragraphs, in the claims and / or in the description and drawings, and in particular the individual features thereof, may be taken independently or in any combination. That is, all embodiments and / or features of any embodiment can be combined in any way and / or combination, unless such features are incompatible. The applicant reserves the right to change any originally filed claim or file any new claim, accordingly, including the right to amend any originally filed claim to depend from and / or incorporate any feature of any other claim although not originally claimed in that manner.

[0083] Having described aspects of the disclosure in detail, it will be apparent that modifications and variations are possible without departing from the scope of aspects of the disclosure as defined in the appended claims. As various changes could be made in the above constructions, products, and methods without departing from the scope of aspects of the disclosure, it is intended that all matter contained in the above description and shown in the accompanying drawings shall be interpreted as illustrative and not in a limiting sense.

Claims

1.A method comprising:receiving, by a large language model (LLM) , an input prompt;converting, by the LLM, the input prompt into a sequence of input tokens;processing, by the LLM, the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, wherein processing the sequence of input tokens using the multi-head attention model comprises allocating key-value cache memory to individual heads of the multi-head attention model; andoutputting, by the LLM, the response to the input prompt.2.The method of claim 1, wherein processing the sequence of input tokens using the multi-head attention model further comprises:assigning importance scores to the individual heads of the multi-head attention model, wherein the key-value cache memory is allocated to the individual heads of the multi-head attention model based on the importance scores assigned to the individual heads of the multi-head attention model.3.The method of claim 2, wherein the importance scores are assigned to the individual heads of the multi-head attention model based on both retrieval attributes and reasoning attributes of the individual heads of the multi-head attention model.4.The method of claim 3, wherein the importance scores are determined using an updated Needle-in-a-Haystack test that incorporates a Retrieval-Reasoning evaluation.5.The method of claim 4, wherein the updated Needle-in-a-Haystack test evaluates both the retrieval attributes and the reasoning attributes of the individual heads of the multi-head attention model.6.The method of claim 3, wherein the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model.7.The method of claim 1, wherein the key-value cache memory is allocated to the individual heads of the multi-head attention model at a head-level of the multi-head attention model rather than a layer-level of the multi-head attention model.8.A computer for supporting a large language model (LLM) , the computer comprising:a processor; anda computer-readable medium storing programming instructions that, when executed by the processor, cause the processor to:receive an input prompt;convert the input prompt into a sequence of input tokens;process the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, wherein the instructions to process the sequence of input tokens using the multi-head attention model include instructions to allocate key-value cache memory to individual heads of the multi-head attention model; andoutput the response to the input prompt.9.The computer of claim 8, wherein the instructions to process the sequence of input tokens using the multi-head attention model further include instructions to:assign importance scores to the individual heads of the multi-head attention model, wherein the key-value cache memory is allocated to the individual heads of the multi-head attention model based on the importance scores assigned to the individual heads of the multi-head attention model.10.The computer of claim 9, wherein the importance scores are assigned to the individual heads of the multi-head attention model based on both retrieval attributes and reasoning attributes of the individual heads of the multi-head attention model.11.The computer of claim 10, wherein the importance scores are determined using an updated Needle-in-a-Haystack test that incorporates a Retrieval-Reasoning evaluation.12.The computer of claim 11, wherein the updated Needle-in-a-Haystack test evaluates both the retrieval attributes and the reasoning attributes of the individual heads of the multi-head attention model.13.The computer of claim 10, wherein the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model.14.The computer of claim 8, wherein the key-value cache memory is allocated to the individual heads of the multi-head attention model at a head-level of the multi-head attention model rather than a layer-level of the multi-head attention model.15.A computer program product for supporting a large language model (LLM) , the computer program product storing programming instructions that, when executed by a processor, cause the processor to:receive an input prompt;convert the input prompt into a sequence of input tokens;process the sequence of input tokens using a multi-head attention model to obtain a sequence of output tokens forming a response to the input prompt, wherein the instructions to process the sequence of input tokens using the multi-head attention model include instructions to allocate key-value cache memory to individual heads of the multi-head attention model; andoutput the response to the input prompt.16.The computer program product of claim 15, wherein the instructions to process the sequence of input tokens using the multi-head attention model further include instructions to:assign importance scores to the individual heads of the multi-head attention model, wherein the key-value cache memory is allocated to the individual heads of the multi-head attention model based on the importance scores assigned to the individual heads of the multi-head attention model.17.The computer program product of claim 16, wherein the importance scores are assigned to the individual heads of the multi-head attention model based on both retrieval attributes and reasoning attributes of the individual heads of the multi-head attention model.18.The computer program product of claim 17, wherein the importance scores are determined using an updated Needle-in-a-Haystack test that incorporates a Retrieval-Reasoning evaluation.19.The computer program product of claim 18, wherein the updated Needle-in-a-Haystack test evaluates both the retrieval attributes and the reasoning attributes of the individual heads of the multi-head attention model.20.The computer program product of claim 17, wherein the reasoning attributes of the individual heads are different than attention scores associated with the multi-head attention model.