State space model accelerator, Mama model calculation method and Mama model calculation system

Abstract

The invention discloses a state space model accelerator, a Mama model calculation method and a Mama model calculation system. The state space model accelerator comprises a reconfigurable calculation unit used for linear calculation and element-by-element calculation; the reconfigurable computing unit comprises a reduction tree and a processing unit array; the processing unit array comprises a plurality of reconfigurable processing units; during linear calculation, each reconfigurable processing unit calculates input data; the reduction tree is in a reduction mode, and reduction tree slices in the reduction tree carry out reduction calculation on calculation results of the multiple corresponding reconfigurable processing units; when element-by-element calculation is carried out, each reconfigurable processing unit carries out element-by-element calculation on input data and outputs calculation results in parallel; the reduction tree is in a non-reduction mode. According to the Mama accelerator, a reduction-replaceable PE array architecture is adopted, and element-by-element operation and linear operation can be efficiently supported at the same time.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence, and in particular to a state space model accelerator, a Mamba model calculation method and a calculation system. Background Art

[0002] State-space models (SSMs), especially Mamba, have demonstrated remarkable capabilities in various fields including language, images, audio, and genomics. Prior to Mamba, large transformer-based language models achieved great success in sequence modeling due to the self-attention mechanism. However, it suffers from quadratic growth in storage and computational complexity as the sequence length increases, and has difficulty processing sequences with long-range dependencies. Compared to other SSM and pioneers such as Transformer, Mamba achieves higher algorithmic accuracy (Mamba-2.8b is more accurate than OPT-6.7b) and more efficient hardware computing efficiency.

[0003] Since Mamba is a newer model architecture and introduces the SSM computing module, which contains a large number of element-by-element operations, most of the current large model accelerators are specially designed for the transformer architecture and are therefore not suitable for Mamba.

[0004] There are the following problems when running Mamba on existing mainstream GPUs:

[0005] (1) Incompatibility between element-by-element operations and Tensor Cores. Linear operations (such as matrix multiplication and convolution) and element-by-element operations are the two main operations in Mamba. However, as the length of the sequence increases, the proportion of time spent on element-by-element operations that do not require reduction increases significantly, resulting in extremely low computational efficiency on GPUs that are mainly based on Tensor Cores that focus on accelerating linear operations.

[0006] (2) Nonlinear functional units have large area overhead. The exponential function and the Sigmoid Linear Unit (SiLU) function are two nonlinear functions in Mamba. The exponential function of the optimized nonlinear functional unit still occupies 30% of the PE area.

[0007] (3) Element-by-element operations have large memory accesses but limited input data sharing. Linear and element-by-element operations in Mamba exhibit large variance in computational intensity (e.g., ≈3 orders of magnitude) and large variance in read / write ratios (e.g., >3 orders of magnitude). Due to the limited data sharing of element-by-element operations, it is ineffective to apply existing methods such as tiling to element-by-element operations. Summary of the invention

[0008] The purpose of the present invention is to overcome the above defects. The present invention provides a state space model accelerator, a Mamba model calculation method and a calculation system. The state space model accelerator accelerates the scheduling balance of element-by-element operations and linear operations in Mamba. To achieve the above purpose, the technical solution of the present invention includes:

[0009] A state space model accelerator, comprising a reconfigurable computing unit for performing linear calculation and element-by-element calculation; the reconfigurable computing unit comprises a reduction tree and a processing unit array; the processing unit array comprises a plurality of reconfigurable processing units;

[0010] When performing linear calculation, each of the reconfigurable processing units calculates the input data; the reduction tree is in reduction mode, and each reduction tree slice inside it performs reduction calculation on the calculation results of the corresponding multiple reconfigurable processing units;

[0011] When performing element-by-element calculation, each of the reconfigurable processing units performs element-by-element calculation on the input data and outputs the calculation results in parallel; the reduction tree is in a non-reduction mode.

[0012] A further improvement of the present invention is that the linear calculation includes matrix multiplication operation; during the matrix multiplication operation, each of the reconfigurable processing units is configured to use a floating-point multiplication unit to calculate the product of two input data; the reduction calculation performed by the reduction tree slice includes the sum of the calculation results of each of the corresponding reconfigurable processing units through multi-level addition calculation.

