A single-request heterogeneous memory optimization method and system for large language model reasoning

Abstract

The application provides a single-request heterogeneous memory optimization method and system for large language model reasoning, which realizes attention head granularity zero software overhead access statistics by sinking key value cache heat collection to memory controller and CXL device controller hardware units; combined with hierarchical storage scheduling and heat decay mechanism, the power law characteristics of key value cache access are matched, high frequency hot data is directed to reside in GPU HBM, and the high speed memory missing rate is reduced. Relying on independent hardware prefetch engine and CXL.cache protocol primitive, three-level priority arbitration is realized to realize hardware-level parallel of cross-layer data transfer and GPU calculation; with the help of HBM double buffering ping-pong pipeline, the prefetch delay is hidden in the QKV projection calculation window, and the idle bandwidth is reused. Compared with the traditional pure software scheme, the scheme eliminates the heat collection overhead, improves the HBM bandwidth utilization and the CXL memory acceleration value, significantly reduces the single-request reasoning delay, and can adapt to the memory characteristics of different models and sequence lengths, and long-term maintain the optimal scheduling efficiency.

Description

Technical Field

[0001] This invention relates to the field of electronic digital data processing technology, and in particular to a single-request heterogeneous memory optimization method and system for large language model inference. Background Technology

[0002] Large Language Models (LLMs), with their superior semantic understanding and generation capabilities, have been widely applied to numerous artificial intelligence tasks such as intelligent dialogue, code generation, and scientific computing. However, the computational and storage overhead of the LLM inference process, especially the autoregressive decoding stage, is extremely large, which has become a key bottleneck restricting its large-scale deployment.

[0003] Currently, LLM inference commonly employs the Transformer architecture, with its core component being the Multi-Head Attention (MHA) mechanism. During autoregressive generation, to avoid repeatedly calculating the key-value (KV) tensors of historical tokens, the system needs to cache historical key-value pairs in high-speed memory, i.e., a KV cache. Taking the current mainstream 70-byte parameter scale model as an example, its inference process requires a GB-level key-value cache. With the continuous growth of model parameter size and sequence length, the storage capacity requirement for the key-value cache expands dramatically.

[0004] In terms of storage, limited by the capacity and cost of on-chip high-bandwidth memory (HBM), the HBM space of a single GPU is typically 80GB to 96GB, which cannot simultaneously accommodate the complete model weights and all key-value caches. Therefore, existing systems generally adopt a heterogeneous memory architecture, which combines GPU HBM, near-end memory extended based on the CXL (Compute Express Link) protocol, and host-side CPU DDR memory to form a multi-layered, large-capacity hybrid memory pool to carry model weights and key-value caches. Under this architecture, the efficient scheduling of key-value caches across heterogeneous memory layers directly determines inference latency and system throughput.

[0005] To address the aforementioned heterogeneous memory scheduling problem, the industry has proposed several solutions. Regarding key-value cache popularity awareness, existing inference frameworks typically maintain access counters at the software layer to identify frequently accessed KV data, thereby guiding data migration between memory layers. However, this approach requires executing several additional CPU or GPU instructions for each memory access to record popularity information, introducing an additional 5% to 15% runtime overhead. Furthermore, the granularity of data collection is limited by the framework's abstraction level, only accurate to the request or model level, and cannot be refined to individual attention heads or cache lines, thus limiting the accuracy of popularity information.

[0006] Regarding data prefetching, to hide the latency of reading data from slow memory, existing solutions propose asynchronous prefetching of key-value caches via software-scheduled DMA (Direct Memory Access). However, such prefetch instructions share the instruction scheduling channel with GPU computation instructions, making it impossible to achieve hardware-level parallelism with GPU computation. Furthermore, the prefetch granularity is limited by the DMA descriptor size, making it difficult to match the logical structure of the key-value cache, resulting in low prefetch efficiency.

[0007] In terms of pipelined execution, LLM single-layer decoding includes two types of operations: computationally intensive QKV projection and memory-intensive Attention computation. Existing inference systems usually execute these operations serially, failing to fully utilize the idle memory bandwidth during QKV projection and the idle GPU computing power during Attention computation, resulting in a significant waste of hardware resources.

[0008] In addition, the CXL protocol itself provides hardware primitives such as non-blocking transactions, cache consistency atomic operations, and MSI-X interrupts in the CXL.cache sub-protocol. However, current mainstream inference frameworks only regard CXL extended memory as a passive storage extension resource and fail to make full use of the scheduling capabilities of the above hardware primitives.

[0009] It is evident that existing technologies urgently need a new solution to achieve accurate awareness, efficient prefetching, and pipelined execution of key-value caches in heterogeneous memory architectures, in order to minimize inference latency and improve memory bandwidth utilization. Summary of the Invention

[0010] In view of this, embodiments of the present invention provide a single-request heterogeneous memory optimization method and system for large language model inference, which solves the problem of low scheduling efficiency of key-value cache in a heterogeneous memory system composed of GPU HBM, CXL extended memory and CPU DDR in large language model inference scenarios.

[0011] One aspect of the present invention provides a single-request heterogeneous memory optimization method for large language model inference, the method comprising the following steps: The system monitors memory access instructions based on the access pattern awareness unit pre-embedded in the server memory controller and CXL device controller. It filters access to the key-value cache by the range of physical addresses accessed in memory, and maps the physical addresses accessed in memory to the corresponding attention heads according to the attention head mapping table. It counts the number of times each attention head is accessed within a set time window and exposes it to the inference engine through the heat register interface. The cache addresses of the key-value caches corresponding to each attention head are divided into a high-speed hot data layer based on HBM, a medium-speed warm data layer based on CXL, and a low-speed cold data layer based on DDR. The inference engine decays the number of times each attention head is accessed according to a set time step interval, and calculates the heat score of each attention head in combination with the weight of the current data layer, and sorts them in descending order. The key-value caches of the first set number of attention heads in the sorted order are marked as prefetch candidates, added to the prefetch command queue, and written to the hardware prefetch engine pre-embedded in the CXL device controller through the MMIO interface. In response to the prefetch trigger signal from the access pattern awareness unit or the prefetch instruction from the inference engine, the hardware prefetch engine implements zero-software-overhead key-value cache prefetching based on hardware primitives of the CXL.cache protocol. A prefetch arbitrator prioritizes the key-value caches corresponding to each attention head in the prefetch command queue based on SLA urgency, heat score, and prefetch benefit estimation. Each clock cycle, the highest-priority cache is selected from the prefetch command queue and prefetched to HBM. Upon completion of prefetching, the inference engine is asynchronously notified via an MSI-X interrupt. In HBM, different cache areas are configured for ping-pong switching, simultaneously used for current computation and background prefetching. During the QKV projection stage of single-layer decoding, the key-value cache required for the next computation is prefetched. During the attention computation stage, the prefetched key-value cache is retrieved without waiting. In some embodiments, filtering access to the key-value cache by the range of physical addresses accessed in memory includes: obtaining the physical address in the memory access instruction and extracting the high-order bits, comparing it with the start and end addresses of a pre-configured key-value cache memory region; if the key-value cache memory region is hit, then performing subsequent access count statistics; if the key-value cache memory region is not hit, then allowing access directly without processing.

[0012] In some embodiments, the attention head mapping table is implemented based on a tri-state content-addressable memory. When allocating the key-value cache, the inference engine writes the mapping relationship between the cache address and the attention head into the attention head mapping table through the MMIO interface. The allocation method includes continuous allocation and paginated allocation. When allocating the key-value cache, the inference engine synchronously records the data level corresponding to the corresponding cache address. The data level includes: the high-speed hot data layer, the medium-speed warm data layer, or the low-speed cold data layer. The method assigns an independent hardware counter to each attention head to form an access frequency counter array, which is used to count the number of times the attention head is accessed, achieving zero software overhead.

[0013] In some embodiments, the inference engine attenuates the number of times each attention head is accessed at a set time step interval, and calculates the popularity score of each attention head in combination with the weight of the current data layer to which it belongs. This includes: reading the number of times each attention head is accessed in the access pattern perception unit, finding the data layer to which the key-value cache of each attention head belongs, and determining the weight of the current data layer to which it belongs; multiplying the number of times all attention heads are accessed by α at a third set number of time steps to achieve attenuation, where α < 1; and multiplying the number of times each attention head is accessed and its corresponding weight of the current data layer to obtain the corresponding popularity score.

