Large language model accelerator system based on three-dimensional NAND flash memory

Abstract

The application discloses a large language model accelerator system based on three-dimensional NAND flash memory, and belongs to the technical field of calculation, estimation or counting. The system comprises three-dimensional NAND flash memory for performing feedforward neural network calculation, an auxiliary calculation unit for performing attention mechanism calculation, a DRAM chip for storing attention mechanism related weights and KV cache, and interconnection resources for realizing data interaction between components. The three-dimensional NAND flash memory comprises a logic chip and a NAND array chip, the logic chip is used for controlling and performing calculation, and the NAND array chip is used for storing weights and participating in calculation. The application further proposes a scheduling method based on KV cache perception, comprising decomposition and dynamic allocation of calculation tasks. Through the collaborative design of the hardware architecture and the scheduling method, the application can alleviate the memory wall bottleneck in large language model calculation, reduce energy consumption and improve overall calculation performance.

Description

Technical Field

[0001] This invention relates to integrated circuits and artificial intelligence computing technology, and in particular discloses a large language model accelerator system based on three-dimensional NAND flash memory computing and its dynamic scheduling method, belonging to the technical field of computing, calculation or counting. Background Technology

[0002] This section is intended to provide background or context for the technical solutions of the invention as set forth in the claims. The description herein does not imply acceptance that it is prior art simply because it is included in this section.

[0003] With the rapid development of natural language processing, large language models (LLMs) such as the GPT, LLaMa, and Deepseek series place extremely high demands on computing power and storage bandwidth. These models typically contain billions to hundreds of billions of parameters, and their training and inference processes involve a large number of matrix-vector multiplication and addition calculations as well as attention mechanism computations. Traditional implementations based on GPUs or general-purpose accelerators face bottlenecks in computing power and storage. For example, although the peak computing power of GPUs can exceed hundreds of TOPS, reaching trillions of operations per second, when the GPU's video memory cannot hold all the model parameters, some data needs to be offloaded to DRAM or even SSDs with relatively low bandwidth. Taking the PCIe 4.0 x4 interface as an example, its bandwidth is only 8GB / s, while the model parameter size can reach tens or even hundreds of GB, resulting in high latency and low energy efficiency.

[0004] In recent years, Processing-In-Memory (PIM) computing technology has gained widespread attention due to its ability to alleviate the memory wall problem in data-intensive applications. PIM integrates computing units within the memory chip, reducing the overhead of moving data to the processor and facilitating low-power computation of large-scale matrix multiplication and addition. However, DRAM, as a volatile medium, requires periodic refreshes to maintain data state during operation, and data is lost upon power failure. Furthermore, DRAM has low storage density, and the capacity of a single chip is insufficient to store large model parameters. Therefore, DRAM-based PIM architectures typically require expensive Through Silicon Via (TSV) technology to bond stacked DRAM layers, making it difficult to reduce power consumption and cost, especially posing challenges for edge applications.

[0005] With the development of 3D NAND flash memory technology, its enormous capacity and internal bandwidth characteristics reveal its potential as a non-volatile storage medium for large language model flash computing. Performing weighted accumulation calculations at the flash array control logic level to directly execute multiply-accumulate operations of feedforward neural networks within the memory holds promise for overcoming performance bottlenecks caused by data transmission and significantly reducing chip area, power consumption, and cost.

[0006] In summary, achieving a complete and efficient inference architecture for large language models with tens or even hundreds of billions of data points mainly faces the following challenges:

[0007] The DRAM-based architecture still needs to be initialized with non-volatile low-speed storage, and traditional storage media can hardly meet the storage capacity and access bandwidth requirements brought about by the increase in the size of large language models.

[0008] The NAND-based architecture has high weight read latency and an imperfect data update strategy between on-chip cache and flash array, which may lead to performance bottlenecks during multi-layer neural network inference.

[0009] The storage requirements and data flow scheduling of attention mechanism computation are not fully optimized, resulting in frequent data movement between auxiliary computing units and storage units;

[0010] The lack of a dynamic scheduling mechanism for computation of multi-layer feedforward neural networks and attention mechanisms makes it impossible to fully utilize the internal parallel computing resources of flash memory, resulting in resource waste.

[0011] To address the aforementioned issues, it is necessary to design a novel large language model accelerator system based on three-dimensional NAND flash memory (In-Flash Computing / Processing-In-Flash, IFC / PIF) and a corresponding dynamic scheduling method to efficiently support attention mechanism computation and feedforward neural network computation, thereby improving the performance and energy efficiency of large language model inference. Summary of the Invention

[0012] The purpose of this invention is to provide a large language model accelerator system based on three-dimensional NAND flash memory and its dynamic scheduling method, so as to solve the technical problems of storage bandwidth bottleneck, insufficient utilization of computing resources and large data migration overhead in the existing technology in feedforward neural network computing and attention mechanism computing, thereby achieving the invention objective of improving the inference performance and energy efficiency of large language models.

[0013] To achieve the above objectives, the present invention proposes the following technical solution:

[0014] A large language model accelerator system based on 3D NAND flash memory includes: DRAM, an auxiliary computing unit, a NAND array chip, a flash memory logic chip, a first interconnect channel, a second interconnect channel, and a third interconnect channel. The DRAM stores attention weights and key and value vectors generated during attention calculation. The auxiliary computing unit reads attention weights from the DRAM, performs attention calculation, writes the key and value vectors generated during attention calculation back to the DRAM, outputs attention activation results, and obtains decoding results or performs the next attention calculation based on the feedforward neural network calculation results. The NAND array chip stores large language model parameters. The flash memory logic chip receives attention activation results to trigger feedforward neural network calculation, reads feedforward neural network weights and the final output projection weights from the large language model parameters stored in the NAND array chip, and writes the feedforward neural network calculation results back to the auxiliary computing unit. The first interconnect channel couples the auxiliary computing unit and the flash memory logic chip. The second interconnect channel couples the DRAM and the auxiliary computing unit. The third interconnect channel couples the flash memory logic chip and the NAND array chip.

[0015] As a further optimization of the large language model accelerator system based on three-dimensional NAND flash memory, the NAND array chip contains at least one cluster of shared read address decoders and partial read address circuits. Each cluster contains at least one storage plane. The smallest set of weights of the large language model that participates in the calculation consecutively is taken as weight atom blocks. Based on the principle of continuous storage location of weight atom blocks, the mapping relationship between weight atom blocks and clusters and the mapping relationship between weight atom blocks and storage planes are established.

[0016] As a further optimization of the large language model accelerator system based on 3D NAND flash memory, the auxiliary computing unit includes: an interface module, a control module, and an attention calculation mechanism computing module; the interface module is used to couple the flash memory logic chip through the first interconnect channel; the control module is used to manage the write-back of key vectors and value vectors to DRAM, adjust the scheduling graph and generate cooperative control signals after real-time sensing of the KV cache, the scheduling graph is a bitmap that instructs the auxiliary computing unit and the flash memory logic chip to cooperate in performing attention calculation tasks, and the cooperative control signals are instructions that instruct the flash memory logic chip to perform a unit attention calculation task; the attention calculation mechanism computing module is used to complete the unit attention calculation task executed by the auxiliary computing unit as instructed by the scheduling graph, and obtain the decoding result or execute the next attention calculation based on the calculation result of the feedforward neural network.

[0017] As a further optimization of the large language model accelerator system based on 3D NAND flash memory, the flash memory logic chip includes: an interface unit, a control unit, an on-chip cache system, a bit error correction subunit, an on-chip computation subunit, and a nonlinear computation subunit. The interface unit is used to couple to the auxiliary computation unit via a first interconnect channel, receive attention activation results and cooperative control signals, and transmit feedforward neural network computation results back to the auxiliary computation unit. The control unit is used to manage access to the NAND array chip and the on-chip cache system, and manage matrix multiplication and addition operations in the unit attention computation task corresponding to the cooperative control signals. The on-chip cache system is used to store the feedforward neural network weights, final output projection weights, and attention weights read from the NAND array chip. The bit error correction subunit is used to process the feedforward neural network weights and final output projection weights... The system performs error correction on the weights and attention weights; the flash memory computing subunit is used to perform matrix-vector multiplication and addition on the attention activation values ​​and the corrected feedforward neural network weights, or to perform matrix-vector multiplication and addition on the output of the nonlinear computing subunit and the corrected feedforward neural network weights, or to perform matrix-vector multiplication and addition on the unit attention calculation task corresponding to the cooperative control signal based on the corrected attention weights, or to perform matrix-vector multiplication and addition on the output of the last layer of the feedforward neural network and the projected weights of the final output after error correction, write the matrix-vector multiplication and addition calculation results into the on-chip cache system, and transmit the final result of the feedforward neural network to the interface unit; the nonlinear computing subunit is used to perform activation or normalization operations on the matrix-vector multiplication and addition calculation results related to the feedforward neural network weights read from the on-chip cache system, and transmit the intermediate result of the feedforward neural network to the interface unit.

[0018] As a further optimization scheme for a large language model accelerator system based on 3D NAND flash memory, the control module adjusts the scheduling graph and generates cooperative control signals after real-time sensing of the KV cache. Specifically, it monitors the total time of the attention calculation mechanism in each step of the attention calculation module and calculates the difference between the total time of two adjacent attention calculation steps. Exceeding the threshold When the bitmap rewrite operation is triggered, based on the difference... and threshold Calculate adjustment force The intensity will be adjusted. The corresponding unit attention calculation task is configured for the NAND array chip.

[0019] As a further optimization scheme for a large language model accelerator system based on 3D NAND flash memory, the threshold is... P is the cluster page cache capacity, and u is the weighted atomic block size. To mitigate the delay in the computational unit's processing of weighted atom block data; the computational adjustment force is... , This is for rounding up.

[0020] As a further optimization scheme for the large language model accelerator system based on three-dimensional NAND flash memory, the computation sub-units in the flash memory adopt a rectangular or non-rectangular array structure based on multiplier-adder primitives; or, adopt an SRAM in-memory computation macro array.

[0021] As a further optimization scheme for a large language model accelerator system based on three-dimensional NAND flash memory, the key vectors and value vectors of different attention heads are mapped to data blocks or channels in DRAM, and the key vectors and value vectors that are frequently accessed by the auxiliary computing unit are mapped to the on-chip cache of the auxiliary computing unit. The key vectors and value vectors for the next attention calculation are prefetched before the previous attention calculation is finished.

