A method, device and storage medium for balancing attention load

Abstract

The application discloses an attention load balancing method and device, equipment and storage medium. Including: obtaining attention calculation data, splitting the attention calculation data, generating a split task, and distributing the split task to each computing unit; determining the task starting index of each computing unit, traversing each split task and giving a linear number, determining the first position parameter of each computing unit according to the task starting index and the linear number; based on the first position parameter, the attention calculation of each split task is executed, and the final result is written into the video memory. By dividing the data, the load of each computing unit can be balanced, and the waste of computing power can be avoided; by determining the first position parameter based on the linear number, the task positioning and scheduling logic can be simplified, and the scheduling complexity is reduced; by executing attention calculation based on the first position parameter, the parallel efficiency of memory access and calculation can be improved, the video memory overhead is reduced, and the ASIC hardware architecture characteristics are fully adapted.

Description

Technical Field

[0001] This invention relates to the field of computer technology, and in particular to an attention load balancing method, apparatus, device, and storage medium. Background Technology

[0002] With the rapid development of generative recommender systems, the demand for variable-length sequence modeling is increasing. The Hierarchical Sequential Transduction Unit (HSTU) operator, with its efficient sequence modeling capabilities, is gradually replacing the traditional Transformer self-attention mechanism and becoming the core operator of generative recommender systems. When processing variable-length sequences, HSTU attention computation requires irregular segmentation of the key / value sequence to adapt to inputs of different lengths. Application-Specific Integrated Circuits (ASICs), with their high parallelism, low power consumption, and dedicated multi-level caches, have become the mainstream architecture for hardware acceleration of HSTU operators. How to achieve efficient computation of variable-length HSTU attention on an ASIC platform is currently a key focus of the industry.

[0003] To address the uneven computational load caused by variable-length sequences, existing technologies typically employ dynamic task allocation algorithms based on min-heaps to achieve load balancing among multiple computing units. These schemes maintain the load status of each computing unit in real time using a min-heap structure, dynamically allocating the split attention computation tasks to computing units with lower loads. They also rely on GPU memory to temporarily store intermediate results and use additional operators to manage tasks, process parameters, and merge results, thereby attempting to ensure balanced computational utilization across all computing units.

[0004] However, existing dynamic task allocation schemes based on min-heaps have significant drawbacks when performing variable-length HSTU attention computations on ASIC architectures: min-heaps require frequent insertion, adjustment, and maintenance, resulting in a scheduling complexity of O(N log K), which leads to significant scheduling overhead in large-scale parallel scenarios; complex task description information such as task length and priority needs to be maintained, making the task management structure cumbersome and dependent on additional operators; intermediate results must be written to GPU memory for temporary storage and merging, increasing GPU memory bandwidth usage, memory access latency, and storage space consumption; and the on-chip multi-level cache and computation unit collaboration features of ASICs are not fully utilized, making it impossible to achieve deep parallelism between computation and memory access, thus limiting overall hardware utilization and computational performance. Summary of the Invention

[0005] This invention provides an attention load balancing method, apparatus, device, and storage medium to solve the technical problems of high scheduling complexity, complex task management, reliance on video memory, and low hardware utilization in existing attention computing under ASIC architecture.

[0006] According to one aspect of the present invention, an attention load balancing method is provided, applied to an application-specific integrated circuit (ASIC) architecture, the method comprising: Acquire attention calculation data, segment the attention calculation data, generate segmentation tasks, and distribute the segmentation tasks evenly to each computing unit; Determine the task start index of each computing unit, traverse each segmented task and assign a linear number to the task, and determine the initial position parameter of each computing unit based on the task start index and the task linear number. Attention calculations are performed on each segmentation task based on the initial position parameters, and the final results are written to the video memory.

[0007] Optionally, the attention computation data is segmented to generate a segmentation task, including: dividing the attention computation into layers according to preset dimensions to obtain layered data, wherein the layered data includes batches, attention heads, query sequences, and key-value sequences; dividing the query sequences into query sub-blocks of fixed length, and dividing the key-value sequences into key-value sub-blocks of fixed length; and generating a segmentation task by using the combination of batches, attention heads, query sub-blocks, and key-value sub-blocks as the smallest computational unit.

[0008] Optionally, after segmenting the attention computation data to generate segmentation tasks, the method further includes: identifying invalid computation blocks introduced by the mask; filtering invalid computation blocks in the segmentation tasks to obtain valid segmentation tasks; and evenly distributing the segmentation tasks to each computation unit, including: evenly distributing the valid segmentation tasks to each computation unit.

[0009] Optionally, each segmentation task is traversed and assigned a linear task number, including: traversing each segmentation task in the order of batch, attention head, query sub-block, and key-value sub-block; and assigning a unique and continuously increasing sequence number to each segmentation task as the task linear number.

[0010] Optionally, the initial position parameters of each computing unit are determined based on the task start index and the task linear number, including: matching the task start index of each computing unit with the task linear number; when the task start index and the task linear number are consistent, the corresponding segmentation task is taken as the target task; and the coordinates of the computing block corresponding to the target task are taken as the initial position parameters of each computing unit.

[0011] Optionally, after using the computation block coordinates corresponding to the target task as the initial position parameters of each computation unit, the method further includes: performing incremental calculations based on the initial position parameters to obtain the next round of position parameters, wherein the next round of position parameters includes the next round of key block coordinates and the next round of query block coordinates.

