Retrieval enhancement generation reasoning acceleration system and method based on in-memory calculation

Abstract

The invention discloses a retrieval enhancement generation reasoning acceleration system and method based on in-memory calculation. According to the system, a decentralized document attention decomposition mechanism is adopted, attention calculation of document key value pairs is distributed to a plurality of in-memory processing devices to be independently executed, and the data transmission amount is reduced through local aggregation. The system architecture comprises a plurality of in-memory computing PIM devices, and each device integrates a PIM unit and a controller side unit to jointly complete a computing task. The scheduler adopts a hybrid scheduling strategy combining static mapping and dynamic triggering, supports multiple parallel modes, and improves the resource utilization rate through an interlaced execution mechanism. According to the method, data transmission overhead and calculation delay in RAG reasoning are remarkably reduced, and the method is suitable for a language model reasoning scene with large-scale knowledge enhancement.

Description

Technical Field

[0001] This invention relates to the fields of artificial intelligence and computer architecture technology, and in particular to a retrieval-augmented generation (RAG) inference acceleration system and method based on in-memory computing. Background Technology

[0002] Large-scale language models have achieved great success in many fields, but their inherent closed-world nature makes them prone to illusions and knowledge obsolescence. Retrieval-enhanced generation (RAG) effectively improves the factual accuracy and timeliness of generated content by combining LLMs with external knowledge bases. However, RAG systems require concatenating long retrieved documents with user queries as input during inference, leading to a sharp increase in the input sequence and resulting in huge computational and communication overhead.

[0003] Existing in-memory computing accelerators are primarily designed for standard LLM inference. Their design, which concentrates attention computation in dedicated computing units, fails to eliminate document key-value (KV) transport. Furthermore, their heterogeneous architecture, which separates GEMM from attention, cannot efficiently handle the GEMM operations common in RAG (Reference Aggregate Inference) where multiple input vectors are multiplied by the same weight matrix. Therefore, current technologies still lack effective solutions that fundamentally address the bottlenecks in data movement and computational efficiency in RAG inference. Summary of the Invention

[0004] In view of the shortcomings of the prior art, the present invention provides a retrieval enhancement generation inference acceleration system and method based on in-memory computing, which aims to solve at least one or more technical problems existing in the prior art.

[0005] To achieve the above objectives, the first aspect of the present invention provides a retrieval enhancement generation inference algorithm based on in-memory computation. Speed ​​system, including: A computing cluster consisting of multiple homogeneous in-memory processing devices interconnected via the CXL protocol, the computing cluster being logically divided into a document attention cluster and a context execution cluster; in, The in-memory processing devices in the document attention cluster are used for distributed storage of pre-computed document key-value pairs and employ a resource-optimized in-memory computing microarchitecture to perform parallel document attention computations; the resource-optimized in-memory computing microarchitecture includes: The buffer copying unit for GEMM operations configures multiple input buffers for the same arithmetic logic unit, enabling multiple input vectors to share the loaded weight matrix. For nonlinear functions, a lookup table is used to segment the linear approximation unit, which completes the calculation of complex functions through one interval lookup and one multiplication-addition operation; The decentralized document attention decomposition module is configured as follows: The global attention computation is decoupled into document attention computation performed in the document attention cluster and query-response attention computation performed in the context execution cluster; Control each in-memory processing device to independently compute local attention results, including local output vectors, local maxima, and local summation terms, based on local document key-value pairs; By using a global normalization step based on the local maximum and local summation terms, the local attention results of each device are fused with the query-response attention results; The load balancing hybrid scheduler is configured as follows: During system initialization, computational tasks are statically mapped according to the model structure, and during runtime, dynamic task scheduling and migration are performed based on the load status of the two clusters.

[0006] Specifically, the in-memory processing device includes a PIM unit and a controller-side unit: The in-memory computing unit, located next to the DRAM memory, includes several multipliers and adders, comparators, and a double-buffered input buffer; it is used to perform basic matrix-vector multiplication and nonlinear function approximation calculations based on lookup tables. The controller-side unit is located in the memory controller. Each memory controller integrates a near-memory unit, which includes an addition tree for aggregating the outputs of all memory computing units within the channel, and a dedicated high-precision Softmax unit for performing precise Softmax calculations that require numerical stability in the attention mechanism; it is used to perform aggregation operations and precise Softmax calculations across memory banks or across devices.

