Data processing method and device, electronic equipment and storage medium

Abstract

The invention relates to the technical field of artificial intelligence chips, and provides a data processing method and device, electronic equipment and a storage medium, and the method comprises the steps: obtaining a to-be-processed data set; segmenting each piece of to-be-processed data in the to-be-processed data set to obtain a plurality of data blocks corresponding to each piece of to-be-processed data; distributing the data processing tasks corresponding to the to-be-processed data to at least one processing unit; and for each processing unit, each data block related to the current data processing task of the processing unit is preloaded to a target memory area, each data block is loaded from the target memory area for data processing, and the target memory area is a non-consistent memory access area. According to the method and the device, the data blocks related to the current data processing task of each processing unit are preloaded to the target memory area, so that bandwidth competition in the data loading process can be effectively avoided, the stability and the high efficiency of data loading are ensured, and the operator performance is improved.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence chip technology, and in particular to a data processing method, device, electronic device and storage medium. Background Art

[0002] With the rapid development of artificial intelligence (AI), the ability of models to process large amounts of data is being impacted by memory and computational complexity. For example, in models based on the Transformer architecture, the memory requirements for attention, one of its core operations, increase quadratically with the sequence length when processing long sequence inputs. This can easily exceed the hardware memory capacity, limiting the model's ability to process long text and large amounts of data.

[0003] To address this issue, the industry has proposed various optimization strategies, with block processing being one effective approach. For example, FlashAttention, a representative technical solution, divides the Q (Query), K (Key), and V (Value) matrices used in the attention calculation into multiple small blocks, then performs calculations on them. This effectively reduces the memory required for a single calculation, enabling the model to process longer input sequences with limited hardware resources.

[0004] However, when calculating by block, each Q block needs to traverse all K and V blocks to complete the calculation, resulting in the same K and V blocks being loaded repeatedly. Moreover, since Q, K, and V blocks are usually stored in the UMA (Uniform Memory Access) area of HBM (High Bandwidth Memory), bandwidth competition problems may easily arise when loading Q, K, and V blocks from HBM UMA, thereby limiting the data transmission speed. The computing unit becomes idle while waiting for data transmission and cannot fully utilize its computing power, which in turn seriously restricts operator performance and model training and inference speed. Summary of the Invention

[0005] The present invention provides a data processing method, device, electronic device and storage medium to solve the defect of related technologies in which repeated loading of data during data processing causes bandwidth competition, thereby affecting operator performance.

[0006] The present invention provides a data processing method, comprising: Acquire a data set to be processed, wherein the data set to be processed includes at least one batch of data to be processed; Segmenting each data to be processed in the set of data to be processed to obtain a plurality of data blocks corresponding to each data to be processed; Allocating data processing tasks corresponding to each piece of data to be processed to at least one processing unit; For each processing unit, each data block related to the current data processing task of the processing unit is preloaded into a target memory area, and each data block is loaded from the target memory area for data processing. The target memory area is a non-uniform memory access area.

[0007] According to a data processing method provided by the present invention, the data processing task includes attention calculation based on each data block. Accordingly, the data to be processed in the data set to be processed is segmented to obtain multiple data blocks corresponding to each data to be processed, including: Convert any data to be processed into a matrix to be processed, wherein the matrix to be processed includes a query matrix, a key matrix and a value matrix; Each matrix in the matrix to be processed is divided to obtain a plurality of query matrix blocks, key matrix blocks and value matrix blocks corresponding to any data to be processed.

[0008] According to a data processing method provided by the present invention, preloading each data block related to the current data processing task of the processing unit into a target memory area, and loading each data block from the target memory area for data processing, includes: Preloading each key matrix block and each value matrix block related to the current data processing task of the processing unit into a target memory area; Traverse each query matrix block related to the current data processing task of the processing unit, preload the currently traversed query matrix block into the target memory area, and load the query matrix block, each key matrix block and each value matrix block from the target memory area to perform data processing.

[0009] According to a data processing method provided by the present invention, the query matrix block, each key matrix block, and each value matrix block are loaded from the target memory area to perform data processing, including: Traversing the key matrix blocks and the value matrix blocks, and loading the query matrix block, the currently traversed key matrix block, and the value matrix block from the target memory area; Perform calculation based on the currently loaded query matrix block and key matrix block to obtain a first calculation result; Calculation is performed based on the first calculation result and the currently loaded value matrix block to obtain a second calculation result, and the second calculation result is accumulated to the output result.

[0010] According to a data processing method provided by the present invention, the step of allocating the target memory area includes: Receive memory allocation request; The memory allocation request is parsed to obtain memory parameters, and memory allocation is performed on an off-chip memory based on the memory parameters to obtain a target memory area, wherein the memory parameters include memory size, memory location, and access mode.

[0011] According to a data processing method provided by the present invention, the data to be processed includes any one of text data, image data, and audio data.

[0012] The present invention also provides a data processing device, comprising: A data acquisition module is used to acquire a set of data to be processed, wherein the set of data to be processed includes at least one batch of data to be processed; A data segmentation module, configured to segment each to-be-processed data in the to-be-processed data set to obtain a plurality of data blocks corresponding to each to-be-processed data; A task allocation module, configured to allocate the data processing tasks corresponding to the data to be processed to at least one processing unit; The data processing module is used to preload each data block related to the current data processing task of each processing unit into a target memory area, and load each data block from the target memory area for data processing, wherein the target memory area is a non-uniform memory access area.

[0013] The present invention also provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and running on the processor, wherein the processor implements any of the above-mentioned data processing methods when executing the computer program.

[0014] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which implements any of the above-mentioned data processing methods when executed by a processor.

[0015] The present invention also provides a computer program product, comprising a computer program, wherein when the computer program is executed by a processor, the computer program implements any of the above-mentioned data processing methods.