[0012] Optionally, attention calculation is performed on each segmentation task based on the initial position parameters, including: reading the current query value and the current key value based on the initial position parameters, calculating the query key matrix multiplication, obtaining the current matrix multiplication result and storing it in the first cache; asynchronously reading the current value, storing it in the first cache, and performing activation calculation on the current matrix multiplication result, obtaining the activation calculation result and storing it in the first cache; asynchronously reading the next key value based on the next key block coordinates, and asynchronously reading the current bias value based on the initial position parameters, adding the current bias value and the activation calculation result to obtain... The biasing result of this round is stored in the first cache; length scaling and masking processing are performed on the biasing result of this round to obtain the masking result of this round, and the masking result of this round is stored in the first cache; the query value of the next round is read asynchronously based on the block coordinates of the next round query; matrix multiplication is performed on the masking result of this round and the value of this round in the first cache to obtain the calculation result of this round; in the next round calculation, the asynchronously read query value of the next round and the key value of the next round are used as the new query value of this round and the key value of this round, and the calculation is repeated until all the partitioning tasks of the current calculation unit are completed.

[0013] Optionally, after obtaining the calculation result of this round, the method further includes: determining whether the calculation group corresponding to the calculation result of this round is complete, wherein the calculation group is all the split tasks corresponding to the same query block divided according to the query offset; if so, the calculation result of this round is directly written to the video memory as the final result; otherwise, the calculation result of this round is written to the second cache as an intermediate result. After all the split tasks are calculated, the intermediate results belonging to the same calculation group in the second cache are merged to obtain the final result and written to the video memory.

[0014] According to another aspect of the present invention, an electronic device is provided, the electronic device comprising: At least one processor; and a memory communicatively connected to the at least one processor; The memory stores a computer program that can be executed by the at least one processor, which enables the at least one processor to perform an attention load balancing method according to any embodiment of the present invention.

[0015] According to another aspect of the present invention, a computer-readable storage medium is provided, the computer-readable storage medium storing computer instructions for causing a processor to execute and implement an attention load balancing method according to any embodiment of the present invention.

[0016] The technical solution of this invention can balance the load of each computing unit and avoid wasting computing power by dividing the attention calculation data and distributing the tasks evenly to each computing unit; by determining the initial position parameter based on linear numbering, the task location and scheduling logic can be simplified and the scheduling complexity can be reduced; by performing attention calculation based on the initial position parameter and outputting the result, the parallel efficiency of memory access and calculation can be improved, the video memory overhead can be reduced, and the characteristics of ASIC hardware architecture can be fully adapted.

[0017] It should be understood that the description in this section is not intended to identify key or essential features of the embodiments of the present invention, nor is it intended to limit the scope of the invention. Other features of the invention will become readily apparent from the following description. Attached Figure Description

[0018] To more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0019] Figure 1 This is a flowchart of an attention load balancing method provided in Embodiment 1 of the present invention; Figure 2 This is a flowchart of another attention load balancing method provided in Embodiment 2 of the present invention; Figure 3 This is a schematic diagram of an attention load balancing device according to Embodiment 3 of the present invention; Figure 4 This is a schematic diagram of the structure of an electronic device that implements an attention load balancing method according to an embodiment of the present invention. Detailed Implementation

[0020] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0021] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover a non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0022] Example 1 Figure 1 This document provides a flowchart of an attention load balancing method according to Embodiment 1 of the present invention. This embodiment can be applied to Application-Specific Integrated Circuit (ASIC) architectures, and is particularly suitable for HSTU attention hardware acceleration scenarios in generative recommendation systems with variable-length sequences. This method can be executed by an attention load balancing device, which can be implemented in hardware and / or software, and can be configured in an ASIC chip controller or a computer controller. Figure 1 As shown, the method includes: S110. Obtain attention calculation data, segment the attention calculation data, generate segmentation tasks, and distribute the segmentation tasks evenly to each calculation unit.

[0023] It should be noted that the attention load balancing method in this embodiment is specifically designed for attention computation of the Hierarchical Sequential Transduction Unit (HSTU). It is an algorithm scheme that distributes the computational load of multiple computing units evenly on the application-specific integrated circuit (ASIC) hardware, eliminates idle waiting, and reduces scheduling and memory access overhead.

[0024] The attention computation data refers to the input data required for HSTU attention computation, including query value (Q), key value (K), data value (V), and bias value (rab). It also includes the actual length information of each sequence, causal mask information, and target mask information. The segmentation task refers to the smallest computational task block obtained by segmenting the attention computation data according to a specified dimension. In this embodiment, the fixed length of the query sub-block and key sub-block can be set to 128 to maximize the 2D computing power utilization of the matching ASIC architecture. The computation unit refers to the 2D matrix multiplication and addition unit within the ASIC, organized in an array, responsible for performing operations such as QK dot product and PV weighted summation. In this embodiment, the number of computation units can be 24.

[0025] Specifically, after obtaining the total number of all the split tasks, all the split tasks can be distributed sequentially and evenly to each computing unit in the form of continuous intervals based on the total number of computing units contained in the ASIC architecture. This ensures that the amount of computing tasks undertaken by each computing unit remains consistent, avoiding the waste of computing power due to uneven load. At the same time, the traditional minimum heap dynamic scheduling method is abandoned, reducing the scheduling complexity from O(N log K) to O(1), which fully adapts to the parallel computing characteristics of ASIC hardware.

[0026] Optionally, the attention computation data is segmented to generate a segmentation task, including: dividing the attention computation into layers according to preset dimensions to obtain layered data, wherein the layered data includes batches, attention heads, query sequences, and key-value sequences; dividing the query sequences into query sub-blocks of fixed length, and dividing the key-value sequences into key-value sub-blocks of fixed length; and generating a segmentation task by using the combination of batches, attention heads, query sub-blocks, and key-value sub-blocks as the smallest computational unit.