[0014] In some embodiments, after calculating the heat score of each attention head in conjunction with the weight of the current data layer and sorting them in descending order, the method further includes: According to preset rules, the sorted attention heads are divided into high-temperature segments matching the high-speed hot data layer, medium-temperature segments matching the medium-speed warm data layer, and low-temperature segments matching the low-speed cold data layer. A migration instruction queue is generated and written to the pre-embedded hardware prefetch engine in the CXL device controller through the MMIO interface to perform data transfer.

[0015] Specifically, the sorted attention heads are divided into a high-temperature segment matching the high-speed hot data layer, a medium-temperature segment matching the medium-speed warm data layer, and a low-temperature segment matching the low-speed cold data layer according to preset rules. This includes marking the first 20% of the sorted attention heads as the high-temperature segment, marking the first 20% to 60% of the attention heads as the medium-temperature segment, and marking the last 40% of the attention heads as the low-temperature segment.

[0016] In some embodiments, the prefetch command queue supports three modes: single prefetch, batch discontinuous data prefetch, and predictive prefetch. The predictive prefetch is based on the access pattern sensing unit storing the number of times each attention head is accessed in a fourth set number of consecutive sampling windows and fitting the trend of change. The key value cache corresponding to the attention head whose access count is continuously higher than a first set threshold and whose trend of change is rising or remains high is included in the prefetch command queue. Based on the prefetch arbitrator, the key-value caches corresponding to each attention head in the prefetch command queue are sorted according to a three-level priority system: SLA urgency, popularity score, and prefetch benefit prediction, including: Based on the time margin before the SLA defaults, the SLA urgency level is assigned according to preset rules. The prefetching benefit is calculated based on the expected hit rate of the key-value cache corresponding to each attention point, the amount of data prefetched, and the transmission latency. The calculation formula is as follows: ; Where V represents the pre-fetched revenue. This indicates the expected hit rate. This indicates the amount of data to be prefetched. This indicates the transmission delay; The SLA urgency, the heat score, and the prefetch benefit are weighted and summed to calculate a comprehensive priority, and the key-value caches corresponding to each attention point in the prefetch command queue are sorted based on the comprehensive priority.

[0017] In some embodiments, the standard for the access pattern sensing unit to generate a prefetch trigger signal is: if the number of times the attention head is accessed exceeds a second set threshold in a second set number of consecutive sampling windows, a prefetch trigger signal is generated. Each time the inference engine accesses the key-value cache, it checks whether the corresponding cache address has been prefetched to HBM and calculates the overall hit rate. If the overall hit rate is less than a third set threshold, the second set threshold is increased by a first ratio to reduce the prefetch coverage, where the first ratio is greater than 1; if the overall hit rate is greater than a fourth set threshold, the second set threshold is decreased by a second ratio to increase the prefetch coverage, where the second ratio is less than 1. Furthermore, if the overall hit rate is greater than the fifth set threshold, a third cache area is added in HBM; if the overall hit rate is less than the sixth set threshold, the two cache areas in HBM are reduced to a single cache area and ping-pong switching is stopped.

[0018] On the other hand, the present invention also provides a single-request heterogeneous memory optimization system for large language model inference. The system combines GPU HBM, GPU near-end memory extended via the CXL protocol, and host-side CPU DDR memory. An access pattern awareness unit is pre-embedded in the server memory controller and the CXL device controller. The access pattern awareness unit includes an address range comparator, a tri-state content-addressable memory for implementing an attention head mapping table, an access frequency counter array, a heat register interface, and a prefetch trigger signal generator. The system is used to execute the computer program or instructions. When the computer program / instructions are executed, the device / system implements the steps of the above method.

[0019] On the other hand, the present invention also provides a computer-readable storage medium having a computer program or instructions stored thereon, which, when executed by a processor, implement the steps of the above-described method.

[0020] On the other hand, the present invention also provides a computer program product, including a computer program or instructions that, when executed by a processor, implement the steps of the above-described method.

[0021] The single-request heterogeneous memory optimization method and system for large language model inference described in this invention adopts a collaborative optimization architecture of hardware awareness, software decision-making, hardware execution, and pipeline interweaving. By sinking the collection of key-value cache access heat to the hardware units of the memory controller and CXL device controller, it realizes zero-software-overhead access frequency statistics at the attention head granularity. Combined with hierarchical storage scheduling strategy and heat decay mechanism, it can accurately match the access power-law characteristics of key-value cache, and direct high-frequency hot data to reside on GPU HBM, effectively reducing high-speed memory miss rate. At the same time, relying on the independently running hardware prefetch engine and the native hardware primitives of the CXL.cache protocol, it combines SLA urgency, heat score, and prefetch benefit to comprehensively evaluate prefetch priority for key-value cache, realizing full hardware-level parallelism of cross-memory level data transfer and GPU computing. Then, by using the HBM double-buffered ping-pong switching pipeline design, the key-value cache prefetch latency is hidden within the QKV projection computing time window, making full use of idle memory bandwidth during computationally intensive stages. Compared to traditional pure software scheduling schemes, this invention eliminates the runtime overhead of heat collection, improves HBM bandwidth utilization and the active acceleration value of CXL extended memory, significantly reduces single-request inference latency, and can adapt to the memory access characteristics of different models and sequence lengths to maintain long-term optimal scheduling efficiency.

[0022] Additional advantages, objects, and features of the invention will be set forth in part in the description which follows, and will also become apparent in part to those skilled in the art upon studying the description, or may be learned by practice of the invention. The objects and other advantages of the invention can be realized and obtained by means of the structures specifically pointed out in the description and drawings.

[0023] Those skilled in the art will understand that the objectives and advantages achievable with the present invention are not limited to those specifically described above, and that the above and other objectives achievable with the present invention will become clearer from the following detailed description. Attached Figure Description

[0024] The accompanying drawings, which are included to provide a further understanding of the invention and form part of this application, are not intended to limit the scope of the invention. In the drawings: Figure 1 This is a flowchart illustrating a single-request heterogeneous memory optimization method for large language model inference according to an embodiment of the present invention.

[0025] Figure 2 This is a diagram of the four-layer linkage scheduling architecture used in the single-request heterogeneous memory optimization method for large language model inference according to an embodiment of the present invention.

[0026] Figure 3This is a server system architecture diagram applicable to the single-request heterogeneous memory optimization method for large language model inference according to an embodiment of the present invention.

[0027] Figure 4 This is a timing diagram of the computation-memory access double-buffered pipeline in a single-request heterogeneous memory optimization method for large language model inference according to an embodiment of the present invention. Detailed Implementation

[0028] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the embodiments and accompanying drawings. Here, the illustrative embodiments and descriptions of this invention are used to explain the invention, but are not intended to limit the invention.

[0029] It should also be noted that, in order to avoid obscuring the invention with unnecessary details, only the structures and / or processing steps closely related to the solution according to the invention are shown in the accompanying drawings, while other details that are not closely related to the invention are omitted.

[0030] It should be emphasized that the term "including / comprises" as used herein refers to the presence of a feature, element, step, or component, but does not exclude the presence or addition of one or more other features, elements, steps, or components.

[0031] During large language model inference, after model weights reside in GPU high-bandwidth memory (HBM), the key-value cache (KV Cache) for a single inference request is inefficiently scheduled in a heterogeneous memory system consisting of remaining HBM space, CXL extended memory, and CPU DDR. Existing pure software solutions rely on software counters for key-value cache heat collection, which introduces 5% to 15% runtime overhead and the granularity can only reach the request or hierarchical level, failing to accurately perceive the access frequency of individual attention heads. In the prefetching stage, the software-scheduled DMA prefetching shares the instruction scheduling channel with GPU computation instructions, making hardware-level parallelism impossible, and the prefetching granularity is limited by the DMA descriptor size, making it difficult to match the logical structure of the key-value cache. In the computation execution stage, computationally intensive QKV projection and memory-intensive Attention computation are executed serially, resulting in alternating idle GPU computing power and memory bandwidth. At the same time, CXL extended memory, which has hardware-level capabilities such as non-blocking transactions and cache consistency, is only treated as a passive storage resource, and its active scheduling potential remains untapped. The aforementioned defects result in GPU computing units frequently being in a state of waiting for memory access to complete, leading to high inference latency, insufficient utilization of heterogeneous memory bandwidth, and difficulty in fully leveraging the performance advantages of multi-level heterogeneous memory architecture.