[0022] As a further optimization of the large language model accelerator system based on 3D NAND flash memory, the second interconnect channel accesses DRAM data blocks in blocks. Based on the timing arbitration of the data block transmission order of the multi-head attention access request, the write-back operation of non-critical data blocks is postponed when the read traffic of the auxiliary computing unit exceeds the limit.

[0023] As a further optimization of the large language model accelerator system based on three-dimensional NAND flash memory, the third interconnect channel uses hybrid bonding technology to couple the flash memory logic chip and the NAND array chip.

[0024] The present invention adopts the above technical solution, and the proposed architecture and dynamic scheduling method have the following significant advantages compared with the prior art.

[0025] (1) The accelerator system proposed in this invention makes full use of the characteristics of three-dimensional NAND flash memory and hybrid bonding, decouples the feedforward neural network computation and attention mechanism computation to the hardware part with efficient division of labor and cooperation, and distributes the bandwidth pressure to the three interconnect channels.

[0026] (2) The accelerator system proposed in this invention performs matrix-vector multiplication and addition directly inside the NAND memory, reducing the overhead of model weight transmission, alleviating the memory wall bottleneck of the traditional von Neumann computing architecture, reducing data migration overhead, and improving energy efficiency.

[0027] (3) The accelerator system proposed in this invention supports pipelined flash memory access and dynamic scheduling, avoiding the decrease in architecture throughput caused by the mismatch between computation latency and flash memory access latency and the increase in computation in long context reasoning scenarios.

[0028] (4) The control functions in the flash memory and auxiliary computing unit of the accelerator system proposed in this invention are both relatively independent and closely coordinated, which improves the versatility and efficiency of the large language model inference solution based on this architecture, while reducing the hardware implementation cost, especially providing an innovative solution for edge inference scenarios.

[0029] In summary, this invention provides an effective solution for efficient hardware acceleration of large language models by combining flash memory computing with auxiliary computing units and introducing a dynamic scheduling mechanism. Attached Figure Description

[0030] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are merely some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort. The drawings in this invention are as follows:

[0031] Figure 1 This is a schematic diagram of a large language model accelerator system based on three-dimensional NAND flash memory provided in one embodiment of the present invention.

[0032] Figure 2 This is a flowchart of a basic scheduling algorithm provided in one embodiment of the present invention.

[0033] Figure 3 This is a schematic diagram of data flow in a conventional large language model inference hardware scheme.

[0034] Figure 4 This is a schematic diagram of the data flow of an inference scheme provided in one embodiment of the present invention.

[0035] Figure 5 This is a schematic diagram of a first interconnection channel and its main functions provided in one embodiment of the present invention.

[0036] Figure 6 This is a schematic diagram of a second interconnection channel and its main functions provided in one embodiment of the present invention.

[0037] Figure 7 This is a schematic diagram of a third interconnection channel and its main functions provided in one embodiment of the present invention.

[0038] Figure 8 This is a schematic diagram of the internal composition of a flash memory logic chip provided in one embodiment of the present invention.

[0039] Figure 9 This is a schematic diagram of the internal composition of an auxiliary computing unit provided in one embodiment of the present invention.

[0040] Figure 10 This is a schematic diagram of a computational subunit based on multiply-accumulate primitives provided in one embodiment of the present invention.

[0041] Figure 11 This is a schematic diagram of a computational subunit based on a stored-macro structure provided in one embodiment of the present invention.

[0042] Figure 12 This is a schematic diagram of flash memory plane cluster partitioning provided in one embodiment of the present invention.

[0043] Figure 13 This is a schematic diagram of the dynamic scheduling mechanism provided in one embodiment of the present invention.

[0044] Figure 14 This is a schematic diagram illustrating a dynamic scheduling mechanism implementation method provided in one embodiment of the present invention.

[0045] Figure 15 This is a comparison chart of inference latency between the architecture provided in one embodiment of the present invention and the GPU solution.

[0046] Figure 16 This is a diagram illustrating the stabilization effect of dynamic scheduling on inference latency in long contexts, provided in one embodiment of the present invention. Detailed Implementation

[0047] To make the objectives, technical solutions, and beneficial effects 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 described embodiments are merely illustrative of the invention and are not intended to limit the invention.

[0048] It should be noted that, in this application, relational terms such as "first" and "second" are used merely to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a system, method, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a system, method, or apparatus.

[0049] The auxiliary computing unit of this invention can be a single processor or a collective term for multiple processing elements. For example, it can be one or more neural processing units (NPUs), application-specific integrated circuits (ASICs), or one or more integrated circuits configured to implement this invention, such as one or more digital signal processors (DSPs), or one or more field-programmable gate arrays (FPGAs). In a specific implementation, as one embodiment, the auxiliary computing unit may include one or more NPUs. Each of these NPUs may contain one or more central processing units (CPUs) or one or more state machine circuits.

[0050] In the embodiments provided by this invention, it should be understood that the disclosed systems, apparatuses, and methods can be implemented in other ways. For example, the described system embodiments are merely illustrative; for instance, the division of units described is only a logical functional division, and in actual implementation, there may be other division methods. For example, the nonlinear calculations in the feedforward neural network computing unit can be reused or implemented using resources in the auxiliary computing unit. Furthermore, multiple units, modules, or components can be combined or integrated into another device, or some features can be ignored or not executed. For example, for low-speed inference applications, the volatile memory chip in the architecture may not be included; instead, the intermediate cache is written to flash memory for management. Another point is that the term "connection" in the following embodiments should be understood as "electrical connection," "communication connection," etc., if the connected circuits, modules, units, etc., transmit electrical signals or data to each other.