[0016] The data processing method, device, electronic device and storage medium provided by the present invention distribute the data processing tasks corresponding to each data to be processed to at least one processing unit, so that each processing unit can process each data to be processed in parallel, thereby improving the overall processing efficiency. For each processing unit, by preloading each data block related to the current data processing task of the processing unit into the target memory area of the corresponding non-uniform memory access, the data loading operations of different processing units are independent of each other, thereby effectively avoiding bandwidth competition during the data loading process and ensuring the stability and efficiency of data transmission. In addition, since bandwidth competition is avoided, each processing unit does not need to be idle during the data processing process due to waiting for data transmission, and can continuously and efficiently perform data processing tasks, and the hardware computing power is fully utilized, which helps to improve operator performance and accelerate the model training and inference speed. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] In order to more clearly illustrate the technical solutions in the present invention or related technologies, the following is a brief introduction to the drawings required for use in the embodiments or related technical descriptions. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0018] Figure 1 It is a schematic diagram of data processing in related art; Figure 2 is a schematic structural diagram of a graphics processor provided by the present invention; Figure 3 It is a flow chart of the data processing method provided by the present invention; Figure 4 It is a schematic diagram of data loading of each processing unit provided by the present invention; Figure 5 1 is a flow chart of attention calculation based on input data preloading provided by the present invention; Figure 6 is a schematic diagram of the performance comparison of operators provided by the present invention; Figure 7 It is a structural diagram of the data processing device provided by the present invention; Figure 8 It is a structural schematic diagram of the electronic device provided by the present invention. DETAILED DESCRIPTION

[0019] To make the objectives, technical solutions, and advantages of the present invention more clear, the technical solutions of the present invention will be clearly and completely described below in conjunction with the accompanying drawings. Obviously, the embodiments described are only some of the embodiments of the present invention, not all of them. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts shall fall within the scope of protection of the present invention.

[0020] For ease of understanding, the technical terms involved in the embodiments of the present invention are explained below.

[0021] Uniform Memory Access (UMA): Each processing unit can access the same memory area. For example, taking the off-chip memory HBM as an example, for the UMA area in HBM, a piece of data is divided into data block 1, data block 2, data block 3, and data block 4 and then stored in the UMA area. Assuming there are processing units 1 to 4, each processing unit can access all data blocks in the UMA area. Specifically, processing unit 1 can access data blocks 1 to 4, processing unit 2 can also access data blocks 1 to 4, and similarly, processing units 3 and 4 can also access data blocks 1 to 4.

[0022] Non-Uniform Memory Access (NUMA): Each processing unit can only access its own memory region and cannot access the memory regions of other processing units. For example, taking the off-chip HBM memory as an example, the NUMA regions in the HBM are assigned to each processing unit. Suppose there are processing units 1 through 4, where the NUMA region corresponding to processing unit 1 is denoted as NUMA1, the NUMA region corresponding to processing unit 2 is denoted as NUMA2, the NUMA region corresponding to processing unit 3 is denoted as NUMA3, and the NUMA region corresponding to processing unit 4 is denoted as NUMA4. Each processing unit can only access data stored in its corresponding NUMA region. Suppose a piece of data is divided into data blocks 1, data blocks 2, data blocks 3, and data blocks 4, respectively, and stored in NUMA1 through NUMA4. In this case, processing unit 1 can only access data block 1 in NUMA1, processing unit 2 can only access data block 2 in NUMA2, processing unit 3 can only access data block 3 in NUMA3, and processing unit 4 can only access data block 4 in NUMA4.

[0023] Amid the rapid development of artificial intelligence (AI), model performance and computational efficiency have become key factors hindering its widespread adoption. Memory limitations and computational complexity are two key challenges. For example, when processing long sequences of input data, the memory requirements of the attention calculation, one of its core operations, increase quadratically with the sequence length. With traditional computing methods, excessively long input data can easily exceed the hardware's memory capacity, resulting in computational failures and severely impacting the model's ability to process large amounts of data.

[0024] To address this dilemma, the industry has proposed various optimization strategies, with block processing being one of the most effective approaches. Technical solutions, such as Flash Attention, perform block operations on the Q, K, and V matrices used in attention calculations. Specifically, this technology partitions the originally large Q, K, and V matrices into multiple smaller sub-matrices, then performs calculations on each block sequentially. This divide-and-conquer strategy effectively reduces the memory required for a single calculation, enabling the model to process longer input sequences with limited hardware resources. This significantly improves the model's ability to process large amounts of data, opening up the possibility of applying AI technology to more complex scenarios.

[0025] While the block processing approach alleviates memory pressure to a certain extent, it still exhibits numerous drawbacks in practical applications, severely hindering further improvements in computing performance. Specifically, during block-based computation, all K and V blocks must be traversed for each Q block. This means that the same K and V blocks are loaded repeatedly when processing different Q blocks. During this process, since Q, K, and V blocks are typically stored in the HBM UMA area, and HBM memory access bandwidth is limited by hardware conditions, particularly the bandwidth of the chip's external bus, bandwidth contention is prone to occur when loading Q, K, and V blocks from the HBM UMA. For example, when performing computational tasks such as matrix multiplication of Q and K, loading both the Q and K blocks from the HBM UMA simultaneously competes for limited memory access bandwidth. This bandwidth contention limits data transfer speeds, causing computing units to idle while waiting for data transfer, unable to fully utilize their computing power. This, in turn, severely impacts operator performance and is considered a bottleneck restricting further improvements in model training and inference speeds.

[0026] Figure 1 It is a schematic diagram of data processing in related technologies, such as Figure 1As shown in Figure 1, AI chips, such as GPUs (Graphics Processing Units), GPGPUs (General-purpose computing on Graphics Processing Units), and TPUs (Tensor Processing Units), include multiple SPCs (Streaming Processor Clusters). These SPCs can be used to execute different data processing tasks in parallel, or to execute different subtasks of the same data processing task in parallel. For example, suppose a device includes 16 SPCs (such as Figure 1 In the process of data processing (such as attention calculation), for each SPC, the input data required is usually stored in the UMA area of HBM (such as SPC0~SPC15). Figure 1 As shown in the figure, HBM0~HBM15), since these memory areas are consistent memory access areas, each SPC can access the data stored in HBM0~HBM15.

