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

Abstract

The invention discloses a large language model accelerator architecture based on a three-dimensional NAND flash memory, and belongs to the technical field of calculation, reckoning or counting. The architecture comprises a three-dimensional NAND flash memory used for executing feedforward neural network calculation; the auxiliary calculation unit is used for executing attention mechanism calculation; the DRAM chip is used for storing attention mechanism related weights and KV cache; and the interconnection resource is used for realizing data interaction among the components. Wherein the three-dimensional NAND flash memory comprises a logic chip and an NAND array chip, the logic chip is used for controlling and executing calculation, and the NAND array chip is used for storing weights and participating in calculation. The invention further provides a scheduling method based on KV cache awareness. The scheduling method comprises decomposition and dynamic allocation of calculation tasks. Through collaborative design of the hardware architecture and the scheduling method, the memory wall bottleneck in large language model calculation can be relieved, the energy consumption is reduced, and the overall calculation performance is improved.

Description

TECHNICAL FIELD

[0001] The present application relates to integrated circuits and artificial intelligence computing technology, and particularly discloses a large language model accelerator architecture based on three-dimensional NAND flash computing and a dynamic scheduling method thereof, and belongs to the technical field of computing, reckoning or counting. BACKGROUND

[0002] This section is intended to provide background or context to the technical solution of the present application stated in the claims. The description herein is not admitted to be prior art merely because it is included in this section.

[0003] With the rapid development of natural language processing, large language models (LLMs) such as GPT series, LLaMa series and Deepseek series have extremely high requirements for computing power and storage bandwidth. Such models usually contain tens of billions to hundreds of billions of parameters, and involve a large number of matrix vector multiplication and attention mechanism calculations in the training and inference process. The traditional implementation based on GPU or general-purpose accelerator has a bottleneck between computing power and storage. For example, although the peak computing power of GPU can exceed hundreds of TOPS, reaching trillion operations per second, when the GPU's video memory cannot accommodate all model parameters, part of the data needs to be offloaded to DRAM or SSD with relatively low bandwidth. For example, with a PCIe 4.0 x4 interface, only 8GB / s, the model parameter amount can reach tens of GB or even hundreds of GB, resulting in high latency and low energy efficiency.

[0004] In recent years, processing-in-memory (PIM) technology has attracted much attention because it can alleviate the memory wall problem in data-intensive applications. PIM integrates computing units inside the memory chip, reducing the overhead of data movement to the processor, which helps to achieve low-energy computing for large-scale matrix multiplication. However, DRAM is a volatile medium that needs to be periodically refreshed during operation to maintain data state, and data will be lost after power-off. In addition, the storage density of DRAM is not high, and the single-chip capacity is difficult to meet the storage of large model parameters. Therefore, the PIM architecture based on DRAM usually needs to use expensive through silicon via (TSV) technology to bond stacked DRAM layers, which makes it difficult to reduce power consumption and cost, especially for end-side applications.

[0005] With the development of three-dimensional NAND flash memory technology, its huge capacity and internal bandwidth characteristics contain the potential of large language model flash computing non-volatile storage carrier. Weighted accumulation calculation is performed at the flash array control logic level to realize the direct execution of the multiplication and addition operation of the feedforward neural network in the memory, which is expected to overcome the performance bottleneck caused by data transmission and greatly reduce the chip area, power consumption and cost.

[0006] In summary, the implementation of tens of billions or even hundreds of billions of complete and efficient large language model inference architectures mainly faces the following challenges:

[0007] The DRAM-based architecture weights still need to be initialized from non-volatile low-speed storage, and traditional storage media are difficult to meet the storage capacity and access bandwidth requirements brought by the increase in the size of large language models;

[0008] The NAND-based architecture weight reading delay is high, and the data update strategy between the on-chip cache and the flash array is not perfect, which may cause performance bottlenecks in the multi-layer neural network inference process;

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

[0010] There is a lack of dynamic scheduling mechanism for multi-layer feedforward neural network and attention mechanism calculation, which cannot fully utilize the flash internal parallel computing resources, causing resource waste.