[0027] Specifically, attention computation-related data can be hierarchically divided according to preset dimensions, splitting the overall data into mutually independent hierarchical data. Hierarchical data specifically includes batches, attention heads, query sequences, and key-value sequences. After hierarchical division, query sequences and key-value sequences are further regularized and divided according to pre-defined fixed lengths, splitting consecutive query sequences into multiple query sub-blocks of equal length, and similarly splitting consecutive key-value sequences into multiple key-value sub-blocks of equal length. After sequence segmentation, the combination of batches, attention heads, query sub-blocks, and key-value sub-blocks serves as the smallest independent computational unit. Each uniquely corresponding combination is encapsulated as an independent segmentation task. All segmentation tasks together constitute a complete set of attention computation tasks. This segmentation method is fully compatible with subsequent invalid computation block filtering, task balancing, linear numbering, and computation-memory access pipeline execution processes, ensuring uniform computational granularity while fully aligning with the parallel execution and multi-level cache access characteristics of ASIC hardware architecture.

[0028] In another specific implementation, besides using a fixed block size to partition by batch, attention header, query sequence, and key-value sequence dimensions, an adaptive block size can also be used for partitioning based on the actual sequence length. The partitioning dimension can be adjusted to prioritize batch dimension, header dimension, or query dimension. The partitioning granularity is not limited to a fixed size and supports various block sizes such as 64, 256, and 512. Invalid calculations can be retained during the partitioning stage and removed in subsequent calculation stages, all of which can generate partitioning tasks that meet load balancing requirements.

[0029] Optionally, after segmenting the attention computation data to generate segmentation tasks, the method further includes: identifying invalid computation blocks introduced by the mask; filtering invalid computation blocks in the segmentation tasks to obtain valid segmentation tasks; and evenly distributing the segmentation tasks to each computation unit, including: evenly distributing the valid segmentation tasks to each computation unit.

[0030] Specifically, based on the causal and target masking rules used in HSTU attention computation, the invalid computation regions covered by the mask during the computation of query and key-value blocks can be identified, thereby locating invalid computation blocks that do not participate in actual computation under the corresponding dimension. Invalid computation blocks cannot produce valid computation results due to mask limitations; if they were to participate in scheduling and execution, it would waste hardware computing power and memory resources. After locating invalid computation blocks, all previously segmented tasks can be screened one by one, removing those belonging to invalid computation blocks from the overall task set, retaining only tasks that can produce valid computation results, thus obtaining valid segmented tasks. This reduces the total amount of subsequent scheduling and computation, improving hardware execution efficiency.

[0031] Furthermore, in the process of evenly distributing the partitioned tasks to each computing unit, the original partitioned tasks are no longer used as the allocation object. Instead, the effective partitioned tasks after filtering out invalid computing blocks are used as the allocation basis. The total number of effective partitioned tasks is counted, and all effective partitioned tasks are sequentially and evenly distributed to each computing unit in a continuous interval according to the number of computing units in the ASIC architecture. This ensures that the effective computing load undertaken by each computing unit is consistent, avoiding uneven load and idle computing power. At the same time, it reduces the complexity of task scheduling and fully adapts to the parallel computing and multi-level caching characteristics of ASIC hardware.

[0032] S120. Determine the task start index of each computing unit, traverse each segmented task and assign a linear number to the task, and determine the initial position parameter of each computing unit based on the task start index and the task linear number.

[0033] The task start index refers to the first number of the continuous task interval handled by each computing unit after the tasks are evenly divided, used to locate which task the unit starts computing from. The task linear numbering refers to assigning consecutive integer numbers to all the smallest divided computing tasks in a fixed order from beginning to end. This fixed order can be based on the batch, attention head, query block, or key / value block order. The initial position parameter refers to the coordinate parameters corresponding to the first task of the computing unit, calculated backwards from the task start index and task linear numbering, used to read the first round of computing data from GPU memory.

[0034] As can be seen, this implementation transforms the complex task management problem into a simple index matching problem, and the scheduling problem into an index matching problem. Using the task's starting index as the search benchmark, it performs location matching within the linearly generated task numbers. Through coordinate conversion, it obtains the initial position parameters corresponding to the first sub-task to be executed in each computing unit, providing a reliable addressing basis for data reading and computation initiation.

[0035] Optionally, each segmentation task is traversed and assigned a linear task number, including: traversing each segmentation task in the order of batch, attention head, query sub-block, and key-value sub-block; and assigning a unique and continuously increasing sequence number to each segmentation task as the task linear number.

[0036] Specifically, all valid segmentation tasks filtered by invalid computation blocks can be traversed sequentially according to the hierarchical order of batch, attention head, query sub-block, and key-value sub-block. During the traversal, a fixed order is maintained: from the first batch to the last batch, from the first attention head to the last attention head within each batch, from the first query sub-block to the last query sub-block within each attention head, and from the first key-value sub-block to the last key-value sub-block within each query sub-block. During this sequential traversal, a unique and continuously increasing integer sequence number is assigned to each independent valid segmentation task. This sequence number increments sequentially from the initial value, without repetition, skipping, or gaps, uniquely identifying the position of the corresponding segmentation task in the overall task sequence, thus forming a linear task numbering system. By constructing this linear task numbering system, the originally multi-dimensional segmentation tasks can be transformed into a one-dimensional linear index structure, simplifying task location and scheduling logic and providing a foundation for subsequently determining the initial position parameters based on the task's starting index.

[0037] Optionally, the initial position parameters of each computing unit are determined based on the task start index and the task linear number, including: matching the task start index of each computing unit with the task linear number; when the task start index and the task linear number are consistent, the corresponding segmentation task is taken as the target task; and the coordinates of the computing block corresponding to the target task are taken as the initial position parameters of each computing unit.

[0038] Specifically, the starting index of each computational unit can be compared and matched one by one with the linear task numbers generated in the previous traversal. The matching process uses the starting index of the task as the search benchmark, and locates the matching number value in the continuously increasing linear task numbers. When a task starting index is found to be equal to a linear task number value, the segmented task corresponding to that number is determined as the target task that the current computational unit needs to execute first. Then, the coordinate information of the computational blocks to which the target task belongs, such as the batch, attention head, query sub-block, and key-value sub-block, is extracted, and the above coordinate information is used as the initial position parameter of the current computational unit.