[0027] Taking attention calculation as an example, assume that the Q matrix is divided into 16 blocks, each of which is processed by an SPC. For each SPC, when performing the attention calculation corresponding to that block, it needs to load the corresponding Q block from the HBM and traverse all K and V blocks and load them simultaneously. During this process, multiple SPCs may access the UMA area of the HBM to load data blocks simultaneously, causing bandwidth contention. This can cause the corresponding SPC to become idle while waiting for data loading and transmission, affecting overall computing efficiency.

[0028] In this regard, the present invention provides a data processing method, which preloads the data blocks related to the current data processing task of each processing unit into the corresponding target memory area of non-uniform memory access, so that the subsequent data loading operations of different processing units can be independent of each other, thereby effectively avoiding bandwidth competition during the data loading process, and thus overcoming the above-mentioned defects.

[0029] It should be noted that the execution entity of the data processing method provided by the present invention can be an artificial intelligence chip such as GPU, GPGPU, TPU, etc. The following takes GPU as an example to briefly introduce the structure of the execution entity of the present invention.

[0030] Figure 2 Schematic diagram of the structure of the graphics processor provided by the present invention, such as Figure 2As shown, the graphics processor 200 includes at least a plurality of stream processor clusters 201 and a video memory 202, wherein each stream processor cluster 201 includes a plurality of computing units 203, and each computing unit 203 includes at least an on-chip cache 204. Here, the video memory 202 can be an off-chip memory HBM or other types of memory. The on-chip cache 204 is a temporary memory with a smaller capacity than the video memory 202, but a faster data exchange speed than the video memory 202. It should be understood that the data to be processed, the data blocks, and the output results of the data processing involved in the various embodiments of the present invention can be stored in the video memory 202. The processing unit involved in the various embodiments of the present invention can be a stream processor cluster 201 or a computing unit 203 in the stream processor cluster 201.

[0031] It can be understood that the present invention plays a positive role in reducing bandwidth competition and resource idleness by optimizing data processing procedures and rationally utilizing hardware architecture characteristics, effectively improving the hardware performance of artificial intelligence chips such as GPUs, and providing more efficient hardware support for large-scale data processing and complex computing tasks.

[0032] Specifically, in GPU architectures, data to be processed is typically stored in video memory (such as the UMA area of HBM). Simultaneous access to this data by multiple processing units (such as a streaming processor cluster or its internal computing units) can easily lead to bandwidth contention, reducing data transmission speeds. However, the present invention uses a preloading mechanism to preload each data block related to each processing unit's current data processing task into its corresponding NUMA area. Subsequent processing units then load data directly from this area. This means that each processing unit's subsequent data loading operations are independent of each other, without having to wait for other processing units to complete data access. This not only avoids bandwidth competition between multiple processing units, but also reduces waiting time during data loading, thereby improving data transmission efficiency. Furthermore, by avoiding bandwidth competition and latency issues during data loading, processing units spend less time waiting for data loading, enabling more continuous data processing operations and reducing idle time for hardware resources. This allows for more efficient utilization of the computing resources of hardware chips such as GPUs, enabling more data to be processed in the same amount of time, thereby improving overall hardware performance and processing efficiency.

[0033] Furthermore, the data processing method provided by the present invention is applicable to any operator computation scenario that requires repeated data loading, such as the Attention operator, matrix multiplication operator, and convolution operator. To facilitate understanding of the technical solution provided by the present invention, the following mainly uses the computation of the Attention operator as an example to introduce the technical solution provided by the present invention.

[0034] Figure 3 It is a flow chart of the data processing method provided by the present invention, such as Figure 3 As shown, the method includes: Step 310: Acquire a data set to be processed, wherein the data set to be processed includes at least one batch of data to be processed.

[0035] Specifically, a data set to be processed refers to a set of data collected and organized to complete a specific data processing task, and it contains all the raw data that needs to be processed. The data set to be processed can be obtained by receiving user input through a user interface, or by reading stored data files from a file system, or by obtaining data from other systems or services through a network interface, or by collecting data in real time through sensors, such as images and audio. It should be understood that the method for obtaining the data set to be processed can be determined based on the actual application scenario, and the embodiments of the present invention do not specifically limit this.

[0036] It's understandable that to optimize data processing efficiency, adapt to hardware resource limitations, and meet specific algorithm requirements, the data in a data set can be divided into one or more batches. This allows for parallel processing of these data sets, fully utilizing computing resources and accelerating model training and inference. Here, a batch of data refers to a subset of the data set and can be considered a set of data in a single processing task.

[0037] It should be noted that the data processing method provided by the embodiment of the present invention can be applied to multiple fields such as text processing, image processing, and speech processing. In different application scenarios, the data to be processed has different physical meanings. For example, in the field of text processing, the data to be processed may be text data. Specifically, in the text generation task, the set of data to be processed may include multiple prompt texts input by the user into the large language model, and a batch of data to be processed is a prompt text input by the user. For another example, in the field of image processing, the data to be processed may be images collected by a camera. For another example, in the field of speech processing, the data to be processed may be audio data collected by a microphone.

[0038] Step 320 : Segment each data to be processed in the data set to be processed to obtain a plurality of data blocks corresponding to each data to be processed.

[0039] Specifically, each piece of data in the data set to be processed can be segmented based on data characteristics, hardware architecture, and computing requirements. The segmentation method and granularity depend on the specific data processing task. For example, if the data to be processed is prompt text entered by the user and the data processing task is Attention calculation, the Attention calculation requires matrix segmentation of the input data to be processed. First, the data to be processed is converted into a query matrix (Q), a key matrix (K), and a value matrix (V). Then, based on the data shape, hardware architecture, and computing requirements, these matrices are segmented to obtain multiple small sub-matrix blocks (such as query matrix blocks, key matrix blocks, and value matrix blocks).