[0032] In view of this, the present invention provides a single-request heterogeneous memory optimization method for large language model inference, such as... Figure 1As shown, the method includes the following steps S101~S103: Step S101: Monitor memory access instructions based on the access pattern awareness unit pre-embedded in the server memory controller and CXL device controller, filter access to key-value cache by the range of memory access physical addresses, map memory access physical addresses to corresponding attention heads according to the attention head mapping table, count the number of times each attention head is accessed within a set time window, and expose it to the inference engine through the heat register interface.

[0033] Step S102: Divide the cache address of the key-value cache corresponding to each attention head into a high-speed hot data layer based on HBM, a medium-speed warm data layer based on CXL, and a low-speed cold data layer based on DDR; the inference engine decays the number of times each attention head is accessed according to a set time step interval, and calculates the heat score of each attention head in combination with the weight of the current data layer and sorts them in descending order; mark the key-value cache of the first set number of attention heads in the sorted order as prefetch candidates, add them to the prefetch command queue, and write them to the hardware prefetch engine pre-embedded in the CXL device controller through the MMIO interface.

[0034] Step S103: In response to the access pattern awareness unit's prefetch trigger signal or the inference engine's prefetch instruction, the hardware prefetch engine implements zero-software-overhead key-value cache prefetching based on the CXL.cache protocol's hardware primitives. The prefetch arbitrator prioritizes the key-value caches corresponding to each attention head in the prefetch command queue based on SLA urgency, heat score, and prefetch benefit estimation. Each clock cycle, the highest-priority cache is selected from the prefetch command queue and prefetched to HBM. After prefetching is complete, the inference engine is asynchronously notified via an MSI-X interrupt. In HBM, different cache areas are configured for ping-pong switching, simultaneously used for current computation and background prefetching. During the QKV projection stage of single-layer decoding, the key-value cache required for the next computation is prefetched. During the attention computation stage, the key-value cache prefetched in the previous step is retrieved without waiting.

[0035] Step S101 moves the hot data collection function of the key-value cache from the traditional software layer to the hardware logic of the memory controller and CXL device controller. It is executed in parallel with the main memory access path in a bypass statistical manner, completely eliminating the runtime overhead caused by the software counter. At the same time, the granularity of hot data collection is refined from the request / level to the individual attention head level, providing an accurate data foundation for subsequent fine-grained hierarchical scheduling.

[0036] Step S101 first uses the pre-embedded address range comparison logic built into the access mode perception unit to pre-configure the start and end physical addresses of the key-value cache memory region. It then performs address filtering on all memory access instructions, performing subsequent statistics only on key-value cache accesses falling within the target range. Accesses to non-KV data are allowed directly, without interfering with the critical path of normal memory access. Furthermore, a hardware-implemented attention head mapping table stores the mapping relationship between physical addresses and attention head identifiers in a tri-state content-addressable memory, supporting both continuous memory allocation and paging allocation modes. The mapping rules are pre-written by the inference engine via the MMIO interface, allowing the accessed physical address to be mapped to the corresponding level ID and attention head ID within a single cycle. Next, each attention head is configured with an independent hardware access counter, automatically accumulating the number of accesses to the corresponding head within a set sliding time window. A saturation counting design avoids overflow, and the entire process runs purely on hardware logic, requiring no additional instructions from the CPU / GPU. Finally, the counting results are exposed to the upper-layer inference engine via the heat register interface using MMIO, while simultaneously outputting a hardware-level prefetch trigger signal to support subsequent automatic prefetching.

[0037] In some embodiments, filtering access to the key-value cache by the range of physical addresses accessed in memory includes: obtaining the physical address in the memory access instruction and extracting the high-order bits, comparing it with the start and end addresses of the pre-configured key-value cache memory region; if the key-value cache memory region is hit, then performing subsequent access count statistics; if the key-value cache memory region is not hit, then allowing access directly without processing.

[0038] Specifically, an address range comparator hardware module is configured with an operating frequency of 1GHz, adapted to a 4KB memory page granularity and a 40-bit physical address space. When the memory controller receives a memory access instruction, the filtering process is completed serially through a three-stage pipeline without blocking the main memory access path: The first-stage pipeline extracts the physical address carried by the instruction and extracts the high-order page address through bitwise operations on the rising edge of the clock to remove interference from page offsets; the second-stage pipeline uses two parallel hardware comparators to simultaneously perform two comparison operations: "page address ≥ start address high-order bits" and "page address ≤ end address high-order bits," and outputs the two comparison results synchronously; the third-stage pipeline performs a logical AND operation on the two comparison results. If the result is true, the access is determined to have hit the KV Cache memory region, a valid hit signal is generated, and the normalized page address is sent to the subsequent attention head mapping module, initiating the access count statistics process; if the result is false, it is determined to be a non-KV Cache access, and a bypass allow signal is directly output without triggering any subsequent heat statistics operations, and the memory access request is executed normally along the original main path of the memory controller. In this embodiment, the complete address filtering process only requires 3 clock cycles, corresponding to a total latency of 3ns, which is within the margin of the inherent latency of memory access and will not increase the critical path time of memory access. At the same time, only the KV Cache access that hits the target region is processed, and there is no additional performance loss for non-KV memory access such as model weights and system computing stack, thus ensuring the general access efficiency of the memory controller.

[0039] In some embodiments, the attention head mapping table is implemented based on a tri-state content-addressable memory (TCAM). The inference engine writes the mapping relationship between cache addresses and attention heads into the attention head mapping table via the MMIO interface when allocating key-value caches. Allocation methods include contiguous allocation and paging allocation. For example, the default configuration depth of the tri-state content-addressable memory is 8192 entries. Each entry has a bit width including 1 valid bit, 40 bits of physical base address, 16 bits of step size field, 8 bits of level ID, and 8 bits of attention head ID, totaling 73 bits, which can fully cover the mainstream large model configuration of 128 layers × 64 heads. The TCAM is integrated into the access pattern awareness unit of the memory controller, with a query latency of one clock cycle, and does not increase the critical path latency of memory access. For different memory management modes, both contiguous allocation mode and paging allocation mode write rules are supported.

[0040] In the continuous allocation mode, when the KV Cache is arranged with continuous physical addresses by layer and by head, the inference engine completes the writing of all entries in one go during the model initialization phase. Each entry writes the KV Cache starting base address, continuous address step size, and corresponding layer ID and attention head ID of the corresponding attention head; during querying, the input physical address is matched with the base address and step size, and if a match is found, the corresponding head identifier is directly output.

[0041] In the pagination allocation mode, the inference engine writes mapping entries page by page during the process of allocating physical memory pages for each attention head. The entries only store the page base address, and the step size field is set to 0. During the query, the high-order page base address of the physical address is directly matched, and the corresponding header identifier is output to adapt to the dynamically growing KV Cache sequence length.

[0042] When allocating key-value caches, the inference engine synchronously records the data level corresponding to the corresponding cache address. The data levels include: high-speed hot data layer, medium-speed warm data layer, or low-speed cold data layer.

[0043] The method employs an independent hardware counter for each attention head, forming an access frequency counter array to count the number of accesses, achieving zero software overhead. Specifically, all hardware counters are implemented using pure digital logic and integrated within the access pattern awareness unit. The counters utilize a saturation counting design, maintaining their maximum value to avoid counting errors caused by overflow and topping. Counter triggering and updating are entirely automated by the hardware logic, requiring no additional instructions from the CPU or GPU: when a memory access hits the KV Cache region after address range filtering and obtains a valid attention head ID through TCAM mapping, the attention head ID is directly used as the address index of the counter array, automatically incrementing the corresponding counter on the falling edge of the clock. The entire counting process executes in parallel with the main memory access path, without occupying processor instruction cycles, thus achieving zero software runtime overhead.

[0044] Step S102 relies on the accurate single attention head heat data collected by hardware, utilizes the power-law distribution characteristics of key-value cache access, evaluates the memory access latency impact of each attention head through weighted quantized heat scores, prioritizes high-frequency access data to the high-speed memory level, and maximizes the benefits of limited high-speed memory space; at the same time, it dynamically eliminates historical hotspots through a heat decay mechanism, so that the hierarchical strategy can adapt to changes in attention access patterns during inference.

