Data storage method, large language model inference decoding acceleration method and system based on hybrid bonding architecture
Patent Information
- Application Number
- CN202610938839.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2026-06-26
- Publication Date
- 2026-09-18
AI Technical Summary
但现有基于混合键合架构的大语言模型加速方案,在适配推测解码时仍存在挑战:混合键合架构采用分布式内存设计,每个DRAM bank仅与对应处理元件(PE)直连,无跨bank直接互联链路,跨bank数据访问需经过全局缓冲区中转,非最优的数据存储策略会带来严重的跨bank访问开销,最终抵消混合键合的带宽优势,使得大语言模型推测解码加速过程中的吞吐量有待提高
1、本发明第一方面提供了一种数据存储方法,用于对大语言模型推测解码加速过程中模型的权重数据进行存储。考虑到混合键合近存计算设备中,每个内存块与对应计算单元之间具有高带宽直连关系,而不同内存块之间的数据访问通常需要经由全局缓冲区或片上互联进行中转,本发明根据大语言模型的Transformer层中各权重矩阵的运算顺序和数据流向,将目标大语言模型和草稿模型中的权重矩阵划分为包括查询权重矩阵、键权重矩阵
、值权重矩阵
、和前馈神经网络层中的第一层权重矩阵
的第一类矩阵和包括注意力层中的输出投影权重矩阵
和前馈神经网络层中的第二层权重矩阵
的第二类矩阵,并分别采用互补的行列分片方式进行存储。具体而言,将第一类矩阵划分为N行M列的第一矩阵块阵列,并将同一列下的矩阵块存储至同一内存块中;将第二类矩阵划分为M行N列的第二矩阵块阵列,并将同一行下的矩阵块存储至同一内存块中。由此,在执行查询、键、值生成和前馈网络第一层计算时,各计算单元能够主要访问本地内存块中的权重块完成局部计算;在执行输出投影和前馈网络第二层计算时,也能够使前一阶段产生的局部中间结果与后续所需权重块在内存块维度上保持对应关系,从而减少跨内存块搬移和全局缓冲区中转次数。该存储方式还能够使各内存块和各内存库承担相对均衡的权重访问负载,提升近存计算单元的并行利用率和内部带宽利用率,降低权重读取延迟及数据搬移能耗,进而提高大语言模型推测解码加速过程中的吞吐量。
Smart Images

