A matrix multiplication acceleration method for a deep learning model small batch inference scene, a storage medium and an equipment

Abstract

The application discloses a matrix multiplication acceleration method for a small-batch inference scene of a deep learning model, a storage medium and equipment, and comprises the following steps: acquiring an activation matrix and a weight matrix input in a current small-batch calculation task in a deep learning model inference process; when a general multiplication instruction is started and called in a calculation kernel of the deep learning model, a data pointer pointing to the weight matrix is bound as a left operand of the general multiplication instruction to read the weight matrix; a data pointer pointing to the activation matrix is bound as a right operand of the general multiplication instruction to read the activation matrix; a tensor core of a GPU calculates a transposed form of the current small-batch calculation task according to the read weight matrix and the activation matrix; and each element in the transposed form of the current small-batch calculation task is subjected to coordinate decoupling and address mapping to realize accelerated calculation of matrix multiplication of the small-batch calculation task.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, specifically to a matrix multiplication acceleration method, storage medium, and device for small-batch inference scenarios of deep learning models. Background Technology

[0002] In recent years, online inference deployment of deep learning models has typically been divided into two stages: pre-filling and decoding. During the decoding stage, the deep learning model generates tokens one by one in an autoregressive manner. At this stage, the batch size or the number of effective tokens input to a specific layer of the deep learning model is usually extremely small. In the underlying matrix multiplication operations, this input processing size or the number of effective tokens directly corresponds to the M-dimensionality of the input activation matrix involved in the computation.

[0003] To ensure extremely low user interaction latency, deep learning models often need to perform matrix multiplication calculations at high frequency with a very small M dimension, especially in hybrid expert architectures. Because the dynamic routing mechanism of hybrid experts allocates tokens to different experts for computation, the actual batch size allocated to a single expert is usually extremely small and highly fragmented, for example, M is only 8 to 32. MegaBlocks is an efficient sparse training technique specifically designed for hybrid expert architectures. This technique reveals the sparse computation challenges faced by hybrid expert architectures when handling dynamic routing. However, in microsecond-level small-batch inference scenarios, its macro-dimensional optimization still cannot resolve the underlying computation alignment conflicts.

[0004] Modern AI acceleration hardware such as GPUs and TPUs, in pursuit of maximum data throughput, employ execution engines designed for intensive, block-based computation in their underlying tensor cores and other computing units. These engines impose strict hardware alignment requirements on the dimensions of the input matrices. For example, the WGMMA instruction in NVIDIA's Hopper architecture mandates that the M dimension of the input activation matrix must be aligned to a multiple of 64, while the N dimension of the weight matrix only requires alignment to a multiple of 8.

[0005] Faced with such hardware constraints, existing GEMM algorithms optimized for standard dense matrices prove extremely rigid when handling fragmented tasks. Stream-K is a task-centric parallel workload decomposition technique. While such techniques significantly improve load balancing of dense matrices through workflow decomposition, their standard computational kernel still heavily relies on the alignment granularity requirements of the underlying hardware. When using such existing techniques to compute mini-batch tasks, such as when M=16 needs to be computed, the system is forced to zero-padded M dimension up to the 64 boundary. This padding results in up to 75% of the computational power being wasted on invalid zero-multiply-accumulate operations, severely reducing hardware utilization.

[0006] Since existing standard computing engines have fixed the scheduling mode of tensor cores and lack a flexible alignment mechanism for microsecond-level extremely small batches, how to break the strict M-dimensional alignment limit of the underlying hardware without increasing additional memory read and write overhead, efficiently map highly fragmented small batch tasks to tensor cores, and achieve lossless release of computing power has become a core problem that urgently needs to be solved in online inference of deep learning models. Summary of the Invention

[0007] To address the problems existing in the prior art, this invention provides a matrix multiplication acceleration method, storage medium, and device for small-batch inference scenarios of deep learning models. It overcomes the serious zero-filling computational waste caused by the forced alignment of the input batch size or the number of valid tokens by hardware instructions during small-batch tensor computation in the prior art. At the same time, it breaks through the extremely high memory access and transportation overhead barrier caused by conventional explicit transpose methods, and realizes a zero-overhead implicit transpose matrix multiplication computation.