[0007] Furthermore, the lookup table piecewise linear approximation unit also includes the following configuration: dividing the input domain of the target nonlinear function into multiple intervals, pre-calculating and storing a set of linear coefficients (a, b) in each interval; after receiving the input data, the PIM unit of the in-memory processing device determines the interval to which the input belongs through its internal comparator array, reads the corresponding coefficients (a, b), and then performs a y = a * x + b operation through its inherent multiply-accumulator to obtain the approximate output value of the nonlinear function.

[0008] Furthermore, the hybrid scheduler also includes the following configuration: In tensor parallel mode, computational requests are alternately distributed to the two clusters. Supports two clusters processing different micro-batches of the same model stage in pipelined parallel mode; Monitor the load status of each cluster in real time and dynamically trigger task migration.

[0009] A second aspect of the present invention provides a retrieval enhancement and inference acceleration method based on in-memory computation, comprising the following steps: (1) The system receives user queries and retrieved documents, tokenizes them and initializes the inference task by the host; according to the document attention decomposition mechanism, the query vector is broadcast to multiple in-memory processing devices in the document attention cluster DAC; (2) Each in-memory processing device independently performs local attention calculation on the document key-value pairs stored locally to generate a compact intermediate representation; the intermediate representation of each device is transmitted to the context execution cluster CEC and fused with the query attention result; (3) The context execution cluster performs subsequent layer calculations of the Transformer to complete the reasoning of the entire generation phase; (4) A hybrid scheduling strategy is adopted to coordinate the above steps (2)-(3), and load balancing between the document attention cluster and the context execution cluster is achieved through dynamic triggering and interleaved execution.

[0010] Furthermore, step (1) specifically includes the following sub-steps: (1.1) System initialization and mapping: The host connects to and initializes multiple PIM devices through the CXL switch, and logically divides them into document attention clusters and context execution clusters; the host-side hybrid scheduler divides according to the selected parallel strategy, and statically maps different layers of the RAG model to specific devices of the document attention cluster and context execution cluster; the parallel strategy includes tensor parallel strategy and pipeline parallel strategy; (1.2) Query processing and distribution: The host receives the user's query and the retrieved document ID, tokenizes the query; then, it sends the query vector sequence to the context execution cluster, and resolves the document ID into a physical address within the document attention cluster, notifying the document attention cluster to prepare for computation.

[0011] Further, step (2) includes the following steps: (3.1) Each PIM device in the document attention cluster computes the attention score between the query vector and the local document key in parallel; (3.2) Each device independently calculates its local maximum value and normalization factor; (3.3) Generate unnormalized local attention output; (3.4) The local attention results of each device are finally fused through global normalization.

[0012] Furthermore, in step (3.4), global normalization is achieved by fusing local attention results using the following formula: (3.4.1) Attention score calculation: Calculate the attention scores for the two clusters respectively. For the query vector q, calculate the attention score with the Key vector in each of the two clusters respectively: The Key vector of the document cluster is K. d The key vector of the context cluster is K. c The corresponding attention scores are denoted as follows: and : , ; (3.4.2) Calculation of local maxima: Calculate the maximum score m in each of the two clusters. d and m c Used to stabilize exponentiation in softmax: , ; (3.4.3) Calculate the local output vector and normalization factor: Let the Value vector of the document cluster be... V d The value vector of the context cluster is V c Based on the corresponding attention score and The two clusters respectively calculate the output vector o d and o c and its corresponding normalization factor and : ; ; (3.4.4) Global normalization and output fusion: First, calculate the global maximum value m, and then calculate the global normalization factor. Finally, the output of the two clusters is... d and o c The output vector o is obtained by aggregation and global normalization. ; ; .

[0013] A third aspect of the invention: an electronic device comprising: One or more processors; Memory, used to store one or more programs; When the one or more programs are executed by the one or more processors, the one or more processors implement the aforementioned retrieval enhancement generation inference acceleration method based on in-memory computation.