[0039] Optionally, after using the computation block coordinates corresponding to the target task as the initial position parameters of each computation unit, the method further includes: performing incremental calculations based on the initial position parameters to obtain the next round of position parameters, wherein the next round of position parameters includes the next round of key block coordinates and the next round of query block coordinates.

[0040] Specifically, after using the coordinates of the computation block corresponding to the target task as the initial position parameters of each computation unit, the coordinate information of the batch, attention head, query sub-block, and key-value sub-block can be incrementally calculated based on the initial position parameters of the current computation unit, combined with the execution flow of HSTU attention computation and the traversal order of the segmentation tasks. Without retracing all tasks, the position information corresponding to the next segmentation task to be executed after the current segmentation task is completed can be directly derived, thus obtaining the position parameters for the next round. The position parameters for the next round specifically include the key block coordinates and query block coordinates required for the next round of computation. Obtaining subsequent position parameters through incremental computation avoids repeated task traversal and coordinate lookup operations, reducing computational overhead. Simultaneously, in conjunction with the asynchronous data transfer engine in the ASIC architecture, pipeline parallelism between the current computation and the next round of data prefetching is achieved, fully utilizing on-chip cache resources, improving the overall parallel efficiency of computation and memory access, and ensuring the continuous and efficient execution of the variable-length HSTU attention computation process.

[0041] S130. Perform attention calculations on each segmentation task based on the initial position parameters, obtain the final result, and write it to the video memory.

[0042] In this context, video memory refers to the three-level storage in an ASIC architecture, used to store input data and final output results. ASICs also include on-chip L1 cache, L2 cache, and an asynchronous data transfer engine. The L1 cache is a private high-speed static random-access memory (SRAM) for each computing unit. The L2 cache is used for data exchange across computing units and temporary storage of intermediate results. The asynchronous data transfer engine is responsible for asynchronously moving data between video memory and on-chip cache, and can be executed in parallel with the computation process.

[0043] Specifically, based on the initial position parameters corresponding to each computing unit, the data required for the first round of computation is moved from video memory to on-chip L1 cache through the data transfer engine in the ASIC architecture. Then, the complete HSTU attention computation operations are executed sequentially, including query value multiplication with key matrix, relative bias superposition, activation function calculation, sequence length scaling, masking, and attention weight multiplication with data value matrix. The computation process adopts a pipelined design that combines computation and memory access in parallel, making full use of the ASIC on-chip multi-level cache to improve data access efficiency. Intermediate temporary results are preferentially stored in on-chip L2 cache rather than video memory temporary space. Intermediate results across computing units within the same query group are merged in L2 cache. After all the split tasks are completed and the intermediate results are merged, the final attention computation result is written to video memory. This achieves variable-length HSTU attention load balancing computation with low scheduling overhead, low memory access pressure, and high hardware utilization.

[0044] The technical solution of this invention can balance the load of each computing unit and avoid wasting computing power by dividing the attention calculation data and distributing the tasks evenly to each computing unit; by determining the initial position parameter based on linear numbering, the task location and scheduling logic can be simplified and the scheduling complexity can be reduced; by performing attention calculation based on the initial position parameter and outputting the result, the parallel efficiency of memory access and calculation can be improved, the video memory overhead can be reduced, and the characteristics of ASIC hardware architecture can be fully adapted.

[0045] Example 2 Figure 2 This is a flowchart of an attention load balancing method provided in Embodiment 2 of the present invention. This embodiment adds a specific process to Embodiment 1, where attention calculations are performed on each segmented task based on the initial position parameters, the final result is obtained, and written to the video memory. The specific content of steps S250-S260 is largely the same as steps S120-S130 in Embodiment 1, and therefore will not be repeated in this embodiment. Figure 2 As shown, the method includes: S210. Obtain attention calculation data, segment the attention calculation data, generate segmentation tasks, and distribute the segmentation tasks evenly to each calculation unit.

[0046] Optionally, the attention computation data is segmented to generate a segmentation task, including: dividing the attention computation into layers according to preset dimensions to obtain layered data, wherein the layered data includes batches, attention heads, query sequences, and key-value sequences; dividing the query sequences into query sub-blocks of fixed length, and dividing the key-value sequences into key-value sub-blocks of fixed length; and generating a segmentation task by using the combination of batches, attention heads, query sub-blocks, and key-value sub-blocks as the smallest computational unit.

[0047] Optionally, after segmenting the attention computation data to generate segmentation tasks, the method further includes: identifying invalid computation blocks introduced by the mask; filtering invalid computation blocks in the segmentation tasks to obtain valid segmentation tasks; and evenly distributing the segmentation tasks to each computation unit, including: evenly distributing the valid segmentation tasks to each computation unit.

[0048] S220. Determine the task start index of each computing unit, traverse each segmented task and assign a linear number to the task, and determine the initial position parameter of each computing unit based on the task start index and the task linear number.

[0049] Optionally, each segmentation task is traversed and assigned a linear task number, including: traversing each segmentation task in the order of batch, attention head, query sub-block, and key-value sub-block; and assigning a unique and continuously increasing sequence number to each segmentation task as the task linear number.

[0050] Optionally, the initial position parameters of each computing unit are determined based on the task start index and the task linear number, including: matching the task start index of each computing unit with the task linear number; when the task start index and the task linear number are consistent, the corresponding segmentation task is taken as the target task; and the coordinates of the computing block corresponding to the target task are taken as the initial position parameters of each computing unit.

[0051] Optionally, after using the computation block coordinates corresponding to the target task as the initial position parameters of each computation unit, the method further includes: performing incremental calculations based on the initial position parameters to obtain the next round of position parameters, wherein the next round of position parameters includes the next round of key block coordinates and the next round of query block coordinates.

[0052] S230. Based on the initial position parameters, read the query value and key value of this round, calculate the query key matrix multiplication, obtain the matrix multiplication result of this round, and store it in the first cache.