[0008] To achieve the above technical objectives, the present invention adopts the following technical solution:

[0009] A matrix multiplication acceleration method for mini-batch inference scenarios in deep learning models includes the following steps: Step S1: Obtain the activation matrix and weight matrix of the current mini-batch computation task input during the deep learning model inference process; Step S2: When the computation kernel of the deep learning model starts and calls the general multiplication instruction, bind the data pointer pointing to the weight matrix as the left operand of the general multiplication instruction to read the weight matrix; bind the data pointer pointing to the activation matrix as the right operand of the general multiplication instruction to read the activation matrix. Step S3: The GPU's tensor core calculates the transpose of the current mini-batch computation task based on the read weight matrix and activation matrix; Step S4: Decouple the coordinates and map the addresses of each element in the transpose of the current mini-batch computation task to complete the matrix multiplication calculation of the mini-batch computation task.

[0010] Furthermore, the current mini-batch computation task in step S1 is represented as follows:

[0011] in, This indicates the current small batch computation task. This represents the activation matrix input to the current mini-batch computation task, with dimension 1. , This indicates the batch size or number of valid tokens input to the current mini-batch computation task. This indicates the network input dimension of the deep learning model for processing the current mini-batch computation task; This represents the weight matrix input to the current mini-batch computation task, with dimension 1. , This represents the network output dimension of the deep learning model when processing the current mini-batch computation task.

[0012] Furthermore, the left operand reads the weight matrix in physical memory in a row-first manner, and the right operand reads the activation matrix in physical memory in a column-first manner.

[0013] Furthermore, if the current mini-batch computation task inputs a batch size or a number of valid tokens... The tensor cores of the GPU do not conform to the network output dimension Requirements, for Continue filling until the requirements are met.

[0014] Furthermore, step S4 includes the following sub-steps: Step S4.1: The computation thread of the deep learning model is responsible for writing the two-dimensional coordinates of each element in the transpose of the current mini-batch computation task. Two-dimensional coordinates of the element at the register level Perform a swap to obtain the swapped coordinates of the elements. ; Step S4.2: The calculation thread swaps the coordinates of the element. The physical write address of the element is determined by the anti-collision shared memory address mapping function, and the write is performed to complete the matrix multiplication calculation for the small batch calculation task.

[0015] Furthermore, the present invention also provides a computer-readable storage medium storing a computer program that enables a computer to execute the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models.

[0016] Furthermore, the present invention also provides an electronic device, including: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the computer program, it implements the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models.

[0017] Compared with the prior art, the present invention has the following beneficial effects: This invention provides a matrix multiplication acceleration method for mini-batch inference scenarios in deep learning models. By transposing the activation and weight matrices input to the mini-batch computation task, the alignment requirement for the M-dimensional matrix multiplication is changed from a coarse-grained M-dimensional alignment requirement to a fine-grained N-dimensional alignment requirement. This significantly reduces zero-padding space and effectively avoids the stepped computational waste caused by forced large-grained alignment. Secondly, this invention achieves near-zero latency in the overall matrix multiplication computation by decoupling the operand positions at the instruction level, effectively breaking through the performance barrier of explicit transposition and realizing flexible dimensional conversion at almost zero cost. Attached Figure Description

[0018] Figure 1 This is a flowchart of the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to the present invention. Figure 2 This is a comparison chart of the operator delay performance of the present invention and existing explicit transpose matrix multiplication methods under different input batch sizes. Detailed Implementation

[0019] The technical solution of the present invention will be further explained and described below with reference to the accompanying drawings.

[0020] In the autoregressive decoding stage of deep learning models, because the model generates content in a token-by-token manner, the batch size of the basic input or the number of effective tokens is inherently small. This is especially true after introducing a hybrid expert MoE architecture, where the dynamic routing mechanism of the hybrid experts allocates tokens to different experts for computation. Therefore, in the actual computation task reaching a single expert's underlying operator, the M-dimensionality of its activation matrix becomes severely fragmented, typically degenerating to extremely small values ​​such as 8 to 32. This extremely small and dynamic M-dimensionality severely conflicts with the alignment dimension (e.g., 64) required by the underlying tensor core for high throughput. This invention cleverly avoids wasting computational power by using a transpose equivalent transformation to map the fragmented dimension, constrained by hardware-mandated alignment, to a dimension that allows for flexible alignment supported by hardware.