[0014] A fourth aspect of the present invention: a computer-readable storage medium having stored thereon computer instructions which, when executed by a processor, implement the steps of the aforementioned in-memory computation-based retrieval enhancement generative inference acceleration method.

[0015] The beneficial technical effects of the present invention are as follows: First, this invention is built upon CXL memory technology and innovatively proposes a decentralized execution paradigm. Compared to existing solutions that only optimize computation or only optimize storage, the system systematically solves the bottleneck of RAG inference through the collaborative design of execution model, hardware architecture, and scheduling strategy. Second, utilizing a decentralized document attention decomposition mechanism, the system decomposes and distributes traditional centralized attention computation across multiple in-memory processing devices, enabling document key-value data to be processed locally at the storage location. Only lightweight intermediate results need to be transmitted, reducing communication volume from the order of document tokens to the order of query / response tokens, fundamentally alleviating interconnect bandwidth pressure. Moreover, through a resource-optimized in-memory computation microarchitecture, the system's PIM unit efficiently supports RAG-specific GEMM operations, which multiply multiple input vectors with the same weight matrix, by "copying only the buffer rather than the entire computation unit," maintaining high throughput while significantly reducing hardware overhead. Furthermore, it employs a piecewise linear approximation method based on lookup tables to implement nonlinear function computation in memory with minimal logic. Finally, based on a load-balanced inter-cluster scheduling strategy, the system effectively smooths the load difference between the document attention cluster and the context execution cluster through a hybrid scheduling approach combining static mapping and dynamic triggering, as well as an interleaved cluster execution mechanism. This reduces device idle time and significantly improves the overall resource utilization and throughput of the system. It provides an efficient and scalable hardware acceleration solution for large-scale, low-latency RAG services. Attached Figure Description

[0016] Figure 1 This is an overall architecture diagram of a preferred embodiment of the retrieval enhancement and generation inference acceleration system based on in-memory computing provided by the present invention; Figure 2 This is a flowchart of a preferred embodiment of the decentralized document attention decomposition mechanism provided by the present invention; Figure 3 This is a detailed diagram of the internal architecture of a PIM device according to a preferred embodiment of the present invention; Figure 4 This is a microarchitecture diagram of a PIM unit according to a preferred embodiment of the present invention; Figure 5This is a schematic diagram of a tensor parallel model mapping according to a preferred embodiment of the present invention; Figure 6 This is a schematic diagram of a pipeline parallel model mapping according to a preferred embodiment of the present invention; Figure 7 This is a timing diagram of the interleaved cluster execution according to a preferred embodiment of the present invention. Detailed Implementation

[0017] The following is a detailed explanation with reference to the accompanying drawings.

[0018] This invention first provides a retrieval enhancement and inference acceleration system based on in-memory computing, comprising: A computing cluster consisting of multiple homogeneous in-memory processing devices interconnected via the CXL protocol, the computing cluster being logically divided into a document attention cluster and a context execution cluster; The in-memory processing devices in the document attention cluster are used for distributed storage of pre-computed document key-value pairs and employ a resource-optimized in-memory computing microarchitecture to perform parallel document attention computations; the resource-optimized in-memory computing microarchitecture includes: The buffer copying unit for GEMM operations configures multiple input buffers for the same arithmetic logic unit, enabling multiple input vectors to share the loaded weight matrix. For nonlinear functions, a lookup table is used to segment the linear approximation unit, which completes the calculation of complex functions through one interval lookup and one multiplication-addition operation; The decentralized document attention decomposition module is configured as follows: The global attention computation is decoupled into document attention computation performed in the document attention cluster and query-response attention computation performed in the context execution cluster; Control each in-memory processing device to independently compute local attention results, including local output vectors, local maxima, and local summation terms, based on local document key-value pairs; By using a global normalization step based on the local maximum and local summation terms, the local attention results of each device are fused with the query-response attention results; The load balancing hybrid scheduler is configured as follows: During system initialization, operators are statically mapped to two clusters based on the model structure (i.e., prior knowledge). Downstream execution is dynamically triggered at runtime based on data availability. The task migration mechanism dynamically migrates overloaded tasks from the context execution cluster to idle devices in the document attention cluster for execution.