[0053] Specifically, firstly, based on the determined initial location parameters, the physical addresses of the current query value Q and key value K corresponding to the current task block in video memory are located. Then, the ASIC's dedicated asynchronous data transfer engine (DTE) is activated, using Direct Memory Access (DMA) to move the query value and key value from external video memory to the independent private first cache L1 of each computing unit. The first cache is a high-speed SRAM, serving only the current computing unit, and features low latency and high throughput. After the query value and key value are loaded into L1 and are in a ready state, the ASIC's built-in 2D matrix multiply-accumulate (MMA) hardware unit is called to perform matrix multiplication of the query value and key value within the first cache, obtaining the matrix multiplication result of the query-key combination. This result is temporarily stored directly in the first cache without being written to external video memory, thereby reducing latency and bandwidth consumption caused by off-chip memory access.

[0054] S240. Asynchronously read the values ​​of this round, store them in the first cache, and perform activation calculation on the matrix multiplication results of this round to obtain the activation calculation results and store them in the first cache.

[0055] Specifically, after completing the current round of query key matrix multiplication, the block storage address of the current round's value V is immediately calculated based on the coordinate information of the current task block. The value is then read from the background by the asynchronous data transfer engine (DTE) and loaded into the first cache. This asynchronous read operation proceeds in parallel with subsequent calculations, preventing idle waiting by the computing unit and maximizing hardware resource utilization. Simultaneously, the matrix multiplication result stored in the first cache is subjected to silu activation calculation according to the calculation logic specified by HSTU attention, completing the nonlinear transformation processing. The result of the activation calculation is then written back to the first cache for storage.

[0056] S250. Asynchronously read the key value of the next round based on the key block coordinates of the next round, and asynchronously read the offset value of the current round based on the initial position parameters. Add the offset value of the current round and the activation calculation result to obtain the offset result of the current round and store it in the first cache.

[0057] Specifically, while performing activation calculations, based on the key block coordinates obtained from previous incremental calculations, the asynchronous prefetching operation for the next round of key values ​​is initiated in advance. The key values ​​required for the next round of calculations are loaded into the first cache, achieving pipeline overlap between data prefetching and the current calculation. Simultaneously, based on the initial position parameters, the storage location of the current round's bias value in video memory is determined. The current round's bias value is asynchronously read into the first cache via DTE. Following the HSTU calculation process, the current round's bias value is added element-wise with the activation calculation result to obtain the current round's bias addition result, which is then stored in the L1 cache.

[0058] S260. Perform length scaling and masking processing on the biasing result of this round to obtain the masking result of this round, and store the masking result of this round into the first cache.

[0059] Specifically, after the biasing operation is completed, the HSTU attention calculation process continues, and the result of this round of biasing is subjected to sequence length scaling, i.e., divided by a fixed parameter scaling_seqlen, to complete numerical normalization. Then, combined with the causal masking and target masking rules, the scaled data is masked, and invalid calculation positions are masked and set to zero to obtain the masking result of this round that meets the requirements of HSTU calculation, and this result is stored in the first cache.

[0060] S270. Asynchronously read the next round of query values ​​based on the block coordinates of the next round of query.

[0061] Specifically, during the intervals of mask processing, the asynchronous reading of the next round of query values ​​is initiated based on the block coordinates of the next round of queries. The query values ​​required for the next round are preloaded into the first cache, further enhancing the parallel efficiency of computation and memory access, and maximizing the utilization of the ASIC on-chip cache and asynchronous transmission capabilities.

[0062] S280. Perform matrix multiplication based on the current round mask result and the current round value in the first cache to obtain the current round calculation result.

[0063] Specifically, after the masking process is completed, the masking result in the first cache is used as the attention weight, and a matrix multiplication operation is performed with the asynchronously loaded value V in the current round to obtain the calculation result of the current task block.

[0064] S290. In the next round of calculation, the next round query value and the next round key value read asynchronously are used as the new current round query value and the current round key value, and the calculation is repeated until all the subdivision tasks of the current calculation unit are completed.

[0065] Specifically, after the current task block is computed, the next round's query value and key value, asynchronously prefetched from the previous stage, are used as the query value and key value for the new round of computation. Following the same execution flow, all operations—query key matrix multiplication, activation computation, bias value superposition, sequence length scaling, masking, and numerical matrix multiplication—are repeated sequentially until all subdivision tasks within the continuous intervals allocated to the current computation unit are completed. The entire computation process strictly adheres to the parallel execution characteristics and on-chip memory access rules of ASIC hardware. All intermediate results are transferred and temporarily stored in the on-chip L1 and L2 caches, completely independent of the GPU workspace for storing temporary data. Simultaneously, through a pipeline design that deeply parallelizes computation and memory access, memory access latency is reduced, hardware computation unit utilization is improved, and ultimately, efficient load-balanced execution of HSTU attention computation in variable-length sequence scenarios is achieved.

[0066] Optionally, after obtaining the calculation result of this round, the method further includes: determining whether the calculation group corresponding to the calculation result of this round is complete, wherein the calculation group is all the split tasks corresponding to the same query block divided according to the query offset; if so, the calculation result of this round is directly written to the video memory as the final result; otherwise, the calculation result of this round is written to the second cache as an intermediate result. After all the split tasks are calculated, the intermediate results belonging to the same calculation group in the second cache are merged to obtain the final result and written to the video memory.