[0021] like Figure 1 This is a flowchart of a matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to the present invention. The matrix multiplication acceleration method includes the following steps: Step S1: Obtain the activation matrix and weight matrix of the current mini-batch computation task input during the deep learning model inference process, wherein the current mini-batch computation task in this invention is represented as:

[0022] in, This indicates the current small batch computation task. This represents the activation matrix input to the current mini-batch computation task, with dimension 1. , This indicates the batch size or number of valid tokens input to the current mini-batch computation task. This indicates the network input dimension of the deep learning model for processing the current mini-batch computation task; This represents the weight matrix input to the current mini-batch computation task, with dimension 1. , This represents the network output dimension of the deep learning model when processing the current mini-batch computation task.

[0023] Step S2: When the deep learning model's computation kernel starts and calls the general multiplication instruction, the data pointer pointing to the weight matrix is ​​bound to the left operand of the general multiplication instruction to read the weight matrix; the data pointer pointing to the activation matrix is ​​bound to the right operand of the general multiplication instruction to read the activation matrix. This ensures that the logical transposition does not incur any additional data handling or memory copying overhead during actual physical execution. Specifically, the left operand reads the weight matrix in physical memory row-major order, and the right operand reads the activation matrix in physical memory column-major order. When the deep learning model reads the activation matrix that was originally stored row-major order... When the right operand position is passed, the hardware forces it to be loaded according to the column-major order. By reading a matrix stored row-majorly by column, mathematically and physically, this is equivalent to directly reading the activation matrix. The transpose of the data. Taking advantage of this read characteristic, the activation matrix is ​​completed the instant the command retrieves the data. And matrix weight matrix The transpose transformation achieves true zero-overhead transpose. Through this process, the originally highly fragmented M dimension, which was constrained by mandatory alignment, becomes a "flexible dimension" in the new task; while the originally larger N dimension is mapped to a "mandatory alignment dimension" in the new task, fundamentally avoiding the need for invalid zero-padding of the small M dimension.

[0024] After performing the transpose, the dimension in the original matrix multiplication calculation becomes N dimensions in the transposed matrix multiplication calculation. This is relevant if the batch size or the number of valid tokens input to the current mini-batch computation task is [not specified]. The tensor cores of the GPU do not conform to the network output dimension Requirements, for Continue filling until the requirements are met.

[0025] Step S3: The GPU's tensor core calculates the transpose of the current mini-batch computation task based on the read weight matrix and activation matrix.

[0026] Step S4: Before writing the transposed form of the current mini-batch computation task back to globally available memory, it is necessary to reverse it back to the original representation of the mini-batch computation task, while ensuring that the original system's anti-memory block conflict mechanism is not violated. Therefore, this invention performs coordinate decoupling and address mapping on each element in the transposed form of the current mini-batch computation task to complete the matrix multiplication calculation of the mini-batch computation task, including the following sub-steps: Step S4.1: The computation thread of the deep learning model is responsible for writing the two-dimensional coordinates of each element in the transpose of the current mini-batch computation task. Two-dimensional coordinates of the element at the register level Perform a swap to obtain the swapped coordinates of the element. ; Step S4.2: Calculate the swapped coordinates of the element using the thread. The physical write address of the element is determined by the conflict-preventing shared memory address mapping function, and the write is performed. When the data falls into physical memory, its arrangement is automatically restored to the original representation of the mini-batch computation task, completing the matrix multiplication calculation of the mini-batch computation task. This process achieves zero-overhead result reversal while perfectly preserving the conflict-free memory access characteristics of the original code.

[0027] Example 1: Accelerating Matrix Multiplication Based on Expert Dimension of DeepSeek-V3 / R1 Architecture This embodiment runs on an accelerator card with a tensor core and supports low-level matrix multiplication and addition instructions, such as an NVIDIA Hopper architecture GPU. The application scenario is the computation of a hybrid expert module feedforward network (FFN) based on the DeepSeek-V3 / R1 architecture. The weight matrix dimension of the target expert in this hybrid expert module feedforward network (FFN) is... , During the autoregressive decoding phase, the system continuously receives small batches of inference requests, of which the number of valid tokens allocated to the expert is... The step size is dynamically varied between 8 and 128.