Figure CN122777052A_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of large language model inference acceleration technology, and more specifically, relates to a data storage method, a large language model inference decoding acceleration method and system based on hybrid bonding architecture. Background Technology
[0002] Generative large language models based on Transformer are widely used in tasks such as dialogue, code generation, and complex reasoning. Their reasoning process is generally divided into a pre-filling stage and an autoregressive decoding stage. The pre-filling stage computes multiple input requests in parallel; the decoding stage generates tokens autoregressively, with each token depending on the output of the previous token. This strict temporal dependency leads to low inference throughput and high latency, becoming a core bottleneck for deploying large language models at edge devices.
[0003] To improve the parallelism and throughput of the decoding stage, speculative decoding technology is widely used. It introduces a lightweight draft language model (DLM) to pre-generate multiple draft tokens, and then the target large language model (TLM) verifies the multiple draft tokens in parallel.
[0004] To accommodate the drastically different computational characteristics of DLM (memory-intensive) and TLM (compute-intensive) operations in speculative decoding, near-memory computing (NMP) heterogeneous architectures have been widely researched and applied. These architectures integrate logic computing units within DRAM, providing high bandwidth support for memory-intensive operators and reducing data transfer overhead. However, existing NMP heterogeneous architectures typically integrate logic computing units onto DRAM chips, resulting in extremely low computational density due to the logic resource constraints of DRAM technology.
[0005] In emerging hybrid-bonded near-memory computing heterogeneous devices, hybrid bonding (HB) technology vertically stacks DRAM chips and logic chips, achieving high-density interconnects through fusion bonding. It allows for customized design of logic computing units within the logic chips, balancing high memory access bandwidth with powerful computing capabilities. However, existing large language model acceleration solutions based on hybrid bonding architectures still face challenges in adapting to speculative decoding: the hybrid bonding architecture employs a distributed memory design, with each DRAM bank directly connected only to its corresponding processing element (PE), lacking direct cross-bank interconnect links. Cross-bank data access requires transit through a global buffer, and this suboptimal data storage strategy leads to significant cross-bank access overhead, ultimately offsetting the bandwidth advantage of hybrid bonding and requiring improved throughput during large language model speculative decoding acceleration. Summary of the Invention
[0006] In view of the above-mentioned defects or improvement needs of the existing technology, the present invention provides a data storage method, a method and system for accelerating large language model inference decoding based on a hybrid bonding architecture, the purpose of which is to improve the throughput in the process of accelerating large language model inference decoding.
[0007] To achieve the above objectives, in a first aspect, the present invention provides a data storage method, comprising: storing key data in the large language model inference decoding acceleration process; the large language model inference decoding acceleration process is executed using a hybrid bonded near-memory computing device; wherein the hybrid bonded near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: multiple near-memory computing chips; the near-memory computing chips include: vertically stacked upper DRAM chips and lower logic chips; the DRAM chips include: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. The key data includes: weight data for the target large language model (TLM) and the draft model (DLM); the weight data for the TLM and DLM are stored in a DRAM chip in the following manner: For each model in TLM and DLM, for each first-class matrix within it and each second-class matrix Both horizontal partitioning by row index and vertical partitioning by column index are performed to divide each type I matrix. Each matrix is divided into an N-row, M-column first matrix block array and each second-class matrix The first matrix block array is divided into M rows and N columns; the size of each matrix block in the first matrix block array does not exceed [a certain value]. The size of each matrix block in the second matrix block array does not exceed [a certain value]. ; Store the matrix block in the j-th row and k-th column of each first matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; store the matrix block in the k-th row and j-th column of each second matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; ; ; and They correspond to as The number of rows and columns; This represents the floor function; and They correspond to as The number of rows and columns; The first type of matrix includes: query weight matrix Key weight matrix Value weight matrix The weight matrix of the first layer in the feedforward neural network. ; The second type of matrix includes: the output projection weight matrix in the attention layer. and the weight matrix of the second layer in the feedforward neural network layer .
[0008] More preferably, the aforementioned key data further includes: the key-value matrices in the TLM and DLM generated during the large language model inference decoding acceleration process; the key-value matrices are stored in the DRAM chip in the following manner: When the number of attention heads in the key-value matrix is less than or equal to M, one attention head in the key-value matrix is considered as one attention head group; when the number of attention heads in the key-value matrix is greater than M, the attention heads in the key-value matrix are grouped to obtain M' attention head groups. ; The key matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip; for each attention head corresponding to the key matrix K in the key-value matrix, its transpose matrix is... Vertical partitioning by column index yields a matrix block array with one row and N columns. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th column is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; The value matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip. For the value matrix V corresponding to each attention head in the key-value matrix, it is horizontally partitioned according to the row index to obtain an N-row, one-column matrix block array. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th row is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; in, ; and Each is a matrix The number of rows and columns; and These are the number of rows and columns of the value matrix V, respectively.
[0009] Secondly, the present invention provides a method for accelerating large language model inference decoding based on a hybrid bonding architecture, which is executed using a hybrid bonding near-memory computing device; the hybrid bonding near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: multiple near-memory computing chips; the near-memory computing chips include: vertically stacked upper DRAM chips and lower logic chips; the DRAM chips include: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. The above-mentioned methods for accelerating speculative decoding include: S1. Read the weight data of the target large language model TLM and draft model DLM from the DRAM chip, perform the pre-filling calculation of the inference request in parallel based on the embedding vector of each inference request, generate the first output token of the TLM corresponding to each inference request as the current header token of the inference request, and store the key-value matrix of the TLM and DLM corresponding to each token in each inference request obtained during the pre-filling calculation into the DRAM chip. S2, One round of speculative decoding operation: Execute the speculative decoding operation of each inference request in parallel based on the current header token of each inference request; S3. Take the last token in the output token of each inference request in the previous round of speculation and decoding as the current header token of the inference request, and repeat S2 for iteration until an output end symbol is generated or the sum of the number of output tokens of each round of speculation and decoding reaches the preset number. The speculative decoding operation for each inference request r includes a draft prediction operation and a verification operation executed sequentially for the inference request r. The draft prediction operation performed on inference request r includes: reading the weight data of DLM from the DRAM chip, the key-value matrix of the draft model corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of DLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing multiple iterative predictions using DLM based on the read data to generate n draft tokens for inference request r, and storing the key-value matrix of DLM corresponding to each generated draft token into the DRAM chip; the input of the first iterative prediction for inference request r in the draft prediction operation of the first round of speculation decoding operation is the current header token of the inference request; the input of subsequent iterative predictions for inference request r in the draft prediction operation of each round of speculation decoding operation is the output of the previous iterative prediction of the request; The verification operation performed on inference request r includes: reading the weight data of TLM from the DRAM chip, the key-value matrix of TLM corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of TLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing a verification of the n draft tokens of inference request r using TLM based on the read data, and using the verified draft tokens as the output tokens of inference request r in the current round of speculation decoding operation; storing the key-value matrix of TLM corresponding to the output token of inference request r in the current round of speculation decoding operation into the DRAM chip, and deleting the key-value matrix of DLM corresponding to each draft token that failed verification from the DRAM chip; In this process, the key data in the large language model inference decoding acceleration process is stored in a DRAM chip using the data storage method provided in the first aspect of the present invention; the key data includes: weight data of TLM and DLM; or; the key data includes: weight data of TLM and DLM, and key-value matrices in TLM and DLM.
[0010] More preferably, the n draft tokens generated for each inference request in the draft prediction operation are in a tree structure; Each memory bank includes a first fixed storage area for storing the key-value matrix of the DLM corresponding to the draft token and a second fixed storage area for storing the key-value matrix of the DLM corresponding to the output token; The draft prediction operation performed on inference request r also includes: When generating the first draft token corresponding to the inference request r, a corresponding logic tree is created. The root node of this logic tree is used to record the address information of the key-value matrix of the DLM corresponding to the current header token of the inference request r in the memory library. When a draft token corresponding to an inference request r is generated, a minimum allocatable unit is obtained from the first fixed storage area of the corresponding memory library and allocated to the key-value matrix of the DLM corresponding to the draft token. A node corresponding to the draft token is established in the corresponding logic tree to record the address information of the key-value matrix of the DLM corresponding to the draft token in the memory library. At the same time, the node corresponding to the draft token in the corresponding logic tree is pointed to its parent node. The verification operations performed on inference request r also include: For the draft token corresponding to the verified inference request r, move the key-value matrix of its corresponding DLM from the first fixed area to the second fixed area of the memory library; The above-mentioned S2 also includes: after the verification operation for all inference requests is completed, releasing all data in the first fixed area of each memory bank.
[0011] More preferably, the computing element includes: a hybrid bonding controller and multiple multiply-accumulate operation units; The hybrid bonding controller is used to access the memory library that is hybrid-bonded to the computing element it resides in; Multiple multiply-accumulate units are used to perform matrix-vector multiplication or matrix multiplication operations on weight data, key-value matrix data or intermediate activation data read from the corresponding memory library.
[0012] More preferably, the hybrid bonded near-memory computing device further includes a neural network processor.
[0013] More preferably, before each round of speculative decoding operation is executed, a corresponding executor is assigned to each operator in that round of speculative decoding operation, including: For each model in the operator set corresponding to this round of speculative decoding operations, compute the operator op and calculate its arithmetic strength. ;when At that time, the near-in-memory computing module is used as the executor of the operator op in this round of speculation decoding operation; when When the neural network processor is used as the executor of the operator op in that round of speculation-decoding operation; when ,and In this case, the device with higher actual computing power in the neural network processor and near-memory computing module when executing operator op is used as the executor for executing operator op in this round of speculation decoding operation; The operator set includes: the fully connected operator FC1 of the TLM, and the attention operator of the TLM corresponding to each inference request r. And the fully connected operator of DLM in the s-th iteration prediction of the draft prediction operation in this round of speculative decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure: the model computation operator set also includes: the DLM attention operator corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure: the model computation operator set also includes: the shared attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. And the independent attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of the speculative decoding operation. ; Arithmetic strength of operator FC1 ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Calculate its arithmetic strength ; in, S represents the number of iterations for each inference request r in the draft prediction operation of this round of speculative decoding. This is the total bandwidth between all memory libraries and their corresponding computing elements in the near-memory computing module. This refers to the external memory access bandwidth of the neural network processor. and B represents the floating-point arithmetic capabilities of the neural network processor and the near-memory computing module, respectively; B represents the number of concurrent inference requests; and the number of tokens. ; and These are the input and output dimensions of the TLM, respectively. The sum of the number of tokens in the inference request r and the number of output tokens of the inference request r obtained from the previous rounds of inference decoding operations; The dimension of the attention head in TLM; the number of draft tokens generated in the s-th iteration prediction when the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure. The value is 1; when the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure, the number of draft tokens generated in the s-th iteration prediction is 1. Let be a preset integer, which satisfies ,and ; and These are the input and output dimensions of the DLM, respectively. This represents the dimension of the attention head in DLM.
[0014] More preferably, the actual computing power of the neural network processor when executing operator op is: ; The actual computing power of the near-memory computing module when executing operator op is: .
[0015] More preferably, in S1 above, a neural network processor is used to perform pre-filling computation of the inference request in parallel based on the embedding vector of each inference request.
[0016] Thirdly, the present invention provides a large language model inference decoding acceleration system based on a hybrid bonding architecture, comprising: a hybrid bonding near-memory computing device; the hybrid bonding near-memory computing device comprising: a near-memory computing module; wherein, the near-memory computing module comprises: multiple near-memory computing chips; the near-memory computing chip comprises: a vertically stacked upper DRAM chip and a lower logic chip; the DRAM chip comprises: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. The aforementioned hybrid-bonded near-memory computing device is used to execute the large language model inference decoding acceleration method provided in the second aspect of the present invention.
[0017] In summary, the above-described technical solutions conceived in this invention can achieve the following beneficial effects: 1. The first aspect of this invention provides a data storage method for storing the weight data of a model during the acceleration of large language model inference decoding. Considering that in hybrid bonded near-memory computing devices, each memory block has a high-bandwidth direct connection with its corresponding computing unit, and data access between different memory blocks usually requires transit through a global buffer or on-chip interconnect, this invention, based on the operation order and data flow of each weight matrix in the Transformer layer of the large language model, divides the weight matrices in the target large language model and the draft model into a system including a query weight matrix. Key weight matrix Value weight matrix The weight matrix of the first layer in the feedforward neural network. The first type of matrix and the output projection weight matrix including the attention layer and the weight matrix of the second layer in the feedforward neural network layer The second type of matrix is stored using complementary row and column partitioning. Specifically, the first type of matrix is divided into an N-row, M-column array of first matrix blocks, with matrix blocks in the same column stored in the same memory block; the second type of matrix is divided into an M-row, N-column array of second matrix blocks, with matrix blocks in the same row stored in the same memory block. Thus, during querying, key-value generation, and the first layer computation of the feedforward network, each computation unit can primarily access the weight blocks in its local memory block to complete local computations. During output projection and the second layer computation of the feedforward network, the local intermediate results generated in the previous stage can maintain a correspondence with the subsequent required weight blocks in the memory block dimension, thereby reducing the number of cross-memory block moves and global buffer transfers. This storage method also enables each memory block and memory library to bear a relatively balanced weight access load, improving the parallel utilization and internal bandwidth utilization of near-memory computation units, reducing weight read latency and data movement energy consumption, and thus improving throughput in the large language model inference decoding acceleration process.
[0018] 2. Further, the data storage method provided by this invention is used to store the key-value matrix generated during the inference decoding acceleration process of a large language model; to store the key matrices corresponding to different attention head groups in the key-value matrix into different memory blocks in a DRAM chip; to store the value matrices corresponding to different attention head groups in the key-value matrix into different memory blocks in a DRAM chip; and simultaneously, according to the operation order and data flow of each weight matrix in the Transformer layer of the large language model, to transpose the key matrix... Vertical partitioning by column index yields a matrix block array with one row and N columns. The value matrix V is horizontally partitioned according to its row index to obtain an N-row, one-column matrix block array. ;Will The matrix block under column j and The matrix block in row j is stored in the j-th memory bank of the corresponding memory block in the DRAM chip. This allows the key-value matrix to be mapped to the memory bank of the DRAM chip using an alternating row and column sharding method. This fully utilizes the characteristic that each memory bank in the near-memory computing chip is hybrid-bonded to the corresponding computing element, enabling the entire single attention head computation process to be completed within a single memory block. This minimizes the access overhead across memory banks and memory blocks, maximizes the internal bandwidth utilization of the hybrid bonding architecture, and further improves the throughput in the large language model inference decoding acceleration process.
[0019] 3. The second aspect of the present invention provides a method for accelerating large language model speculation decoding based on a hybrid bonding architecture. The key data in the speculation decoding acceleration process is stored in a DRAM chip using the data storage method provided in the first aspect of the present invention, thereby improving the throughput in the large language model speculation decoding acceleration process.
[0020] 4. Furthermore, the language model inference decoding acceleration method provided by this invention proposes a tree-structured key-value memory management mechanism for the scenario in which n draft tokens generated for each inference request in the draft prediction operation are in a tree structure. Based on the logic tree, the method realizes the cascading invalidation of key-value cache and the reuse of fixed block memory, eliminates the memory fragmentation problem caused by the abandoned branches of the tree structure, improves the device memory utilization rate, and reduces the key-value cache allocation latency, which is especially suitable for edge devices.
[0021] 5. Furthermore, the language model speculation decoding acceleration method provided by the present invention, in the scenario where the n draft tokens generated for each inference request in the draft prediction operation are in a tree structure, takes into account that the number of draft tokens generated in each round increases under the tree structure. At this time, the computing element is set to include multiple multiply-accumulators to better adapt to the high computational intensity characteristics of tree structure speculation decoding.
[0022] 6. Furthermore, the language model speculation decoding acceleration method provided by the present invention provides an arithmetic strength-aware adaptive operator scheduling strategy, which analyzes the running characteristics of the operators to be executed in the current round, counts the arithmetic strength of the operators to be executed, and allocates different operators to neural network processors or near-memory computing modules for execution, so that they can run in the optimal execution hardware, adapt to the dynamic load characteristics of speculation decoding, and greatly improve the utilization of hardware resources. Attached Figure Description
[0023] Figure 1 This is a schematic diagram of the near-memory computing chip provided in an embodiment of the present invention; Figure 2A schematic diagram illustrating the block division method of the weight matrices of TLM and DLM provided in an embodiment of the present invention; Figure 3 This is a schematic diagram of the block division method of the key-value matrix provided in an embodiment of the present invention; Figure 4 This is a schematic diagram of the near-memory processing execution flow provided in an embodiment of the present invention; Figure 5 A schematic diagram of the structure of the large language model inference decoding acceleration system provided in an embodiment of the present invention; Figure 6 This is a schematic diagram of the microarchitecture of the computing unit provided in an embodiment of the present invention; Figure 7 A schematic diagram of the microarchitecture of a computing element provided in an embodiment of the present invention; Figure 8 This is a schematic diagram of an arithmetic strength-aware operator scheduling strategy provided in an embodiment of the present invention. Detailed Implementation
[0024] 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 accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention. Furthermore, the technical features involved in the various embodiments of this invention described below can be combined with each other as long as they do not conflict with each other.
[0025] To achieve the above objectives, in a first aspect, the present invention provides a data storage method, comprising: storing key data in the large language model inference decoding acceleration process; the large language model inference decoding acceleration process is executed using a hybrid-bonded near-memory computing device; wherein the hybrid-bonded near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: multiple near-memory computing chips; as shown... Figure 1 As shown, the near-memory computing chip includes: a vertically stacked upper DRAM chip and a lower logic chip; the DRAM chip includes: M memory blocks (i.e., memory channels). Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory libraries in the corresponding memory blocks; each memory library is hybrid-bonded to its corresponding computing element; preferably, each memory library is hybrid-bonded to its corresponding computing element via Cu-Cu hybrid bonding.
[0026] Key data includes: weight data for the target large language model (TLM) and the draft model (DLM); such as Figure 2 As shown, the weight data for TLM and DLM are stored in the DRAM chip in the following manner: For each model in TLM and DLM, for each first-class matrix within it and each second-class matrix Both horizontal partitioning by row index and vertical partitioning by column index are performed to divide each type I matrix. Each matrix is divided into an N-row, M-column first matrix block array and each second-class matrix The first matrix block array is divided into M rows and N columns; the size of each matrix block in the first matrix block array does not exceed [a certain value]. The size of each matrix block in the second matrix block array does not exceed [a certain value]. ; Store the matrix block in the j-th row and k-th column of each first matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; store the matrix block in the k-th row and j-th column of each second matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; ; ; and They correspond to as The number of rows and columns; This represents the floor function; and They correspond to as The number of rows and columns; The first type of matrix includes: query weight matrix Key weight matrix Value weight matrix The weight matrix of the first layer in the feedforward neural network. ; The second type of matrix includes: the output projection weight matrix in the attention layer. and the weight matrix of the second layer in the feedforward neural network layer .
[0027] Preferably, in one optional implementation, the aforementioned key data further includes: the key-value matrix (i.e., KV-Cache) in the TLM and DLM generated during the large language model inference decoding acceleration process; such as Figure 3 As shown, the key-value matrix is stored in the DRAM chip in the following manner: When the number of attention heads in the key-value matrix is less than or equal to M, one attention head in the key-value matrix is considered as one attention head group; when the number of attention heads in the key-value matrix is greater than M, the attention heads in the key-value matrix are grouped to obtain M' attention head groups. ; The key matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip; for each attention head corresponding to the key matrix K in the key-value matrix, its transpose matrix is... Vertical partitioning by column index yields a matrix block array with one row and N columns. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th column is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; The value matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip. For the value matrix V corresponding to each attention head in the key-value matrix, it is horizontally partitioned according to the row index to obtain an N-row, one-column matrix block array. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th row is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; in, ; and Each is a matrix The number of rows and columns; and These are the number of rows and columns of the value matrix V, respectively.
[0028] Secondly, the present invention provides a method for accelerating large language model inference decoding based on a hybrid bonding architecture, which is executed using a hybrid bonding near-memory computing device; the hybrid bonding near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: multiple near-memory computing chips; the near-memory computing chips include: vertically stacked upper DRAM chips and lower logic chips; the DRAM chips include: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory libraries in the corresponding memory blocks; each memory library is hybrid-bonded to its corresponding computing element; preferably, each memory library is hybrid-bonded to its corresponding computing element via Cu-Cu hybrid bonding.
[0029] The above-mentioned methods for accelerating speculative decoding include: S1. Read the weight data of the target large language model TLM and draft model DLM from the DRAM chip, perform the pre-filling calculation of the inference request in parallel based on the embedding vector of each inference request, generate the first output token of the TLM corresponding to each inference request as the current header token of the inference request, and store the key-value matrix of the TLM and DLM corresponding to each token in each inference request obtained during the pre-filling calculation into the DRAM chip. S2, One round of speculative decoding operation: Execute the speculative decoding operation of each inference request in parallel based on the current header token of each inference request; S3. Take the last token in the output token of each inference request in the previous round of speculation and decoding as the current header token of the inference request, and repeat S2 for iteration until an output end symbol is generated or the sum of the number of output tokens of each round of speculation and decoding reaches the preset number. The speculative decoding operation for each inference request r includes a draft prediction operation and a verification operation executed sequentially for the inference request r. The draft prediction operation performed on inference request r includes: reading the weight data of DLM from the DRAM chip, the key-value matrix of the draft model corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of DLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing multiple iterative predictions using DLM based on the read data to generate n draft tokens for inference request r, and storing the key-value matrix of DLM corresponding to each generated draft token into the DRAM chip; the input of the first iterative prediction for inference request r in the draft prediction operation of the first round of speculation decoding operation is the current header token of the inference request; the input of subsequent iterative predictions for inference request r in the draft prediction operation of each round of speculation decoding operation is the output of the previous iterative prediction of the request; The verification operation performed on inference request r includes: reading the weight data of TLM from the DRAM chip, the key-value matrix of TLM corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of TLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing a verification of the n draft tokens of inference request r using TLM based on the read data, and using the verified draft tokens as the output tokens of inference request r in the current round of speculation decoding operation; storing the key-value matrix of TLM corresponding to the output token of inference request r in the current round of speculation decoding operation into the DRAM chip, and deleting the key-value matrix of DLM corresponding to each draft token that failed verification from the DRAM chip; In this process, the key data in the large language model inference decoding acceleration process is stored in a DRAM chip using the data storage method provided in the first aspect of the present invention; the key data includes: weight data of TLM and DLM; or; the key data includes: weight data of TLM and DLM, and key-value matrices in TLM and DLM.
[0030] like Figure 4 The diagram shows the near-memory processing execution flow.
[0031] Preferably, in one optional implementation, the n draft tokens generated for each inference request in the draft prediction operation are in a tree structure; Each memory bank includes a first fixed storage area for storing the key-value matrix of the DLM corresponding to the draft token and a second fixed storage area for storing the key-value matrix of the DLM corresponding to the output token; The draft prediction operation performed on inference request r also includes: When generating the first draft token corresponding to the inference request r, a corresponding logic tree is created. The root node of this logic tree is used to record the address information of the key-value matrix of the DLM corresponding to the current header token of the inference request r in the memory library. When a draft token corresponding to an inference request r is generated, a minimum allocatable unit is obtained from the first fixed storage area of the corresponding memory library and allocated to the key-value matrix of the DLM corresponding to the draft token. A node corresponding to the draft token is established in the corresponding logic tree to record the address information of the key-value matrix of the DLM corresponding to the draft token in the memory library. At the same time, the node corresponding to the draft token in the corresponding logic tree is pointed to its parent node. The verification operations performed on inference request r also include: For the draft token corresponding to the verified inference request r, move the key-value matrix of its corresponding DLM from the first fixed area to the second fixed area of the memory library; The above-mentioned S2 also includes: after the verification operation for all inference requests is completed, releasing all data in the first fixed area of each memory bank.
[0032] Considering that the number of draft tokens generated in each round increases under a tree structure, preferably, the above-mentioned computing elements include: a hybrid bonding controller and multiple multiply-accumulate operation units; The hybrid bonding controller is used to access the memory library that is hybrid-bonded to the computing element it resides in; Multiple multiply-accumulate units are used to perform matrix-vector multiplication or matrix multiplication operations on weight data, key-value matrix data or intermediate activation data read from the corresponding memory library.
[0033] Preferably, in an optional implementation, the hybrid bonded near-memory computing device further includes a neural network processor.
[0034] In one optional implementation, before each round of speculative decoding operation is executed, a corresponding executor is assigned to each operator in that round of speculative decoding operation, including: For each model in the operator set corresponding to this round of speculative decoding operations, compute the operator op and calculate its arithmetic strength. ;when At that time, the near-in-memory computing module is used as the executor of the operator op in this round of speculation decoding operation; when When the neural network processor is used as the executor of the operator op in that round of speculation-decoding operation; when ,and In this case, the device with higher actual computing power in the neural network processor and near-memory computing module when executing operator op is used as the executor for executing operator op in this round of speculation decoding operation; The operator set includes: the fully connected operator FC1 of the TLM, and the attention operator of the TLM corresponding to each inference request r. And the fully connected operator of DLM in the s-th iteration prediction of the draft prediction operation in this round of speculative decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure: the model computation operator set also includes: the DLM attention operator corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure: the model computation operator set also includes: the shared attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. And the independent attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of the speculative decoding operation. ; Arithmetic strength of operator FC1 ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Calculate its arithmetic strength ; in, S represents the number of iterations for each inference request r in the draft prediction operation of this round of speculative decoding. This is the total bandwidth between all memory libraries and their corresponding computing elements in the near-memory computing module. This refers to the external memory access bandwidth of the neural network processor. and B represents the floating-point arithmetic capabilities of the neural network processor and the near-memory computing module, respectively; B represents the number of concurrent inference requests; and the number of tokens. ; and These are the input and output dimensions of the TLM, respectively. The sum of the number of tokens in the inference request r and the number of output tokens of the inference request r obtained from the previous rounds of inference decoding operations; The dimension of the attention head in TLM; the number of draft tokens generated in the s-th iteration prediction when the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure. The value is 1; when the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure, the number of draft tokens generated in the s-th iteration prediction is 1. Let be a preset integer, which satisfies ,and ; and These are the input and output dimensions of the DLM, respectively. This represents the dimension of the attention head in DLM.
[0035] In one optional implementation, the actual computing power of the neural network processor when executing operator op is: ; The actual computing power of the near-memory computing module when executing operator op is: .
[0036] Preferably, in an alternative implementation, in S1 above, a neural network processor is used to perform pre-filling computation of the inference request in parallel based on the embedding vector of each inference request.
[0037] It should be noted that the above reasoning request can be one or more of text, images, audio, and video, and there is no limitation here.
[0038] To further illustrate the large language model inference decoding acceleration provided by the present invention, a detailed description is given below with reference to a specific embodiment 1: This embodiment provides a method for accelerating large language model speculation decoding based on a hybrid bonding architecture, employing a hybrid-bonded near-memory computing device to perform pre-filling computation and speculation decoding operations. In this embodiment, the hybrid-bonded near-memory computing device includes a host, a neural network processor (NPU), and a near-memory computing module; the near-memory computing module includes multiple near-memory computing chips (i.e., HB-NMP chips). The near-memory computing chips are vertically stacked from an upper DRAM wafer and a lower logic wafer via Cu-Cu hybrid bonding. The DRAM chip includes M memory blocks; Each memory block is divided into N memory banks. The logic chip includes M computing units (PUs) corresponding to M memory blocks; each computing unit includes N computing elements (PEs) corresponding to N memory banks in the corresponding memory blocks; each memory bank and its corresponding computing element are connected by Cu-Cu hybrid bonding.
[0039] In this embodiment, the computing unit also includes a global buffer, a dedicated function unit, and a router; The global buffer is used to cache intermediate results generated by multiple computing elements; the dedicated function unit includes one or more of the following: reducer, accumulator, divider, and exponentiation unit, used to perform corresponding calculations (such as reduction, softmax, normalization, or activation function calculation); the router is used to transmit intermediate results or perform aggregate communication operations between different computing units.
[0040] In this embodiment, the computing element includes: a hybrid bonding controller, an input buffer, a weight buffer, an output buffer, and multiple multiply-accumulate operation units; The hybrid bonding controller is used to access the memory bank that is hybrid-bonded to the computing element; Multiple multiply-accumulate units are used to perform matrix-vector multiplication or matrix multiplication operations on weight data, key-value matrix data, or intermediate activation data read from the corresponding memory library.
[0041] The above-mentioned method for accelerating large language model inference decoding may include the following steps: A1. Model Deployment and Data Pre-configuration Phase: The host loads the model parameters of the target large language model TLM and the draft model DLM (i.e., draft language model), and maps the model weight matrix and key-value cache to multiple HB-NMP chips according to the preset communication optimization data mapping strategy.
[0042] Specifically, for the weight matrix used for QKV generation and feedforward network expansion, a mapping method is adopted that first shards by memory channel (i.e. memory block) dimension and then shards by bank dimension; for the weight matrix used for attention output projection and feedforward network back projection, a corresponding sharding method is adopted, that is, first perform row partitioning in the computation unit dimension and then perform column partitioning in the bank dimension.
[0043] For key-value caches, allocation is based on attention heads or attention head groups, with each attention head fixedly assigned to a corresponding memory channel of the near-memory computing chip. Within each memory channel, the key matrix KT is column-partitioned along the bank dimension, and the value matrix V is row-partitioned along the bank dimension. Key-value caches within the same channel are stored in an interleaved manner across all banks of that channel, maintaining load balancing across banks as the sequence length increases. Based on this, key-value caches corresponding to the same attention head can be mapped to the same channel or the same processing unit's associated storage area as much as possible, reducing cross-bank access overhead.
[0044] Simultaneously, a fixed-capacity tree-structured key-value cache pool is pre-allocated. This is achieved by using a fixed region (referred to as the first fixed cache region) within the HB-NMP chip's memory library as the tree-structured key-value cache pool (i.e., the Tree-KV pool). The Tree-KV pool is divided into multiple minimum allocatable units (MAUs) to store the key-value cache corresponding to draft tokens during the speculation decoding process. A corresponding logical tree is also established to record the parent-child topology relationship of draft tokens and the address information of the corresponding MAUs. The total number of minimum allocatable units matches the maximum number of draft tokens that the target large language model can process in a single parallel verification.
[0045] Generally, a logical tree can be established in the host or near-memory computing module. Preferably, in this embodiment, the logical tree is established in the host.
[0046] A2. Pre-filling stage: The embedding vectors of the inference requests are input into the target large language model (TLM) for pre-filling computation. The embedding vectors of the inference requests are obtained as follows: the host receives user inference requests and encodes all requests in parallel to obtain the corresponding embedding vectors.
[0047] In this stage, the TLM performs forward computation on the input embedding vector to generate the first output token, and simultaneously constructs an initial key-value cache corresponding to the input sequence. The initial key-value cache is written into the storage space of multiple HB-NMP chips according to the data mapping method in step A1, serving as the context cache basis for subsequent rounds of speculative decoding operations.
[0048] A3. One round of speculative decoding operation: Operator dynamic scheduling phase: Before each round of speculative decoding operation, a corresponding executor is assigned to each operator in that round of speculative decoding operation. Specifically, the running characteristics of the operators to be executed in the current round are analyzed, the arithmetic strength of the operators to be executed is statistically analyzed, and different operators are assigned to the NPU or HB-NMP chip for execution according to the arithmetic strength-aware adaptive operator scheduling strategy, and the corresponding working mode is switched.
[0049] The hybrid-bonded near-memory computing device supports two operating modes: normal storage and near-memory processing. In normal storage mode, all computing elements in the HB-NMP chip remain inactive, and the neural network processor accesses each DRAM bank via an external interface. In near-memory processing mode, computing elements access each DRAM bank in parallel via the corresponding hybrid-bonding controller, with only one mode enabled at a time to avoid row buffer interference in the DRAM banks.
[0050] Specifically, based on the current batch size (i.e. the number of concurrent inference requests), the number of draft tokens, the current sequence length, and the model dimension parameters, the arithmetic strength of the fully connected operator and the attention operator is calculated; then, combined with the computing power of the NPU and the external memory access bandwidth, the near-memory computing power of the HB-NMP chip and the internal bank-level bandwidth, the expected execution efficiency of each operator to be executed on different hardware units is determined.
[0051] Based on the roofline model, the actual achievable computing power of the operator in different hardware units is calculated: The actual computing power of the neural network processor when executing the operator op is: ; The actual computing power of the near-memory computing module when executing operator op is: .
[0052] By comparing the actual achievable computing power of operators on two types of hardware units, operators are dynamically allocated to hardware units with higher computing power for execution. Specifically: For operators that are memory-restricted in the near-memory computing module (i.e., satisfying...) Operators that can fully utilize the high internal bandwidth of the near-memory computing module are therefore assigned to the near-memory computing module for execution; conversely, operators that are computationally constrained on the NPU (i.e., satisfying...) are assigned to the near-memory computing module for execution. If an operator is in a computationally constrained state on the near-memory computing module (i.e., satisfies...), it is mapped to the NPU for execution to leverage the NPU's superior computing power; At the same time, it is in a memory-restricted state on the NPU (i.e., it meets the following conditions). When hybrid-bonded near-memory computing devices are used, they will directly compare... and The value will map the operator to a device with higher performance.
[0053] Parallelization strategies can be summarized into three scenarios: when the total number of draft tokens is small, both the fully connected operator and the attention operator are allocated to the near-memory computing module; when the number of draft tokens per request is moderate, the fully connected operator is allocated to the neural network processor, and the attention operator is allocated to the near-memory computing module; when the number of draft tokens per request is large, the fully connected operator and the shared attention operator are allocated to the neural network processor, and the independent attention operator is allocated to the hybrid bonded near-memory computing system.
[0054] For bandwidth-sensitive operators, they are scheduled to be executed on the HB-NMP chip; for computationally intensive operators, they are scheduled to be executed on the NPU; for operators near different resource bottleneck boundaries, the execution location is selected based on the estimated execution performance results.
[0055] Once the execution location is determined, the corresponding HB-NMP chip is controlled to operate in normal storage mode or near-memory processing mode to support heterogeneous execution in the subsequent verification stage.
[0056] Draft tree prediction stage: Using the output token generated in the previous stage or the valid token obtained in the previous verification stage as the current input, the draft model DLM is called to perform draft prediction.
[0057] In this embodiment, DLM performs S rounds of iterative prediction. In each round, n candidate draft tokens are generated for the current layer node, constructing a draft token tree containing multi-path candidate sequences. .
[0058] When generating each draft token, a minimum allocatable unit (MAU) is allocated from the first fixed cache area for each draft token to store the corresponding key-value cache generated by the draft token during the DLM or verification process; at the same time, a node corresponding to the draft token is established in the logical tree, recording the MAU address of the corresponding draft token, the parent node identifier and its topological relationship in the token tree, so as to track the topological dependency relationship through the logical tree.
[0059] For multiple candidate sequences with a shared prefix, the key-value cache of the shared prefix portion can be associated with the same prefix node, thus avoiding duplicate storage.
[0060] In this embodiment, the process of generating n draft tokens in a tree structure for inference request r includes: Using the current header token of reasoning request r as the root node, construct the draft token tree corresponding to the reasoning request r; the root node does not include the above n draft tokens; Suppose that the draft prediction operation includes S iterations of prediction, and the s-th iteration corresponds to the s-th level node of the draft token tree, s=1,2,…,S; In the first iteration of prediction, the current head token is input into DLM to obtain one or more candidate tokens, and at least one token is selected from the candidate tokens as the child node of the root node; In the s-th iteration prediction, when s>1, the s-th iteration prediction will be... Each node in layer 1 serves as a parent node, and based on the token corresponding to each parent node and the historical tokens on the path from the root node to that parent node, DLM is called to predict the next token, resulting in one or more candidate tokens for each parent node; then at least one candidate token is selected for each parent node as its child node. For each selected candidate token, a node corresponding to the candidate token is created in the draft token tree, and the pointing relationship between the node and its parent node is recorded, such that the path from the root node to any node represents a draft token prefix, and the path from the root node to any leaf node represents a draft token sequence to be verified; the same prefix tokens in different draft token sequences are shared through the common ancestor node in the draft token tree. When the number of non-root nodes in the draft token tree reaches n, or the number of iterations reaches S, or the preset termination condition is met, the draft prediction operation stops, resulting in a tree structure composed of n draft tokens.
[0061] Parallel validation phase of the target model: The draft token tree, which is constructed during the draft tree prediction phase, is input into the target large language model TLM, and the TLM performs parallel verification on multiple draft tokens in the draft token tree.
[0062] During verification, the TLM performs forward computation on each candidate branch based on the topological relationship of the token tree, and compares whether the candidate tokens at each level are consistent with the actual output of the TLM.
[0063] For a given branch, the following checks are performed sequentially starting from the root node: when the current draft token matches the TLM output, the draft token is accepted, and the verification of its successor nodes continues; when the first inconsistent draft token appears, the draft token is rejected, and the acceptance check of subsequent nodes of that branch is stopped.
[0064] A4. Iterative Update and Result Output Stage: The draft token verified by TLM for each request and the correction token given by TLM in the current round are concatenated into a complete generated result and output. The last valid token (i.e. the output token) is used as the input for the next round of speculation decoding operation.
[0065] Repeat step A3 until an end symbol is generated or the preset maximum output length is reached. Specifically, in one optional implementation, when the input sequence length (i.e., the number of tokens in the inference request) is... At that time, the maximum output length can be set to The maximum output length is 3072 tokens. For example, when the input sequence length is 1024 tokens, the maximum output length is 3072 tokens.
[0066] After decoding is complete, the host concatenates the valid tokens generated in all rounds to obtain the final output and returns it to the user.
[0067] After the verification operations for all inference requests are completed, the key-value cache update and storage reclamation phase is performed: After the target model is validated, different processing is performed on the nodes in the draft token tree that are accepted (validated successfully) and rejected (validated unsuccessfully).
[0068] For an accepted draft token, its corresponding MAU is directly incorporated into the storage link or mapping relationship corresponding to the main sequence key-value cache (that is, the key-value matrix of its corresponding DLM is moved from the first fixed area to the second fixed area in the corresponding memory library for storing the key-value matrix of the DLM corresponding to the output token), so that this part of the key-value cache is transformed into a valid context cache that can be reused in the subsequent decoding process without the need for additional data copying.
[0069] For rejected draft tokens, the node and all its descendant nodes are identified based on the topological relationship recorded in the logical tree. All corresponding minimum allocable units are released in batches for repeated allocation in subsequent rounds. The released minimum allocable units are then reintegrated into the tree-structured key-value cache pool, completing closed-loop capacity management.
[0070] During the release process, only the storage resources corresponding to the rejected branches are reclaimed, while the key-value caches corresponding to the accepted branches and their common prefixes are retained to reduce storage fragmentation and redundant allocation overhead.
[0071] In this embodiment, by pre-mapping model parameters and key-value caches to multiple HB-NMP chips, and combining tree structure draft generation, arithmetic strength-aware heterogeneous scheduling, and a logic tree-based Tree-KV pool management mechanism, the parallel verification efficiency of the speculative decoding stage can be improved, cross-bank data access overhead can be reduced, and the storage management burden during the tree structure draft branch recycling process can be reduced, while ensuring the consistency of the target model output.
[0072] In summary, this embodiment employs hybrid bonding technology, customizing the computing unit in the logic chip to enhance near-memory computing power and adapt to the high computational intensity characteristics of tree-structured speculative decoding. This invention utilizes a communication-optimized data mapping strategy to map the weight matrix and key-value cache to the channels and banks of the DRAM chip using an alternating row and column partitioning method, minimizing cross-bank access overhead. An arithmetic strength-aware operator scheduler is designed to dynamically calculate the arithmetic strength of operators based on the parallelism of tree-structured tokens, allocating operators to neural network processors or near-memory computing modules for execution. Simultaneously, a tree-structured key-value cache pool management mechanism is introduced, implementing cascaded invalidation and memory reuse of the key-value cache based on the logic tree, eliminating memory fragmentation caused by discarded branches. This embodiment fully leverages the high bandwidth and high computing power advantages of the hybrid bonding architecture, significantly improving the inference throughput and energy efficiency of speculative decoding, and is particularly suitable for the low-latency deployment requirements of large language models at the edge.
[0073] Thirdly, the present invention provides a large language model speculation decoding acceleration system based on a hybrid bonding architecture, comprising: a hybrid bonding near-memory computing device; wherein, the hybrid bonding near-memory computing device comprises: a near-memory computing module; wherein, the near-memory computing module comprises: multiple near-memory computing chips; the near-memory computing chip comprises: a vertically stacked upper DRAM chip and a lower logic chip; the DRAM chip comprises: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. The aforementioned hybrid-bonded near-memory computing device is used to execute the large language model inference decoding acceleration method provided in the second aspect of the present invention.
[0074] The related technical solutions are the same as the large language model inference decoding acceleration method provided in the second aspect of this invention, and will not be described in detail here.
[0075] To further illustrate the large language model inference decoding acceleration system provided by the present invention, a detailed description is given below with reference to a specific embodiment 2: The large language model inference decoding acceleration system provided in this embodiment is used to execute the large language model inference decoding acceleration method provided in Embodiment 1, such as... Figure 5 As shown, it includes a host computer, a neural network processor (NPU), a near-memory computing module including multiple HB-NMP chips, and a target large language model (TLM) and a draft model (DLM) running on the hardware.
[0076] The host is connected to the NPU and HB-NMP chips via a high-speed interconnect bus. The host is responsible for receiving user inference requests, performing request encoding and embedding generation (i.e., encoding each inference request to obtain the embedding vector of each inference request), scheduling the entire inference process, and managing the construction of the token tree, mask generation, and token acceptance / rejection logic.
[0077] The Neural Processing Unit (NPU) is used to execute computationally intensive operators. It adopts a TPU architecture and has multiple 128×128 systolic arrays and SIMD vector processing units built in. With a main frequency of 1GHz, it is used to execute computationally intensive FC operators and shared attention operators, and is also responsible for the working mode switching, task distribution and result aggregation of the HB-NMP chip.
[0078] The near-memory computing module executes bandwidth-sensitive operators and also serves as the main memory system for the neural network processor. Each HB-NMP chip is vertically stacked from an upper-layer DRAM chip to a lower-layer logic chip via Cu-Cu hybrid bonding, wherein: A single DRAM chip is divided into 8 memory channels (i.e. memory blocks), each memory block contains 16 memory banks, each bank has a capacity of 32MB, and the total capacity of a single DRAM chip is 4GB (in this embodiment there are 4 DRAM chips, and the total capacity of the near-memory computing module is 16GB), providing ultra-high internal memory access bandwidth; The logic chip is divided into 8 computing units (PUs), which correspond one-to-one with the 8 memory channels in the DRAM; for example Figure 6 As shown, each PU contains 16 computing elements (PEs), which are paired one-to-one with 16 memory banks within the channel; as... Figure 7 As shown, each PE has 8 built-in multiply-accumulate units (GEMV units in this embodiment), which can perform multi-way multiply-accumulate operations in a single cycle. All computing elements within the computing unit can execute computing tasks synchronously and in parallel, providing aggregated bandwidth and parallel computing power support for near-memory processing.
[0079] Each PU is also configured with a global buffer, a dedicated function unit (SFU), and a router. The global buffer is an on-chip storage resource shared by all computing elements within the computing unit, used to cache intermediate computation results during operator execution. The dedicated function unit (SFU) is used to perform nonlinear operator computations other than matrix multiplication and accumulation. It has a built-in reducer, accumulator, divider, and exponentiation unit, supporting hardware acceleration for functions such as softmax, SiLU, and GeLU. The router integrates a local switching switch to achieve efficient data routing and interaction between different processing units; specifically, the router implements efficient data routing between PUs and supports aggregated communication operations such as all-reduce.
[0080] The aforementioned computing elements are the basic computational execution cores within the computing unit. Each computing element is exclusively paired with a DRAM bank within its corresponding DRAM channel, and the two are interconnected via dedicated hybrid-bonded I / O and a hybrid-bonded controller. Each computing element integrates multiple general-purpose matrix-vector multiplication units, enabling multi-way multiplication and accumulation operations to be completed within a single clock cycle, providing fundamental computational power support for the core matrix operations of fully connected operators and attention operators. The operating state of the computing elements is linked to the operating mode of the HB-NMP chip. In near-memory processing mode, the computing elements access their paired DRAM banks in parallel through a dedicated hybrid-bonded controller, performing memory-intensive operator computations locally. In normal storage mode, all computing elements remain inactive, without interfering with DRAM bank access. The hardware design of the computing elements eliminates the on-chip network module, and their dedicated execution flow does not involve communication between different computing elements within the same computing unit, saving chip area for expanding computational logic and increasing computational density.
[0081] The system has three core functional modules built in: The communication-optimized data mapping module is used to map model weights and key-value caches to the near-memory computing module using a specific partitioning strategy. Specifically, it performs communication-optimized sharding mapping between the weight matrix and the key-value cache to complete the layout optimization of model parameters in the HB-NMP chip.
[0082] Arithmetic strength-aware operator scheduling module: This module dynamically allocates operators to the optimal hardware unit for execution based on their arithmetic strength at runtime. Specifically, it collects the parallelism information of the current token tree in real time and calculates the arithmetic strength of each operator, such as... Figure 8 As shown, the dynamic mapping from the operator to the NPU or HB-NMP chip is completed based on the roofline model; The tree-structured key-value cache pool management module is used to allocate storage units for tokens and combine a logical tree structure to allocate and reclaim key-value caches for draft tokens in order to eliminate memory fragmentation. Specifically, it maintains a fixed-capacity tree-structured key-value cache pool and logical tree to achieve fast allocation, cascading release, and fragmentation-free management of key caches.
[0083] In this embodiment, the system can flexibly expand the number of HB-NMP chips according to the model size, supporting the edge deployment of different mainstream large language models; through pipelined parallelism of the model layer, different Transformer layers are mapped to different HB-NMP chips, further improving the inference parallelism.
[0084] This embodiment is based on a heterogeneous accelerator designed with a hybrid bonding architecture. The computing elements are customized in the logic chip design, which breaks through the computing power limitation of traditional DRAM-PIM. It can fully adapt to the high computing intensity characteristics of tree structure speculation and decoding, while taking into account the high bandwidth requirements of memory-intensive operators, and achieve synergistic optimization of computing power and bandwidth.
[0085] Those skilled in the art will readily understand that the above are merely preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.
Claims
1. A data storage method, characterized in that, include: Store key data in the process of accelerating the inference and decoding of large language models; The large language model inference decoding acceleration process is performed using a hybrid bonded near-memory computing device; Hybrid bonding A near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: a plurality of near-memory computing chips; the near-memory computing chip includes: a vertically stacked upper DRAM chip and a lower logic chip; the DRAM chip includes: M memory blocks; ; each memory block is divided into N memory banks; ; the logic chip includes M computing units corresponding to the M memory blocks; the computing unit includes: N computing elements corresponding to the N memory banks in the corresponding memory block; each memory bank is connected to the corresponding computing element by hybrid bonding; The key data includes: weight data for the target large language model (TLM) and the draft model (DLM); the weight data for the TLM and DLM are stored in a DRAM chip in the following manner: For each model in TLM and DLM, for each first-class matrix within it and each second-class matrix Both horizontal partitioning by row index and vertical partitioning by column index are performed to divide each type I matrix. Each matrix is divided into an N-row, M-column first matrix block array and each second-class matrix The first matrix block array is divided into M rows and N columns; the size of each matrix block in the first matrix block array does not exceed [a certain value]. The size of each matrix block in the second matrix block array does not exceed [a certain value]. ; Store the matrix block in the j-th row and k-th column of each first matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; store the matrix block in the k-th row and j-th column of each second matrix block array into the j-th memory bank of the k-th memory block of the DRAM chip; ; ; and They correspond to as The number of rows and columns; This represents the floor function; and They correspond to as The number of rows and columns; The first type of matrix includes: query weight matrix Key weight matrix Value weight matrix The weight matrix of the first layer in the feedforward neural network. ; The second type of matrix includes: the output projection weight matrix in the attention layer. and the weight matrix of the second layer in the feedforward neural network layer .
2. The data storage method according to claim 1, characterized in that, Key data also includes: key-value matrices in the TLM and DLM generated during the large language model inference decoding acceleration process; the key-value matrices are stored in the DRAM chip in the following manner: When the number of attention heads in the key-value matrix is less than or equal to M, one attention head in the key-value matrix is considered as one attention head group; when the number of attention heads in the key-value matrix is greater than M, the attention heads in the key-value matrix are grouped to obtain M' attention head groups. ; The key matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip; for each attention head corresponding to the key matrix K in the key-value matrix, its transpose matrix is... Vertical partitioning by column index yields a matrix block array with one row and N columns. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th column is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; The value matrices corresponding to different attention head groups in the key-value matrix are stored in different memory blocks in the DRAM chip. For the value matrix V corresponding to each attention head in the key-value matrix, it is horizontally partitioned according to the row index to obtain an N-row, one-column matrix block array. ; The size of each matrix block in the array does not exceed [a certain value]. ;Will The matrix block under the j-th row is stored in the j-th memory bank of the corresponding memory block in the DRAM chip; in, and Each is a matrix The number of rows and columns; and These are the number of rows and columns of the value matrix V, respectively.
3. A method for accelerating large language model inference decoding based on a hybrid bonding architecture, characterized in that, Execution is performed using a hybrid bonded near-memory computing device; Hybrid bonding The near-memory computing device includes: a near-memory computing module; the near-memory computing module includes: multiple near-memory computing chips; the near-memory computing chips include: vertically stacked upper DRAM chips and lower logic chips; the DRAM chips include: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. Speculative decoding acceleration methods include: S1. Read the weight data of the target large language model TLM and draft model DLM from the DRAM chip, perform the pre-filling calculation of the inference request in parallel based on the embedding vector of each inference request, generate the first output token of the TLM corresponding to each inference request as the current header token of the inference request, and store the key-value matrix of the TLM and DLM corresponding to each token in each inference request obtained during the pre-filling calculation into the DRAM chip. S2, One round of speculative decoding operation: Execute the speculative decoding operation of each inference request in parallel based on the current header token of each inference request; S3. Take the last token in the output token of each inference request in the previous round of speculation and decoding as the current header token of the inference request, and repeat S2 for iteration until an output end symbol is generated or the sum of the number of output tokens of each round of speculation and decoding reaches the preset number. The speculative decoding operation for each inference request r includes a draft prediction operation and a verification operation executed sequentially for the inference request r. The draft prediction operation performed on inference request r includes: reading the weight data of DLM from the DRAM chip, the key-value matrix of the draft model corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of DLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing multiple iterative predictions using DLM based on the read data to generate n draft tokens for inference request r, and storing the key-value matrix of DLM corresponding to each generated draft token into the DRAM chip; the input of the first iterative prediction for inference request r in the draft prediction operation of the first round of speculation decoding operation is the current header token of the inference request; the input of subsequent iterative predictions for inference request r in the draft prediction operation of each round of speculation decoding operation is the output of the previous iterative prediction of the request; The verification operation performed on inference request r includes: reading the weight data of TLM from the DRAM chip, the key-value matrix of TLM corresponding to each token in inference request r obtained during the pre-filling calculation process, and the key-value matrix of TLM corresponding to the output token of inference request r obtained in each round of speculation decoding operation before the current round of speculation decoding operation; then performing a verification of the n draft tokens of inference request r using TLM based on the read data, and using the verified draft tokens as the output tokens of inference request r in the current round of speculation decoding operation; storing the key-value matrix of TLM corresponding to the output token of inference request r in the current round of speculation decoding operation into the DRAM chip, and deleting the key-value matrix of DLM corresponding to each draft token that failed verification from the DRAM chip; The key data in the large language model inference decoding acceleration process is stored in a DRAM chip using the data storage method described in claim 1 or 2; the key data includes: weight data of TLM and DLM; or; the key data includes: weight data of TLM and DLM, and key-value matrices in TLM and DLM.
4. The method for accelerating large language model inference decoding according to claim 3, characterized in that, In the draft prediction operation, the n draft tokens generated for each inference request are arranged in a tree structure. Each memory bank includes a first fixed storage area for storing the key-value matrix of the DLM corresponding to the draft token and a second fixed storage area for storing the key-value matrix of the DLM corresponding to the output token; The draft prediction operation performed on inference request r also includes: When generating the first draft token corresponding to the inference request r, a corresponding logic tree is created. The root node of this logic tree is used to record the address information of the key-value matrix of the DLM corresponding to the current header token of the inference request r in the memory library. When a draft token corresponding to an inference request r is generated, a minimum allocatable unit is obtained from the first fixed storage area of the corresponding memory library and allocated to the key-value matrix of the DLM corresponding to the draft token. A node corresponding to the draft token is established in the corresponding logic tree to record the address information of the key-value matrix of the DLM corresponding to the draft token in the memory library. At the same time, the node corresponding to the draft token in the corresponding logic tree is pointed to its parent node. The verification operations performed on inference request r also include: For the draft token corresponding to the verified inference request r, move the key-value matrix of its corresponding DLM from the first fixed area to the second fixed area of the memory library; S2 also includes: after the verification operation for all inference requests is completed, releasing all data in the first fixed area of each memory bank.
5. The method for accelerating large language model inference decoding according to claim 4, characterized in that, The computing element includes: a hybrid bonding controller and multiple multiply-accumulate operation units; The hybrid bonding controller is used to access the memory library that is hybrid-bonded to the computing element in which it resides; The plurality of multiply-accumulate operation units are used to perform matrix-vector multiplication or matrix multiplication operations on the weight data, key-value matrix data or intermediate activation data read from the corresponding memory library.
6. The method for accelerating large language model inference decoding according to any one of claims 3-5, characterized in that, Hybrid-bonded near-memory computing devices also include: neural network processors.
7. The method for accelerating large language model inference decoding according to claim 6, characterized in that, Also includes: Before each round of speculative decoding operations, a corresponding executor is assigned to each operator in that round of speculative decoding operations, including: For each model in the operator set corresponding to this round of speculative decoding operations, compute the operator op and calculate its arithmetic strength. ;when At that time, the near-in-memory computing module is used as the executor of the operator op in this round of speculation decoding operation; when When the neural network processor is used as the executor of the operator op in that round of speculation-decoding operation; when ,and In this case, the device with higher actual computing power in the neural network processor and near-memory computing module when executing operator op is used as the executor for executing operator op in this round of speculation decoding operation; The operator set includes: the fully connected operator FC1 of the TLM, and the attention operator of the TLM corresponding to each inference request r. And the fully connected operator of DLM in the s-th iteration prediction of the draft prediction operation in this round of speculative decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure: the model computation operator set also includes: the DLM attention operator corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. ; When the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure: the model computation operator set also includes: the shared attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of this round of speculation decoding operation. And the independent attention operator of DLM in the s-th iteration prediction corresponding to each inference request r in the draft prediction operation of the speculative decoding operation. ; Arithmetic strength of operator FC1 ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Arithmetic strength ; Operator Calculate its arithmetic strength ; in, S represents the number of iterations for each inference request r in the draft prediction operation of this round of speculative decoding. This is the total bandwidth between all memory libraries and their corresponding computing elements in the near-memory computing module. This refers to the external memory access bandwidth of the neural network processor. and B represents the floating-point arithmetic capabilities of the neural network processor and the near-memory computing module, respectively; B represents the number of concurrent inference requests; and the number of tokens. ; and These are the input and output dimensions of the TLM, respectively. The sum of the number of tokens in the inference request r and the number of output tokens of the inference request r obtained from the previous rounds of inference decoding operations; The dimension of the attention head in TLM; the number of draft tokens generated in the s-th iteration prediction when the n draft tokens generated for each inference request r in the draft prediction operation are in a non-tree structure. The value is 1; when the n draft tokens generated for each inference request r in the draft prediction operation are in a tree structure, the number of draft tokens generated in the s-th iteration prediction is 1. Let be a preset integer, which satisfies ,and ; and These are the input and output dimensions of the DLM, respectively. This represents the dimension of the attention head in DLM.
8. The method for accelerating large language model inference decoding according to claim 7, characterized in that, The actual computing power of a neural network processor when executing an operator (op) is: ; The actual computing power of the near-memory computing module when executing operator op is: .
9. The method for accelerating large language model inference decoding according to claim 6, characterized in that, In S1, a neural network processor is used to perform pre-filling computations for each inference request in parallel based on the embedding vector of that inference request.
10. A large language model inference decoding acceleration system based on a hybrid bonding architecture, characterized in that, include: Hybrid-bonded near-memory computing devices; Hybrid bonding The near-memory computing device includes: a near-memory computing module; wherein, the near-memory computing module includes: multiple near-memory computing chips; the near-memory computing chips include: vertically stacked upper DRAM chips and lower logic chips; the DRAM chips include: M memory blocks; Each memory block is divided into N memory banks; The logic chip includes M computing units corresponding to M memory blocks; each computing unit includes N computing elements corresponding to N memory banks in the corresponding memory blocks; each memory bank is co-bonded to its corresponding computing element. A hybrid-bonded near-memory computing device is used to execute the large language model speculation decoding acceleration method of any one of claims 3-9.