[0011] To solve the above problems, it is necessary to design a new type of three-dimensional NAND flash memory-based In-Flash Computing / Processing-In-Flash (IFC / PIF) large language model accelerator architecture and a dynamic scheduling method to efficiently support attention mechanism calculation and feedforward neural network calculation and improve the performance and energy efficiency of large language model inference. SUMMARY

[0012] The purpose of the present application is to provide a three-dimensional NAND flash memory-based large language model accelerator architecture and its dynamic scheduling method to solve the technical problems of storage bandwidth bottleneck, insufficient utilization of computing resources and large data movement overhead in the prior art in feedforward neural network calculation and attention mechanism calculation, thereby achieving the purpose of improving the performance and energy efficiency of large language model inference.

[0013] To achieve the above purpose, the present application proposes the following technical solutions:

[0014] The large language model accelerator architecture based on three-dimensional NAND flash memory includes: a DRAM, an auxiliary computing unit, a NAND array chip, a flash logic chip, a first interconnection channel, a second interconnection channel, and a third interconnection channel; the DRAM is used to store attention weights and key vectors and value vectors generated in the attention calculation process; the auxiliary computing unit is used to read the attention weights from the DRAM, perform attention calculation, write the key vectors and value vectors generated in the attention calculation process back to the DRAM, output the attention activation result, and obtain the decoding result according to the feedforward neural network calculation result or perform the next step of attention calculation; the NAND array chip is used to store large language model parameters; the flash logic chip is used to receive the attention activation result to trigger the feedforward neural network calculation, read the feedforward neural network weights and the last output projection weights from the large language model parameters stored in the NAND array chip, and write the feedforward neural network calculation result back to the auxiliary computing unit; the first interconnection channel is used to couple the auxiliary computing unit and the flash logic chip; the second interconnection channel is used to couple the DRAM and the auxiliary computing unit; and the third interconnection channel is used to couple the flash logic chip and the NAND array chip.

[0015] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the NAND array chip contains at least one cluster sharing a read address decoder and part of a read address circuit, each cluster contains at least one storage plane, the smallest set of large language model weights continuously participating in calculation is taken as a weight atomic block, and a mapping relationship of the weight atomic block to the cluster and a mapping relationship of the weight atomic block to the storage plane are established according to the principle of continuous storage position of the weight atomic block.

[0016] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the auxiliary computing unit includes: an interface module, a control module, and an attention calculation mechanism calculation module; the interface module is used to couple the flash logic chip through the first interconnection channel; the control module is used to manage the operation of writing the key vectors and value vectors back to the DRAM, adjust the scheduling diagram in real time after the KV cache, and generate a cooperative control signal, the scheduling diagram is a bitmap indicating that the auxiliary computing unit and the flash logic chip cooperatively perform an attention calculation task, and the cooperative control signal is an instruction indicating that the flash logic chip performs a unit attention calculation task; the attention calculation mechanism calculation module is used to complete the unit attention calculation task indicated by the scheduling diagram for the auxiliary computing unit to perform, obtain the decoding result according to the feedforward neural network calculation result, or perform the next step of attention calculation.

[0017] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the flash logic chip comprises: an interface unit, a control unit, an on-chip cache system, a bit error correction subunit, a flash-in-computation subunit, and a nonlinear computation subunit; the interface unit is used to couple the auxiliary computation unit through a first interconnection channel, receive attention activation results and collaborative control signals, and transmit feedforward neural network computation results to the auxiliary computation unit; the control unit is used to manage NAND array chip access and on-chip cache system access, and manage matrix multiplication and addition operations in unit attention computation tasks corresponding to the collaborative control signals; the on-chip cache system is used to store 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 correct the feedforward neural network weights, the final output projection weights, and the attention weights; the flash-in-computation subunit is used to perform matrix vector multiplication and addition computation on the attention activation values and the corrected feedforward neural network weights, or perform matrix vector multiplication and addition computation on the nonlinear computation subunit output results and the corrected feedforward neural network weights, or perform matrix multiplication and addition operations in unit attention computation tasks corresponding to the corrected attention weights according to the collaborative control signals, or perform matrix vector multiplication and addition computation on the last layer feedforward neural network output and the corrected final output projection weights, write the matrix vector multiplication and addition computation results into the on-chip cache system, and transmit the final feedforward neural network results to the interface unit; and the nonlinear computation subunit is used to perform activation or normalization operation on the matrix vector multiplication and addition computation results related to the feedforward neural network weights read from the on-chip cache system, and transmit the intermediate feedforward neural network results to the interface unit.