[0067] The computation group is divided according to the query offset, referring to the set of all key-value splitting tasks corresponding to the same query block. All tasks within the same query group jointly complete the complete attention calculation of a query block. Only when all tasks within the same group have completed their calculations can a valid and continuous output result be obtained. During the judgment process, the query block identifier, query offset position, and the total number of pre-allocated tasks for the query group can be used to determine whether the current calculation result is the last calculation result for that query group, thus determining whether the current computation group is complete. If the judgment result indicates that the current computation group is complete, it means that all key-value splitting tasks corresponding to the query block have been executed within the current computation unit. There is no task splitting across computation units, and there is no need to splice and accumulate intermediate results from multiple chips. In this case, the current calculation result can be directly used as the final calculation result for the query block and written to the video memory via ASIC's DTE using Direct Memory Access (DMA), completing the output process of the current computation group. If the determination result indicates that the current computation group is incomplete, it means that the splitting task corresponding to the query block has been distributed to multiple different computation units by the sequential equalization algorithm. The result of this round of computation is only a part of the entire computation group and needs to be merged with the intermediate results of the same group output by other computation units to form a complete result. At this time, the result of this round of computation needs to be written as an intermediate result into the second cache on the ASIC chip. This cache is used to realize data interaction and temporary result storage between different computation units. It can replace the video memory workspace that must be used in the traditional solution, effectively reducing video memory bandwidth occupation and memory access latency. After all the splitting tasks of the current computation unit and all other computation units have been completed, the system will perform a global synchronization operation to ensure that the intermediate results of all computation units have been stably written into the second cache. Then, according to the query offset, all intermediate results in the second cache are grouped and collected, and all start-type intermediate results, intermediate-type intermediate results, and end-type intermediate results belonging to the same computation group are identified. The computation unit where the first task of the computation group is located starts the merging computation operation, and the intermediate results of all the segments in the same group are accumulated and integrated in sequence to form a continuous and complete final computation result. Finally, the merged final result is written to the video memory through the data transfer engine (DTE). The above processing method makes full use of the hardware characteristics of the ASIC on-chip multi-level cache, combined with the sequential task allocation mechanism, simplifies the scheduling logic of merging results across computing units, and ensures the integrity and correctness of HSTU attention calculation results in variable-length sequence scenarios without relying on additional temporary video memory space, while improving overall computing efficiency and hardware resource utilization.

[0068] The technical solution of this invention effectively reduces off-chip memory access overhead by relying on a first cache to cache computational data. Through multi-stage asynchronous data reading and synchronous execution of computation, parallel computation and memory access pipelines are achieved, hiding data read latency. Scaling and masking processes are strictly followed to ensure computational accuracy, and all tasks are completed iteratively using prefetched data. By determining the integrity of computational groups and selecting between direct writing to video memory or merging into a second cache, the result output logic is simplified, avoiding the use of temporary video memory space, reducing memory access overhead, and improving result merging efficiency.

[0069] Example 3 Figure 3 This is a schematic diagram of an attention load balancing device provided in Embodiment 3 of the present invention. Figure 3 As shown, the device includes a task segmentation and allocation module 310, which is used to: acquire attention calculation data, segment the attention calculation data, generate segmented tasks, and distribute the segmented tasks evenly to each computing unit; The initial position parameter determination module 320 is used to: determine the task start index of each computing unit, traverse each segmented task and assign a linear number to the task, and determine the initial position parameter of each computing unit based on the task start index and the task linear number. Attention calculation module 330 is used to: perform attention calculation on each segmented task based on the initial position parameters, obtain the final result and write it to the video memory.

[0070] Optionally, the task segmentation and allocation module 310 is specifically used to: divide the attention calculation into layers according to a preset dimension to obtain layer data, wherein the layer data includes batches, attention heads, query sequences, and key-value sequences; divide the query sequences into query sub-blocks of a fixed length, and divide the key-value sequences into key-value sub-blocks of a fixed length; and generate segmented tasks by using the combination of batches, attention heads, query sub-blocks, and key-value sub-blocks as the smallest calculation unit.

[0071] Optionally, the apparatus further includes: an invalid computation block removal module, used to segment the attention computation data, generate segmentation tasks, and then determine the invalid computation blocks introduced by the mask; filter the invalid computation blocks in the segmentation tasks to obtain valid segmentation tasks; and distribute the segmentation tasks evenly to each computation unit, including: distributing the valid segmentation tasks evenly to each computation unit.

[0072] Optionally, the initial position parameter determination module 320 specifically includes: a task linear number determination unit, used to: sequentially traverse each segmented task according to the order of batch, attention head, query sub-block, and key-value sub-block; and assign a unique and continuously increasing sequence number to each segmented task as the task linear number.

[0073] Optionally, the initial position parameter determination module 320 specifically includes: an initial position parameter determination unit, used to: match the task start index of each computing unit with the task linear number; when the task start index matches the task linear number, take the corresponding segmented task as the target task; and take the computing block coordinates corresponding to the target task as the initial position parameter of each computing unit.

[0074] Optionally, the device further includes: a position increment calculation module, used to perform incremental calculation based on the initial position parameters after using the calculation block coordinates corresponding to the target task as the initial position parameters of each calculation unit to obtain the next round of position parameters, wherein the next round of position parameters includes the next round of key block coordinates and the next round of query block coordinates.

[0075] Optionally, the attention calculation module 330 specifically includes: a current round calculation result generation unit, used for: reading the current round query value and current round key value based on the initial position parameters, calculating the query key matrix multiplication, obtaining the current round matrix multiplication result and storing it in the first cache; asynchronously reading the current round value, storing it in the first cache, and performing activation calculation on the current round matrix multiplication result to obtain the activation calculation result and storing it in the first cache; asynchronously reading the next round key value based on the next round key block coordinates, and asynchronously reading the current round bias value based on the initial position parameters, adding the current round bias value and the activation calculation result. The process involves obtaining the biased result for the current round and storing it in the first cache; performing length scaling and masking on the biased result to obtain the masked result for the current round, and storing the masked result in the first cache; asynchronously reading the query value for the next round based on the block coordinates of the next round query; performing matrix multiplication on the masked result and the value for the current round in the first cache to obtain the calculation result for the current round; and repeating the calculation in the next round by using the asynchronously read query value and key value for the next round as the new query value and key value for the current round, until all the partitioning tasks of the current calculation unit have been completed.