[0028] The matrix multiplication acceleration method in this embodiment is based on Let's take an example to illustrate.

[0029] (1) Dimension matching and decision-making: The system scheduler captures the current mini-batch computation task of the Hybrid Expert Module Feedforward Network (FFN) and inputs the activation matrix. The dimension is Weight matrix The dimension is The underlying general-purpose multiplication instruction mandates that the left operand... Dimensions must be aligned to 64, right operands Dimensional alignment is supported with a step size of 8. Determine the current... The alignment requirement of 64 is not met. If the native implementation is used, it would need to be forcibly padded to 64, resulting in a large amount of invalid computation. Therefore, the scheduler decides to trigger the hardware adaptation transpose mapping mechanism of this invention.

[0030] (2) Logical shape reconstruction: The system mathematically reconstructs the small-batch computation task into a transpose equivalent form. At this point, the left operand becomes the weight matrix. The transpose of the form, with dimensions of The right operand becomes the activation matrix. The transpose of the form, with dimensions of After reconstruction, the new task's M' = 2048 perfectly satisfies the 64-fold multiple alignment requirement, and the new task's N' = 16 satisfies the hardware-supported 8-fold alignment requirement, logically avoiding the need for forced zero padding. Therefore, when the computation kernel starts, no explicit memory read / write functions are started for transposition; the thread block directly points to the activation matrix. The pointer is bound to the right operand descriptor of the general multiplication instruction, pointing to the weight matrix. The pointer is bound to the left operand descriptor of the general multiplication instruction. Since hardware typically expects the right operand to be column-majorly laid out, the activation matrix, which is normally stored in row-major order in physical memory... The input matrix is ​​automatically reinterpreted and read column by column by the hardware. This underlying hardware read characteristic allows the system to implicitly transpose the input matrix within zero cycles.

[0031] (3) Tensor core computation: The GPU's tensor core is calculated according to the reconstructed, unfilled compact dimension. , Perform high-throughput multiply-accumulate operations and slice the calculated results into registers.

[0032] (4) Coordinate Decoupling, Reversal, and Write-back: In the final stage, the result slice in the register needs to be written to shared memory. The system has a pre-defined Swizzle addressing function to prevent memory block conflicts. Each thread obtains the two-dimensional coordinates of the elements in its result slice to be written. Then, first, coordinate decoupling and flipping are performed, changing them to... Then, the flipped coordinates are passed to the Swizzle addressing function to generate a physical address and written. When the data is written to physical memory, its arrangement is automatically restored to the original requirements. The layout of small-batch computing tasks not only retains the conflict-free memory access characteristics of the original code, but also completes the result inversion with zero overhead.

[0033] To verify the effectiveness of this embodiment, the number of valid input tokens was tested. Perform parameter scan tests with a step size of 8 in the range of 8 to 128: Due to the mandatory 64-bit alignment, the execution time of the native standard computing library is... It remains constant between 64 and 64, and At step 128, a step-like jump occurs; however, this invention reduces the alignment granularity to a step size of 8, effectively avoiding this waste of large-granularity alignment.

[0034] like Figure 2 If the conventional method of explicit transposition and overlaying the native computing library is used, the high overhead of explicit transposition memory copy will cause the total latency to surge to 7.93 times to 14.75 times that of the native baseline. However, since the present invention completely eliminates the physical transposition cost, the total computation latency is strictly controlled to 0.98 times to 1.01 times that of the ideal pure computation time.

[0035] This embodiment, based on expert-level experiments using the DeepSeek-V3 / R1 architecture, demonstrates that the present invention successfully translates theoretical computational savings into actual performance improvements. Compared to native standard computing libraries, the execution time of the present invention is reduced to 0.60 to 1.07 times the baseline, with an average latency reduction of 17.26% and a maximum latency reduction of up to 40%. Although in When the dimension is exactly a multiple of 64, the native baseline has a slight advantage because it does not require padding and has a coarser block granularity. However, in most irregular and highly fragmented small-batch inference scenarios, this invention achieves significant acceleration and provides flexible support for irregular shapes.