[0013] A further improvement of the present invention is that during the linear calculation process, the on-chip buffer is configured as a read buffer; the memory access handler loads the input data of the linear calculation from the global memory to fill the read buffer; the reconfigurable computing unit reads the required data from the read buffer as input data, and writes the calculation result back to the global memory.

[0014] A further improvement of the present invention is that when performing element-by-element calculation, each of the reconfigurable processing units outputs the calculation results in parallel to the P2D buffer.

[0015] A further improvement of the present invention is that the element-by-element calculation includes element-by-element multiplication and element-by-element addition of matrices.

[0016] A further improvement of the present invention is that the element-by-element calculation includes element-by-element exponential operation, and each of the reconfigurable processing units uses a fast biased exponential algorithm to perform element-by-element exponential operation.

[0017] A further improvement of the present invention is that the element-by-element calculation includes element-by-element SiLU operation; each of the reconfigurable processing units uses a piecewise approximate SiLU algorithm to perform element-by-element SiLU operation.

[0018] A further improvement of the present invention is that it also includes an instruction processing unit; the instruction processing unit is used to load instructions from the global memory and control the reconfigurable computing unit according to the instructions.

[0019] The present invention also provides a Mamba model calculation method, which is implemented using the above-mentioned state space model accelerator; it is characterized in that the SSM calculation module of the Mamba model includes a process of multiple cyclic updates of the state h, during which the input variables that participate in the update of the state h multiple times are cached in an on-chip buffer to reduce the number of memory accesses.

[0020] The present invention also provides a computing system, which includes the above-mentioned state space model accelerator.

[0021] The beneficial effects of the present invention include:

[0022] 1. The state-space model accelerator reduces the replaceable PE array architecture, which can efficiently support both element-by-element operations and linear operations;

[0023] 2. The approximate algorithms of SiLU and e^x in Mamba and the corresponding hardware architecture can perform element-by-element calculations efficiently;

[0024] 3. Different buffer management strategies for different operation types in Mamba; can reduce the number of memory accesses, minimize idle computing resources, and thus speed up the Mamba computing process. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] Figure 1 The calculation process of the Mamba model;

[0026] Figure 2 is the ratio distribution of linear calculation to element-by-element calculation in the Mamba model;

[0027] Figure 3 is an architecture diagram of the state space model accelerator in this embodiment;

[0028] Figure 4 is a schematic diagram of a reconfigurable processing unit in this embodiment;

[0029] Figure 5 Schematic diagram of the fast biased index algorithm;

[0030] Figure 6 Schematic diagram of intra-operation and inter-operation cache management strategies;

[0031] Figure 7 This is a performance comparison chart of the state space model accelerator of this embodiment and the prior art. DETAILED DESCRIPTION

[0032] The following describes the embodiments of the present invention by specific examples, and those skilled in the art can easily understand other advantages and effects of the present invention from the contents disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and the details in this specification can also be modified or changed in various ways based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that the following embodiments and features in the embodiments can be combined with each other without conflict.

[0033] like Figure 3 , Figure 4 As shown, an embodiment of the present invention includes a state space model accelerator. This embodiment takes the calculation of the Mamba model as an example to introduce the state space model accelerator. The Mamba model includes linear operations, element operations, and a small number of activation functions. The difference between element-by-element operations and linear operations is whether reduction is performed. The reconfigurable computing unit (RCU) of this embodiment can perform element-by-element operations based on the processing unit array (PE array) of the Tensor Core by disabling the reduction tree.