[0076] Optionally, the attention calculation module 330 further includes: an intermediate result merging unit, used to: after obtaining the calculation result of the current round, determine whether the calculation group corresponding to the calculation result of the current round is complete, wherein the calculation group is all the split tasks corresponding to the same query block divided according to the query offset; if so, write the calculation result of the current round directly as the final result into the video memory; otherwise, write the calculation result of the current round as an intermediate result into the second cache, and after all the split tasks have been calculated, merge the intermediate results belonging to the same calculation group in the second cache to obtain the final result and write it into the video memory.

[0077] The technical solution of this invention can balance the load of each computing unit and avoid wasting computing power by dividing the attention calculation data and distributing the tasks evenly to each computing unit; by determining the initial position parameter based on linear numbering, the task location and scheduling logic can be simplified and the scheduling complexity can be reduced; by performing attention calculation based on the initial position parameter and outputting the result, the parallel efficiency of memory access and calculation can be improved, the video memory overhead can be reduced, and the characteristics of ASIC hardware architecture can be fully adapted.

[0078] The attention load balancing device provided in this embodiment of the invention can execute an attention load balancing method provided in any embodiment of the invention, and has the corresponding functional modules and beneficial effects of the method.

[0079] Example 4 Figure 4 A schematic diagram of an electronic device 10 that can be used to implement embodiments of the present invention is shown. The electronic device is intended to represent various forms of digital computers, such as laptop computers, desktop computers, workstations, personal digital assistants, servers, blade servers, mainframe computers, and other suitable computers. The electronic device can also represent various forms of mobile devices, such as personal digital assistants, cellular phones, smartphones, wearable devices (e.g., helmets, glasses, watches, etc.), and other similar computing devices. The components shown herein, their connections and relationships, and their functions are merely illustrative and are not intended to limit the implementation of the invention described and / or claimed herein.

[0080] like Figure 4 As shown, the electronic device 10 includes at least one processor 11 and a memory, such as a read-only memory (ROM) or random access memory (RAM), communicatively connected to the at least one processor 11. The memory stores computer programs executable by the at least one processor. The processor 11 can perform various appropriate actions and processes based on the computer program stored in the ROM 12 or loaded into the RAM 13 from the storage unit 18. The RAM 13 can also store various programs and data required for the operation of the electronic device 10. The processor 11, ROM 12, and RAM 13 are interconnected via a bus 14. Input / output (I / O) interfaces are also connected to the bus 14.

[0081] Multiple components in electronic device 10 are connected to I / O interface 15, including: input unit 16, such as keyboard, mouse, etc.; output unit 17, such as various types of displays, speakers, etc.; storage unit 18, such as disk, optical disk, etc.; and communication unit 19, such as network card, modem, wireless transceiver, etc. Communication unit 19 allows electronic device 10 to exchange information / data with other devices through computer networks such as the Internet and / or various telecommunications networks.

[0082] Processor 11 can be a variety of general-purpose and / or special-purpose processing components with processing and computing capabilities. Some examples of processor 11 include, but are not limited to, a central processing unit (CPU), a graphics processing unit (GPU), various special-purpose artificial intelligence (AI) computing chips, various processors running machine learning model algorithms, a digital signal processor (DSP), and any suitable processor, controller, microcontroller, etc. Processor 11 performs the various methods and processes described above, such as an attention load balancing method.

[0083] In some embodiments, an attention load balancing method may be implemented as a computer program tangibly contained in a computer-readable storage medium, such as storage unit 18. In some embodiments, part or all of the computer program may be loaded and / or installed on electronic device 10 via ROM 12 and / or communication unit 19. When the computer program is loaded into RAM 13 and executed by processor 11, one or more steps of the attention load balancing method described above may be performed. Alternatively, in other embodiments, processor 11 may be configured to perform an attention load balancing method by any other suitable means (e.g., by means of firmware).

[0084] Various embodiments of the systems and techniques described above herein can be implemented in digital electronic circuit systems, integrated circuit systems, field-programmable gate arrays (FPGAs), application-specific integrated circuits (ASICs), application-specific standard products (ASSPs), systems-on-a-chip (SoCs), payload-programmable logic devices (CPLDs), computer hardware, firmware, software, and / or combinations thereof. These various embodiments may include implementations in one or more computer programs that can be executed and / or interpreted on a programmable system including at least one programmable processor, which may be a dedicated or general-purpose programmable processor, capable of receiving data and instructions from a storage system, at least one input device, and at least one output device, and transmitting data and instructions to the storage system, the at least one input device, and the at least one output device.

[0085] Computer programs used to implement the methods of the present invention may be written in any combination of one or more programming languages. These computer programs may be provided to a processor of a general-purpose computer, a special-purpose computer, or other programmable data processing device, such that when executed by the processor, the computer programs cause the functions / operations specified in the flowcharts and / or block diagrams to be performed. The computer programs may be executed entirely on a machine, partially on a machine, or as a standalone software package, partially on a machine and partially on a remote machine, or entirely on a remote machine or server.

[0086] In the context of this invention, a computer-readable storage medium can be a tangible medium that may contain or store a computer program for use by or in conjunction with an instruction execution system, apparatus, or device. A computer-readable storage medium may include, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, apparatus, or devices, or any suitable combination thereof. Alternatively, a computer-readable storage medium may be a machine-readable signal medium. More specific examples of machine-readable 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, portable compact disk read-only memory (CD-ROM), optical storage devices, magnetic storage devices, or any suitable combination thereof.

[0087] To provide interaction with a user, the systems and techniques described herein can be implemented on an electronic device having: a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user; and a keyboard and pointing device (e.g., a mouse or trackball) through which the user provides input to the electronic device. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback (e.g., visual feedback, auditory feedback, or tactile feedback); and input from the user can be received in any form (including sound input, voice input, or tactile input).