[0019] Based on the above system, such as Figure 2As shown, the present invention also provides a retrieval enhancement generation inference acceleration method based on in-memory computing, comprising the following steps: (1) The system receives the user's query and retrieved documents, which are tokenized by the host and the reasoning task is initialized; (2) Based on the document attention decomposition mechanism, the query vector is broadcast to multiple in-memory processing devices in the Document Attention Cluster (DAC); (3) Each in-memory processing device independently performs local attention calculations on the document key-value pairs stored locally to generate a compact intermediate representation; (4) Transmit the intermediate representations of each device to the Context Execution Cluster (CEC) and fuse them with the query attention results; (5) The context execution cluster performs subsequent layer calculations of the Transformer model to complete the reasoning of the entire generation stage; (6) A hybrid scheduling strategy is adopted to coordinate the above steps (3) to (5), and load balancing between the document attention cluster and the context execution cluster is achieved through dynamic triggering and interleaved execution.

[0020] Specifically, the system proposed in this invention utilizes the high-speed communication protocol CXL (Compute Express Link) to construct a computing cluster composed of multiple homogeneous in-memory computing PIM devices, and logically divides it into a document attention cluster and a context execution cluster to support decentralized execution. The document attention cluster, composed of a first group of in-memory processing devices, is used to distribute and store pre-computed document key-value pairs and perform parallel document attention computation. The context execution cluster, composed of a second group of in-memory processing devices, is used to perform subsequent layer computations of the Transformer model.

[0021] For the decentralized document attention decomposition module, the core of this invention lies in decoupling global attention computation from centralized processing into document attention computation executed in a parallel document attention cluster and query-response attention computation executed in a context execution cluster. By utilizing the associative law of attention scores and the composability of Softmax, each device in the document attention cluster can independently compute its local results (including output vector, local maxima, and normalization factor), and then fuse them with the query-response attention results through a numerically stable global normalization step. This mechanism ensures that, under the premise of mathematical equivalence, document key-value data does not need to leave its storage device, thereby achieving an order-of-magnitude reduction in communication volume.

[0022] The global normalization described above achieves the fusion of local attention results through the following calculation process: First, attention score calculation: Calculate the attention score for each of the two clusters. For the query vector q, calculate the attention score with the key vectors in both clusters: the key vector for the document cluster is K. d The key vector of the context cluster is K. c The corresponding attention scores are denoted as follows: and : , ; Secondly, local maximum calculation: calculate the maximum score m in each of the two clusters. d and m c Used to stabilize exponentiation in softmax: , ; Then, calculate the local output vector and the normalization factor: Let V be the value vector of the document cluster. d The value vector of the context cluster is V. c Based on the corresponding attention score s d and s c The two clusters respectively calculate the output vector o d and o c and its corresponding normalization factor l d and l c : , ; , ; Finally, global normalization and output fusion: First, calculate the global maximum value m, then calculate the global normalization factor l, and finally fuse the outputs o of the two clusters. d and o c The output vector o is obtained by aggregation and global normalization. ; ; .

[0023] This invention presents a resource-optimized in-memory computing microarchitecture designed to efficiently support Retrieval Enhancement Generation (RAG) inference primitives within strict area and power consumption constraints of the in-memory computing unit. For the buffer copying unit in GEMM operations (multiplying multiple input vectors with the same weight matrix), the in-memory computing unit (PIM) employs a "buffer copying" strategy. This involves configuring multiple input buffers for the same set of arithmetic logic units, allowing multiple input vectors to share the loaded weight matrix, avoiding redundant memory accesses, and achieving parallelism at low cost. For the lookup table piecewise linear approximation unit for nonlinear functions, a "lookup table piecewise linear approximation" method is used. The input range of the nonlinear function is divided into multiple intervals, each storing a set of linear approximation parameters. This transforms complex function calculations into a single interval lookup and a single multiply-add operation, achieving high-energy-efficiency approximation computation with minimal hardware overhead.