[0051] In their research on neural network inference accelerators, the inventors discovered that while existing widely discussed in-memory computing solutions can alleviate the "memory wall" problem between storage and computation to some extent, they still have the following shortcomings: First, the development of large language models places higher demands on storage capacity, bandwidth, operator types, and data scheduling, which existing architectures do not adequately support; Second, due to the autoregressive decoding nature of large language models, for traditional solutions that separate storage and computation, each generated token requires traversing and accessing weight data, resulting in massive data transfer between the storage medium and the computation unit; Third, there is a lack of effective dynamic coordination between feedforward neural network computation and attention mechanism computation, leading to increased computational load on the attention mechanism and decreased architecture throughput under long contextual inference.

[0052] To this end, the inventors proposed a large language model accelerator system based on three-dimensional NAND flash memory, and a dynamic scheduling method based on KV cache awareness. By executing the feedforward neural network calculation, which occupies the main part of the model parameters, inside the flash memory logic chip, the attention mechanism calculation is completed in the auxiliary computing unit. Combined with high-bandwidth DRAM and multi-layer interconnect design, the system achieves efficient acceleration of the large language model inference process.

[0053] like Figure 1 As shown, this embodiment provides a large language model accelerator system based on three-dimensional NAND flash memory, including three-dimensional NAND flash memory, DRAM chips, auxiliary computing units, and interconnect resources. The three-dimensional NAND flash memory is composed of NAND array chips and flash memory logic chips. The NAND array chips are used to store all model parameters and participate in matrix-vector multiplication and addition operations. All model parameters include: attention weight parameters, feedforward neural network weight parameters, final output projection weights used to map hidden states to vocabulary probability distributions, and scaling coefficients and bias parameters used in layer normalization operations, etc., representing large language model parameters. The DRAM chips are used to store attention weights and intermediate data. The flash memory logic chips are used to load relevant weights from the NAND array chips and perform feedforward neural network calculations and final output projection calculations. The auxiliary computing unit loads attention weights from the DRAM and undertakes the attention mechanism calculation task. The interconnect resources enable high-speed communication between the above components. The interconnect resources include: a first interconnect channel for data and cooperative control signal transmission between the auxiliary computing unit and the flash memory logic chips; a second interconnect channel for data transmission between the DRAM chips and the auxiliary computing unit; and a third interconnect channel for data transmission between the flash memory logic chips and the NAND array chips.

[0054] like Figure 2As shown, this invention proposes a basic scheduling algorithm flow, including the following steps: S1, firstly, the LLM model parameters are programmed into the NAND array chip, and the flash memory logic chip controls the loading of feedforward neural network calculation parameters via the third interconnect channel; S2, then the attention weights are loaded into the DRAM, and the auxiliary computing unit controls the loading of attention weights via the second interconnect channel for calculation preparation; S3, the auxiliary computing unit then completes the attention calculation, including generating query vectors, key vectors, and value vectors based on the weights, calculating the attention score, and performing multi-head output projection to obtain the attention activation result. The generated key vectors and value vectors are stored in the DRAM chip; S4, the auxiliary computing unit transmits the attention activation result to the flash memory logic chip via the first interconnect channel to trigger its... The process involves performing feedforward neural network calculations. Finally, depending on whether it's the output of the last layer of the feedforward neural network, step S5 is executed. If it is the output of the last layer, the final result is generated based on the final output projection weights read from the NAND array chip. If it's not the output of the last layer, the activation result of the current output of the feedforward neural network is used as an intermediate result. The final result or intermediate result is transmitted to the auxiliary computing unit via the first interconnect channel, and the intermediate result participates in the next attention calculation. Step S6 involves the auxiliary computing unit applying a nonlinear transformation to the final result. Based on the nonlinear transformation result, the next input vector is obtained. If the output termination condition is met, the next input vector is output as the decoding result; otherwise, the next attention calculation is performed based on the next input vector. This process is executed cyclically under a dynamic scheduling mechanism until a complete output sequence is generated.

[0055] like Figure 3 and Figure 4 As shown, the data flow under the inference scheme of this invention has significant advantages in terms of storage bandwidth utilization, energy efficiency, and latency compared to conventional large model inference hardware solutions such as GPU inference. By performing multiply-accumulate operations directly within the flash array, this invention reduces weight transfer overhead, thereby enabling it to exhibit throughput exceeding that of traditional von Neumann architectures in large language model inference.

[0056] In the accelerator architecture of this invention, interconnect resources handle the data and control signal transmission between flash memory logic chips, NAND array chips, auxiliary computing units, and DRAM chips. Addressing the characteristics of large data traffic, high bandwidth requirements, and complex control paths in large language model inference, this invention proposes a separate interconnect channel design to decouple and distribute bandwidth requirements.

[0057] 1. First Interconnection Channel