[0034] The reconfigurable computing unit (RCU) of this embodiment is a reduction replaceable processing unit array. The replaceable processing unit (PE) array supports configurability. The reduction tree (Reduction Tree) consists of 16 slices and runs in two modes: reduction mode for linear operations and non-reduction mode for element-by-element operations. In the reduction mode (such as Figure 4 As shown in the lower left corner, for a 16-to-1 reduction tree slice, the output of the 16×1 processing unit array is fed into a reduction tree slice that uses multi-level addition to calculate the sum. In non-reduction mode, the output of the 16×16 processing unit array directly skips the reduction. In addition, the last level addition of each reduction tree slice supports three inputs to accumulate partial results of linear operations.

[0035] In this embodiment, the reconfigurable computing unit (RCU) includes a 16×16 processing unit array and a reduction tree, as shown in the left figure below. Each reconfigurable processing unit (RPE) is configured to support three main calculations and three data paths: a shift path for exponential functions, a segment path for SiLU functions, and a normal path for addition or multiplication. The normal path contains a floating-point multiplier, a floating-point adder, and a multiplexer unit to handle element-by-element multiplication and addition calculations. Figure 4 As shown, the reconfigurable processing unit includes the following working modes:

[0036] (1) Matrix multiplication (MM-RCU). The reduction tree in the RCU and the floating-point multiplication unit in the RPE are enabled. Therefore, for a matrix multiplication operation of two matrices of size 16×16, the results calculated by the 16×16 multiplier array are then passed through the reduction tree to produce 16 final results. This process is then repeated a total of 16 times to obtain the complete result of the 16×16 matrix multiplication. To support the accumulation of partial sums, an additional adder is added to the last level of the reduction tree.

[0037] (2) Element-by-element operation (EW-RCU). The reduction tree is disabled, the P2D buffer is enabled, and the floating-point multiplier or addition unit in the RPE is activated. For the element-by-element multiplication operation of two 16×16 matrices, the result calculated by the 16×16 RPE array maintains the same size of 16×16 and is output in parallel to the P2D buffer; the P2D buffer is used to cache data for subsequent calculations or write data to the on-chip buffer.

[0038] (3) Exponential operation (EXP-RCU). The reduction tree is disabled. The floating-point multiplier, addition unit, and exponential shift unit in the RPE are enabled. Therefore, for a 16×16 matrix that performs an exponential function operation, the RCU first performs element-by-element multiplication using the multiplier and then performs element-by-element addition using the adder. Subsequently, the exponential shift unit performs logical operations, shift operations, and bias operations to obtain the final output.

[0039] (4) Activation function (SiLU-RCU). The reduction tree is disabled. The floating-point multipliers or adders, range detectors, and constant output units in the RPE are enabled, while the reduction tree is disabled. For a 16×16 input matrix, each input is first differentiated according to the range, and then each input is processed using constant output or normal element-by-element calculation depending on the range.

[0040] In this embodiment, the reusable nonlinear functional unit of the reconfigurable processing unit array is used to quickly perform element-by-element calculations, which includes the following characteristics:

[0041] By analyzing the input range of nonlinear functions in Mamba. The input of the exponential function is mostly between -7 and 0, especially the value slightly less than 0, while the input range of the SiLU function is from -5 to 4. On the one hand, the present embodiment can only approximate the nonlinear function within these ranges to focus on preventing the loss of precision. On the other hand, the exponential function and SiLU are similar to element-by-element operations except that each value is scaled. Therefore, in order to avoid increasing the area overhead and preventing the loss of precision, the present embodiment only approximates these nonlinear functions within a small range and decomposes them into several element-by-element operations to reuse the alternative processing unit array architecture described above.

[0042] (1) Fast biased exponential algorithm: Given the special distribution observed in the input of the exponential function, namely the outer product of Δ and A, this embodiment uses a set of data points x = -7 / n, n = 1, 2, ..., 200, where the density increases as they approach zero, to evaluate the deviation of the approximate calculation from the original exponential value. Therefore, this embodiment modifies the fast exp algorithm to adapt to the specific data range and adds a bias at the end to improve accuracy. The exponential function algorithm includes three main steps: linearly convert the input x to x', whose expression is x' = ax + b; Figure 5 As shown, x' is multiplied by 2^23 and then corrected to an unsigned integer; x is treated as a floating point number and a bias c is added. The purpose of the fast exp algorithm is to put x / ln2 into the exponent position of e^x so that e^x =

[0043] 2^(x / ln2). Floating point numbers do not have an intuitive representation like unsigned integers, so they cannot be used directly as the exponent bits of e^x, so it is necessary to convert to uint. In addition, a bias is required to compensate for the conversion loss. In actual calculations, the exponent of floating point numbers is usually subtracted by 127. Now the coefficient a=1 / ln(2) and the term b=127+c can be determined, and c is used as the final bias, which effectively reduces the average loss caused by the specific data distribution. This embodiment only adds an exponent shift unit and reuses other element-by-element units to achieve fast exponential operation with a bias.

[0044] (2) Piecewise approximation SiLU algorithm. The calculation formula of the SiLU function is SiLU(x)=x / (1+e^(-x)), which involves not only exponentiation but also division. In this embodiment, the SiLU function is segmented in this interval for approximation, and a 4-segment piecewise function is used to approximate the SiLU function, which is represented by the following figure: In each interval, further piecewise approximation is performed according to the distribution of the data.

[0045]

[0046] Where: f(x) is the SiLU function, and x is the input of the SiLU function.

[0047] like Figure 6 As shown, in this embodiment, different buffer management strategies for different operation types are adjusted to maximize data sharing within and between operations. Specifically, it includes:

[0048] (1) Intra-operation management. For linear operations, such as Figure 6As shown in the upper part, the entire on-chip buffer is configured as a read buffer. The memory access handler loads linear inputs from the global memory to fill the buffer. The computing unit then sequentially reads the required data from the buffer, performs the calculation, and writes the calculation results back to the global memory. For Element-wise 2 operations, even if the read / write ratio is relatively low, it can be effectively regarded as a matrix multiplication with two dimensions reduced to 1, so data needs to be shared between the inputs. Therefore, in this embodiment, a small area is reserved in the on-chip buffer to store them. Element-wise 1 operation refers to the element-by-element addition or multiplication of two matrices, including SiLU and exponential operations. Element-wise 2 refers to the outer product of two vectors or matrices, which is essentially the Cartesian product of different dimensions of two tensors to increase the dimension. There is room for input data reuse in Element-wise 2.

[0049] (2) Operation room management. Figure 6 As shown in the middle, Element-wise 1 does not reuse data when computing input data. The basic approach involves reading data from global memory and writing it back directly after computing. For a single computation, data reuse optimizations such as blocking are not feasible. For Element-wise 1 and Element-wise 2 operations, in order to maximize buffer utilization, this embodiment mainly optimizes data sharing between adjacent element-wise operations, such as Figure 6 Shown at bottom.

[0050] like Figure 1 As shown, Figure 1 The right side shows the calculation process of the SSM calculation module of the Mamba model; during the calculation process, the input X undergoes three matrix multiplications to obtain Δ, B, and C, Δ is multiplied element by element by A to obtain ΔA, and Δ is multiplied element by element by B to obtain ΔB. The SSM calculation module includes the process of updating the state h multiple times in a loop.

[0051] In the process of obtaining state h each time, ΔA, ΔB, X and the previous state h are used as input variables, and the new state h is obtained by adding the product of ΔA and h and the product of ΔB and x.

[0052] In the process of continuously updating the state h, h needs to be repeatedly read and written L times. In addition, the corresponding ΔA, ΔB and x need to be repeatedly read L times. Therefore, for the element-by-element operations concentrated in the SSM process, this embodiment caches the above direct results into the buffer. Through the operation-level cache management strategy, the number of memory accesses is minimized to the maximum extent, and the idle computing resources are minimized, thereby accelerating the Mamba computing process.

[0053] The state space model accelerator of this embodiment further includes an instruction processing unit (InstructionProcessing); the instruction processing unit is used to load instructions from the global memory and control the reconfigurable computing unit according to the instructions.

[0054] An embodiment of the present invention also includes a Mamba model calculation method, which is implemented using the above-mentioned state space model accelerator. In the SSM calculation module of the Mamba model, the above-mentioned inter-operation management is used to cache input variables that participate in the state h update multiple times in an on-chip buffer to reduce the number of memory accesses.

[0055] The embodiment of the present invention also includes a computing system, which includes the above-mentioned state space model accelerator. The state space model accelerator can be integrated into the computing system in the form of a board or a chip on board, or can be integrated into a system on chip as an IP core.

[0056] like Figure 7 As shown, compared with Mamba-CPU and Mamba-GPU, MARCA in this embodiment achieves acceleration of 463.32x / 11.66x and 194.26x / 4.93x respectively, and the energy efficiency is improved by a maximum of 9761.42x / 242.52x and an average of 3415.55x / 42.49x respectively.

[0057] As shown in Table 1, the total power consumption and area of ​​the design of this embodiment are only 10.44 watts and 221.88 square millimeters, respectively. For the on-chip buffer, this embodiment uses eDRAM to reduce area and energy consumption. In terms of computational accuracy, this embodiment uses 32-bit fixed-point representation, which is sufficient to maintain the accuracy of Mamba reasoning. The following table provides a detailed breakdown of area and power consumption. The on-chip buffer consumes most of the power consumption (>60%) and area (about 80%) to support more memory accesses and data sharing for element-by-element operations. The compute engine consumes 38.67% of the power consumption and 20.57% of the area to perform linear, element-by-element, exponential, and SiLU calculations. Other components use less power and area due to the simple implementation of normalization and instruction fetch and decode.

[0058] Table-1

[0059]

[0060] The above embodiments are merely illustrative of the principles and effects of the present invention, and are not intended to limit the present invention. Anyone familiar with the art may modify or alter the above embodiments without departing from the spirit and scope of the present invention. Therefore, all equivalent modifications or alterations made by a person of ordinary skill in the art without departing from the spirit and technical concept disclosed by the present invention shall still be covered by the claims of the present invention.

Claims

1. A state space model accelerator, characterized in that: It includes a reconfigurable computing unit for performing linear calculation and element-by-element calculation; the reconfigurable computing unit includes a reduction tree and a processing unit array; the processing unit array includes a plurality of reconfigurable processing units; When performing linear calculation, each of the reconfigurable processing units calculates the input data; the reduction tree is in reduction mode, and each reduction tree slice inside it performs reduction calculation on the calculation results of the corresponding multiple reconfigurable processing units; When performing element-by-element calculation, each of the reconfigurable processing units performs element-by-element calculation on the input data and outputs the calculation results in parallel; the reduction tree is in a non-reduction mode.

2. A state space model accelerator according to claim 1, characterized in that: The linear calculation includes matrix multiplication operation; during the matrix multiplication operation, each of the reconfigurable processing units is configured to use a floating-point multiplication unit to calculate the product of two input data; the reduction calculation performed by the reduction tree slice includes the sum of the calculation results of each of the reconfigurable processing units corresponding to the multi-level addition calculation.

3. A state space model accelerator according to claim 2, characterized in that: During the linear calculation process, the on-chip buffer is configured as a read buffer; the memory access handler loads the input data of the linear calculation from the global memory to fill the read buffer; the reconfigurable computing unit reads the required data from the read buffer as input data, and writes the calculation result back to the global memory.

4. A state space model accelerator according to claim 1, characterized in that: When performing element-by-element calculation, each of the reconfigurable processing units outputs the calculation results in parallel to the P2D buffer.

5. A state space model accelerator according to claim 1, characterized in that: Element-by-element computations include element-by-element multiplication and element-by-element addition of matrices.

6. A state space model accelerator according to claim 1, characterized in that: The element-by-element calculation includes element-by-element exponential operation, and each of the reconfigurable processing units uses a fast biased exponential algorithm to perform element-by-element exponential operation.

7. A state space model accelerator according to claim 1, characterized in that: The element-by-element calculation includes element-by-element SiLU operation; each of the reconfigurable processing units adopts a piecewise approximate SiLU algorithm to perform element-by-element SiLU operation.

8. A state space model accelerator according to claim 1, characterized in that: It also includes an instruction processing unit; the instruction processing unit is used to load instructions from the global memory and control the reconfigurable computing unit according to the instructions.

9. A Mamba model calculation method, implemented by using the state space model accelerator described in any one of claims 1 to 8; characterized in that: The SSM calculation module of the Mamba model includes a process of updating the state h multiple times in a loop. During this process, the input variables that participate in the update of the state h multiple times are cached in an on-chip buffer to reduce the number of memory accesses.

10. A computing system, characterized in that: The method comprises the state space model accelerator as claimed in any one of claims 1 to 8.

Images