[0040] It's understandable that after data segmentation, the multiple data blocks corresponding to each piece of data to be processed are smaller data units obtained by segmenting the original data in some way. In the example of Attention calculation, after each piece of data to be processed (such as the prompt text entered by the user) is converted into the Q, K, and V matrices, the query matrix block, key matrix block, and value matrix block obtained by further segmentation are the multiple data blocks corresponding to the data to be processed.

[0041] Step 330: Allocate the data processing tasks corresponding to the data to be processed to at least one processing unit.

[0042] It should be noted that the data processing task corresponding to each piece of data to be processed refers to the specific calculation or operation that needs to be performed on the data to be processed. In the fields of artificial intelligence and machine learning, these tasks can include attention calculation, matrix multiplication, convolution operations, etc. For example, if the data to be processed is the prompt text input by the user into a large language model, and the data processing task is attention calculation, the data processing task corresponding to each piece of data to be processed is to perform attention calculation on the prompt text.

[0043] Specifically, before processing each batch of data in a data set, the data processing tasks corresponding to each batch of data need to be assigned to one or more processing units so that these tasks can be executed in parallel. Task allocation can be achieved through appropriate task scheduling and allocation algorithms, which take into account factors such as the number of processing units on the hardware, the load on each processing unit, and the computing power. The specific allocation method can be static (determined during program compilation) or dynamic (adjusted based on real-time conditions during program runtime).

[0044] For example, if the data set to be processed includes multiple prompt texts, the attention calculation task corresponding to each prompt text can be assigned to a processing unit (such as SPC) for processing. For example, assuming that the data set to be processed includes 8 prompt texts, each prompt text corresponds to an attention calculation task, namely Task 0 to Task 7, and there are 4 idle processing units on the hardware, namely SPC0, SPC1, SPC2 and SPC3, then Task 0 to Task 3 can be assigned to SPC0 to SPC3 for processing respectively. After SPC0 to SPC3 complete the processing of Task 0 to Task 3, Task 4 to Task 7 can be assigned to SPC0 to SPC3 for further processing.

[0045] Step 340: For each processing unit, preload each data block related to the current data processing task of the processing unit into a target memory area, and load each data block from the target memory area for data processing. The target memory area is a non-uniform memory access area.

[0046] It should be noted that the target memory area refers to the memory area used to store data blocks related to the processing unit's current data processing task. In this embodiment of the present invention, the target memory area is designated as a non-uniform memory access (NUMA) area. This means that each processing unit can only access data in its corresponding NUMA area and cannot access data in the NUMA areas corresponding to other processing units. This design effectively avoids bandwidth contention during data loading.

[0047] Specifically, for each processing unit, before processing the data to be processed, each data block related to the current data processing task of the processing unit can be pre-loaded into the target memory area to prepare the relevant data for the subsequent data processing process.

[0048] Here, for each processing unit, the data blocks associated with that unit's current data processing task are the data blocks that the processing unit needs to access and process when performing its task. For example, in an Attention computation, if a processing unit is responsible for processing a query matrix block (Q block), then the data blocks associated with that processing unit's current data processing task include the query matrix block and all key matrix blocks (K block) and value matrix blocks (V block) that need to be computed with it.

[0049] It's understood that after the data to be processed is partitioned, the corresponding data blocks are typically stored in the UMA area of the HBM. For each processing unit, after determining the data blocks relevant to that unit's current data processing task, these data blocks are preloaded from the HBM UMA area into the target memory area (i.e., NUMA area) corresponding to that processing unit for subsequent data loading and processing. The data preloading process involves memory copying or data transfer operations, which can be implemented through hardware instructions or software algorithms.

[0050] Once the relevant data blocks are preloaded into the target memory area, the processing unit can load the corresponding data blocks from that area for data processing. For example, in the example of Attention calculation, the processing unit will traverse all query matrix blocks related to the current data processing task and, for each query matrix block, load the query matrix block and its corresponding key matrix block and value matrix block from the target memory area, and then perform the attention calculation. This process is repeated until all query matrix blocks have been processed.

[0051] Figure 4 This is a schematic diagram of data loading for each processing unit provided by the present invention, such as Figure 4 As shown in the figure, assume that the GPU includes 16 processing units, namely SPC0 to SPC15. Each SPC corresponds to a target memory region, for example, SPC0 corresponds to NUMA0, SPC1 corresponds to NUMA1, and so on. Assume that the data set to be processed includes 16 batches of prompt text, and each SPC is responsible for processing the attention calculation task corresponding to the prompt text batch. For each SPC, the prompt text to be processed can be converted into Q, K, and V matrices. These matrices are then split and each matrix block is stored in the UMA region of the HBM. Taking SPC0 as an example, before performing attention calculation on the Q, K, and V matrices of the prompt text, all blocks of the K and V matrices can be preloaded from the UMA region of the HBM into NUMA0. Then, all Q blocks are traversed and the currently traversed block Qi is preloaded into NUMA0. Subsequently, the corresponding Q, K, and V blocks can be loaded from NUMA0 to perform the attention calculation. It should be understood that in an embodiment of the present invention, by preloading the Q, K, and V blocks into the NUMA area before performing the attention calculation, the bandwidth competition problem in the subsequent data loading process can be effectively avoided, thereby ensuring that the data loading process can be carried out stably and efficiently, thereby improving the overall computing efficiency.

[0052] The method provided by the embodiment of the present invention distributes the data processing tasks corresponding to each data to be processed to at least one processing unit, so that each processing unit can process each data to be processed in parallel, thereby improving the overall processing efficiency. For each processing unit, by preloading each data block related to the current data processing task of the processing unit into the target memory area of the corresponding non-uniform memory access, the data loading operations of different processing units are independent of each other, thereby effectively avoiding bandwidth competition and ensuring the stability and efficiency of data transmission. In addition, since bandwidth competition is avoided, each processing unit does not need to be idle during the data processing process due to waiting for data transmission, and can continuously and efficiently perform data processing tasks, and the hardware computing power is fully utilized, which helps to improve operator performance and accelerate the model training and inference speed.

