Multi-head attention task accelerator mapping configuration generation method and system
By constructing memory overhead constraints and using reinforcement learning algorithms to search in a reduced mapping space, the mapping configuration for multi-head attention tasks is optimized, solving the problems of high memory access overhead and low efficiency in traditional methods, and achieving efficient mapping configuration generation and accelerator optimization.
Patent Information
- Application Number
- CN202511992788.3
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-12-26
- Publication Date
- 2026-05-05
AI Technical Summary
When traditional multi-head attention tasks are deployed on spatial accelerators, the limited on-chip memory capacity leads to frequent memory access overhead and difficulty in finding the optimal mapping configuration. Existing technologies use reinforcement learning methods to spend time filtering suitable configurations in a huge mapping space, which reduces efficiency.
By constructing shared memory overhead and register memory overhead constraints, eliminating invalid block factors, and using a pre-defined reinforcement learning algorithm to search for the optimal configuration in the reduced mapping space, the reinforcement learning algorithm is optimized by combining computation latency, memory access latency, and power consumption.
It improves the efficiency of generating optimal mapping configurations, reduces memory access overhead and computation time, and optimizes the operating efficiency of the space accelerator.
Smart Images
Figure CN121979560A_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of electronic information technology, and in particular relates to a method and system for mapping and configuring multi-head attention task accelerators. Background Technology
[0002] Traditional approaches employ an operator mapping strategy to deploy multi-head attention tasks on space accelerators with multi-level memory hierarchies. Under this strategy, each operator in the multi-head attention task reads the input tensor from the off-chip memory of the space accelerator, transfers the input tensor to on-chip memory for computation, and then writes the computation result back to off-chip memory. However, due to the limited capacity of on-chip memory, intermediate results between operators need to be written back to off-chip memory for frequent read and write operations in subsequent computations, resulting in a large amount of unnecessary memory access overhead.
[0003] To address the aforementioned issues, existing technologies utilize a block-based mechanism to decompose tensors into smaller blocks adapted to on-chip memory. This allows intermediate variables to be temporarily stored in on-chip memory and directly reused by downstream operators, effectively reducing access to off-chip memory. However, existing technologies involving block-based mapping configurations rely on reinforcement learning methods to search an unpruned and extremely large mapping space. Consequently, significant time is required to select suitable mapping configurations, making it difficult to find the optimal configuration and drastically reducing the search efficiency for the optimal mapping configuration. Summary of the Invention
[0004] The present invention aims to provide a method and system for generating mapping configurations for multi-head attention task accelerators, in order to solve the above-mentioned technical problems and improve the generation efficiency of optimal mapping configurations.
[0005] To address the aforementioned technical problems, this invention provides a method for generating a multi-head attention task accelerator mapping configuration, comprising the following steps: Obtain the operator-shared parameters, head dimension, and key sequence length for the multi-head attention task to be accelerated; Obtain the total shared memory capacity and total register capacity of the target space accelerator; Based on the operator common parameters, the head dimension, and the key sequence length, the shared memory overhead and register memory overhead are constructed. Based on the shared memory overhead and the total shared memory capacity, a shared memory overhead constraint is constructed; Based on the register memory overhead and the total register capacity, a register memory overhead constraint is constructed; Set the block factor corresponding to each task parameter in the common parameters of the operators to the preset minimum value of the block factor, and perform a constraint range step on each task parameter to obtain the range of the block factor corresponding to each task parameter. Based on the range of the block factor, an initial mapping space is obtained, and based on the initial mapping space, a search is performed under a preset reinforcement learning algorithm to generate the optimal mapping configuration for the target space accelerator. The constraint range step includes: Substitute the block factor corresponding to the task parameters other than the current task parameters in the common parameters of the operators into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size; Based on the maximum block size and the current task parameters, generate the block factor range corresponding to the current task parameters.
[0006] In the above scheme, the operator shared parameters represent the task parameters that the block factors in the multi-head attention task must be consistent across all operators. Based on the operator shared parameters, head dimension, and key sequence length, shared memory overhead and register memory overhead are constructed. Shared memory overhead represents the size of the shared memory space occupied by the block size corresponding to a set of block factors in the accelerator, and register memory overhead represents the size of the register space occupied by the block size corresponding to a set of block factors in the accelerator. Subsequently, based on the block factor range obtained by solving the formed shared memory overhead constraints and register memory overhead constraints, this scheme can eliminate invalid block factor values that will cause memory overflow, ensuring that each candidate mapping configuration in the initial mapping space will not cause memory overflow. This eliminates the need for time-consuming searches for invalid mapping configurations that cause memory overflow when using reinforcement learning algorithms, and enables faster and more stable acquisition of the optimal mapping configuration of the target space accelerator, improving the generation efficiency of the optimal mapping configuration.
[0007] Further, the construction of shared memory overhead and register memory overhead based on the operator common parameters, the head dimension, and the key sequence length includes: constructing a first operator shared memory overhead and a second operator shared memory overhead based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator common parameters; constructing a first operator register memory overhead based on the key sequence length and the batch size, number of heads, and query sequence length in the operator common parameters; constructing a second operator register memory overhead based on the head dimension and the batch size, number of heads, and query sequence length in the operator common parameters; constructing shared memory overhead based on the first operator shared memory overhead and the second operator shared memory overhead; and constructing register memory overhead based on the first operator register memory overhead and the second operator register memory overhead.
[0008] Further, the step of constructing the first operator shared memory overhead and the second operator shared memory overhead based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator common parameters includes: constructing the first operator input tensor memory overhead based on the head dimension and the batch size, number of heads, and query sequence length in the operator common parameters; constructing the first operator weight tensor memory overhead based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters; constructing the second operator input tensor memory overhead based on the key sequence length and the batch size, number of heads, and query sequence length in the operator common parameters; constructing the second operator weight tensor memory overhead based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters; constructing the first operator shared memory overhead based on the first operator input tensor memory overhead and the first operator weight tensor memory overhead; and constructing the second operator shared memory overhead based on the second operator input tensor memory overhead and the second operator weight tensor memory overhead.
[0009] The register memory overhead and shared memory overhead constructed by the above scheme can accurately characterize the actual storage occupancy of tensors of different operators in the space accelerator at each layer, so as to ensure that the estimation of space accelerator hardware resources is more accurate in the subsequent search for mapping configurations. This enables more accurate and efficient screening of feasible configurations that meet register capacity constraints, thereby improving the generation efficiency of the optimal mapping configuration.
[0010] Further, in the step of obtaining an initial mapping space based on the block factor range, and performing a search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator, the initial mapping space includes the block factor range and a preset parallel factor range. Obtaining the parallel factor range includes: obtaining the number of computing cores of the target space accelerator and constructing parallel computing overhead constraints based on the number of computing cores; obtaining the parallel factors corresponding to the operator common parameters, the head dimension, and the key sequence length; and obtaining the parallel factor range based on the parallel factors and the parallel computing overhead constraints.
[0011] Further, in the step of obtaining an initial mapping space based on the block factor range, and performing a search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator, the step of performing a search under the preset reinforcement learning algorithm based on the initial mapping space to obtain the optimal mapping configuration for the target space accelerator includes: performing a factor decision step based on each target factor range in the initial mapping space to obtain the optimal mapping configuration for the target space accelerator; the factor decision step includes: performing a search under the preset reinforcement learning algorithm based on the current target factor range to obtain a definite value corresponding to the current target factor range, and marking the current target factor range as decided. Based on the determined value, under the preset factor range update algorithm, the second factor range corresponding to the target factor range not marked as decided in the initial mapping space is obtained; the initial mapping space is updated according to the determined value corresponding to the target factor range marked as decided in the initial mapping space and the second factor range; if it is confirmed that there is a target factor range not marked as decided in the initial mapping space, the factor decision step is re-executed based on a target factor range not marked as decided in the initial mapping space; otherwise, the factor decision step is stopped, and the optimal mapping configuration of the target space accelerator is generated according to the determined value corresponding to each target factor range in the initial mapping space.
[0012] Further, obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the preset factor range update algorithm includes: confirming that the current target factor is a block factor, then obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the shared memory overhead constraint and the register memory overhead constraint; otherwise, obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the parallel computing overhead constraint.
[0013] In the above scheme, after determining the definite value of the target factor corresponding to each target factor range, the remaining unmarked target factor ranges are immediately updated with this definite value as the second factor range. This second factor range eliminates the invalid parts of the target factor range that conflict with the definite value, so the second factor range will be smaller than the original target factor range. Subsequently, the initial mapping space is updated based on the second factor range and the definite values of the target factors corresponding to the target factor ranges marked as decided. Therefore, after each definite value is obtained, the remaining unmarked target factor ranges will gradually shrink. Compared with the traditional scheme that searches in an unpruned and extremely large mapping space, this scheme uses a preset reinforcement learning algorithm to search for definite values in this continuously shrinking target factor range, which can greatly reduce the time required to generate the optimal mapping configuration and improve the generation efficiency of the optimal mapping configuration.
[0014] Further, after obtaining the initial mapping space based on the block factor range, and performing a search based on the initial mapping space under a preset reinforcement learning algorithm to generate the optimal mapping configuration for the target space accelerator, the method further includes: obtaining the computation latency, memory access latency, and memory access power consumption of the target space accelerator after executing the multi-head attention task based on the optimal mapping configuration; generating a reward value based on the computation latency, the memory access latency, and the memory access power consumption; and optimizing the reinforcement learning algorithm based on the reward value under a preset optimization algorithm.
[0015] In the above scheme, the optimal mapping configuration is applied to the computation latency, memory access latency, and memory power consumption of the space accelerator. The reinforcement learning algorithm is then optimized based on the reward value generated from the computation latency, memory access latency, and memory power consumption, so that the next optimal mapping configuration will have lower computation latency, memory access latency, and memory power consumption, thereby improving the operating efficiency of the space accelerator.
[0016] Further, the step of substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the shared memory overhead constraint and the register memory overhead constraint to obtain the maximum block size includes: substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the shared memory overhead constraint to obtain the first maximum block size; substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the register memory overhead constraint to obtain the second maximum block size; confirming that the first maximum block size is less than the second maximum block size, then using the first maximum block size as the maximum block size; otherwise, using the second maximum block size as the maximum block size.
[0017] Further, generating the block factor range corresponding to the current task parameters based on the maximum block size and the current task parameters includes: calculating the minimum block factor based on the ratio of the current task parameters and the maximum block size; calculating the maximum block factor based on the ratio of the current task parameters under the preset minimum block size; and generating the block factor range corresponding to the current task parameters based on the minimum block factor and the maximum block factor.
[0018] This invention also provides a multi-head attention task accelerator mapping configuration generation system for implementing any of the multi-head attention task accelerator mapping configuration generation methods described above, comprising: a task parameter acquisition module for acquiring the operator common parameters, head dimension, and key sequence length of the multi-head attention task to be accelerated; an accelerator parameter acquisition module for acquiring the total shared memory capacity and total register capacity of the target space accelerator; a memory overhead construction module for constructing the shared memory overhead and register memory overhead based on the operator common parameters, the head dimension, and the key sequence length; a shared memory overhead constraint construction module for constructing shared memory overhead constraints based on the shared memory overhead and the total shared memory capacity; and a register overhead constraint construction module for constructing register memory overhead constraints based on the register memory overhead and the total register capacity. The module for obtaining the block factor range is used to set the block factor corresponding to each task parameter in the operator common parameters to a preset minimum value of the block factor, and to perform a constraint range step on each task parameter to obtain the block factor range corresponding to each task parameter; the module for generating the optimal mapping configuration is used to obtain an initial mapping space based on the block factor range, and to search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator; wherein, the constraint range step includes: substituting the block factors corresponding to task parameters other than the current task parameter in the operator common parameters into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size; and generating the block factor range corresponding to the current task parameter based on the maximum block size and the current task parameter.
[0019] In the above scheme, the range of block factors obtained by solving the shared memory overhead constraints and register memory overhead constraints can eliminate invalid block factor values that will cause memory overflow, ensuring that each candidate mapping configuration factor in the initial mapping space will not cause memory overflow. This eliminates the need to spend time searching for invalid mapping configurations that cause memory overflow when using reinforcement learning algorithms in the future, and can obtain the optimal mapping configuration of the target space accelerator quickly and stably, thereby improving the generation efficiency of the optimal mapping configuration. Attached Figure Description
[0020] Figure 1A flowchart illustrating the steps of a method for generating a mapping configuration for a multi-head attention task accelerator, as provided in an embodiment of the present invention. Figure 2 This is a schematic diagram of a multi-head attention task accelerator mapping configuration generation system provided in an embodiment of the present invention. Detailed Implementation
[0021] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. 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 are within the scope of protection of the present invention.
[0022] Please see Figure 1 This embodiment provides a method for generating a mapping configuration for a multi-head attention task accelerator, including the following steps: Step S1: Obtain the operator-shared parameters, head dimension, and key sequence length for the multi-head attention task to be accelerated; Step S2: Obtain the total shared memory capacity and total register capacity of the target space accelerator; Step S3: Based on the operator common parameters, the head dimension, and the key sequence length, construct the shared memory overhead and register memory overhead; Step S4: Based on the shared memory overhead and the total shared memory capacity, construct shared memory overhead constraints; Step S5: Based on the register memory overhead and the total register capacity, construct register memory overhead constraints; Step S6: Set the block factor corresponding to each task parameter in the common parameters of the operators to the preset minimum value of the block factor, and perform the constraint range step on each task parameter to obtain the range of the block factor corresponding to each task parameter. Step S7: Based on the block factor range, obtain the initial mapping space, and search under the preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration of the target space accelerator; The constraint range step includes: Substitute the block factor corresponding to the task parameters other than the current task parameters in the common parameters of the operators into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size; Based on the maximum block size and the current task parameters, generate the block factor range corresponding to the current task parameters.
[0023] In the above embodiments, the operator shared parameters represent the task parameters that the block factor in the multi-head attention task must be consistent across all operators. Based on the operator shared parameters, head dimension, and key sequence length, shared memory overhead and register memory overhead are constructed. Shared memory overhead represents the size of the shared memory space occupied by the block size corresponding to a set of block factors in the accelerator, and register memory overhead represents the size of the register space occupied by the block size corresponding to a set of block factors in the accelerator. Subsequently, based on the block factor range obtained by solving the formed shared memory overhead constraints and register memory overhead constraints, invalid block factor values that would cause memory overflow can be eliminated, ensuring that each candidate mapping configuration in the initial mapping space will not cause memory overflow. This eliminates the need for time-consuming searches for invalid mapping configurations that cause memory overflow when using reinforcement learning algorithms, and allows for faster and more stable acquisition of the optimal mapping configuration of the target space accelerator, improving the generation efficiency of the optimal mapping configuration.
[0024] Further, the construction of shared memory overhead and register memory overhead based on the operator common parameters, the head dimension, and the key sequence length includes: constructing a first operator shared memory overhead and a second operator shared memory overhead based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator common parameters; constructing a first operator register memory overhead based on the key sequence length and the batch size, number of heads, and query sequence length in the operator common parameters; constructing a second operator register memory overhead based on the head dimension and the batch size, number of heads, and query sequence length in the operator common parameters; constructing shared memory overhead based on the first operator shared memory overhead and the second operator shared memory overhead; and constructing register memory overhead based on the first operator register memory overhead and the second operator register memory overhead.
[0025] In one embodiment, the space accelerator has three memory levels: global memory, shared memory, and registers. After operator-to-operator partitioning in the multi-head attention task, the input tensors and weight tensors of the operators are stored in the shared memory, while the output tensors are stored in registers. Therefore, this embodiment utilizes the query sequence length corresponding to the input tensor Q. The length of the key sequence corresponding to the weight tensor K The feature dimension of the tensor in each attention head, i.e., the head dimension D, the number of samples processed simultaneously, i.e., the batch size B, and the number of heads working in parallel in the attention mechanism H, are used to construct the first operator register memory overhead corresponding to operator BMM1 in the multi-head attention task. ,in, , , B1 is the partitioning factor corresponding to batch size B, B0 is the block size obtained by partitioning batch size B under partitioning factor B1, H1 is the partitioning factor corresponding to number of heads H, and H0 is the block size obtained by partitioning number of heads H under partitioning factor H1. 1 represents the length of the query sequence. The corresponding block factor, 0 represents the length of the query sequence. In block factor The block size obtained by the first block division; and the memory overhead of the second operator register corresponding to the construction operator BMM2 in this embodiment is Furthermore, in this embodiment, the head dimension D is not divided into blocks; therefore, the block factor corresponding to the head dimension D is... And the corresponding block size .
[0026] Further, the step of constructing the first operator shared memory overhead and the second operator shared memory overhead based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator common parameters includes: constructing the first operator input tensor memory overhead based on the head dimension and the batch size, number of heads, and query sequence length in the operator common parameters; constructing the first operator weight tensor memory overhead based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters; constructing the second operator input tensor memory overhead based on the key sequence length and the batch size, number of heads, and query sequence length in the operator common parameters; constructing the second operator weight tensor memory overhead based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters; constructing the first operator shared memory overhead based on the first operator input tensor memory overhead and the first operator weight tensor memory overhead; and constructing the second operator shared memory overhead based on the second operator input tensor memory overhead and the second operator weight tensor memory overhead.
[0027] The register memory overhead and shared memory overhead constructed in the above embodiments can accurately characterize the actual storage occupancy of tensors of different operators in the space accelerator at each layer, so as to ensure that the estimation of space accelerator hardware resources is more accurate in the subsequent search for mapping configurations, thereby enabling more accurate and efficient screening of feasible configurations that meet register capacity constraints, and thus improving the generation efficiency of the optimal mapping configuration.
[0028] In the above embodiments, the memory overhead for constructing the first operator input tensor corresponding to operator BMM1 in the multi-head attention task is: In the multi-head attention task, the memory overhead of the first operator weight tensor corresponding to operator BMM1 is... In multi-head attention tasks, the memory overhead of the second operator input tensor corresponding to operator BMM2 is... In multi-head attention tasks, the memory overhead of the second operator weight tensor corresponding to operator BMM2 is... .
[0029] In the above embodiments, based on shared memory overhead and total shared memory capacity Constructing shared memory overhead constraints is and That is Based on register memory overhead and total register capacity The memory overhead constraint for constructing registers is: and .
[0030] Further, in the step of obtaining an initial mapping space based on the block factor range, and performing a search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator, the initial mapping space includes the block factor range and a preset parallel factor range. Obtaining the parallel factor range includes: obtaining the number of computing cores of the target space accelerator and constructing parallel computing overhead constraints based on the number of computing cores; obtaining the parallel factors corresponding to the operator common parameters, the head dimension, and the key sequence length; and obtaining the parallel factor range based on the parallel factors and the parallel computing overhead constraints.
[0031] It should be noted that the parallel factor refers to the parameter used to control the scale of parallel execution of multi-head attention computing tasks on space accelerators.
[0032] In the above embodiments, the parallelism factor among operators in the multi-head attention task is affected by the number of computational cores of the space accelerator. Constraints. In the above embodiments, based on the number of computing cores. Construct parallel computing overhead constraints as ,in, For the common parameters of the operators, query sequence length Corresponding parallel factor The parallelism factor corresponding to batch size B and To obtain the parallelism factor corresponding to the number of heads H, The length of the key sequence Corresponding parallel factor and Let D be the parallel factor corresponding to the head dimension, and based on the parallel factor and the parallel computing overhead constraint, the range of the parallel factor is obtained.
[0033] Further, in the step of obtaining an initial mapping space based on the block factor range, and performing a search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator, the step of performing a search under the preset reinforcement learning algorithm based on the initial mapping space to obtain the optimal mapping configuration for the target space accelerator includes: performing a factor decision step based on each target factor range in the initial mapping space to obtain the optimal mapping configuration for the target space accelerator; the factor decision step includes: performing a search under the preset reinforcement learning algorithm based on the current target factor range to obtain a definite value corresponding to the current target factor range, and marking the current target factor range as decided. Based on the determined value, under the preset factor range update algorithm, the second factor range corresponding to the target factor range not marked as decided in the initial mapping space is obtained; the initial mapping space is updated according to the determined value corresponding to the target factor range marked as decided in the initial mapping space and the second factor range; if it is confirmed that there is a target factor range not marked as decided in the initial mapping space, the factor decision step is re-executed based on a target factor range not marked as decided in the initial mapping space; otherwise, the factor decision step is stopped, and the optimal mapping configuration of the target space accelerator is generated according to the determined value corresponding to each target factor range in the initial mapping space.
[0034] Further, obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the preset factor range update algorithm includes: confirming that the current target factor is a block factor, then obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the shared memory overhead constraint and the register memory overhead constraint; otherwise, obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value under the parallel computing overhead constraint.
[0035] In the above embodiments, after determining the determined value of the target factor corresponding to each target factor range, the remaining unmarked target factor ranges are immediately updated to a second factor range using the determined value. This second factor range eliminates the invalid parts of the target factor range that conflict with the determined value, so the second factor range is smaller than the original target factor range. Subsequently, the initial mapping space is updated based on the second factor range and the determined values of the target factors corresponding to the target factor ranges marked as decided. Therefore, after each determined value is obtained, the remaining unmarked target factor ranges will gradually shrink. Compared with the traditional scheme that searches in an unpruned and extremely large mapping space, this embodiment uses a preset reinforcement learning algorithm to search for determined values in this continuously shrinking target factor range, which can greatly reduce the time required to generate the optimal mapping configuration and improve the generation efficiency of the optimal mapping configuration.
[0036] Further, after obtaining the initial mapping space based on the block factor range, and performing a search based on the initial mapping space under a preset reinforcement learning algorithm to generate the optimal mapping configuration for the target space accelerator, the method further includes: obtaining the computation latency, memory access latency, and memory access power consumption of the target space accelerator after executing the multi-head attention task based on the optimal mapping configuration; generating a reward value based on the computation latency, the memory access latency, and the memory access power consumption; and optimizing the reinforcement learning algorithm based on the reward value under a preset optimization algorithm.
[0037] In the above embodiments, the computation latency, memory access latency, and memory access power consumption of the optimal mapping configuration are obtained and applied to the space accelerator. The reinforcement learning algorithm is then optimized based on the reward value generated from the computation latency, memory access latency, and memory access power consumption, so that the computation latency, memory access latency, and memory access power consumption of the next optimal mapping configuration are even lower, thereby improving the operating efficiency of the space accelerator.
[0038] It should be noted that the total latency in multi-head attention tasks is the sum of the latencies of each operator, where the latency of each operator includes computation latency and memory access latency.
[0039] In one embodiment, the preset optimization algorithm is the REINFORCE algorithm and the Adam optimizer. For each operator, the computation delay is... ,in The number of operations performed by the operator. This represents the number of computing cores used. Memory access latency, on the other hand, represents the maximum memory access latency across all memory levels. , among which, the The memory access latency at the memory level is The number of memory accesses is Memory bandwidth is Once the computation latency and memory access latency are obtained, the latency of each operator can be calculated. Total latency of multi-head attention mechanism ,in This refers to all operators in the multi-head attention mechanism. Furthermore, in this embodiment, the memory access power consumption is... This is the weighted sum of memory accesses at each memory level in the space accelerator.
[0040] Further, the step of substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the shared memory overhead constraint and the register memory overhead constraint to obtain the maximum block size includes: substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the shared memory overhead constraint to obtain the first maximum block size; substituting the block factors corresponding to task parameters other than the current task parameters in the operator's shared parameters into the register memory overhead constraint to obtain the second maximum block size; confirming that the first maximum block size is less than the second maximum block size, then using the first maximum block size as the maximum block size; otherwise, using the second maximum block size as the maximum block size.
[0041] Further, generating the block factor range corresponding to the current task parameters based on the maximum block size and the current task parameters includes: calculating the minimum block factor based on the ratio of the current task parameters and the maximum block size; calculating the maximum block factor based on the ratio of the current task parameters under the preset minimum block size; and generating the block factor range corresponding to the current task parameters based on the minimum block factor and the maximum block factor.
[0042] In one embodiment, the minimum preset partitioning factor is set to 1. When the current task parameter is the number of heads H, other task parameters, i.e., the query sequence length, are... The block factor corresponding to the batch size B is set to 1. Therefore, under the shared memory overhead constraint, according to Get the maximum size of the first block ,in accordance with Obtain the maximum size of the second block. The smaller of the two maximum block sizes is taken as the maximum block size for the current task parameter header number H, and then based on... This gives the range of the partitioning factor. This is an example, not a limitation. If the number of heads H=16, the maximum corresponding partitioning size is... Then the range of the partitioning factor should be able to accommodate the partition size. ≤8, meaning the block factor range is [2, 4, 8, 16].
[0043] Example, not limitation, mapping the BERT model for the multi-head attention task onto the NVDLA accelerator to obtain the query sequence length. Batch size 1. Number of heads Key sequence length and head dimension Total shared memory capacity of NVDLA accelerator Total register capacity Calculate the number of cores One. Key sequence length and head dimension No blocks will be used, meaning it will always remain , number of heads and query sequence length block factor and Maintain consistency across all operators. When performing inter-operator block partitioning, allocate blocks for the first operator BMM1 and the second operator BMM2 respectively. The shared memory capacity. For BMM 1, the shared memory overhead of the first operator. First operator register memory overhead For BMM 2, the second operator has shared memory overhead. Second operator register memory overhead When performing intra-operator parallelism, the required number of computational cores... .
[0044] Subsequently, the maximum block size corresponding to each block factor of each operator is calculated; for the first operator BMM 1, the number of heads... The corresponding maximum block size Query sequence length The corresponding maximum block size For the second operator BMM 2, the number of heads The corresponding maximum block size Query sequence length The corresponding maximum block size And calculate the minimum block factor. For BMM 1, the number of heads The corresponding minimum block factor and query sequence length The corresponding minimum block factor For BMM 2, head count The corresponding minimum block factor and query sequence length The corresponding minimum block factor This example also yields the maximum parallel factor for each parallel factor. , , , Based on the minimum block factor and the maximum parallel factor mentioned above, invalid block factor values that could lead to memory overflow can be removed, ensuring that the target factor corresponding to each target factor range in the initial mapping space will not cause memory overflow.
[0045] Please see Figure 2 This embodiment also provides a multi-head attention task accelerator mapping configuration generation system to implement any of the multi-head attention task accelerator mapping configuration generation methods described above. The system includes: a task parameter acquisition module for acquiring the operator common parameters, head dimension, and key sequence length of the multi-head attention task to be accelerated; an accelerator parameter acquisition module for acquiring the total shared memory capacity and total register capacity of the target space accelerator; a memory overhead construction module for constructing shared memory overhead and register memory overhead based on the operator common parameters, head dimension, and key sequence length; a shared memory overhead constraint construction module for constructing shared memory overhead constraints based on the shared memory overhead and the total shared memory capacity; and a register overhead constraint construction module for constructing register memory overhead constraints based on the register memory overhead and the total register capacity. The system comprises: a block factor range acquisition module, used to set the block factor corresponding to each task parameter in the operator common parameters to a preset minimum value of the block factor, and to perform a constraint range step on each task parameter to obtain the block factor range corresponding to each task parameter; and an optimal mapping configuration generation module, used to obtain an initial mapping space based on the block factor range, and to search under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator; wherein, the constraint range step includes: substituting the block factors corresponding to task parameters other than the current task parameter in the operator common parameters into the shared memory overhead constraint and the register memory overhead constraint for solving to obtain the maximum block size; and generating the block factor range corresponding to the current task parameter based on the maximum block size and the current task parameter.
[0046] In the above embodiments, the range of block factors obtained by solving the shared memory overhead constraints and register memory overhead constraints can eliminate invalid block factor values that would cause memory overflow, ensuring that each candidate mapping configuration factor in the initial mapping space will not cause memory overflow. This eliminates the need for time-consuming searches for invalid mapping configurations that would cause memory overflow when using reinforcement learning algorithms to search, and enables the optimal mapping configuration of the target space accelerator to be obtained quickly and stably, thereby improving the generation efficiency of the optimal mapping configuration.
[0047] The above description represents the preferred embodiments of the present invention. It should be noted that those skilled in the art can make various improvements and modifications without departing from the principles of the present invention, and these improvements and modifications are also considered to be within the scope of protection of the present invention.
Claims
1. A method for generating mapping configurations for a multi-head attention task accelerator, characterized in that, Includes the following steps: Obtain the operator-shared parameters, head dimension, and key sequence length for the multi-head attention task to be accelerated; Obtain the total shared memory capacity and total register capacity of the target space accelerator; Based on the operator common parameters, the head dimension, and the key sequence length, the shared memory overhead and register memory overhead are constructed. Based on the shared memory overhead and the total shared memory capacity, a shared memory overhead constraint is constructed; Based on the register memory overhead and the total register capacity, a register memory overhead constraint is constructed; Set the block factor corresponding to each task parameter in the common parameters of the operators to the preset minimum value of the block factor, and perform a constraint range step on each task parameter to obtain the range of the block factor corresponding to each task parameter. Based on the range of the block factor, an initial mapping space is obtained, and based on the initial mapping space, a search is performed under a preset reinforcement learning algorithm to generate the optimal mapping configuration for the target space accelerator. The constraint range step includes: Substitute the block factor corresponding to the task parameters other than the current task parameters in the common parameters of the operators into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size; Based on the maximum block size and the current task parameters, generate the block factor range corresponding to the current task parameters.
2. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 1, characterized in that, The construction of shared memory overhead and register memory overhead based on the operator common parameters, the head dimension, and the key sequence length includes: Based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator common parameters, the first operator shared memory overhead and the second operator shared memory overhead are constructed. Based on the key sequence length and the batch size, number of heads, and query sequence length in the common parameters of the operator, the memory overhead of the first operator register is constructed. The second operator register memory overhead is constructed based on the head dimension and the batch size, number of heads, and query sequence length in the operator common parameters; Based on the shared memory overhead of the first operator and the shared memory overhead of the second operator, construct the shared memory overhead; The register memory overhead is constructed based on the memory overhead of the first operator register and the memory overhead of the second operator register.
3. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 2, characterized in that, The construction of the first operator shared memory overhead and the second operator shared memory overhead based on the head dimension, the key sequence length, and the batch size, number of heads, and query sequence length in the operator shared parameters includes: Based on the head dimension and the batch size, number of heads and query sequence length in the common parameters of the operator, the memory overhead of the first operator input tensor is constructed. Based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters, construct the memory overhead of the first operator weight tensor; Based on the key sequence length and the batch size, number of heads, and query sequence length in the common parameters of the operator, construct the memory overhead of the second operator input tensor; Based on the key sequence length, the head dimension, and the batch size and number of heads in the operator common parameters, construct the memory overhead of the second operator weight tensor; Based on the memory overhead of the first operator input tensor and the memory overhead of the first operator weight tensor, construct the shared memory overhead of the first operator; Based on the memory overhead of the second operator input tensor and the memory overhead of the second operator weight tensor, the shared memory overhead of the second operator is constructed.
4. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 1, characterized in that, In the step of obtaining an initial mapping space based on the block factor range, and searching under a preset reinforcement learning algorithm based on the initial mapping space to generate the optimal mapping configuration for the target space accelerator, the initial mapping space includes the block factor range and a preset parallel factor range. The acquisition of the parallel factor range includes: Obtain the number of computing cores in the target space accelerator, and construct parallel computing overhead constraints based on the number of computing cores; Obtain the common parameters of the operators, the parallel factor corresponding to the head dimension and the key sequence length; Based on the parallel factor and the parallel computing overhead constraint, the range of the parallel factor is obtained.
5. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 4, characterized in that, In the step of obtaining an initial mapping space based on the block factor range, and then searching based on the initial mapping space under a preset reinforcement learning algorithm to generate the optimal mapping configuration for the target space accelerator, the step of searching based on the initial mapping space under a preset reinforcement learning algorithm to obtain the optimal mapping configuration for the target space accelerator includes: Based on each target factor range in the initial mapping space, a factor decision step is performed to obtain the optimal mapping configuration of the target space accelerator. The factor decision-making steps include: Based on the current target factor range, a search is performed under the preset reinforcement learning algorithm to obtain the definite value corresponding to the current target factor range, and the current target factor range is marked as decided. Based on the determined value, under the preset factor range update algorithm, the second factor range corresponding to the target factor range that is not marked as decided in the initial mapping space is obtained; The initial mapping space is updated based on the determined values corresponding to the target factor ranges marked as decided in the initial mapping space and the second factor range; If it is confirmed that there is an unmarked target factor range in the initial mapping space, the factor decision-making step is re-executed based on an unmarked target factor range in the initial mapping space; otherwise, the factor decision-making step is stopped, and the optimal mapping configuration of the target space accelerator is generated based on the determined value corresponding to each target factor range in the initial mapping space.
6. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 5, characterized in that, The step of obtaining the second factor range corresponding to the target factor range not marked as decided in the initial mapping space based on the determined value and under the preset factor range update algorithm includes: If the current target factor is confirmed to be a block factor, then based on the determined value under the shared memory overhead constraint and the register memory overhead constraint, the second factor range corresponding to the target factor range not marked as decided in the initial mapping space is obtained; otherwise, based on the determined value under the parallel computing overhead constraint, the second factor range corresponding to the target factor range not marked as decided in the initial mapping space is obtained.
7. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 1, characterized in that, After obtaining the initial mapping space based on the block factor range, and generating the optimal mapping configuration for the target space accelerator by searching under a preset reinforcement learning algorithm based on the initial mapping space, the method further includes: Obtain the computation latency, memory access latency, and memory access power consumption of the target space accelerator after executing the multi-head attention task according to the optimal mapping configuration; A reward value is generated based on the computation latency, the memory access latency, and the memory access power consumption. The reinforcement learning algorithm is optimized based on the reward value and a preset optimization algorithm.
8. The method for generating a mapping configuration for a multi-head attention task accelerator according to claim 1, characterized in that, The step of substituting the block factor corresponding to the task parameters other than the current task parameters in the shared parameters of the operators into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size includes: Substitute the block factor corresponding to the task parameters other than the current task parameters in the shared parameters of the operators into the shared memory overhead constraint to solve for the maximum value of the first block size; Substitute the block factor corresponding to the task parameters other than the current task parameters in the common parameters of the operators into the register memory overhead constraint to solve for the maximum value of the second block size; If the maximum value of the first block size is confirmed to be less than the maximum value of the second block size, then the maximum value of the first block size is taken as the maximum block size; otherwise, the maximum value of the second block size is taken as the maximum block size.
9. A method for generating a mapping configuration for a multi-head attention task accelerator according to claim 8, characterized in that, The step of generating the block factor range corresponding to the current task parameters based on the maximum block size and the current task parameters includes: The minimum value of the block factor is calculated based on the ratio of the current task parameters and the maximum value of the block size; The maximum value of the block factor is obtained by calculating the ratio based on the current task parameters at the minimum preset block size; Based on the minimum and maximum values of the partitioning factors, the range of partitioning factors corresponding to the current task parameters is generated.
10. A multi-head attention task accelerator mapping configuration generation system, characterized in that, A method for generating a multi-head attention task accelerator mapping configuration as described in any one of claims 1 to 9, comprising: The task parameter acquisition module is used to obtain the operator-shared parameters, head dimension, and key sequence length of the multi-head attention task to be accelerated; The accelerator parameter acquisition module is used to obtain the total shared memory capacity and total register capacity of the target space accelerator. A memory overhead construction module is used to construct shared memory overhead and register memory overhead based on the operator common parameters, the head dimension, and the key sequence length; A shared memory overhead constraint construction module is used to construct shared memory overhead constraints based on the shared memory overhead and the total shared memory capacity; The register overhead constraint construction module is used to construct register memory overhead constraints based on the register memory overhead and the total register capacity; The block factor range acquisition module is used to set the block factor corresponding to each task parameter in the common parameters of the operators to a preset minimum value of the block factor, and to perform a constraint range step on each task parameter to obtain the block factor range corresponding to each task parameter. The optimal mapping configuration generation module is used to obtain an initial mapping space based on the block factor range, and to search for the optimal mapping configuration of the target space accelerator based on the initial mapping space under a preset reinforcement learning algorithm. The constraint range step includes: Substitute the block factor corresponding to the task parameters other than the current task parameters in the common parameters of the operators into the shared memory overhead constraint and the register memory overhead constraint to solve for the maximum block size; Based on the maximum block size and the current task parameters, generate the block factor range corresponding to the current task parameters.