[0045] Specifically, step S102 constructs a three-level heterogeneous memory tiered architecture. The GPU HBM is set as a high-speed hot data layer to carry key-value caches with high-frequency access, the CXL extended memory is set as a medium-speed warm data layer to carry medium-frequency access data, and the CPU DDR main memory is set as a low-speed cold data layer to carry low-frequency access data. Next, the inference engine periodically reads the access count of each attention head collected by the hardware, and calculates the heat score of each attention head by weighting it with the weight coefficient of the memory layer where the data is currently located. The higher the score, the greater the impact of the access of that head on the overall inference latency. All attention heads are sorted in descending order of heat score, and the first set number of heads with the highest ranking are selected as promotion candidates. The remaining available capacity of the target memory layer HBM is combined to generate a matching key-value cache migration instruction, which is written to the prefetch command queue and sent to the hardware prefetch engine through the MMIO interface. The software is only responsible for decision-making and does not perform actual data transfer to avoid occupying GPU computing resources. Finally, all hardware access counters are decayed according to the set time step interval to gradually reduce the weight of historical access records, avoid early hot data occupying high-speed memory space for a long time, and ensure the dynamic adaptability of the tiered strategy.

[0046] In some embodiments, the inference engine decays the number of times each attention head is accessed at a set time step interval, and calculates the popularity score of each attention head by combining the current data layer weight. This includes: reading the number of times each attention head is accessed in the access pattern awareness unit, finding the data layer to which the key-value cache of each attention head belongs, and determining the current data layer weight; multiplying the number of times all attention heads are accessed by α at intervals of a third set number of time steps to achieve decay, where α < 1. For example, decay is performed on all counters every T steps, and the calculation formula is: count - k = count × α, α = 0.9. The number of times each attention head is accessed and its corresponding current data layer weight are multiplied to obtain the corresponding popularity score.

[0047] The hierarchical weights are set based on the access latency and access speed of different memory layers. These weights are configurable parameters, written to the configuration file when the inference engine starts. For example, the weight coefficient for GPU HBM corresponding to the high-speed hot data layer is 4, the weight coefficient for CXL extended memory corresponding to the medium-speed warm data layer is 2, and the weight coefficient for CPU DDR corresponding to the low-speed cold data layer is 1. The single attention head's heat score = the number of accesses in the current cycle of that head × the weight coefficient of its current data layer. A higher value indicates a greater total memory access latency cost incurred by the attention head's KV Cache, and a more significant latency optimization benefit from migrating it to higher-speed memory.

[0048] In some embodiments, after calculating the heat score of each attention head based on the weight of the current data layer and sorting them in descending order, the method further includes: dividing the sorted attention heads into a high-heat segment matching high-speed hot data layers, a medium-heat segment matching medium-speed warm data layers, and a low-heat segment matching low-speed cold data layers according to preset rules. Specifically, the top 20% of the sorted attention heads are marked as high-heat segments, the top 20% to 60% of the attention heads are marked as medium-heat segments, and the bottom 40% of the attention heads are marked as low-heat segments. A migration instruction queue is generated and written to the pre-embedded hardware prefetch engine in the CXL device controller via MMIO to perform data transfer.

[0049] Step S103 utilizes a hardware prefetch engine independent of the GPU computing unit, leveraging the non-blocking transaction characteristics of the CXL.cache protocol to achieve asynchronous data transfer across memory levels, enabling the prefetch operation to run completely in parallel with GPU computation. Simultaneously, by utilizing the temporal alternation of computationally intensive QKV projection and memory-intensive Attention computation in single-layer decoding of large language models, and through the HBM double-buffered ping-pong switching mechanism, the key-value cache prefetch process is completely hidden within the computation time window of QKV projection, eliminating memory access waits during the Attention stage, ultimately achieving complete hiding of prefetch latency and full utilization of hardware resources.

[0050] Specifically, step S103 employs a dual-trigger startup mechanism. Prefetching can be initiated directly by the hardware prefetch trigger signal output by the access mode awareness unit in step S101, achieving zero-latency response to hotspot access. Alternatively, it can be initiated by a prefetch instruction issued by the inference engine, supporting global hierarchical scheduling. Furthermore, a prefetch arbiter sorts tasks in the prefetch command queue according to three priorities: SLA urgency, popularity score, and estimated prefetch benefit. Priority is quickly compared using bit concatenation, and the highest-priority task is selected for execution each clock cycle, prioritizing prefetch tasks with high real-time performance, high access popularity, and high transmission benefits. The hardware prefetch engine performs non-blocking data transmission based on the native hardware primitives of the CXL.cache protocol. After issuing a prefetch request, it immediately processes the next queue task. Data transmission and writing are completed asynchronously in the background. After prefetching is complete, the inference engine is asynchronously notified via the MSI-X interrupt mechanism, without occupying GPU SM computing resources or instruction scheduling channels throughout the process. Furthermore, two equal-sized buffers are set up in the GPU HBM to achieve ping-pong switching. During the QKV projection calculation stage of single-layer decoding, the hardware prefetch engine synchronously prefetches the key-value cache required for the next calculation into the backup buffer. When entering the Attention calculation stage, the ready KV data is read directly from the currently active buffer without waiting for memory transfer. After the single-step calculation is completed, the roles of the two buffers are switched through hardware pointer swapping, without data copying, thus completely hiding the prefetch latency within the time window of the computationally intensive stage. At the same time, a built-in hardware miss fallback path is provided. If the prefetch is not completed on time, it will automatically switch to reading directly from CXL memory to ensure the correctness of the calculation.

[0051] In some embodiments, the prefetch command queue supports three modes: single prefetch, batch discontinuous data prefetch, and predictive prefetch.

[0052] Predictive prefetching is based on the access pattern awareness unit storing the access counts of each attention head within a fourth set number of consecutive sampling windows, and fitting the trend. Attention heads whose access counts are continuously higher than a first set threshold and whose trend is increasing or remains high are cached in the prefetch command queue. For example, every 32 decoding steps constitute a counting sampling window, and a decision is made every 4 sampling windows, with the first set threshold set to 80. The number of accesses of each attention head within 4 consecutive windows W is counted. Taking 4 attention heads currently residing in the CXL warm data layer as an example, the specific details are as follows: Attention point A: W1=85 times, W2=92 times, W3=105 times, W4=118 times. The number of visits to all four windows is higher than the threshold of 80 times, and the number increases window by window, showing a continuous upward trend, which meets the prediction prefetching conditions.

[0053] Attention head B: W1=95 times, W2=98 times, W3=92 times, W4=97 times. The number of visits to all four windows is higher than the threshold. The maximum fluctuation of adjacent windows is about 6%, which is within 10%, indicating stable high-level visits, meeting the prediction and prefetching conditions.

[0054] Attention head C: W1=70 times, W2=88 times, W3=95 times, W4=102 times. The number of visits to the first window is lower than the threshold, which does not meet the requirement of continuously exceeding the threshold, and therefore it is not included in the prediction prefetch.

[0055] Attention head D: W1=110 times, W2=95 times, W3=82 times, W4=70 times. The number of visits shows a continuous downward trend, and the last window has fallen below the threshold, indicating that the popularity is declining, and therefore it is not included in the prediction prefetching.

[0056] In some embodiments, the key-value caches corresponding to each attention head in the prefetch command queue are sorted according to three levels of priority: SLA urgency, popularity score, and prefetch benefit prediction, based on the prefetch arbitrator, including steps S1031~S1033: Step S1031: Assign an SLA urgency level according to the time margin before the request defaults, based on preset rules.

[0057] Specifically, when the inference engine issues a prefetch command to the prefetch command queue (PCQ) of the hardware prefetch engine (HPC), it first obtains the difference between the SLA (Service Level Agreement) deadline of the current inference request and the current system time. This difference is the time margin before the request defaults on the SLA. Subsequently, the time margin is quantified and assigned a value according to a preset mapping rule. The mapping rule is as follows: when the time margin is greater than 100ms, it is determined to be a safe state, and the SLA urgency is assigned a low priority value, such as 0~3; when the time margin is between 100ms and 20ms, it is determined to be a warning state, and the SLA urgency is assigned a medium priority value, such as 4~11; when the time margin is less than 20ms, it is determined to be a critical state, and the SLA urgency is assigned a high priority value, such as 12~15, where the larger the value, the higher the urgency.

[0058] Step S1032: Calculate the prefetching benefit based on the expected hit rate of the key-value cache corresponding to each attention point, the amount of data prefetched, and the transmission latency. The calculation formula is as follows: ; Where V represents the pre-collected revenue. Indicates the expected hit rate. Indicates the amount of data to be prefetched. This indicates transmission delay.

[0059] The expected hit rate is obtained by adjusting the historical hit rate. If the current popularity score of the header is high and it has already resided in HBM, the probability of it being accessed again is high, and the expected hit rate is adjusted upward, such as by multiplying by a coefficient of 1.1 to 1.3. If the popularity score of the header continues to decline and it resides in a non-HBM data layer, the expected hit rate is adjusted downward, such as by multiplying by a coefficient of 0.7 to 0.9. Through the above adjustments, the expected hit rate reflects the recent trend of the header's access pattern.