[0053] Based on the above embodiment, the data processing task includes attention calculation based on each data block. Accordingly, step 320 specifically includes: Step 321: convert any data to be processed into a matrix to be processed, where the matrix to be processed includes a query matrix, a key matrix, and a value matrix.

[0054] Specifically, in attention calculation tasks, any data to be processed (such as prompt text input by the user into a large language model) must first be converted into a matrix to be processed through linear transformation operations or embedding operations for subsequent calculation and processing. For example, taking text data as an example, word embedding technology can be used to convert each word or phrase in the text into a high-dimensional vector, and then these vectors are organized into a matrix form. Specifically, in models based on the Transformer architecture, these vectors are typically organized into a query matrix (Q), a key matrix (K), and a value matrix (V) for subsequent attention calculations.

[0055] It should be understood that the matrix to be processed refers to the data in matrix form after conversion for subsequent calculations. In embodiments of the present invention, the matrix to be processed may include a query matrix (Q), a key matrix (K), and a value matrix (V). These matrices contain all the information required for attention calculations and can efficiently calculate the associations between data through matrix operations.

[0056] Step 322 : Slice each matrix in the matrix to be processed to obtain a plurality of query matrix blocks, key matrix blocks, and value matrix blocks corresponding to any data to be processed.

[0057] Specifically, the process of partitioning each matrix (e.g., Q, K, and V) typically involves splitting the large matrix into multiple smaller sub-matrices. The partitioning method and granularity depend on the specific data shape, hardware architecture, and computational requirements.

[0058] For example, using Attention computation as an example, the Q, K, and V matrices can be split along a certain dimension (such as the sequence length dimension) to generate multiple query matrix blocks, key matrix blocks, and value matrix blocks. The size of each matrix block can be determined based on the computing power and memory capacity of the processing unit to ensure efficient parallel computing.

[0059] It can be understood that the multiple query matrix blocks, key matrix blocks, and value matrix blocks corresponding to any data to be processed refer to the multiple smaller sub-matrix blocks obtained by converting the original data to be processed into the matrix to be processed and then performing a slicing operation. These matrix blocks are the basic units for subsequent data processing tasks. They are preloaded into the target memory area and traversed and calculated by the processing unit to obtain the final processing results.

[0060] Based on any of the above embodiments, step 340 specifically includes: Step 341 : preload each key matrix block and each value matrix block related to the current data processing task of the processing unit into a target memory area.

[0061] It should be noted that step 341 preloads the key matrix blocks and value matrix blocks associated with the processing unit's current data processing task into the target memory area (i.e., the NUMA area in the HBM). This process is automatically performed by the GPU's memory management system, ensuring that the processing unit can quickly access the required data blocks when performing data processing.

[0062] Specifically, when a processing unit (such as an SPC) is assigned a data processing task, the GPU identifies all key matrix blocks and value matrix blocks related to the task and preloads them from the HBM's UMA (uniform memory access) area into the target memory area (NUMA area) corresponding to the processing unit. This allows the processing unit to access these data blocks with lower latency and higher bandwidth, thereby improving data processing efficiency.

[0063] Step 342, traverse each query matrix block related to the current data processing task of the processing unit, preload the currently traversed query matrix block into the target memory area, and load the query matrix block, each key matrix block and each value matrix block from the target memory area to perform data processing.

[0064] Specifically, for each processing unit, after preloading all key matrix blocks (i.e., K blocks) and value matrix blocks (i.e., V blocks) related to its current data processing task into the target memory area, it can traverse the query matrix blocks (i.e., Q blocks) related to the current data processing task of the processing unit, so as to preload the currently traversed Q blocks into the target memory area, and load the Q blocks, related K blocks and V blocks from the target memory area for data processing.

[0065] Specifically, the processing unit will traverse all query matrix blocks related to the current data processing task, and for each query matrix block, subsequent processing steps will be performed. During the traversal process, for the currently traversed query matrix block, the GPU will preload it from the UMA area of the HBM to the NUMA area corresponding to the processing unit. Subsequently, the processing unit will load the currently traversed query matrix block and all key matrix blocks and value matrix blocks that have been preloaded into the target memory area from the target memory area, and perform specific data processing tasks based on these matrix blocks, such as attention calculations. The results of each calculation are accumulated to the final output result, which can be achieved through an accumulator or a similar data structure. This process is repeated until all query matrix blocks have been processed, and the final data processing result can be obtained.

[0066] Based on any of the above embodiments, in step 342, the step of loading the query matrix block, each key matrix block, and each value matrix block from the target memory area and performing data processing includes: Step 3421: traverse the key matrix blocks and the value matrix blocks, and load the query matrix block, the currently traversed key matrix block, and the value matrix block from the target memory area.

[0067] It's important to note that this step involves traversing each key matrix block (K block) and each value matrix block (V block) and loading the currently traversed query matrix block (denoted as Qi), key matrix block (denoted as Ki), and value matrix block (denoted as Vi) from the target memory area (i.e., the NUMA area in the HBM). This is a critical step in the Attention calculation, ensuring that the processing unit has access to all necessary data blocks during calculation.

[0068] Specifically, for each query matrix block, a processing unit (such as an SPC) traverses all key and value matrix blocks. The processing unit can maintain a pointer or index to traverse all key and value matrix blocks. For each key and value matrix block combination (Ki, Vi), the processing unit loads them from the target memory area and simultaneously loads the currently traversed query matrix block Qi. Once these matrix blocks are loaded into the processing unit, subsequent calculations can be performed.

[0069] Step 3422: Perform calculation based on the currently loaded query matrix block and key matrix block to obtain a first calculation result.

[0070] It should be noted that this step is the process of calculating the first calculation result based on the currently loaded query matrix block (Qi) and key matrix block (Ki). In the attention calculation, this step usually involves a matrix multiplication operation to calculate the similarity (or attention score) between the Q block and the K block.