[0058] like Figure 5As shown, the first interconnect channel is used to transmit data and control signals between the flash memory logic chip and the auxiliary computing unit, supporting bidirectional transmission and broadcast modes. During the attention mechanism calculation process, this channel can broadcast the attention activation results generated by the auxiliary computing unit to the flash memory logic chip, and return the intermediate activation values ​​or final result data output by the feedforward neural network to the auxiliary computing unit during the result collection phase. By introducing a broadcast and collection mechanism in the channel, redundant data movement is avoided, and bus utilization is improved.

[0059] 2. Second Interconnection Channel

[0060] like Figure 6 As shown, the second interconnect channel is used to transfer attention mechanism-related weights and key-value (KV) caches between the auxiliary computing unit and the DRAM chip. This channel employs a high-bandwidth interface, such as LPDDR series or a customized wide bus interface, to support massively parallel read and write operations. Since the attention mechanism's access to the KV cache is frequent and irregular, this invention introduces flow control and cache mapping strategies in the second interconnect channel to reduce access conflicts while ensuring bandwidth.

[0061] The implementation details of the flow control and cache mapping strategy for the second interconnection channel are as follows:

[0062] The second interconnect channel is primarily used for data interaction between the auxiliary computing unit and the DRAM chip, especially for the key vectors, value vectors, and their caches required for the attention mechanism computation. Since the multi-head attention mechanism needs to repeatedly read large-scale KV caches during inference, and the computation of query vectors exhibits strong sequence dependencies, this channel may experience access congestion and bandwidth waste. Therefore, this invention proposes the following two optimization strategies:

[0063] 1. Flow control strategy

[0064] Block access: The KV cache is divided into fixed-size data blocks, and each access is transmitted in blocks, avoiding bus overload caused by small-granularity random access.

[0065] Priority arbitration: When multiple heads simultaneously initiate access requests, different priorities are set based on the urgency of the requests, such as whether the current query vector is waiting for the return of key / value vector data, and the arbitrator dynamically schedules the data transmission order.

[0066] Bandwidth balancing: When the read traffic of attention computation is too high, the scheduler can temporarily postpone the write-back operation of some non-critical data blocks to ensure the real-time performance of read operations.

[0067] 2. Cache mapping strategy

[0068] Separate mapping: KV caches with different attention heads are mapped to independent banks or channels in DRAM, reducing conflicts within the same bank and improving parallelism.

[0069] Hot data residency: For frequently accessed KV cache data, partially map it to a small-capacity on-chip cache within the auxiliary computing unit to avoid frequent cross-channel transmission.

[0070] Timing alignment: Prefetch the key and value vector data required for query vector calculation, and load the data required for the next step before the previous layer calculation is finished, thereby reducing waiting latency.

[0071] Through the aforementioned flow control and cache mapping strategies, the second interconnect channel can support attention mechanism computation under conditions of high bandwidth, low conflict, and low latency: it reduces bus conflicts and data blocking when multi-head attention is in parallel; it maintains efficient access and updates of the KV cache in long context reasoning scenarios; and it improves the energy efficiency and stability of the interaction between the auxiliary computing unit and DRAM.

[0072] 3. Third Interconnection Channel

[0073] like Figure 7 As shown, the third interconnect channel is used during operation to transmit feedforward neural network weights, final output projection weights, and attention weights between the flash memory logic chip and the NAND array chip. To reduce latency and power consumption, this channel uses hybrid bonding technology to achieve physical connections, thereby realizing extremely short interconnect paths on-chip. The third interconnect channel has high parallelism and low latency characteristics, enabling pipelined updates of weight data in the on-chip cache system while the flash memory logic chip performs computations, achieving overlap between weight loading and feedforward computation.

[0074] 4. Channel Coordination Mechanism

[0075] During the inference process of a large language model, three interconnected channels work in parallel:

[0076] The first channel is mainly used for control and result interaction; the second channel is mainly responsible for high-bandwidth weight / cache transmission; the third channel is mainly used to realize high parallelism and low latency connection between flash logic and storage array.

[0077] This hierarchical interconnection scheme avoids the congestion problem of traditional single-bus architecture, ensures effective collaborative execution between attention mechanism and feedforward computing, thereby improving overall throughput and energy efficiency.

[0078] The system of this invention is designed in accordance with the principle that "feedforward computation is completed on the flash memory side, and the attention mechanism is completed on the auxiliary computing unit side." To this end, the flash memory logic chip and the auxiliary computing unit are each configured with corresponding functional sub-units or sub-modules, and are closely coordinated through interconnect channels.