[0060] Step S1033: Calculate the overall priority by weighted summation of SLA urgency, popularity score and prefetching benefit, and sort the key-value caches corresponding to each attention head in the prefetching command queue based on the overall priority.

[0061] In some embodiments, the access pattern awareness unit generates a prefetch trigger signal based on the following criteria: if the number of times an attention head is accessed exceeds a second set threshold for a consecutive second set number of sampling windows, a prefetch trigger signal is generated. Further, each time the inference engine accesses the key-value cache, it checks whether the corresponding cache address has been prefetched into the HBM and calculates the overall hit rate. If the overall hit rate is less than a third set threshold, the second set threshold is increased by a first proportion to reduce the prefetch coverage, where the first proportion is greater than 1; if the overall hit rate is greater than a fourth set threshold, the second set threshold is decreased by a second proportion to increase the prefetch coverage, where the second proportion is less than 1. Furthermore, if the overall hit rate is greater than a fifth set threshold, a third cache is added to the HBM; if the overall hit rate is less than a sixth set threshold, the two caches in the HBM are reduced to a single cache, and ping-pong switching is stopped.

[0062] In one specific embodiment, the prefetch trigger signal (PTS) generation standard of the Access Pattern Aware Unit (KAPSU) is set as follows: when the number of times a certain attention point is accessed exceeds a set threshold θ = 100 times in three consecutive sampling windows, the PTS signal is set, directly triggering the hardware prefetch engine to start the prefetch operation without the intervention of the inference engine software; at the same time, the inference engine checks whether the corresponding address has been prefetched to HBM each time it accesses the KV Cache, and calculates the overall hit rate accordingly. A threshold evaluation is performed every 100 steps: when the overall hit rate is lower than 50%, it indicates that the prefetch is too aggressive, and the system increases the threshold θ by 1.2 times, i.e., θ←θ×1.2, to reduce unnecessary prefetch operations and save memory bandwidth; when the overall hit rate is higher than 85%, it indicates that the prefetch is too conservative, and the system decreases the threshold θ by 0.8 times, i.e., θ←θ×0.8, to increase the prefetch coverage and improve the HBM hit rate, while the value range of θ is constrained to [50, To prevent extreme adjustments within the 500] range, when the overall hit rate is above 80% and HBM space is ample, the system adds a third buffer in the HBM, expanding the double-buffered pipeline to a triple-buffered pipeline, allowing for longer prefetch lead to further improve parallelism. When the overall hit rate is below 50% or HBM space is tight, the system downgrades from triple-buffered or double-buffered to single-buffered mode, freeing up HBM space for other hot data. Switching between all buffering modes is completed through the MMIO configuration register, with a switching latency of less than 100ns, without interrupting inference services. Through the above threshold adaptation and pipeline depth dynamic adjustment mechanism, the system can maintain the overall hit rate within the optimized range of 70%~80% under different workloads and access modes, achieving a dynamic balance between prefetch accuracy and hardware resource utilization.

[0063] On the other hand, the present invention also provides a single-request heterogeneous memory optimization system for large language model inference. The system combines GPU HBM, GPU near-end memory extended via the CXL protocol, and host-side CPU DDR memory. An access pattern awareness unit is pre-embedded in the server memory controller and the CXL device controller. The access pattern awareness unit includes an address range comparator, a tri-state content-addressable memory for implementing an attention head mapping table, an access frequency counter array, a heat register interface, and a prefetch trigger signal generator. The system is used to execute the computer program or instructions. When the computer program / instructions are executed, the device / system implements the steps of the above method.

[0064] On the other hand, the present invention also provides a computer-readable storage medium having a computer program or instructions stored thereon, which, when executed by a processor, implement the steps of the above-described method.

[0065] On the other hand, the present invention also provides a computer program product, including a computer program or instructions that, when executed by a processor, implement the steps of the above-described method.

[0066] The present invention will now be described with reference to a specific embodiment: This embodiment provides a single-request heterogeneous memory optimization method and system for large language model inference. It constructs a five-layer linkage closed-loop optimization system including hardware awareness, heat decision, hardware prefetching, pipeline execution and effect feedback. For the running efficiency of a single inference request in heterogeneous memory, it comprehensively eliminates software overhead from four links: heat collection, prefetch triggering, data transfer and computation interleaving, and realizes hardware-level end-to-end optimization.

[0067] This embodiment includes the following five steps, refer to... Figure 2 This is applied to server systems that include GPU High Bandwidth Memory (HBM), CXL extended memory devices, and CPU main memory. For example... Figure 3 As shown, the server system consists of three parts: GPU, CXL 2.0 / 3.0 device, and CPU. The inference data is divided into three levels of storage according to access frequency: 96GB of HBM on the GPU side stores L1 layer hot data, which is calculated by the GPU using QKV / Attention. 1TB of CXL extended memory on the CXL device side stores L2 layer warm data; and 512GB of DDR5 main memory on the CPU side stores L3 layer cold data. To achieve intelligent tiering and efficient prefetching, the system deploys access pattern awareness units (KAPSU-G, KAPSU-C, and KAPSU-D) at three tiers: KAPSU-G for GPU HBM heat detection, KAPSU-C for CXL memory heat detection, and KAPSU-D for DDR heat detection. KAPSU-G and KAPSU-C transmit heat information to the HPC hardware prefetch engine via PTS signals, allowing the HPC to prefetch warm / hot data into HBM. The CXL.cache protocol engine handles CrdUp, WrInv, and non-blocking transactions, ensuring cache consistency. The CPU-side inference engine is globally responsible for heat aggregation, tiering decisions, and QoS quota management, and controls the CXL.cache protocol engine through MMIO configuration. The GPU-side memory controller handles address decoding and access scheduling. This hardware-level heat detection and prefetching mechanism enables heterogeneous memory collaborative optimization between HBM, CXL extended memory, and DDR5 main memory in single-request scenarios.

[0068] Step 1: Embed a KV Cache Access Pattern Aware Unit (KAPSU) in the memory controller. A dedicated KV Cache access pattern aware unit (KAPSU) is embedded in both the server's memory controller (MC) and CXL device controller. The KAPSU's hardware structure comprises the following five sub-modules: (1) Address Range Comparator (ARC): Configures the start and end addresses [base, limit] of the KV Cache memory region, and only counts memory accesses falling within this range to avoid interfering with normal access to non-KV Cache data. When the GPU or CPU executes a memory access instruction, a physical address phys_addr is generated. ARC needs to determine whether the physical address hits the configured KV Cache memory region. If it does, it continues to record the access heat; if not, it allows the access directly without any additional processing. A three-stage pipelined comparison logic is adopted: the first stage extracts the high-order bits of the physical address (addr[63:12]), the second stage compares base_addr ≤ phys_addr and phys_addr ≤ limit_addr in parallel, and the third stage performs a logical AND operation on the results of the two comparisons in the second stage. If both results are true, the physical address is hit and used to generate a normalized address that can be used for subsequent lookups. At the same time, a bitwise AND operation is performed with addr_mask. Total latency: 3 clock cycles, 3ns for an example of 1GHz, which does not affect the critical path of memory access.

[0069] (2) Attention Head Map (HMT): Maps physical addresses to (layer_id, head_id) tuples. HMT is a mapping table implemented in TCAM hardware. The software informs in advance which physical address range belongs to which attention head through MMIO. When memory access occurs, HMT can retrieve (layer_id, head_id) in a single cycle for subsequent heat counting.

[0070] When allocating KV Cache memory, the inference engine writes the mapping relationship to the HMT via the MMIO (Memory-Mapped I / O) interface. For continuous allocation scenarios, a linear mapping of base_address + stride is used; for paginated allocation scenarios, such as vLLM PagedAttention, the inference engine dynamically writes the mapping entry when allocating each page. The HMT query method is divided into two types depending on the allocation mode: 1) If the entry is in continuous allocation mode: check if phys_addr is within the range of [addr_base, base_address + stride - 1]; 2) If it is in paginated allocation mode: directly compare whether the high-order bits of phys_addr after removing the page offset are equal to the base_address of the entry. HMT is implemented using TCAM (Ternary Content-Addressable Memory), with a default depth of 8192 entries. It supports 128 layers × 64 heads = 8192 attention heads. The width of each entry is: valid(1) + addr_base(40) + stride(16) + layer_id(8) + head_id(8) = 73 bits. Query method: Input phys_addr, perform parallel matching via TCAM, and output (layer_id, head_id), completed in a single cycle.