[0071] Specifically, the processing unit performs matrix multiplication on the query matrix block Qi and the key matrix block Ki to obtain an intermediate result. This intermediate result is then processed through other operations (such as the Softmax function) to obtain a first calculation result. Here, the first calculation result is used to represent the correlation or attention weight between the query matrix block Qi and the key matrix block Ki.

[0072] Step 3423: Perform calculation based on the first calculation result and the currently loaded value matrix block to obtain a second calculation result, and add the second calculation result to the output result.

[0073] It should be noted that this step is the process of performing a calculation based on the first calculation result and the currently loaded value matrix block (Vi) to obtain the second calculation result, and then accumulating the second calculation result to the output result. In attention calculation, this step usually involves a weighted sum operation, which is used to weight the V blocks according to the attention score to obtain the final output vector.

[0074] Specifically, the processing unit uses the first calculation result (i.e., the attention weight) to perform a weighted summation operation on the value matrix block Vi to obtain a second calculation result. Here, the second calculation result is used to represent the information of the value matrix block weighted according to the attention weight. The processing unit then adds this second calculation result to the output result. Here, the output result is a cumulative result that contains the final information after all query matrix blocks are processed. In the example of attention calculation, the output result can be a set of weighted value matrices, which contains the results of all query matrix blocks and the corresponding key matrix blocks and value matrix blocks.

[0075] Based on any of the above embodiments, the step of allocating the target memory area includes: Receive memory allocation request; The memory allocation request is parsed to obtain memory parameters, and memory allocation is performed on an off-chip memory based on the memory parameters to obtain a target memory area, wherein the memory parameters include memory size, memory location, and access mode.

[0076] Specifically, before executing a data processing task, a corresponding target memory area can be pre-allocated to each processing unit. During the target memory area allocation process, the device side (such as the GPU) first receives a memory allocation request from the host side (such as the CPU). This request includes information about the memory to be allocated and specific requirements for the memory allocation.

[0077] Specifically, the GPU communicates with the host computer via a system bus, such as PCIe (Peripheral Component Interconnect Express), to receive memory allocation requests from the CPU or driver. Upon receiving a memory allocation request, the GPU parses it to identify key information in the request and obtain memory parameters such as the size, amount, location, and access mode of the memory to be allocated.

[0078] As you can understand, memory parameters refer to a series of key factors that need to be considered when allocating memory. For example, memory parameters can include memory size, memory location, and access mode. Memory size refers to the amount of memory space to be allocated, typically measured in bytes. Memory location refers to the specific location in the HBM where the memory is allocated. Access mode refers to how data blocks are accessed in memory. Different access modes may affect data access efficiency. For example, in a NUMA (Non-Uniform Memory Access) architecture, selecting the appropriate access mode can optimize data access latency and bandwidth.

[0079] Based on the parsed memory parameters, the GPU will allocate memory on the HBM. This can be achieved through the following steps: First, based on the memory size specified in the request, the GPU will reserve the corresponding space in the HBM. Then, based on the memory location requirements and the HBM hardware architecture, the GPU will select a suitable location for memory allocation. Next, based on the access mode requirements, the GPU will configure the corresponding memory access control mechanisms, which can ensure that the data can meet the latency and bandwidth requirements when accessed. After completing the above steps, the GPU can successfully allocate the required memory space for the target memory area on the HBM. This space will be used to store data blocks related to the current data processing task of the processing unit and support the efficient execution of subsequent data processing tasks.

[0080] Based on any of the above embodiments, the data to be processed includes any one of text data, image data, and audio data.

[0081] It should be noted that the data processing method provided by the present invention can be applied to multiple fields such as text processing, image processing, and speech processing. In different application scenarios, the data to be processed has different physical meanings.

[0082] For example, in the field of text processing, the data to be processed can be text data. Taking the application scenario of a large language model as an example, the prompt text entered by the user can be used as the data to be processed. When applying the attention mechanism to this data, it is first converted into Q, K, and V matrices through a linear transformation. Q, K, and V are then partitioned into different blocks. Subsequently, the processing unit preloads all K and V blocks into their corresponding target memory areas and then begins traversing all Q blocks. For each Q block, the processing unit preloads it into the target memory area. After preloading the data blocks, the processing unit begins traversing the K and V blocks, reading a Q, K, and V block from the target memory area during each traversal and performing attention calculations. After completing the traversal of all K and V blocks, processing the next Q block begins. Finally, after processing all Q blocks, the processing unit outputs the corresponding text feature representation. It should be understood that in order to reduce data handling, the processing unit will store the intermediate results corresponding to each K and V block in shared memory or registers until all K and V blocks are calculated and then store the results in HBM.

[0083] For example, in the field of image processing, the data to be processed can be images captured by a camera. After being segmented, preloaded and calculated by the processing unit, the corresponding image features can be output for subsequent image processing tasks, such as image classification, target detection, etc.

[0084] For example, in the field of speech processing, the data to be processed may be audio data collected by a microphone. After being segmented, preloaded and calculated by the processing unit, the corresponding audio features can be output for subsequent audio processing tasks, such as speech recognition, speech translation, etc.

[0085] The method provided by the embodiment of the present invention effectively avoids bandwidth competition and improves hardware computing power utilization by dividing the data to be processed and distributing it to processing units, and preloading related data blocks into the target memory area (non-uniform memory access area), thereby providing an efficient and reliable solution for data processing tasks in different fields.

[0086] Based on any of the above embodiments, Figure 5 This is a flow chart of the attention calculation based on input data preloading provided by the present invention, such as Figure 5 As shown, the method includes: Step S1, receiving a temporary memory (i.e., NUMA) allocation request sent by the host, and performing NUMA memory allocation according to the request; Step S2: For each processing unit (such as SPC), preload the complete K and V matrices (i.e., including all K blocks and all V blocks) from the HBM to the NUMA corresponding to the processing unit; Step S3, traverse all Q blocks and preload the currently traversed Q block (i.e., Qi) from HBM to NUMA; Step S4: traverse all K and V blocks, and load the current Qi and the currently traversed K and V blocks (i.e., Ki and Vi) from NUMA to perform FlashAttention calculation; Step S5, repeat the above steps S3 and S4 until all Q blocks are traversed.