[0024] A hybrid load balancing scheduler utilizes inter-cluster scheduling strategies to address the load imbalance issue between clusters caused by decentralized execution. The scheduler employs a hybrid approach combining static mapping and dynamic triggering: during system initialization, operators are statically mapped to the document attention cluster and the context execution cluster based on the model structure; during runtime, downstream execution is dynamically triggered based on data availability, reducing waiting latency. Furthermore, through an interleaved cluster execution mechanism, requests are alternately allocated to the document attention cluster and the context execution cluster under tensor parallelism, while under pipeline parallelism, both clusters are allowed to process different micro-batches of the same stage. It also supports dynamically migrating some tasks from the context execution cluster to idle document attention clusters for execution, thereby maximizing device utilization.

[0025] Example 1 For the accelerated generation and inference process of retrieval enhancement based on in-memory computation, the transmission bottleneck of massive document key-value pairs and the computational inefficiency of GEMM operations involving multiple input vectors multiplied by the same weight matrix are key factors limiting system performance. This invention provides a decentralized in-memory computation architecture built on the CXL protocol, aiming to combine the large capacity advantage of CXL memory with the high bandwidth advantage of in-memory computation to achieve efficient RAG inference.

[0026] Specifically, the following are the specific implementation steps for the Meridian system to perform RAG inference: Step 1: System initialization and mapping. See also Figure 1 The host connects to and initializes multiple PIM devices via a CXL switch, logically dividing them into document attention clusters and context execution clusters. The scheduler on the host side (the central processing unit (CPU) and its operating system environment) determines the parallelism strategy based on the selected strategy, such as... Figure 5 The tensor parallel strategy shown may divide the model layers according to... Figure 6 The pipeline parallel strategy shown is used to partition the RAG model, statically mapping different layers of the RAG model to specific devices of the document attention cluster and the context execution cluster.

[0027] Step 2: Query Processing and Distribution. The host receives the user's query and the retrieved document ID, and tokenizes the query. Then, it sends the query vector sequence to the context execution cluster, resolves the document ID to a physical address within the document attention cluster, and notifies the document attention cluster to prepare for computation.

[0028] Step 3: Decentralized document attention calculation (see [link]) Figure 2 For each query vector q generated in the decoding step, each PIM device in the document attention cluster computes the attention score of q and its local document key in parallel. Each device independently computes its local maximum value and generates an unnormalized output and a normalization factor. After aggregating these intermediate results within the device, each device sends them to the context execution cluster.

[0029] Step 4: Global Fusion and Context Generation. After receiving the intermediate results from the document attention cluster, the context execution cluster performs global normalization and fusion with the attention results from its own query, obtaining the final attention output. This output is then processed through layers such as FFN to generate the next token and update the KV cache in the context execution cluster.

[0030] Step 5: Dynamic scheduling and load balancing. See also Figure 5 , Figure 6 and Figure 7 Throughout the process, the scheduler monitors the status of each cluster. Once the document attention cluster completes the document attention for the current batch, it immediately and dynamically triggers the context execution cluster to begin fusion computation. Simultaneously, if the document attention cluster is idle while the context execution cluster is busy, the scheduler uses an interleaved cluster execution mechanism to allocate the next inference request or a portion of the context execution cluster's computational tasks to the document attention cluster, ensuring high cluster efficiency.

[0031] Furthermore, a resource-optimized in-memory computing microarchitecture is provided, which can be used to jointly implement or execute a retrieval enhancement generation inference acceleration method based on in-memory computing as provided in Example 1.

[0032] Specifically, see Figure 3 and Figure 4 The in-memory computing microarchitecture provided in this embodiment may include at least PIM unit design and controller-side unit design.

[0033] According to a preferred embodiment, the PIM unit is configured as follows: See Figure 4Each PIM unit is integrated alongside the DRAM memory bank and contains 16 FP16 multipliers and adders, 16 FP16 comparators, and four double-buffered 4KB input buffers. When performing a GEMM operation that multiplies multiple input vectors with the same weight matrix, the controller writes the multiple input vectors into different buffers via commands. Then, through a sequence of control signals, it schedules the shared arithmetic logic unit to read data from each buffer sequentially and perform calculations with the same weight matrix, thereby achieving single loading and multiple reuse of weights.