[0079] In such Figure 8 In the flash memory logic chip shown, the control unit manages access to the NAND array chip and the on-chip cache system, and manages matrix-vector multiplication and addition operations in the attention calculation task corresponding to the cooperative control signal. The flash memory computing subunit is tightly coupled to the NAND array, and the feedforward neural network weights, final output projection weights, and attention weights read from the NAND array are cached in the on-chip cache system. The bit error correction subunit is used to correct errors in the feedforward neural network weights, final output projection weights, and attention weights. The flash memory computing subunit is used to perform matrix-vector multiplication and addition operations between the attention activation value and the corrected feedforward neural network weights, or to perform matrix-vector multiplication and addition operations between the output result of the nonlinear calculation subunit and the corrected feedforward neural network weights, or to perform unit attention calculation operations corresponding to the cooperative control signal based on the corrected attention weights. In force calculation tasks, matrix-vector multiplication and addition operations are performed on the output of the last layer of the feedforward neural network and the projected weights of the final output after error correction. The result of the matrix-vector multiplication and addition of the attention activation value and the feedforward neural network weights is entered into the nonlinear calculation subunit via the on-chip cache system for activation or normalization. The result obtained by the nonlinear calculation subunit is then subjected to matrix-vector multiplication and addition with the feedforward neural network weights, thus completing one round of feedforward neural network calculation within the memory. The output of the current layer of the feedforward neural network is fed back to the auxiliary calculation unit as an intermediate result. For the output of the last layer of the feedforward neural network, the projected weights of the final output after error correction are multiplied and added to the output of the last layer of the feedforward neural network to obtain the final result of the feedforward neural network. The final result and the intermediate result of the feedforward neural network are fed back to the auxiliary calculation unit. To ensure efficient loading and updating of weight data, the on-chip cache system has a page buffer at the end, followed by a data buffer, and a global calculation buffer at the top. The data transfer latency within the logic chip is matched with the calculation latency, realizing a pipelined weight transfer and calculation process and avoiding data waiting caused by traditional external access.

[0080] In such Figure 9In the auxiliary computing unit shown, each submodule constitutes a complete processing procedure around the attention mechanism: the computation submodule is responsible for generating queries, keys, and values ​​and calculating attention scores; then, the activation function submodule applies attention weights obtained from Softmax, and the computation submodule adjusts the value matrix or vector according to these attention weights and integrates the multi-head output results. The auxiliary control submodule, based on the structural details of different large language models, regulates the timing of the rotation position encoding submodule, normalization submodule, residual connection submodule, and sampling submodule to ensure that the computation process remains consistent with the actual model. Furthermore, the auxiliary computing unit not only performs these operations but also manages the KV cache through its internal DRAM control submodule, dynamically allocating bandwidth and computing resources.

[0081] More importantly, the two control components—the in-flash control subunit and the auxiliary control submodule—are not independent but rather form a cross-chip scheduling relationship through interconnect channels. This allows attention activation results to be seamlessly transmitted to the flash memory side to trigger feedforward computation, while the intermediate activation values ​​output by the feedforward can be directly returned to the auxiliary computation unit to participate in the next attention computation, forming a closed-loop inference process. Simultaneously, the awareness information of the KV cache by the auxiliary control submodule can be synchronized to the in-flash control subunit, enabling dynamic collaboration between different computational parts of the architecture.

[0082] The computational sub-units within flash memory can be implemented in two ways: one is based on an array of multiplier and adder primitives, and the other is based on an array of in-memory computation macros. Both can perform vector dot product operations, but they differ in area efficiency, power consumption, and access latency. This invention allows for flexible selection based on the specific application scenario.

[0083] In one specific embodiment, the computational sub-units within the flash memory preferably employ a rectangular or non-rectangular array structure based on multiplier-adder primitives, such as... Figure 10 As shown. For example, an 8×2 primitive array can be constructed, with each computational primitive containing 16 multipliers for parallel computation of the product of 16 inputs and weights: The 16 products are then combined into a scalar output using an addition tree. This structure offers advantages such as high computational accuracy and stable latency, facilitating seamless integration with nonlinear sub-units and making it suitable for small- to medium-scale inference tasks. In practical applications, the array size can be expanded to 8×8, 16×8, or larger to meet the computational demands of large-scale model inference.

[0084] In another embodiment, the in-flash computing subunit can employ, as shown below: Figure 11The diagram shows an SRAM-CIM macroarray. Each macroarray stores weights in the SRAM array and performs vector multiplication and addition operations through in-memory computation. This architecture is more suitable for large-scale parallel expansion, such as through 32×32 or larger macroarrays, significantly improving bandwidth utilization while reducing energy consumption and area overhead. Its advantages lie in its high energy efficiency and large-scale concurrency capabilities, making it particularly suitable for batch processing of long-context inference or large-parameter models.

[0085] Therefore, the flash computing sub-unit design of the present invention can support scenarios that prioritize accuracy and low latency, as well as scenarios that prioritize high energy efficiency and large-scale parallelism. Developers can choose different implementation methods according to the actual chip design goals.

[0086] like Figure 12 As shown, multiple planes in a NAND array chip can be divided into clusters. Planes within each cluster share a read address decoder and part of the read address circuitry, thus supporting an intra-cluster cooperative access mechanism. Simultaneously, the storage location indices of different parts of the model weights are adjusted according to the cluster and plane numbers to improve data access efficiency. This design can improve parallelism and reduce access latency.

[0087] In one specific embodiment, the placement of model weights follows the principle of "weights that need to be continuously involved in the calculation should also be stored continuously within and between clusters," or the continuity principle. The corresponding implementation is as follows:

[0088] The weights are divided into several atomic blocks, denoted as set . Each block is a minimum contiguous storage unit, such as several bit rows or pages, which facilitates the contiguous allocation of indexes within and between clusters.

[0089] The NAND array chip has K clusters, and each cluster has P planes, so the total number of planes is KP. The clusters are numbered as follows: The plane within the cluster is numbered as The physical block space is indexed by block location as follows: Here, S represents the number of atomic blocks that can be placed in each plane; if the capacity is the same, then This represents the total number of blocks. Define the mapping:

[0090] ,

[0091] The mapping relationship from atomic blocks to clusters and planes derived according to the continuity principle is as follows:

[0092] .

[0093] It is understood that the method disclosed in this embodiment can be modified to obtain other variations by simple substitutions and the addition of restrictions. For example, the weighting of atomic blocks can be performed row by row, column by column, or sub-block by sub-block; the weighting can be implemented in hardware, software, or a combination of both.

[0094] like Figure 13 As shown, this invention further proposes a dynamic scheduling mechanism that allocates tasks based on a scheduling graph. After initialization, the auxiliary control submodule can rewrite the scheduling graph and collaborate with the flash memory control subunit to dynamically allocate attention mechanism-related computational tasks to stabilize the architecture's throughput. Specifically, the scheduling graph is a bitmap that instructs the auxiliary computing unit and the flash memory logic chip to collaboratively execute attention computation tasks. One implementation of this method is as follows: Figure 14 As shown.

[0095] In one embodiment of the present invention, the KV cache-aware scheduling mechanism uses a bitmap M as a scheduling graph to indicate the execution location of unit attention computation tasks. When a bitmap element is "1", it indicates that the corresponding attention computation task is executed by an auxiliary computing unit; when it is "0", it indicates that the task is executed in an external flash memory logic chip. The cooperative control signal is an instruction that instructs the flash memory logic chip to execute the unit attention computation task. The schedulable unit attention computation tasks include querying, key and value vector generation, and multi-head output projection.

[0096] The auxiliary control submodule sets a counter during the autoregressive decoding stage to monitor the total execution time of each attention mechanism step, including querying, key and value vector generation, attention weight calculation, and multi-head output projection. Let the total execution time of the nth decoding step be... The (n-1)th step is When detected:

[0097]

[0098] At that time, a bitmap rewrite operation is triggered. The threshold time is specified. Defined as:

[0099]

[0100] Where P represents the cluster page cache capacity, which is equal to the product of the number of planes within the cluster and the size of the plane page cache, and u represents the weighted atomic block size. This indicates the delay in the auxiliary computing unit's processing of atomic block data.

[0101] Preferably, the adjustment force is defined as:

[0102]

[0103] That is, the time increment ΔT and the threshold time The ratio is rounded up. This value, r, determines the number of tasks that need to be offloaded to the flash memory logic chip for computation. During the bitmap rewriting process, the scheduling logic will adjust the intensity. The corresponding computation task is set to "0". This part of the task will be loaded by the flash memory logic chip from the NAND array chip and matrix-vector multiplication and addition calculations will be performed to offload part of the task. For example, the part before the last non-zero element in the bitmap will be loaded. Position 0. It should be noted that, thanks to the architectural features of this invention, the weights required for calculating queries, key and value vector generation, and multi-head output projection still exist in the flash memory. Therefore, implementing this process only involves changes to the scheduling graph and communication between the auxiliary control submodule and the control subunit within the flash memory, without involving the transmission of any weight data on the first interconnect channel.

[0104] This strategy allows for dynamic adjustment of task allocation between the auxiliary computing unit and the flash memory logic chip based on runtime performance fluctuations, avoiding excessive cache pressure or latency accumulation. This improves system energy efficiency and resource utilization while maintaining decoding speed. It should be understood that this scheduling mechanism can be implemented layer-by-layer or block-by-block based on the model, or iteratively implemented step-by-step autoregressive decoding, to adapt to different model sizes, hardware configurations, and throughput fluctuations. Furthermore, dynamic scheduling can be implemented in hardware, software, or a combination of both.

[0105] like Figure 15 As shown, the system of this invention demonstrates performance improvements compared to the GPU solution under different computing resource configurations. The detailed computing configurations for "Small Computing Configuration," "Medium Computing Configuration," and "Large Computing Configuration" are as follows: all computing resources utilize an array consisting of 16 multipliers and 1 adder tree primitive, with array sizes of 8×2, 12×2, and 16×2, respectively; and the auxiliary computing unit array size is 4×2 for all configurations. This means that for pre-padded / decoded sequences of different lengths, the total latency for inference under the above configurations is lower than that of the GPU-DRAM and GPU-SSD solutions. The GPU-DRAM and GPU-SSD solutions refer to the GPU as the computing unit, with model weights stored in DRAM and SSD, respectively. This indicates that the architecture of this invention can run inference for large language models at high speed.

[0106] like Figure 16 As shown, the KV cache-aware dynamic scheduling mechanism of this invention can effectively suppress the increase in per-Token generation latency and maintain the stability of inference rate even as the context length of the LLaMa3-8B model continuously increases. This effect is particularly suitable for application scenarios that are sensitive to overall inference latency, such as dialogue systems and real-time inference services.

[0107] In summary, this invention achieves high-efficiency acceleration of large language models by combining flash memory computation of 3D NAND flash memory with attention mechanism computation of auxiliary computing units and introducing a dynamic scheduling mechanism. Its overall design method balances high throughput and low power consumption, making it particularly suitable for deployment of large language model inference at the edge.

[0108] While the embodiments disclosed in this invention are as described above, the above content is merely for the purpose of facilitating understanding of this invention and is not intended to limit this application. Any person skilled in the art may make any modifications and changes to the form and details of the implementation without departing from the spirit and scope disclosed in this invention; however, the scope of patent protection of this invention shall still be determined by the scope defined in the appended claims.

Claims