[0087] Figure 6 is a schematic diagram of the performance comparison of the operators provided by the present invention, such as Figure 6 As shown in the figure, the vertical axis represents the clock cycles required for the Attention operator to execute. Fewer clock cycles indicate better operator performance. The horizontal axis represents the number of batches of data to be processed. Each batch of data is assigned to an SPC for processing. For example, when the batch size is 8, it indicates that 8 batches of data are assigned to 8 SPCs for processing.

[0088] Figure 6 The yellow column indicates no preloading of the data to be processed. By comparing the clock cycles required for different batches, we can see that the operator performance does not improve when using 16 SPCs compared to using 1 SPC. Normally, the time spent processing 16 batches of data using 16 SPCs should be similar to that of processing a single batch using 1 SPC. However, as can be seen from the figure, when no preloading of the data to be processed is performed, the time difference between using 16 SPCs and using 1 SPC is quite large (1.47 times), indicating that the operator performance is not meeting expectations.

[0089] The gray column indicates that only the K and V matrices are preloaded, not the Q blocks, and the storage area used for the K and V preloads is not NUMA. In this case, it can be seen that using 16 SPCs to process 16 batches of data takes 1.33 times longer than using 1 SPC to process 1 batch of data, indicating a 1.33-fold performance difference for the operator.

[0090] The blue column indicates that the preloading operation is performed on Q, K, and V, and the preloaded storage area is the NUMA area. Figure 6As can be seen in the figure, the performance of using 16 SPCs to process 16 batches of data is roughly the same as that of using a single SPC to process a single batch of data (only a factor of 1.14 difference). This means that by preloading each data block related to each processing unit's current data processing task into its corresponding target memory area for non-uniform memory access, the embodiments of the present invention can make the data loading operations of different processing units independent of each other, effectively avoiding bandwidth contention during the data loading process, ensuring the stability and efficiency of data transmission, and thus helping to improve operator performance.

[0091] The data processing device provided by the present invention is described below. The data processing device described below and the data processing method described above can be referenced to each other.

[0092] Based on any of the above embodiments, Figure 7 Schematic diagram of the structure of the data processing device provided by the present invention, such as Figure 7 As shown, the device includes: The data acquisition module 710 is used to acquire a set of data to be processed, wherein the set of data to be processed includes at least one batch of data to be processed; The data segmentation module 720 is configured to segment each to-be-processed data in the to-be-processed data set to obtain a plurality of data blocks corresponding to each to-be-processed data; A task allocation module 730 is used to allocate the data processing tasks corresponding to the data to be processed to at least one processing unit; The data processing module 740 is used to preload each data block related to the current data processing task of each processing unit into a target memory area, and load each data block from the target memory area for data processing. The target memory area is a non-uniform memory access area.

[0093] The device provided by the embodiment of the present invention distributes the data processing tasks corresponding to each data to be processed to at least one processing unit, so that each processing unit can process each data to be processed in parallel, thereby improving the overall processing efficiency. For each processing unit, by preloading each data block related to the current data processing task of the processing unit into the target memory area of the corresponding non-uniform memory access, the data loading operations of different processing units are independent of each other, thereby effectively avoiding bandwidth competition and ensuring the stability and efficiency of data transmission. In addition, since bandwidth competition is avoided, each processing unit does not need to be idle during the data processing process due to waiting for data transmission, and can continuously and efficiently perform data processing tasks, and the hardware computing power is fully utilized, which helps to improve operator performance and accelerate the model training and inference speed.

[0094] Based on any of the above embodiments, the data processing task includes attention calculation based on each data block. Accordingly, the data segmentation module 720 is specifically configured to: Convert any data to be processed into a matrix to be processed, wherein the matrix to be processed includes a query matrix, a key matrix and a value matrix; Each matrix in the matrix to be processed is divided to obtain a plurality of query matrix blocks, key matrix blocks and value matrix blocks corresponding to any data to be processed.

[0095] Based on any of the above embodiments, the data processing module 740 includes: a preloading submodule, configured to preload each key matrix block and each value matrix block related to a current data processing task of the processing unit into a target memory area; The processing submodule is used to traverse the query matrix blocks related to the current data processing task of the processing unit, preload the currently traversed query matrix block into the target memory area, and load the query matrix block, key matrix blocks and value matrix blocks from the target memory area to perform data processing.

[0096] Based on any of the above embodiments, the processing submodule is specifically configured to: Traversing the key matrix blocks and the value matrix blocks, and loading the query matrix block, the currently traversed key matrix block, and the value matrix block from the target memory area; Perform calculation based on the currently loaded query matrix block and key matrix block to obtain a first calculation result; Calculation is performed based on the first calculation result and the currently loaded value matrix block to obtain a second calculation result, and the second calculation result is accumulated to the output result.

[0097] Based on any of the above embodiments, the device further includes a memory allocation module, wherein the memory allocation module is configured to: Receive memory allocation request; The memory allocation request is parsed to obtain memory parameters, and memory allocation is performed on an off-chip memory based on the memory parameters to obtain a target memory area, wherein the memory parameters include memory size, memory location, and access mode.

[0098] Based on any of the above embodiments, the data to be processed includes any one of text data, image data, and audio data.

[0099] Figure 8 An example of a physical structure diagram of an electronic device is shown below. Figure 8As shown, the electronic device may include: a processor 810, a communication interface 820, a memory 830 and a communication bus 840, wherein the processor 810, the communication interface 820 and the memory 830 communicate with each other via the communication bus 840. The processor 810 may call the logic instructions in the memory 830 to execute a data processing method, which includes: obtaining a set of data to be processed, wherein the set of data to be processed includes at least one batch of data to be processed; dividing each data to be processed in the set of data to be processed to obtain multiple data blocks corresponding to each data to be processed; allocating the data processing task corresponding to each data to be processed to at least one processing unit; for each processing unit, preloading each data block related to the current data processing task of the processing unit into a target memory area, and loading each data block from the target memory area for data processing, wherein the target memory area is a non-uniform memory access area.