[0036] In one technical solution of the present invention, a computer-readable storage medium is also provided, storing a computer program that enables a computer to execute the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models.

[0037] In one technical solution of the present invention, an electronic device is also provided, including: a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, it implements the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models.

[0038] In the embodiments disclosed in this application, a computer storage medium may be a tangible medium that may contain or store programs for use by or in conjunction with an instruction execution system, apparatus, or device. The computer storage medium may include, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, apparatus, or devices, or any suitable combination of the foregoing. More specific examples of computer storage media include electrical connections based on one or more wires, portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fibers, and portable compact disc read-only memory (CD). ROM), optical storage devices, magnetic storage devices, or any suitable combination of the foregoing.

[0039] Those skilled in the art will recognize that the units and algorithm steps of the various examples described in conjunction with the embodiments disclosed in this application can be implemented in electronic hardware or a combination of computer software and electronic hardware. Whether these functions are implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this application.

[0040] The above are merely preferred embodiments of the present invention. The scope of protection of the present invention is not limited to the above embodiments. All technical solutions falling within the scope of the present invention's concept are within the scope of protection of the present invention. It should be noted that for those skilled in the art, any improvements and modifications made without departing from the principles of the present invention should be considered within the scope of protection of the present invention.

Claims

1. A matrix multiplication acceleration method for small-batch inference scenarios in deep learning models, characterized in that, Includes the following steps: Step S1: Obtain the activation matrix and weight matrix of the current mini-batch computation task input during the deep learning model inference process; Step S2: When the computation kernel of the deep learning model starts and calls the general multiplication instruction, bind the data pointer pointing to the weight matrix as the left operand of the general multiplication instruction to read the weight matrix; bind the data pointer pointing to the activation matrix as the right operand of the general multiplication instruction to read the activation matrix. Step S3: The GPU's tensor core calculates the transpose of the current mini-batch computation task based on the read weight matrix and activation matrix; Step S4: Decouple the coordinates and map the addresses of each element in the transpose of the current mini-batch computation task to complete the matrix multiplication calculation of the mini-batch computation task.

2. The matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to claim 1, characterized in that, The current mini-batch computation task in step S1 is represented as follows: in, This indicates the current small batch computation task. This represents the activation matrix input to the current mini-batch computation task, with dimension 1. , This indicates the batch size or number of valid tokens input to the current mini-batch computation task. This indicates the network input dimension of the deep learning model for processing the current mini-batch computation task; This represents the weight matrix input to the current mini-batch computation task, with dimension 1. , This represents the network output dimension of the deep learning model when processing the current mini-batch computation task.

3. The matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to claim 1, characterized in that, The left operand reads the weight matrix in physical memory in a row-first manner, and the right operand reads the activation matrix in physical memory in a column-first manner.

4. The matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to claim 2, characterized in that, If the current small-batch computation task inputs the batch size or the number of valid tokens... The tensor cores of the GPU do not conform to the network output dimension Requirements, for Continue filling until the requirements are met.

5. The matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models according to claim 1, characterized in that, Step S4 includes the following sub-steps: Step S4.1: The computation thread of the deep learning model is responsible for writing the two-dimensional coordinates of each element in the transpose of the current mini-batch computation task. Two-dimensional coordinates of the element at the register level Perform a swap to obtain the swapped coordinates of the elements. ; Step S4.2: The calculation thread swaps the coordinates of the element. The physical write address of the element is determined by the anti-collision shared memory address mapping function, and the write is performed to complete the matrix multiplication calculation for the small batch calculation task.

6. A computer-readable storage medium storing a computer program, characterized in that, The computer program causes the computer to execute the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models as described in any one of claims 1-5.

7. An electronic device, characterized in that, include: The memory, the processor, and the computer program stored in the memory and executable on the processor, wherein when the processor executes the computer program, it implements the matrix multiplication acceleration method for mini-batch inference scenarios of deep learning models as described in any one of claims 1-5.

Images