1. A large language model accelerator system based on three-dimensional NAND flash memory, characterized in that, include: DRAM is used to store attention weights, as well as key and value vectors generated during attention calculation. The auxiliary computing unit is used to read attention weights from DRAM, perform attention calculations, write the key vectors and value vectors generated during the attention calculation process back to DRAM, output the attention activation results, and obtain the decoding results or execute the next attention calculation based on the feedforward neural network calculation results. NAND array chips are used to store large language model parameters. The flash memory logic chip receives the attention activation results to trigger feedforward neural network computation. It reads the feedforward neural network weights and the final output projection weights from the large language model parameters stored in the NAND array chip, and writes the feedforward neural network computation results back to the auxiliary computing unit. The first interconnect channel is used to couple the auxiliary computing unit and the flash memory logic chip. A second interconnect channel is used to couple the DRAM and the auxiliary computing unit, and, The third interconnect channel is used to couple the flash memory logic chip and the NAND array chip; The auxiliary calculation unit includes: The interface module is used to couple the flash memory logic chip via the first interconnect channel. The control module manages the write-back of key vectors and value vectors to DRAM, adjusts the scheduling graph and generates cooperative control signals after real-time sensing of the KV cache. The scheduling graph is a bitmap instructing the auxiliary computing unit and the flash memory logic chip to collaboratively execute attention computation tasks. The cooperative control signals are instructions instructing the flash memory logic chip to execute unit attention computation tasks. The attention calculation mechanism module is used to complete the unit attention calculation task executed by the auxiliary calculation unit as instructed by the scheduling graph, and to obtain the decoding result or execute the next attention calculation based on the calculation result of the feedforward neural network.

2. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 1, characterized in that, The NAND array chip includes at least one cluster of shared read address decoders and partial read address circuits. Each cluster contains at least one storage plane. The smallest set of weights of a large language model that participates in the calculation consecutively is taken as a weight atom block. Based on the principle of continuous storage location of weight atom blocks, a mapping relationship between weight atom blocks and clusters and a mapping relationship between weight atom blocks and storage planes are established.

3. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 1, characterized in that, The flash memory logic chip includes: The interface unit is used to couple to the auxiliary computing unit through the first interconnection channel, receive attention activation results and collaborative control signals, and transmit feedforward neural network calculation results back to the auxiliary computing unit. The control unit is used to manage access to NAND array chips and on-chip cache architecture, and to manage matrix-vector multiplication and addition operations in unit attention computing tasks corresponding to the cooperative control signals. The on-chip cache system is used to store the feedforward neural network weights, the final output projection weights, and the attention weights read from the NAND array chip; The bit error correction subunit is used to correct errors in the feedforward neural network weights, the final output projection weights, and the attention weights. The flash memory-based computation subunit is used to perform matrix-vector multiplication and addition on the attention activation values ​​and the corrected feedforward neural network weights, or to perform matrix-vector multiplication and addition on the output of the nonlinear computation subunit and the corrected feedforward neural network weights, or to perform matrix multiplication and addition in the unit attention calculation task corresponding to the cooperative control signal based on the corrected attention weights, or to perform matrix-vector multiplication and addition on the output of the last layer of the feedforward neural network and the projected weights of the final output after error correction. The matrix-vector multiplication and addition results are written to the on-chip cache system, and the final result of the feedforward neural network is transmitted to the interface unit; and... The nonlinear computation subunit is used to perform activation or normalization operations on the matrix-vector multiplication and addition calculation results related to the feedforward neural network weights read from the on-chip cache system, and to transmit the intermediate results of the feedforward neural network to the interface unit.

4. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 1, characterized in that, The control module senses the KV cache in real time, adjusts the scheduling diagram, and generates cooperative control signals. Specifically, it monitors the total time for the attention calculation mechanism to execute each step of the attention calculation, and calculates the difference between the total times of two adjacent attention calculation steps. Exceeding the threshold When the bitmap rewrite operation is triggered, based on the difference... and threshold Calculate adjustment force The intensity will be adjusted. The corresponding unit attention calculation task is configured for the NAND array chip.

5. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 4, characterized in that, The threshold is P is the cluster page cache capacity, and u is the weighted atomic block size. To mitigate the delay in the computational unit's processing of weighted atom block data; the computational adjustment force is... , This is for rounding up.

6. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 3, characterized in that, The in-flash computing sub-unit adopts a rectangular or non-rectangular array structure based on multiplier-adder primitives; or, adopts an SRAM in-memory computing macro array.

7. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 1, characterized in that, The key vectors and value vectors of different attention heads are mapped to data blocks or channels in DRAM, and the key vectors and value vectors that are frequently accessed by the auxiliary computing unit are mapped to the on-chip cache of the auxiliary computing unit. The key vectors and value vectors for the next attention calculation are prefetched before the previous attention calculation is finished.

8. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 7, characterized in that, The second interconnect channel accesses DRAM data blocks in blocks, and postpones the write-back operation of non-critical data blocks when the read traffic of the auxiliary computing unit exceeds the limit, based on the transmission order of the data blocks arbitrated by the multi-head attention access request.

9. The large language model accelerator system based on three-dimensional NAND flash memory according to claim 1, characterized in that, The third interconnect channel uses hybrid bonding technology to couple the flash memory logic chip and the NAND array chip.

Images