[0100] Furthermore, the logic instructions in the aforementioned memory 830 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the portion that contributes to the relevant art, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product, stored in a storage medium, includes instructions for enabling a computer device (which can be a personal computer, server, or network device, etc.) to perform all or part of the steps of the methods described in various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as a USB flash drive, a mobile hard drive, a read-only memory (ROM), a random access memory (RAM), a magnetic disk, or an optical disk.

[0101] On the other hand, the present invention also provides a computer program product, which includes a computer program, which can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the data processing methods provided by the above methods, which include: obtaining a set of data to be processed, wherein the set of data to be processed includes at least one batch of data to be processed; dividing each data to be processed in the set of data to be processed to obtain multiple data blocks corresponding to each data to be processed; allocating the data processing tasks corresponding to each data to be processed to at least one processing unit; for each processing unit, preloading each data block related to the current data processing task of the processing unit into a target memory area, and loading each data block from the target memory area for data processing, wherein the target memory area is a non-uniform memory access area.

[0102] On the other hand, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, is implemented to execute the data processing method provided by the above-mentioned methods, the method comprising: obtaining a set of data to be processed, the set of data to be processed including at least one batch of data to be processed; dividing each data to be processed in the set of data to be processed to obtain multiple data blocks corresponding to each data to be processed; allocating the data processing tasks corresponding to each data to be processed to at least one processing unit; for each processing unit, preloading each data block related to the current data processing task of the processing unit into a target memory area, and loading each data block from the target memory area for data processing, the target memory area being a non-uniform memory access area.

[0103] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, i.e., they may be located in one location or distributed across multiple network units. Some or all of the modules may be selected based on actual needs to achieve the objectives of the present embodiment. Persons of ordinary skill in the art will be able to understand and implement the present invention without inventive effort.

[0104] Through the description of the above embodiments, those skilled in the art will clearly understand that each embodiment can be implemented using software plus a necessary general-purpose hardware platform, or of course, hardware. Based on this understanding, the essence of the above technical solution, or the portion that contributes to the relevant technology, can be embodied in the form of a software product. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a magnetic disk, or an optical disk, and includes a number of instructions for causing a computer device (such as a personal computer, server, or network device) to execute the methods described in each embodiment or certain portions of the embodiments.

[0105] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present invention.

Claims

1. A data processing method, characterized in that: include: Acquire a data set to be processed, wherein the data set to be processed includes at least one batch of data to be processed; Segmenting each data to be processed in the set of data to be processed to obtain a plurality of data blocks corresponding to each data to be processed; Allocating data processing tasks corresponding to each piece of data to be processed to at least one processing unit; For each processing unit, each data block related to the current data processing task of the processing unit is preloaded into a target memory area, and each data block is loaded from the target memory area for data processing. The target memory area is a non-uniform memory access area.

2. The data processing method according to claim 1, wherein: The data processing task includes an attention calculation based on each data block. Accordingly, the data to be processed in the data set to be processed is segmented to obtain a plurality of data blocks corresponding to each data to be processed, including: Convert any data to be processed into a matrix to be processed, wherein the matrix to be processed includes a query matrix, a key matrix and a value matrix; Each matrix in the matrix to be processed is divided to obtain a plurality of query matrix blocks, key matrix blocks and value matrix blocks corresponding to any data to be processed.

3. The data processing method according to claim 2, characterized in that: The step of preloading each data block related to the current data processing task of the processing unit into a target memory area, and loading each data block from the target memory area for data processing, comprises: Preloading each key matrix block and each value matrix block related to the current data processing task of the processing unit into a target memory area; Traverse each query matrix block related to the current data processing task of the processing unit, preload the currently traversed query matrix block into the target memory area, and load the query matrix block, each key matrix block and each value matrix block from the target memory area to perform data processing.

4. The data processing method according to claim 3, characterized in that: The step of loading the query matrix block, each key matrix block, and each value matrix block from the target memory area and performing data processing includes: Traversing the key matrix blocks and the value matrix blocks, and loading the query matrix block, the currently traversed key matrix block, and the value matrix block from the target memory area; Perform calculation based on the currently loaded query matrix block and key matrix block to obtain a first calculation result; Calculation is performed based on the first calculation result and the currently loaded value matrix block to obtain a second calculation result, and the second calculation result is accumulated to the output result.

5. The data processing method according to any one of claims 1 to 4, characterized in that: The step of allocating the target memory area includes: Receive memory allocation request; The memory allocation request is parsed to obtain memory parameters, and memory allocation is performed on an off-chip memory based on the memory parameters to obtain a target memory area, wherein the memory parameters include memory size, memory location, and access mode.

6. The data processing method according to any one of claims 1 to 4, characterized in that: The data to be processed includes any one of text data, image data, and audio data.

7. A data processing device, characterized in that: include: A data acquisition module is used to acquire a set of data to be processed, wherein the set of data to be processed includes at least one batch of data to be processed; A data segmentation module, configured to segment each to-be-processed data in the to-be-processed data set to obtain a plurality of data blocks corresponding to each to-be-processed data; A task allocation module, configured to allocate the data processing tasks corresponding to the data to be processed to at least one processing unit; The data processing module is used to preload each data block related to the current data processing task of each processing unit into a target memory area, and load each data block from the target memory area for data processing, wherein the target memory area is a non-uniform memory access area.

8. An electronic device comprising a memory, a processor, and a computer program stored in the memory and running on the processor, characterized in that: When the processor executes the computer program, the data processing method according to any one of claims 1 to 6 is implemented.

9. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the data processing method according to any one of claims 1 to 6 is implemented.

10. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, the data processing method according to any one of claims 1 to 6 is implemented.

Images