[0071] (3) Access Frequency Counter Array (AFC): A hardware counter is allocated to each attention head, with a default width of 16 bits. It automatically increments by 1 each time the corresponding head's KV Cache is accessed. It employs a saturation design, maintaining its maximum value even when overflowing, resulting in zero software overhead. Timing: On the rising edge of the clock, the HMT outputs head_id. The high level of the clock uses head_id as the index address of the AFC. On the falling edge of the clock, the corresponding counter executes count = min(count + 1, MAX_VAL). Throughput: KV Cache access count is processed once per clock cycle. Power Consumption Estimation: 16K 16-bit counters, updating one at a time, with dynamic power consumption less than 1mW.

[0072] (4) Heat Register Interface (HRI): The AFC count value is exposed to the inference engine through MMIO. The inference engine reads the heat data once every N steps for hierarchical decision-making.

[0073] (5) Prefetch Trigger Signal (PTS): When the access frequency of a specific attention head exceeds the threshold θ for W consecutive sampling windows, the PTS signal is set, directly triggering the hardware prefetch engine in the CXL device controller without the need for inference engine software intervention. The PTS is implemented using a finite state machine, with states including idle, counting, and trigger. Default parameters: W=3, T=1024, θ=100. Trigger delay: less than 2 clock cycles from the time the count exceeds the limit to the time the PTS is set.

[0074] Workflow: When the inference engine starts, ARC and HMT are configured via MMIO; thereafter, all memory accesses to the KV Cache are automatically intercepted and counted by KAPSU without consuming any GPU / CPU computing resources; the inference engine reads heat data via HRI every N steps; when AFC detects that the access frequency of a certain head exceeds the threshold, PTS directly triggers hardware prefetching.

[0075] Step 2: The inference engine performs hot aggregation and dynamic tiering of the KV cache. After reading KAPSU's popularity data, the inference engine performs the following hierarchical decision: (1) Read the access count value maintained by the access frequency counter array AFC in KAPSU for each attention head, i.e., the per-head count value, and store it in the AFC hardware counter.

[0076] (2) Calculate the heat score for each attention head: The score(head) is calculated as Σ(access_count[head][tier] × tier_weight[tier]), where access_count[head][tier] refers to the number of times the head is accessed at a certain tier. The inference engine actively records which tier each head is currently in, or KAPSU counts them separately by tier. tier_weight[tier] is the weight coefficient for each tier, where HBM has a weight of 4, CXL has a weight of 2, and DDR has a weight of 1. A higher weight indicates a faster tier. If a head is frequently accessed in HBM (meaning the data is already in HBM and used frequently), it will have a high score, indicating it should remain in HBM. If it is accessed in DDR, a high score indicates it is "hot" and should be placed in HBM. Ultimately, a higher score indicates a greater total access cost contributed by the head, and therefore it should be promoted to a faster tier.

[0077] (3) Sort by popularity score in descending order and mark the KV Cache of the top K attention heads as candidates for promotion.

[0078] (4) Check the remaining capacity of the target memory level and generate a prefetch command queue.

[0079] (5) The migration instructions are written to the prefetch command queue (PCQ) of the Hardware Prefetch Controller (HPC) via MMIO, and the actual data transfer is performed by the hardware. There may be multiple prefetch commands to be processed in the PCQ, and the arbitrator decides which one to execute first.

[0080] (6) Every T steps, perform decay on all counters: count = count × α, α = 0.9, to prevent historical hot data from occupying high-level memory for a long time. If decay is not performed, historically hot heads will remain high indefinitely, occupying HBM even if they are no longer accessed. After decay, old access records are gradually forgotten, and new access patterns can be reflected in the popularity score more quickly. The decayed counter value is still an integer, ensuring that updating the hardware counter is simple.

[0081] Hierarchical Mapping: The inference engine dynamically determines the memory level at which the KV Cache of each attention head should reside based on its heat score. HBM space is limited, so heads cannot be moved to HBM indefinitely. The inference engine checks how much free space is left in HBM, determines how many promotion candidate heads can actually be migrated based on the remaining capacity, and generates a series of migration instructions: reading the KV Cache of certain heads from the current level and writing them to the target level. The software does not move data itself; instead, it writes the migration commands to the HPC's PCQ, which is run independently of the GPU by the HPC hardware prefetch engine. It uses DMA or CXL protocols to perform the actual data transfer, allowing the software to continue with other tasks after issuing the commands.

[0082] In other embodiments, the storage locations of all attention heads corresponding to the KV Cache are migrated as a whole according to the heat score, as shown in Table 1 below.

[0083] Table 1 Step 3: The Hardware Prefetch Engine (HPC) executes the CXL.cache protocol for prefetching. A hardware prefetching engine (HPC) is embedded in the CXL device controller, utilizing the hardware primitives of the CXL.cache protocol to achieve zero-software-overhead KV cache prefetching. The HPC hardware architecture consists of the following four sub-modules: (1) Prefetch Command Queue (PCQ): Configurable depth (default 64 entries), supports three prefetch operations: Single Prefetch PREFETCH_KV(addr, dst_tier): Prefetches KV Cache data from the current tier to the target tier from the specified address.