[0088] The systems and technologies described herein can be implemented in computing systems that include backend components (e.g., as data servers), or middleware components (e.g., application servers), or frontend components (e.g., user computers with graphical user interfaces or web browsers through which users can interact with implementations of the systems and technologies described herein), or any combination of such backend, middleware, or frontend components. The components of the system can be interconnected via digital data communication of any form or medium (e.g., communication networks). Examples of communication networks include local area networks (LANs), wide area networks (WANs), blockchain networks, and the Internet.

[0089] A computing system can include clients and servers. Clients and servers are generally located far apart and typically interact through communication networks. The client-server relationship is created by computer programs running on the respective computers and having a client-server relationship with each other. The server can be a cloud server, also known as a cloud computing server or cloud host, which is a hosting product within the cloud computing service system. It addresses the shortcomings of traditional physical hosts and VPS services, such as high management difficulty and weak business scalability.

[0090] It should be understood that the various forms of processes shown above can be used, with steps reordered, added, or deleted. For example, the steps described in this invention can be executed in parallel, sequentially, or in different orders, as long as the desired result of the technical solution of this invention can be achieved, and this is not limited herein.

[0091] The specific embodiments described above do not constitute a limitation on the scope of protection of this invention. Those skilled in the art should understand that various modifications, combinations, sub-combinations, and substitutions can be made according to design requirements and other factors. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of this invention should be included within the scope of protection of this invention.

Claims

1. An attention load balancing method, characterized in that, Applications in Application-Specific Integrated Circuit (ASIC) architectures include: Acquire attention calculation data, segment the attention calculation data, generate segmentation tasks, and distribute the segmentation tasks evenly to each computing unit; Determine the task start index of each computing unit, traverse each segmented task and assign a linear number to the task, and determine the initial position parameter of each computing unit based on the task start index and the task linear number. Attention calculations are performed on each segmentation task based on the initial position parameters to obtain the final result, which is then written to the video memory.

2. The method according to claim 1, characterized in that, The step of segmenting the attention calculation data to generate a segmentation task includes: The attention calculation is divided into layers according to a preset dimension to obtain layer data, wherein the layer data includes batch, attention head, query sequence and key value sequence; The query sequence is divided into query sub-blocks of fixed length, and the key value sequence is divided into key value sub-blocks of fixed length; The minimum computational unit is a combination of batches, attention heads, query sub-blocks, and key-value sub-blocks to generate a segmentation task.

3. The method according to claim 1, characterized in that, After segmenting the attention calculation data to generate segmentation tasks, the method further includes: Identify invalid computation blocks introduced by the mask; Invalid computation blocks in the segmented task are filtered out to obtain valid segmented tasks; The step of evenly distributing the segmented tasks to each computing unit includes: The effective partitioning tasks are evenly distributed to each computing unit.

4. The method according to claim 1, characterized in that, The step of traversing each segmented task and assigning it a linear number includes: In the order of batch, attention head, query sub-block, key-value sub-block, traverse each segmented task in sequence; Each segmentation task is assigned a unique and continuously increasing sequence number, which serves as the linear numbering of the task.

5. The method according to claim 4, characterized in that, The step of determining the initial position parameters of each computing unit based on the task start index and the task linear number includes: Match the task start index of each computing unit with the task linear number; When the starting index of the task matches the linear number of the task, the corresponding split task is taken as the target task; The coordinates of the computation block corresponding to the target task are used as the initial position parameters of each computation unit.

6. The method according to claim 5, characterized in that, After using the coordinates of the computation block corresponding to the target task as the initial position parameters of each computation unit, the method further includes: Incremental calculations are performed based on the initial position parameters to obtain the next round of position parameters, wherein the next round of position parameters includes the next round of key block coordinates and the next round of query block coordinates.

7. The method according to claim 6, characterized in that, The attention calculation for each segmentation task based on the initial position parameters includes: Based on the initial position parameters, the query value and key value of this round are read, the query key matrix multiplication is calculated, the matrix multiplication result of this round is obtained and stored in the first cache; Asynchronously read the values ​​of the current round, store them in the first cache, and perform activation calculation on the matrix multiplication results of the current round to obtain the activation calculation results and store them in the first cache; The key value of the next round is read asynchronously based on the key block coordinates of the next round, and the offset value of the current round is read asynchronously based on the initial position parameter. The offset value of the current round and the activation calculation result are added together to obtain the offset result of the current round and stored in the first cache. Perform length scaling and masking processing on the current biasing result to obtain the current masking result, and store the current masking result in the first cache; The next round query value is read asynchronously based on the block coordinates of the next round query. The result of this round of calculation is obtained by performing matrix multiplication based on the mask result and the value of this round in the first cache. In the next round of calculation, the next round's query value and key value read asynchronously are used as the new query value and key value for the current round, and the calculation is repeated until all the subdivision tasks of the current calculation unit are completed.

8. The method according to claim 7, characterized in that, After obtaining the results of this round of calculations, the method further includes: Determine whether the calculation group corresponding to the calculation result of this round is complete, wherein the calculation group is all the split tasks corresponding to the same query block divided according to the query offset; If so, write the result of this round of calculation directly into the video memory as the final result; Otherwise, the calculation result of this round is written as an intermediate result to the second cache. After all the split tasks are completed, the intermediate results belonging to the same calculation group in the second cache are merged to obtain the final result and written to the video memory.

9. An electronic device, characterized in that, The electronic device includes: At least one processor; and a memory communicatively connected to the at least one processor; The memory stores a computer program that can be executed by the at least one processor, the computer program being executed by the at least one processor to enable the at least one processor to perform the method of any one of claims 1-8.

10. A computer storage medium, characterized in that, The computer storage medium stores computer instructions that are used to cause a processor to execute the method of any one of claims 1-8.

Images