[0034] According to a preferred embodiment, the nonlinear function approximation method is configured such that the coefficients (a, b) of each interval are pre-calculated and stored in a LUT with a capacity of 16. After receiving the input data, the PIM unit determines the interval to which the input value belongs in parallel using its comparator array, indexes out the corresponding a and b, and then uses its inherent multiply-accumulator to complete the calculation of y = a * x + b in one cycle to obtain an approximate value of the nonlinear function (such as GELU).

[0035] According to a preferred embodiment, the controller-side unit is configured as follows: See Figure 3 Each PIM controller integrates a near-memory unit, which contains an addition tree to aggregate the outputs of all in-memory computing units within the channel, and a dedicated high-precision Softmax unit for performing numerically stable, precise Softmax (high-precision complex function) calculations required in the attention mechanism. The CXL controller integrates multiple RISC-V cores for performing global result aggregation and lightweight control logic across channels and devices.

[0036] In addition, the present invention also provides an electronic device, comprising: One or more processors; Memory, used to store one or more programs; When the one or more programs are executed by the one or more processors, the one or more processors implement the aforementioned retrieval enhancement generation inference acceleration method based on in-memory computation.

[0037] A computer-readable storage medium having computer instructions stored thereon, which, when executed by a processor, implement the steps of the retrieval enhancement generative inference acceleration method based on in-memory computation.

[0038] It should be noted that the specific embodiments described above are exemplary. Those skilled in the art can devise various solutions inspired by the disclosure of this invention, and these solutions all fall within the scope of this invention and its protection. Those skilled in the art should understand that this specification and its accompanying drawings are illustrative and not intended to limit the scope of the claims. The scope of protection of this invention is defined by the claims and their equivalents. This specification contains multiple inventive concepts; terms such as "preferredly," "according to a preferred embodiment," or "optionally" indicate that the corresponding paragraph discloses an independent concept. The applicant reserves the right to file divisional applications based on each inventive concept.

Claims

1. A retrieval enhancement and generative reasoning acceleration system based on in-memory computing, characterized in that, include: A computing cluster consisting of multiple homogeneous in-memory processing devices interconnected via the CXL protocol, the computing cluster being logically divided into a document attention cluster and a context execution cluster; in, The in-memory processing devices in the document attention cluster are used for distributed storage of pre-computed document key-value pairs and employ a resource-optimized in-memory computing microarchitecture to perform parallel document attention computations; the resource-optimized in-memory computing microarchitecture includes: The buffer copying unit for GEMM operations configures multiple input buffers for the same arithmetic logic unit, enabling multiple input vectors to share the loaded weight matrix. For nonlinear functions, a lookup table is used to segment the linear approximation unit, which completes the calculation of complex functions through one interval lookup and one multiplication-addition operation; The decentralized document attention decomposition module is configured as follows: The global attention computation is decoupled into document attention computation performed in the document attention cluster and query-response attention computation performed in the context execution cluster; Control each in-memory processing device to independently compute local attention results, including local output vectors, local maxima, and local summation terms, based on local document key-value pairs; By using a global normalization step based on the local maximum and local summation terms, the local attention results of each device are fused with the query-response attention results; The load balancing hybrid scheduler is configured as follows: During system initialization, computational tasks are statically mapped according to the model structure, and during runtime, dynamic task scheduling and migration are performed based on the load status of the two clusters.

2. The system according to claim 1, characterized in that, The in-memory processing device includes an in-memory computing unit and a controller-side unit: The in-memory computing unit, located next to the DRAM memory, includes several multipliers and adders, comparators, and a double-buffered input buffer; it is used to perform basic matrix-vector multiplication and nonlinear function approximation calculations based on lookup tables. The controller-side unit is located in the memory controller. Each memory controller integrates a near-memory unit, which includes an addition tree for aggregating the outputs of all memory computing units within the channel, and a dedicated high-precision Softmax unit for performing precise Softmax calculations that require numerical stability in the attention mechanism; it is used to perform aggregation operations and precise Softmax calculations across memory banks or across devices.