[0018] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the control module adjusts the scheduling diagram and generates collaborative control signals in real time after sensing the KV cache, specifically: monitoring the total time of the attention computation mechanism computation module to perform each step of attention computation, and triggering a bitmap rewriting operation when the difference between the total times of adjacent two steps of attention computation exceeds a threshold value . The adjustment strength is calculated according to the difference and the threshold value . The unit attention computation task corresponding to the adjustment strength is configured to the NAND array chip.

[0019] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the threshold value is , P is the cluster page cache capacity, u is the weight atomic block size, is the delay of the auxiliary computation unit processing weight atomic block data; and the calculation adjustment strength is . A ceiling function operation.

[0020] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the in-flash computing subunit adopts a rectangular or non-rectangular array structure based on a multiplier-adder cell; or, an SRAM storage and computing integrated macro array.

[0021] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the key vectors and value vectors of different attention heads are respectively mapped to the data blocks or channels of the DRAM, the frequently accessed key vectors and value vectors of the auxiliary computing unit are mapped to the on-chip cache of the auxiliary computing unit, and the key vectors and value vectors of the next step of attention calculation are prefetched when the previous step of attention calculation has not been completed.

[0022] As a further optimization scheme of the large language model accelerator architecture based on three-dimensional NAND flash memory, the second interconnection channel accesses the DRAM data block in blocks, arbitrates the transmission order of the data block according to the time sequence of the multi-head attention access request, and delays the write-back operation of the non-critical data block when the read traffic of the auxiliary computing unit exceeds the limit.

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

[0024] The technical scheme adopted by the present application has the following significant beneficial effects compared with the prior art.

[0025] (1) The accelerator architecture proposed in the present application fully utilizes the characteristics of three-dimensional NAND flash memory and hybrid bonding, decouples the feedforward neural network calculation and attention mechanism calculation into efficient hardware parts that work together, and distributes the bandwidth pressure to three interconnection channels.

[0026] (2) The accelerator architecture proposed in the present application directly performs matrix-vector multiplication and addition inside the NAND memory, reduces the model weight transmission overhead, alleviates the memory wall bottleneck of the traditional von Neumann computing architecture, reduces the data movement overhead, and improves the energy efficiency ratio.

[0027] (3) The accelerator architecture proposed in the present application supports pipelined flash access and dynamic scheduling, avoids the mismatch between calculation delay and flash access delay, and the decrease in architecture throughput caused by increased calculation amount in long context reasoning scenarios.

[0028] (4) The accelerator architecture proposed in the application, the control functions in the flash memory and the auxiliary computing unit are relatively independent and closely cooperate, which improves the universality and efficiency of the large language model inference solution based on the architecture, reduces the hardware implementation cost, and especially provides an innovative solution for the end-side inference scene.

[0029] In summary, the application provides an effective solution for efficient hardware acceleration of large language models by combining storage flash computing with auxiliary computing units and introducing a dynamic scheduling mechanism. BRIEF DESCRIPTION OF DRAWINGS

[0030] In order to more clearly illustrate the technical solutions in the embodiments of the application or the prior art, the following will briefly introduce the drawings needed to be used in the embodiments or the prior art description. Obviously, the drawings in the following description are only some embodiments of the application, and for those skilled in the art, other drawings can also be obtained from these drawings without creative labor. The drawings in the application are:

[0031] Figure 1 The figure shows the large language model accelerator architecture based on three-dimensional NAND flash memory provided in an embodiment of the application.

[0032] Figure 2 The figure shows the basic scheduling algorithm flowchart provided in an embodiment of the application.

[0033] Figure 3 The figure shows the data flow diagram of the conventional large language model inference hardware solution.

[0034] Figure 4 The figure shows the inference solution data flow diagram provided in an embodiment of the application.

[0035] Figure 5 The figure shows the first interconnection channel and its main function diagram provided in an embodiment of the application.

[0036] Figure 6 The figure shows the second interconnection channel and its main function diagram provided in an embodiment of the application.

[0037] Figure 7 The figure shows the third interconnection channel and its main function diagram provided in an embodiment of the application.

[0038] Figure 8 The figure shows the internal composition diagram of the flash logic chip provided in an embodiment of the application.

[0039] Figure 9 The figure shows the internal composition diagram of the auxiliary computing unit provided in an embodiment of the application.

[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 an architecture, 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 architecture, 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 architecture, apparatus, and method can be implemented in other ways. For example, the described architectural 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., have electrical signal or data transmission with 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 address this, the inventors proposed a large language model accelerator architecture based on three-dimensional NAND flash memory, along with a dynamic scheduling method based on KV cache awareness. This method executes feedforward neural network calculations, which constitute the main part of the model parameters, within the flash memory logic chip, while completing attention mechanism calculations in the auxiliary computing unit. Combined with high-bandwidth DRAM and multi-layer interconnect design, this approach achieves efficient acceleration of the large language model inference process.

[0053] like Figure 1 As shown, this embodiment provides a large language model accelerator architecture 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 consists of NAND array chips and flash memory logic chips. The NAND array chips store all model parameters and participate in matrix-vector multiplication and addition operations. These 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 and bias parameters used in layer normalization operations, among other large language model parameters. The DRAM chips store attention weights and intermediate data. The flash memory logic chips load relevant weights from the NAND array chips and perform feedforward neural network calculations and final output projection calculations. The auxiliary computing units load attention weights from the DRAM and perform attention mechanism calculations. The interconnect resources enable high-speed communication between the above components. These interconnect resources include: a first interconnect channel for data and cooperative control signal transmission between the auxiliary computing units and the flash memory logic chips; a second interconnect channel for data transmission between the DRAM chips and the auxiliary computing units; 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 and ensures effective collaborative execution between attention mechanism and feedforward computing, thereby improving overall throughput and energy efficiency.

[0078] The architecture of this invention follows the principle that "feedforward computation is performed on the flash memory side, and the attention mechanism is performed 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 subunits or submodules, 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 architecture 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-padding / decoding 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 architecture 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 attention activation results, and obtain decoding results or execute the next attention calculation based on the calculation results of the feedforward neural network. NAND array chips are used to store large language model parameters; The flash memory logic chip is used to receive the attention activation result to trigger the feedforward neural network calculation, read the feedforward neural network weights and the final output projection weights from the large language model parameters stored in the NAND array chip, and write the feedforward neural network calculation result back to the auxiliary calculation unit. The first interconnection channel is used to couple the auxiliary computing unit and the flash memory logic chip; The 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.

2. The large language model accelerator architecture 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 architecture based on three-dimensional NAND flash memory according to claim 2, characterized in that, The auxiliary calculation unit includes: An interface module is used to couple flash memory logic chips via a 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.

4. The large language model accelerator architecture based on three-dimensional NAND flash memory according to claim 3, 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.

5. The large language model accelerator architecture based on three-dimensional NAND flash memory according to claim 3, 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 A bitmap rewrite operation is triggered at the time, 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.

6. The large language model accelerator architecture based on three-dimensional NAND flash memory according to claim 3, 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.

7. The large language model accelerator architecture based on three-dimensional NAND flash memory according to claim 4, 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.

8. The large language model accelerator architecture 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.

9. The large language model accelerator architecture based on three-dimensional NAND flash memory according to claim 8, 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.

10. The large language model accelerator architecture 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