[0084] Batch prefetch of discontinuous data (BATCH_PREFETCH(addr_list, dst_tier): Batch prefetch, transfers multiple discontinuous KV cache data blocks in a single command.

[0085] Predictive Prefetch (head_id, step_offset): Based on KAPSU's heat data, predict and prefetch the KV data block that is most likely to be accessed next.

[0086] (2) Heat-address conversion logic: Convert the head_id output by KAPSU into the corresponding physical address list.

[0087] (3) Prefetch Arbiter: Prioritized according to the following criteria: ① SLA urgency; ② heat score; ③ Prefetch benefit estimate, benefit = expected_hit_rate × data_size / transfer_latency. Overall priority = urgency × 1,000,000 + heat_score × 1,000 + benefit. The highest priority command is selected from the PCQ for execution each clock cycle.

[0088] (4) Completion notification module: The inference engine is notified asynchronously via the MSI-X interrupt mechanism to complete the prefetching.

[0089] HPC utilizes the CXL.cache protocol for non-blocking transactional prefetching: ① HPC sends a CrdUp request to the CXL Host; ② The CXL Host returns Ack + Data; ③ HPC writes the data to the target memory level; ④ HPC updates its internal state table; ⑤ HPC notifies the inference engine via an MSI-X interrupt. After step ① is sent, HPC immediately processes the next PCQ command, while steps ② and ⑤ are completed asynchronously in the background. The entire process does not consume any computing resources from the GPU SM. Prefetching operations and GPU computation can be completely parallelized without waiting for data returns.

[0090] Prefetch triggering can be achieved in two ways: hardware automatic triggering and software instruction triggering. Hardware automatic triggering is achieved by directly triggering HPC with the PTS signal of KAPSU, without the need for software intervention; software instruction triggering is achieved by the inference engine writing prefetch commands to PCQ via MMIO.

[0091] Step 4: Execution of the compute-memory access double-buffered pipeline A pipeline integrating double buffering and hardware prefetching is implemented in the inference engine. Leveraging the alternating time characteristics of computationally intensive QKV projection and memory-intensive Attention computation in single-layer decoding of LLM inference, hardware-level parallelism between heterogeneous memory prefetching and GPU computation is achieved. Each layer of decoding in LLM inference includes two main operations: QKV projection and Attention computation. In traditional solutions, Attention computation requires reading the KV cache from CXL or DDR, which introduces a latency of hundreds of nanoseconds to microseconds, directly adding to the total latency. The double buffering approach utilizes the QKV projection period to move the KV cache needed for the next step from slow memory to a backup buffer in HBM. This ensures that by the time Attention computation occurs, the data is already waiting in HBM, virtually eliminating the read latency. The two buffers alternately act as the currently used buffer and the background prefetch buffer. The specific process is as follows: (1) Buffer allocation: In HBM, allocate two K data buffers, Buffer A and Buffer B, of equal size. The capacity of each buffer is the amount of KV data required for a single layer and a single step.

[0092] (2) Production line timing: Phase 1: When the GPU SM performs QKV projection, the HPC simultaneously prefetches the KV data of step N+1 from the CXL extended memory into Buffer B; the Attention of step N has not yet started, and Buffer A stores the KV data of step N.

[0093] Phase 2: When the GPU SM performs Attention calculations, it uses the KV from step N, that is, it reads directly from Buffer A without waiting.

[0094] Phase 3: After the current step is completed, the roles of Buffer A and Buffer B are switched by pointer swapping based on the ping-pong pattern. No data copying is required; data from Buffer B doesn't need to be copied to Buffer A. Only the pointers of the two buffers are swapped. This pointer swapping is a hardware-level operation with a latency of less than 1ns, almost zero. Now Buffer B becomes the buffer for the current step, and Buffer A is cleared for the next round of prefetching. In this way, the roles are immediately reversed after each step, and the pipeline continues. The data needed for Attention calculation is already in the HBM's buffer and does not need to wait to be read from CXL. Therefore, the prefetch latency is completely hidden in the QKV projection calculation time. This allows for longer prefetch lead times and improves parallelism.

[0095] (3) Miss rollback: If a prefetch misses, the system automatically rolls back to reading directly from the CXL extended memory, increasing the CXL access latency by approximately 150ns, but this does not affect the accuracy of the calculation. The rollback path is implemented through a hardware multiplexer (MUX) without the need for software intervention.

[0096] (4) Orthogonality with FlashAttention: This pipeline mechanism is orthogonal to FlashAttention and can be superimposed. FlashAttention optimizes the data reuse of the GPU's on-chip SRAM to reduce the number of HBM accesses; this scheme optimizes the data pipeline between heterogeneous memory layers to maximize the HBM hit rate. Both can be enabled simultaneously, and their effects are superimposed.

[0097] For example, such as Figure 4 As shown, the system allocates two equally sized key-value (KV) data buffers, Buffer A and Buffer B, in HBM. During the time interval t0 to t1, the GPU SM performs the Nth step of QKV projection and Attention calculation, directly reading the Nth step KV data from the ready Buffer A. Simultaneously, the HPC hardware prefetch engine prefetches the N+1th step KV data from the CXL extended memory to Buffer B via the CXL.cache protocol. Buffer B is in a prefetching state, and the CXL device performs the corresponding data transfer. At time t1, Buffer A and Buffer B switch roles through hardware pointer exchange without data copying. Buffer B immediately carries the ready KV data of the N+1th step, while Buffer A is freed up for the next round of prefetching. During the time interval t1 to t3, the GPU SM performs the N+1th step of QKV projection and Attention calculation, reading KV data from Buffer B without waiting. The HPC synchronously prefetches the N+2th step KV data to Buffer A. After the CXL device completes the transfer between t2 and t3, Buffer A enters a ready state. At time t3, the pointer switch is performed again, and Buffer... Buffer A becomes the currently active buffer carrying the data for step N+2. GPUSM performs the calculation for step N+2 and reads data from Buffer A during the time period t3 to t5. HPC simultaneously prefetches the KV data for step N+3 into Buffer B. The CXL device completes the transfer between t4 and t5. After that, the QKV projection for step N+3 continues to advance on the switched buffer. The entire pipeline repeats in this way, so that the CXL memory access latency is completely hidden within the calculation time window of QKV projection. Prefetching and transfer and GPU computation achieve hardware-level parallelism.

[0098] Step 5: Prefetching effect feedback and adaptive parameter tuning To achieve adaptive optimization of system parameters, a prefetching effect feedback mechanism is introduced to form a closed loop. Step 5 automatically adjusts the heat threshold, pipeline depth, and layering strategy by real-time statistical analysis of the prefetching hit rate, enabling the system to adapt to different access patterns and maintain a high hit rate and high efficiency over the long term without manual intervention. In this way, the entire five-layer architecture achieves closed-loop optimization of hardware acceleration and software adaptation. Specifically, this includes: (1) Hit Rate Statistics: The inference engine maintains a per-head prefetch hit rate statistics table, which is obtained by checking whether each KVCache access hits the data prefetched by HPC. Each time the inference engine accesses KV Cache data, it checks whether the address has been prefetched to the target level: if hit, hit_count is incremented by 1; if miss, miss_count is incremented by 1. Hit rate = hit_count / (hit_count + miss_count).

[0099] (2) KAPSU threshold adaptation: Threshold evaluation is performed every S steps, with S=100 by default; the current overall hit rate is calculated as: overall_hit_rate=Σ(hit_count[head]) / Σ(hit_count[head]+ miss_count[head]); if overall_hit_rate<0.5: prefetching is too aggressive, θ←θ×1.2 to reduce unnecessary prefetching; if overall_hit_rate>0.85: prefetching is too conservative, θ←θ × 0.8 to increase prefetching coverage; the value range of θ is constrained to [50, 500] to prevent extreme adjustments.

[0100] (3) Adaptive pipeline depth: When the prefetch hit rate is greater than 80% and there is sufficient HBM space, it can be expanded into a triple-buffered pipeline, allowing for longer prefetch lead and improving parallelism.

[0101] When the prefetch hit rate is less than 50% or HBM space is tight, it degrades to single-buffered mode to release HBM space; the buffered mode switching is implemented through the MMIO configuration register, and the switching latency is less than 100ns.

[0102] (4) Tiered Thresholds: Different trigger thresholds can be set for different memory tiers. The prefetch threshold θ_hbm_cxl for HBM→CX can be lower than the prefetch threshold θ_cxl_ddr for CXL→DDR, because HBM space is precious, and more aggressive prefetching of hot data to HBM is more beneficial; while the space from DDR→CXL is relatively abundant, the threshold θ_cxl_ddr can be lower to allow for more aggressive prefetching. This allows for independent configuration of thresholds for each tier, enabling fine-grained management.

[0103] (5) Feedback path: All adaptive adjustment instructions are written to the configuration registers of KAPSU and HPC via MMIO, and the hardware automatically updates the parameters without restarting the inference engine.

[0104] Corresponding to the above method, the present invention also provides an apparatus / system including a computer device, the computer device including a processor and a memory, the memory storing computer instructions, the processor executing the computer instructions stored in the memory, and when the computer instructions are executed by the processor, the apparatus / system performs the steps of the method as described above.

[0105] This invention also provides a computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements the steps of the aforementioned edge computing server deployment method. The computer-readable storage medium can be a tangible storage medium, such as random access memory (RAM), main memory, read-only memory (ROM), electrically programmable ROM, electrically erasable programmable ROM, registers, floppy disks, hard disks, removable storage disks, CD-ROMs, or any other form of storage medium known in the art.

[0106] In summary, the single-request heterogeneous memory optimization method and system for large language model inference described in this invention addresses the heterogeneous memory scheduling scenario for single-request inference of large language models. Through a hardware-software collaborative optimization architecture that integrates hardware awareness, software decision-making, hardware execution, and pipeline interweaving, it achieves systematic performance improvement from four dimensions: memory access statistics, data scheduling, prefetching execution, and resource reuse.

[0107] First, by bringing the KV Cache access frequency collection down to the hardware units of the memory controller and CXL device controller, zero-software-overhead access frequency statistics with attention head granularity are achieved, completely eliminating the runtime performance loss caused by traditional software-layer counters and providing accurate data support for fine-grained memory scheduling.

[0108] Secondly, by combining a tiered storage scheduling strategy with a hotness decay mechanism, it can accurately match the access idempotency characteristics of the KV Cache, directing high-frequency hot data to reside in the GPU HBM, effectively reducing the high-speed memory access miss rate, and maximizing the utilization efficiency of limited high-speed memory space.

[0109] Third, relying on the independently running hardware prefetch engine and the native hardware primitives of the CXL.cache protocol, combined with a three-level arbitration mechanism of SLA urgency, heat score, and prefetch benefits, it achieves complete hardware-level parallelism of cross-memory level data transfer and GPU computing, upgrading CXL extended memory from a passive storage carrier to an active acceleration resource.

[0110] Fourth, by leveraging the HBM double-buffered ping-pong switching pipeline design, the KV cache prefetch latency is completely hidden within the computation time window of the QKV projection, efficiently reusing idle memory bandwidth during computationally intensive phases and significantly improving HBM bandwidth utilization. Compared to traditional pure software scheduling schemes, this invention can significantly reduce single-request inference latency while adapting to memory access characteristics of different model structures and sequence lengths, ensuring the system maintains optimal scheduling efficiency over the long term.

[0111] Those skilled in the art will understand that the exemplary components, systems, and methods described in conjunction with the embodiments disclosed herein can be implemented in hardware, software, or a combination of both. Whether implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this invention. When implemented in hardware, it can be, for example, electronic circuits, application-specific integrated circuits (ASICs), appropriate firmware, plug-ins, function cards, etc. When implemented in software, the elements of this invention are programs or code segments used to perform the desired tasks. The programs or code segments can be stored in a machine-readable medium or transmitted over a transmission medium or communication link via data signals carried in a carrier wave.

[0112] It should be clarified that the present invention is not limited to the specific configurations and processes described above and shown in the figures. For the sake of brevity, detailed descriptions of known methods are omitted here. In the above embodiments, several specific steps are described and shown as examples. However, the method process of the present invention is not limited to the specific steps described and shown. Those skilled in the art can make various changes, modifications, and additions, or change the order of steps, after understanding the spirit of the present invention.

[0113] In this invention, features described and / or illustrated for one embodiment may be used in the same or similar manner in one or more other embodiments, and / or combined with or in place of features of other embodiments.

[0114] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. For those skilled in the art, various modifications and variations of the embodiments of the present invention are possible. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A single-request heterogeneous memory optimization method for large language model inference, characterized in that, The method includes the following steps: The system monitors memory access instructions based on the access pattern awareness unit pre-embedded in the server memory controller and CXL device controller. It filters access to the key-value cache by the range of physical addresses accessed in memory, and maps the physical addresses accessed in memory to the corresponding attention heads according to the attention head mapping table. It counts the number of times each attention head is accessed within a set time window and exposes it to the inference engine through the heat register interface. The cache addresses of the key-value caches corresponding to each attention head are divided into a high-speed hot data layer based on HBM, a medium-speed warm data layer based on CXL, and a low-speed cold data layer based on DDR. The inference engine decays the number of times each attention head is accessed according to a set time step interval, and calculates the heat score of each attention head in combination with the weight of the current data layer, and sorts them in descending order. The key-value caches of the first set number of attention heads in the sorted order are marked as prefetch candidates, added to the prefetch command queue, and written to the hardware prefetch engine pre-embedded in the CXL device controller through the MMIO interface. In response to the prefetch trigger signal from the access pattern awareness unit or the prefetch instruction from the inference engine, the hardware prefetch engine implements zero-software-overhead key-value cache prefetching based on hardware primitives of the CXL.cache protocol. A prefetch arbitrator prioritizes the key-value caches corresponding to each attention head in the prefetch command queue based on SLA urgency, heat score, and prefetch benefit estimation. Each clock cycle, the highest-priority cache is selected from the prefetch command queue and prefetched to HBM. Upon completion of prefetching, the inference engine is asynchronously notified via an MSI-X interrupt. In HBM, different cache areas are configured for ping-pong switching, simultaneously used for current computation and background prefetching. During the QKV projection stage of single-layer decoding, the key-value cache required for the next computation is prefetched. During the attention computation stage, the prefetched key-value cache is retrieved without waiting.

2. The single-request heterogeneous memory optimization method for large language model inference according to claim 1, characterized in that, Access to the key-value cache is filtered by the range of physical addresses accessed in memory, including: obtaining the physical address in the memory access instruction and extracting the high-order bits, comparing it with the start and end addresses of the pre-configured key-value cache memory region; if the key-value cache memory region is hit, then the subsequent access count is performed; if the key-value cache memory region is not hit, then the access is allowed directly without further processing.

3. The single-request heterogeneous memory optimization method for large language model inference according to claim 1, characterized in that, The attention head mapping table is implemented based on a tri-state content-addressable memory. When allocating the key-value cache, the inference engine writes the mapping relationship between the cache address and the attention head into the attention head mapping table through the MMIO interface. The allocation methods include contiguous allocation and paging allocation. When allocating the key-value cache, the inference engine synchronously records the data level corresponding to the corresponding cache address. The data level includes: the high-speed hot data layer, the medium-speed warm data layer, or the low-speed cold data layer. The method assigns an independent hardware counter to each attention head to form an access frequency counter array, which is used to count the number of times the attention head is accessed, achieving zero software overhead.

4. The single-request heterogeneous memory optimization method for large language model inference according to claim 3, characterized in that, The inference engine decays the number of times each attention head is accessed according to a set time step interval, and calculates the popularity score of each attention head in combination with the weight of the current data layer. This includes: reading the number of times each attention head is accessed in the access pattern perception unit, finding the data layer to which the key-value cache of each attention head belongs, and determining the weight of the current data layer; multiplying the number of times all attention heads are accessed by α at intervals of a third set number of time steps to achieve decay, where α < 1; and multiplying the number of times each attention head is accessed and its corresponding weight of the current data layer to obtain the corresponding popularity score.

5. The single-request heterogeneous memory optimization method for large language model inference according to claim 1, characterized in that, After calculating the heat score of each attention head based on the weight of the current data layer and sorting them in descending order, the method further includes: According to preset rules, the sorted attention heads are divided into high-temperature segments matching the high-speed hot data layer, medium-temperature segments matching the medium-speed warm data layer, and low-temperature segments matching the low-speed cold data layer. A migration instruction queue is generated and written to the pre-embedded hardware prefetch engine in the CXL device controller through the MMIO interface to perform data transfer. Specifically, the sorted attention heads are divided into a high-temperature segment matching the high-speed hot data layer, a medium-temperature segment matching the medium-speed warm data layer, and a low-temperature segment matching the low-speed cold data layer according to preset rules. This includes marking the first 20% of the sorted attention heads as the high-temperature segment, marking the first 20% to 60% of the attention heads as the medium-temperature segment, and marking the last 40% of the attention heads as the low-temperature segment.

6. The single-request heterogeneous memory optimization method for large language model inference according to claim 1, characterized in that, The prefetch command queue supports three modes: single prefetch, batch discontinuous data prefetch, and predictive prefetch. The predictive prefetch is based on the access pattern sensing unit storing the number of times each attention head is accessed in a fourth set number of consecutive sampling windows and fitting the trend of change. The key value cache corresponding to the attention head whose access count is continuously higher than a first set threshold and whose trend of change is rising or remains high is included in the prefetch command queue. Based on the prefetch arbitrator, the key-value caches corresponding to each attention head in the prefetch command queue are sorted according to a three-level priority system: SLA urgency, popularity score, and prefetch benefit prediction, including: Based on the time margin before the SLA defaults, the SLA urgency level is assigned according to preset rules. The prefetching benefit is calculated based on the expected hit rate of the key-value cache corresponding to each attention point, the amount of data prefetched, and the transmission latency. The calculation formula is as follows: ； Where V represents the pre-fetched revenue. This indicates the expected hit rate. This indicates the amount of data to be prefetched. This indicates the transmission delay; The SLA urgency, the heat score, and the prefetch benefit are weighted and summed to calculate a comprehensive priority, and the key-value caches corresponding to each attention point in the prefetch command queue are sorted based on the comprehensive priority.

7. The single-request heterogeneous memory optimization method for large language model inference according to claim 1, characterized in that, The standard for the access pattern sensing unit to generate a prefetch trigger signal is: if the number of times the attention head is accessed exceeds a second set threshold in a second set number of consecutive sampling windows, a prefetch trigger signal is generated. Each time the inference engine accesses the key-value cache, it checks whether the corresponding cache address has been prefetched to HBM and calculates the overall hit rate. If the overall hit rate is less than a third set threshold, the second set threshold is increased by a first ratio to reduce the prefetch coverage, where the first ratio is greater than 1; if the overall hit rate is greater than a fourth set threshold, the second set threshold is decreased by a second ratio to increase the prefetch coverage, where the second ratio is less than 1. Furthermore, if the overall hit rate is greater than the fifth set threshold, a third cache area is added in HBM; if the overall hit rate is less than the sixth set threshold, the two cache areas in HBM are reduced to a single cache area and ping-pong switching is stopped.

8. A single-request heterogeneous memory optimization system for large language model inference, characterized in that, The system combines GPU HBM, GPU near-end memory extended via the CXL protocol, and host-side CPU DDR memory. An access pattern awareness unit is pre-embedded in the server memory controller and the CXL device controller. The access pattern awareness unit includes an address range comparator, a tri-state content-addressable memory for implementing an attention head mapping table, an access frequency counter array, a hot register interface, and a prefetch trigger signal generator. The system is used to execute the computer program or instructions. When the computer program / instructions are executed, the device / system implements the steps of the method as described in any one of claims 1 to 7.

9. A computer-readable storage medium having a computer program or instructions stored thereon, characterized in that, When the computer program or instructions are executed by a processor, they implement the steps of the method as described in any one of claims 1 to 7.

10. A computer program product, comprising a computer program or instructions, characterized in that, When the computer program or instructions are executed by a processor, they implement the steps of the method according to any one of claims 1 to 7.

Images