3. The system according to claim 1, characterized in that, The lookup table piecewise linear approximation unit further includes the following configuration: dividing the input domain of the target nonlinear function into multiple intervals, pre-calculating and storing a set of linear coefficients (a, b) in each interval; after receiving the input data, the PIM unit of the in-memory processing device determines the interval to which the input belongs through its internal comparator array, reads the corresponding coefficients (a, b), and then performs a y = a * x + b operation through its inherent multiply-accumulator to obtain the approximate output value of the nonlinear function.

4. The system according to claim 1, characterized in that, The hybrid scheduler also includes the following configurations: In tensor parallel mode, computational requests are alternately distributed to the two clusters. Supports two clusters processing different micro-batches of the same model stage in pipelined parallel mode; Monitor the load status of each cluster in real time and dynamically trigger task migration.

5. The acceleration method for a retrieval-enhanced generative reasoning acceleration system based on in-memory computation according to any one of claims 1-4, characterized in that, Includes the following steps: (1) The system receives user queries and retrieved documents, tokenizes them and initializes the inference task by the host; according to the document attention decomposition mechanism, the query vector is broadcast to multiple in-memory processing devices in the document attention cluster DAC; (2) Each in-memory processing device independently performs local attention calculation on the document key-value pairs stored locally to generate a compact intermediate representation; the intermediate representation of each device is transmitted to the context execution cluster CEC and fused with the query attention result; (3) The context execution cluster performs subsequent layer calculations of the Transformer to complete the reasoning of the entire generation phase; (4) A hybrid scheduling strategy is adopted to coordinate the above steps (2)-(3), and load balancing between the document attention cluster and the context execution cluster is achieved through dynamic triggering and interleaved execution.

6. The method according to claim 5, characterized in that, Step (1) specifically includes the following sub-steps: (1.1) System initialization and mapping: The host connects to and initializes multiple PIM devices through the CXL switch, and logically divides them into document attention clusters and context execution clusters; the host-side hybrid scheduler divides according to the selected parallel strategy, and statically maps different layers of the RAG model to specific devices of the document attention cluster and context execution cluster; the parallel strategy includes tensor parallel strategy and pipeline parallel strategy; (1.2) Query processing and distribution: The host receives the user's query and the retrieved document ID, and tokenizes the query; Subsequently, the query vector sequence is sent to the context execution cluster, and the document ID is resolved into a physical address within the document attention cluster, notifying the document attention cluster to prepare for computation.

7. The method according to claim 5, characterized in that, Step (2) includes the following steps: (3.1) Each PIM device in the document attention cluster computes the attention score between the query vector and the local document key in parallel; (3.2) Each device independently calculates its local maximum value and normalization factor; (3.3) Generate unnormalized local attention output; (3.4) The local attention results of each device are finally fused through global normalization.

8. The method according to claim 7, characterized in that, In step (3.4), global normalization is achieved by fusing local attention results using the following formula: (3.4.1) Attention score calculation: Calculate the attention scores for the two clusters respectively. For the query vector q, calculate the attention score with the Key vector in each of the two clusters respectively: The Key vector of the document cluster is K. d The key vector of the context cluster is K. c The corresponding attention scores are denoted as follows: and : , ； (3.4.2) Calculation of local maxima: Calculate the maximum score m in each of the two clusters. d and m c Used to stabilize exponentiation in softmax: ， ； (3.4.3) Calculate the local output vector and normalization factor: Let the Value vector of the document cluster be... V d The value vector of the context cluster is V c Based on the corresponding attention score and The two clusters respectively calculate the output vector o d and o c and its corresponding normalization factor and : ； ； (3.4.4) Global normalization and output fusion: First, calculate the global maximum value m, and then calculate the global normalization factor. Finally, the output of the two clusters is... d and o c The output vector o is obtained by aggregation and global normalization. ； ； 。 9. An electronic device, characterized in that, include: One or more processors; Memory, used to store one or more programs; When the one or more programs are executed by the one or more processors, the one or more processors implement a retrieval enhancement generation inference acceleration method based on in-memory computation as described in any one of claims 5-8.

10. A computer-readable storage medium storing computer instructions thereon, characterized in that, When executed by the processor, this instruction implements the steps of the retrieval enhancement generation RAG inference acceleration method based on in-memory computation as described in any one of claims 5-8.

Images