NPU layout method and device oriented to Attention operator generalization

Abstract

The invention discloses an NPU (Network Processing Unit) layout method and device for Attention operator generalization, which are applied to a compiler of a large model in an NPU hardware system, and the method comprises the following steps: determining whether the numhead of a Query matrix in an Attention operator is equal to the numkeyvalue head of a Key-Value matrix; if yes, a preset first mapping scheme is adopted to generate a corresponding first layout scheme set; if not, adopting a preset second mapping scheme to generate a corresponding second layout scheme set; obtaining all possible NPU layout schemes based on the first and / or second layout scheme set, and then selecting an optimal NPU layout scheme according to an optimal principle that the overall parameter filling amount is minimum; the method has the beneficial effects that the execution efficiency of the Attention operator can be improved to the greatest extent, and the method has good generalization.

Description

Technical Field

[0001] The present invention relates to the technical field of hardware accelerators, and in particular to an NPU layout method and device for generalization of an Attention operator. Background Art

[0002] The Neural Processing Unit (NPU) is a hardware accelerator specifically designed to accelerate artificial intelligence (AI) and deep learning tasks.

[0003] In NPU hardware systems, a global memory architecture is often used. However, this memory solution cannot fully utilize the performance advantages of the NPU architecture, leading to the emergence of distributed memory solutions.

[0004] Meanwhile, in traditional CPU or GPU hardware systems, the input shape of the Query matrix is (batch_size, num_heads, M, head_dim), while the input shape of the Key and Value matrices is (batch_size, num_key_value_heads, M, head_dim). In the NPU hardware architecture, the execution of the Attention operator needs to be distributed across multiple computing units (tiles), and the minimum data parallel granularity requirements of each tile must be met.

[0005] However, in NPU hardware systems based on distributed memory, implementing the Attention operator more efficiently and flexibly is a significant technical challenge for large-scale model inference. This is primarily due to the difference in NPU hardware data layout compared to traditional CPUs or GPUs. Furthermore, the Attention operator includes multiple variants, such as Multi-Head Attention (MHA), Group-Query Attention (GQA), and Multi-Query Attention (MQA).

[0006] Therefore, there is an urgent need for a generalizable and efficient NPU data layout method to better adapt to various types of attention operators. Summary of the Invention

[0007] In view of the technical deficiencies of the prior art, the embodiment of the present invention aims to provide an NPU layout method and device for the generalization of the Attention operator.

[0008] To achieve the above objectives, in a first aspect, an embodiment of the present invention provides an NPU layout method for generalizing an Attention operator, which is applied to a compiler for a large model in an NPU hardware system, wherein the NPU hardware includes computing units with multiple parallel cores. The method includes:

[0009] Determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal;

[0010] If they are equal, a preset first mapping scheme is used to generate a corresponding first layout scheme set;

[0011] If they are not equal, a preset second mapping scheme is used to generate a corresponding second layout scheme set;

[0012] Based on the first and / or second layout scheme sets, all possible NPU layout schemes are obtained, and then the optimal NPU layout scheme is selected according to the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

[0013] As a preferred implementation of the present application, the first mapping scheme includes a preset normal scheme, a Heads_to_dim scheme, and a Dim_to_heads scheme;

[0014] The normal solution is applicable when the number of heads of Query, Key and Value exactly matches the specified tile_num;

[0015] The Heads_to_dim solution is applicable when num_heads is greater than the specified tile_num. By splitting the num_heads dimension, the num_heads of the query matrix are decomposed so that one tile can contain data of multiple heads at the same time to meet the mapping requirements.

[0016] The Dim_to_heads function is applicable when num_heads is less than the specified tile_num. By splitting the K' dimension of the Query matrix, the complete head_dim data of a single head is distributed to multiple tiles, thereby completing the mapping from num_heads to tile_num. Here, K'=head_dim / / K16.

[0017] As a specific implementation of the present application, the second mapping scheme includes using the preset Heads_to_dim, dim_to_Heads and Pad_per_group_heads schemes for the Query matrix; and using the preset Heads_to_dim, Dim_to_heads and Repeat_kv_heads schemes for the Key-Value matrix;

[0018] The Pad_per_group_heads scheme is used to increase the number of heads in a group by padding a certain number of heads in the group of the Query matrix, thereby achieving the mapping from num_heads to tile_num;

[0019] The Repeat_kv_heads scheme is used to distribute the Key-Value matrix to more tiles by repeating the number of heads of each Key-Value matrix, thereby ensuring that more tiles contain valid data.

[0020] As a specific implementation of the present application, the Pad_per_group_heads solution specifically includes:

[0021] Used when num_heads is not equal to num_key_value_heads, by dividing the heads of the Query matrix into num_key_value_heads groups, that is, each group has num_key_value_groups heads; where num_key_value_groups = num_heads / / num_key_value_heads; this method mainly fills a certain number of heads in the Query matrix group, thereby changing num_key_value_groups; and updates the number of heads of the Query matrix by num_heads = group_factor * num_key_value_heads, where group_factor is the group coefficient representing the number of heads in each group in the Query matrix.

[0022] As a specific implementation method of the present application, in the Repeat_kv_heads scheme, for the Key-Value matrix, the number of heads of the Key-Value matrix is updated by num_key_value_heads = rep_factor * num_key_value_heads; wherein rep_factor is the repeat coefficient, which indicates the number of times each head of the Key-Value matrix is replicated.

[0023] As a preferred implementation of the present application, the method further includes:

[0024] When both the Head_to_dim solution and the Dim_to_heads solution fail to complete the mapping, the padding solution is enabled to ensure that the required tile_num is finally reached by padding the data.

[0025] In a second aspect, an embodiment of the present invention further provides an NPU layout device for generalization of an Attention operator, which is applied to a compiler of a large model in an NPU hardware system, wherein the NPU hardware includes a computing unit with multiple parallel cores, and the device includes:

[0026] The judgment module is used to determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal;

[0027] Mapping module for:

[0028] If they are equal, a preset first mapping scheme is used to generate a corresponding first layout scheme set;

[0029] If they are not equal, a preset second mapping scheme is used to generate a corresponding second layout scheme set;

[0030] A processing module is used to obtain all possible NPU layout schemes based on the first and / or second layout scheme sets, and then select the optimal NPU layout scheme based on the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

[0031] The technical solution provided by the embodiment of the present invention uses the dimensions of the heads of the Query, Key, and Value matrices as the mapping dimensions of the tiles, and uses multiple schemes to achieve efficient mapping of the num_heads of the Query and the num_key_value_heads of the Key-Value matrix to any specified tile_num. These methods, combined with the characteristics of various components of the NPU hardware, can maximize the execution efficiency of the Attention operator; and fully consider that when the Attention operator is executed on the NPU, the diverse parameter combinations of the Attention operator are mainly determined by num_heads and num_key_value_heads, so that Attention operators of different parameter scales can run efficiently, thereby having good generalization and improving the versatility of the Attention operator in the NPU hardware system; it also solves the problem of flexible deployment of the existing Attention operator of large models with arbitrary specifications and parameters in the NPU hardware system based on the distributed memory solution. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] In order to more clearly illustrate the specific implementation of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for the specific implementation or the description of the prior art.

[0033] Figure 1 This is a flowchart of an NPU layout method for generalization of the Attention operator provided by an embodiment of the present invention;

[0034] Figure 2 This is a structural block diagram of another NPU layout device for the generalization of the Attention operator provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0035] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of them. All other embodiments obtained by ordinary technicians in this field based on the embodiments of the present invention without making any creative efforts shall fall within the scope of protection of the present invention.

[0036] It will be understood that when used in this specification and the appended claims, the terms “comprises” and “comprising” indicate the presence of described features, integers, steps, operations, elements and / or components, but do not preclude the presence or addition of one or more other features, integers, steps, operations, elements, components and / or groups thereof.

[0037] Throughout this specification, references to "one embodiment," "an embodiment," "an example," or "an example" mean that a particular feature, structure, or characteristic described in connection with the embodiment or example is included in at least one embodiment of the present invention. Thus, the appearances of the phrases "in one embodiment," "in an embodiment," "an example," or "an example" in various places throughout this specification are not necessarily all referring to the same embodiment or example. Furthermore, the particular features, structures, or characteristics may be combined in any suitable combinations and / or subcombinations in one or more embodiments or examples.

[0038] Attention operator: In this paper, it specifically refers to the operator that processes the attention mechanism in Transformer.

[0039] MHA: Multi-HeadAttention, a multi-head attention mechanism, is an important attention mechanism in deep learning.

[0040] GQA: Group-Query Attention, a group attention mechanism, is an improvement and optimization of MHA.

[0041] MQA: Multi-Query Attention, a multi-query attention mechanism, is also an improvement and optimization of MHA.

[0042] Tile: The smallest parallel computing unit in the NPU hardware system.

[0043] Query: In this invention, it specifically refers to the Query matrix in the Attention operator.

[0044] Key: In this invention, it specifically refers to the Key matrix in the Attention operator.

[0045] Value: In this invention, it specifically refers to the Value matrix in the Attention operator.

[0046] Num_heads: In this invention, it refers to the number of heads of the Query matrix in the Attention operator.

[0047] Num_key_value_heads: In this invention, it refers to the number of heads of the Key and Value matrices in the Attention operator.

[0048] Head_dim: In this invention, it refers to the data length of each head of the Query, Key, and Value matrices in the Attention operator.

[0049] M: In this invention, it refers to the length of the input token when the large model is inferred. Token: A unit in the text, which can be a word, a punctuation mark, a number

[0050] K16: In the present invention, it means that the data is 16-aligned.

[0051] Padding: refers to filling.

[0052] It should be noted that, unless otherwise specified, the technical terms in this embodiment have the common meanings understood in the relevant technical field.

[0053] The inventive concept of this invention is that the NPU accelerates operator execution through parallel computing. It typically contains a certain number of parallel cores, such as 32 or 64 tiles. In actual task scenarios, it may be necessary to fully utilize all cores or only use a specified number of cores based on demand. This efficient and flexible resource management approach significantly improves the operating efficiency of the Attention operator and is the key to breaking through the operator execution bottleneck.

[0054] The parameters of the Attention operator are mainly determined by the number of heads in the query matrix (num_heads) and the number of heads in the key-value matrix (num_key_value_heads), and large models of different scales may have diverse parameter combinations. When executing the Attention operator on the NPU, it is necessary to fully consider these complex and arbitrarily changing parameter configurations. In response to this, the method of the present invention optimizes the data layout so that Attention operators of different parameter scales can run efficiently, thereby having good adaptability; through this design, the present invention can not only solve the efficiency and flexibility problems of executing the Attention operator in the existing NPU hardware, but also support multiple variants and parameter combinations, realizing an innovative data layout solution that is both practical and scalable.

[0055] Please refer to Figure 1 An embodiment of the present invention provides an NPU layout method for generalization of the Attention operator, which is applied to a compiler of a large model in an NPU hardware system, wherein the NPU hardware includes a computing unit with multiple parallel cores. The method includes:

[0056] S101, determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal;

[0057] S102, if they are equal, using a preset first mapping scheme to generate a corresponding first layout scheme set;

[0058] S103, if they are not equal, using a preset second mapping scheme to generate a corresponding second layout scheme set;

[0059] S104, based on the first and / or second layout scheme sets, all possible NPU layout schemes are obtained, and then the optimal NPU layout scheme is selected according to the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

[0060] For ease of description, in a specific embodiment, the calculation unit is represented by a tile, and / / is an arithmetic operator representing integer division.

[0061] It should be noted that under the NPU hardware architecture, the execution of the Attention operator needs to be distributed to multiple computing units (tiles) and needs to meet the minimum data parallel granularity requirements of each tile, such as 16 or 32. The present invention sets the minimum parallelism to 16 (hereinafter referred to as K16), that is, the minimum data block size that needs to be allocated to each tile is K16 data. In order to adapt to this parallel execution, the original input data needs to be rearranged in a K16-aligned manner. When the data is distributed to tile_num tiles, the overall data must ensure the tile_num*K16 alignment requirements.

[0062] Specifically, the original shape of the query matrix (num_heads, M, head_dim) must be converted to the NPU-adapted shape (num_heads, M, head_dim / / K16, K16). Similarly, the original shapes of the key and value matrices (num_key_value_heads, M, head_dim) must be converted to the adapted shape (num_key_value_heads, M, head_dim / / K16, K16). Here, head_dim / / K16 is represented by K'.

[0063] In this embodiment, the first mapping scheme includes a preset normal scheme, a Heads_to_dim scheme, and a Dim_to_heads scheme;

[0064] The normal solution is suitable for use when the number of heads of Query, Key, and Value completely matches the specified tile_num, so that no processing is required.

[0065] The Heads_to_dim solution is applicable when num_heads is greater than the specified tile_num. By splitting the num_heads dimension, the num_heads of the query matrix are decomposed so that one tile can contain data of multiple heads at the same time to meet the mapping requirements.

[0066] Specifically, taking the Query matrix as an example, the Key-Vaule matrix is also applicable. First, convert the Query matrix shape to a similar NPU layout (K16 alignment) (num_heads, head_dim / / K16, M, K16); where K' = head_dim / / K16;

[0067] Because num_heads is greater than tile_num, num_heads needs to be split. Then num_heads can be expressed by the following expression: num_heads = (num_heads / / h2d_factor)*h2d_factor, where h2d_factor represents the split coefficient of num_heads in the case of heads_to_dim, and num_heads / / h2d_factor represents the value assigned to the tile_num dimension under the current split coefficient;

[0068] ① If num_heads / / h2d_factor is not equal to tile_num, then you need to use the Padding solution to fill it to tile_num; assuming the number of heads to be filled is padding_num;

[0069] That is, num_heads / / h2d_factor+padding_num=tile_num. In this case, the NPU layout scheme is (num_heads / / h2d_factor+padding_num,h2d_factor,K',M,K16), and the split coefficient h2d_factor is merged into the K' dimension. In this case, the final NPU layout of the Query and Key-Value matrix is (num_heads / / h2d_factor+padding_num,h2d_factor*K',M,K16), and the total number of parameters that need to be filled is: padding_num*h2d_factor*K'*K16.

[0070] ② If num_heads / / h2d_factor equals tile_num, no padding is required, satisfying the principle of minimizing overall parameter padding. Therefore, the optimal split coefficient for the heads_to_dim solution is h2d_factor = num_heads / / tile_num. With this optimal split coefficient, the NPU layout becomes (tile_num, h2d_factor, K', M, K16). The h2d_factor split coefficient is merged into the K' dimension, resulting in the final NPU layout for the query and key-value matrices being (tile_num, h2d_factor * K', M, K16). The total number of parameters required for padding is 0.

[0071] The Dim_to_heads function is applicable when num_heads is less than the specified tile_num. By splitting the K' dimension of the Query matrix, the complete head_dim data of a single head is distributed to multiple tiles, thereby completing the mapping from num_heads to tile_num. Here, K'=head_dim / / K16.

[0072] Specifically, taking the Query matrix as an example, the Key-Vaule matrix is also applicable. First, convert the Query matrix shape to a K16-aligned (num_heads, K', M, K16) similar to the NPU layout, where K' = head_dim / / K16. Because num_heads is less than tile_num, K' needs to be split. K' can be expressed as follows: K' = (K' / / d2h_factor)*d2h_factor, where d2h_factor represents the K' split coefficient under the dim_to_heads scheme, and num_heads*(K' / / d2h_factor) represents the value assigned to the tile_num dimension under the current split coefficient.

[0073] ① If num_heads*(K' / / d2h_factor) is not equal to tile_num, the Padding scheme needs to be used to fill it to tile_num; assuming that the number of heads to be filled is padding_num, then tile_num = num_heads*(K' / / d2h_factor)+padding_num, then in this case the final NPU layout scheme is (num_heads*(K' / / d2h_factor)+padding_num,d2h_factor,M,K16), and the overall parameter amount that needs to be filled is: padding_num*d2h_factor*M*K16;

[0074] ② If num_heads*(K' / / d2h_factor) is equal to tile_num, no padding is required; if the principle of minimum overall parameter padding is met, the final NPU layout solution is (num_heads*(K' / / d2h_factor), d2h_factor, M, K16), where tile_num = num_heads*(K' / / d2h_factor), and the overall parameter amount that needs to be filled is: 0.

[0075] The second mapping scheme includes using the preset Heads_to_dim, dim_to_Heads and Pad_per_group_heads scheme for the Query matrix; and using the preset Heads_to_dim, Dim_to_heads and Repeat_kv_heads scheme for the Key-Value matrix.

[0076] The Pad_per_group_heads scheme is used to increase the number of heads in a group by padding a certain number of heads in the group of the Query matrix, thereby achieving the mapping from num_heads to tile_num;

[0077] Specifically, for the case where num_heads is not equal to num_key_value_heads, num_key_value_groups = num_heads / / num_key_value_heads. The heads of the Query matrix are divided into num_key_value_heads groups, that is, each group has num_key_value_groups heads. This method mainly pads a certain number of heads in the Query matrix group, thereby changing num_key_value_groups.

[0078] We update the number of heads in the Query matrix using the formula num_heads = group_factor * num_key_value_heads, where group_factor is the group coefficient representing the number of heads in each group in the Query matrix. The initial value is num_key_value_groups, and then it is increased by 1 each time until the maximum limit max_num_heads is reached. After each update of num_heads using group_factor, the heads_to_dim and dim_to_heads schemes are used to update the current Query layout scheme, and the heads_to_dim, dim_to_heads, and repeat_kv_heads schemes are used to update the Key-Value layout scheme. Finally, the layout schemes for the Query and Key-Value matrices for all cases are obtained.

[0079] The Repeat_kv_heads scheme is used to distribute the Key-Value matrix to more tiles by repeating the number of heads of each Key-Value matrix, thereby ensuring that more tiles contain valid data.

[0080] Specifically, for the case where num_heads is not equal to num_key_value_heads, where num_key_value_groups = num_heads / / num_key_value_heads. For the Key-Value matrix, we update the number of heads of the Key-Value matrix by num_key_value_heads = rep_factor * num_key_value_heads, where rep_factor is the repeat coefficient, which represents the number of times each head of the Key-Value matrix is replicated, and is a combination of all factors from the original num_key_value_groups. Each time a new num_key_value_heads is obtained, the Dim_to_heads scheme is applied to the Key-Value matrix to obtain all possible NPU layout schemes.

[0081] By using the dimensions of the query, key, and value matrix heads as the tile mapping dimensions and employing the various methods described above, we achieve efficient mapping of the query's num_heads and the key-value matrix's num_key_value_heads to any specified tile_num. These methods, combined with the characteristics of various NPU hardware components, maximize the execution efficiency of the Attention operator and enhance its versatility within the NPU hardware system. For different query matrix num_heads and key-value matrix num_key_value_heads parameters, the final layout solution adopted may be one of these or a combination of multiple solutions, so the actual layout solution may vary depending on the specific scenario.

[0082] Furthermore, in another embodiment, based on the above technical solution, the method further includes:

[0083] When both the Head_to_dim solution and the Dim_to_heads solution fail to complete the mapping, the padding solution is enabled to ensure that the required tile_num is finally reached by padding the data.

[0084] That is, when the aforementioned methods cannot efficiently map num_heads or num_key_value_heads to the specified tile_num, this solution will be adopted to ensure that the required tile_num is finally reached by padding the data.

[0085] The above solution uses the dimensions of the Query, Key, and Value matrix heads as the tile mapping dimensions, and uses multiple solutions to achieve efficient mapping of the Query's num_heads and the Key-Value matrix's num_key_value_heads to any specified tile_num. These methods, combined with the characteristics of various NPU hardware components, can maximize the execution efficiency of the Attention operator. They also fully consider that when executing the Attention operator on the NPU, the diverse parameter combinations of the Attention operator are mainly determined by num_heads and num_key_value_heads, so that Attention operators of different parameter scales can run efficiently, thereby having good generalization and improving the versatility of the Attention operator in the NPU hardware system. It also solves the problem of flexible deployment of the Attention operator for existing large models of arbitrary specifications and parameters in the NPU hardware system based on the distributed memory solution.

[0086] Based on the same inventive concept, the embodiment of the present invention also provides an NPU layout device for the generalization of the Attention operator, referring to Figure 2 , a compiler for a large model in an NPU hardware system, wherein the NPU hardware includes a computing unit with multiple parallel cores, the apparatus comprising:

[0087] The judgment module is used to determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal;

[0088] Mapping module for:

[0089] If they are equal, a preset first mapping scheme is used to generate a corresponding first layout scheme set;

[0090] If they are not equal, a preset second mapping scheme is used to generate a corresponding second layout scheme set;

[0091] A processing module is used to obtain all possible NPU layout schemes based on the first and / or second layout scheme sets, and then select the optimal NPU layout scheme based on the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

[0092] In this embodiment, the first mapping scheme includes a preset normal scheme, a Heads_to_dim scheme, and a Dim_to_heads scheme;

[0093] The normal solution is applicable when the number of heads of Query, Key and Value exactly matches the specified tile_num;

[0094] The Heads_to_dim solution is applicable when num_heads is greater than the specified tile_num. By splitting the num_heads dimension, the num_heads of the query matrix are decomposed so that one tile can contain data of multiple heads at the same time to meet the mapping requirements.

[0095] The Dim_to_heads function is applicable when num_heads is less than the specified tile_num. By splitting the K' dimension of the Query matrix, the complete head_dim data of a single head is distributed to multiple tiles, thereby completing the mapping from num_heads to tile_num. Here, K'=head_dim / / K16.

[0096] The second mapping scheme includes using the preset Heads_to_dim, dim_to_Heads and Pad_per_group_heads schemes for the Query matrix; and using the preset Heads_to_dim, Dim_to_heads and Repeat_kv_heads schemes for the Key-Value matrix;

[0097] The Pad_per_group_heads scheme is used to increase the number of heads in a group by padding a certain number of heads in the group of the Query matrix, thereby achieving the mapping from num_heads to tile_num;

[0098] The Repeat_kv_heads scheme is used to distribute the Key-Value matrix to more tiles by repeating the number of heads of each Key-Value matrix, thereby ensuring that more tiles contain valid data.

[0099] Furthermore, based on the above technical solution, the processing module is further configured to:

[0100] When both the Head_to_dim solution and the Dim_to_heads solution fail to complete the mapping, the padding solution is enabled to ensure that the required tile_num is finally reached by padding the data.

[0101] It should be noted that for a more specific description of the workflow of the device embodiment, please refer to the aforementioned method embodiment part, which will not be repeated here.

[0102] The entire solution uses the dimensions of the Query, Key, and Value matrix heads as the tile mapping dimensions and uses multiple schemes to achieve efficient mapping of the Query's num_heads and the Key-Value matrix's num_key_value_heads to any specified tile_num. These methods, combined with the characteristics of various NPU hardware components, can maximize the execution efficiency of the Attention operator. They also fully consider that when executing the Attention operator on the NPU, the diverse parameter combinations of the Attention operator are mainly determined by num_heads and num_key_value_heads, allowing Attention operators with different parameter sizes to run efficiently, thereby having good generalization and improving the versatility of the Attention operator in the NPU hardware system.

[0103] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any person skilled in the art can easily conceive of various equivalent modifications or substitutions within the technical scope disclosed in the present invention, and such modifications or substitutions are intended to be within the scope of protection of the present invention. Therefore, the scope of protection of the present invention shall be subject to the scope of protection of the claims.

Claims

1. An NPU layout method for the generalization of the Attention operator, characterized by: A compiler for a large model in an NPU hardware system, wherein the NPU hardware includes a computing unit with multiple parallel cores, the method comprising: Determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal; If they are equal, a preset first mapping scheme is used to generate a corresponding first layout scheme set; If they are not equal, a preset second mapping scheme is used to generate a corresponding second layout scheme set; Based on the first and / or second layout scheme sets, all possible NPU layout schemes are obtained, and then the optimal NPU layout scheme is selected according to the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

2. The NPU layout method for generalization of the Attention operator according to claim 1, characterized in that: The first mapping scheme includes a preset normal scheme, a Heads_to_dim scheme, and a Dim_to_heads scheme; The normal solution is applicable when the number of heads of Query, Key and Value exactly matches the specified tile_num; The Heads_to_dim solution is applicable when num_heads is greater than the specified tile_num. By splitting the num_heads dimension, the num_heads of the query matrix are decomposed so that one tile can contain data of multiple heads at the same time to meet the mapping requirements. The Dim_to_heads function is applicable when num_heads is less than the specified tile_num. By splitting the K' dimension of the Query matrix, the complete head_dim data of a single head is distributed to multiple tiles, thereby completing the mapping from num_heads to tile_num. Here, K'=head_dim / / K16.

3. The NPU layout method for Attention operator generalization according to claim 2, characterized in that: The second mapping scheme includes using the preset Heads_to_dim, dim_to_Heads and Pad_per_group_heads schemes for the Query matrix; and using the preset Heads_to_dim, Dim_to_heads and Repeat_kv_heads schemes for the Key-Value matrix; The Pad_per_group_heads scheme is used to increase the number of heads in a group by padding a certain number of heads in the group of the Query matrix, thereby achieving the mapping from num_heads to tile_num; The Repeat_kv_heads scheme is used to distribute the Key-Value matrix to more tiles by repeating the number of heads of each Key-Value matrix, thereby ensuring that more tiles contain valid data.

4. The NPU layout method for generalization of the Attention operator according to claim 3, characterized in that: The Pad_per_group_heads solution specifically includes: Used when num_heads is not equal to num_key_value_heads, by dividing the heads of the Query matrix into num_key_value_heads groups, that is, each group has num_key_value_groups heads; where num_key_value_groups = num_heads / / num_key_value_heads; this method mainly fills a certain number of heads in the Query matrix group, thereby changing num_key_value_groups; and updates the number of heads of the Query matrix by num_heads = group_factor * num_key_value_heads, where group_factor is the group coefficient representing the number of heads in each group in the Query matrix.

5. The NPU layout method for Attention operator generalization according to claim 3, characterized in that: In the Repeat_kv_heads scheme, for the Key-Value matrix, the number of heads of the Key-Value matrix is updated by num_key_value_heads = rep_factor * num_key_value_heads; wherein rep_factor is the repeat coefficient, which indicates the number of times each head of the Key-Value matrix is replicated.

6. The NPU layout method for generalization of the Attention operator according to any one of claims 2 to 5, characterized in that: The method further comprises: When both the Head_to_dim solution and the Dim_to_heads solution fail to complete the mapping, the padding solution is enabled to ensure that the required tile_num is finally reached by padding the data.

7. An NPU layout device for generalization of Attention operator, characterized in that: A compiler for a large model in an NPU hardware system, wherein the NPU hardware includes a computing unit with multiple parallel cores, the apparatus comprising: The judgment module is used to determine whether the num_heads of the Query matrix and the num_key_value_heads of the Key-Value matrix in the Attention operator are equal; Mapping module for: If they are equal, a preset first mapping scheme is used to generate a corresponding first layout scheme set; If they are not equal, a preset second mapping scheme is used to generate a corresponding second layout scheme set; A processing module is used to obtain all possible NPU layout schemes based on the first and / or second layout scheme sets, and then select the optimal NPU layout scheme based on the optimal principle of minimizing the overall parameter filling amount, so as to reasonably map num_heads or num_key_value_heads to the specified tile_num, thereby achieving the optimal data parallel layout; wherein tile_num represents the number of computing units.

8. The NPU layout device for Attention operator generalization according to claim 7, characterized in that: The first mapping scheme includes a preset normal scheme, a Heads_to_dim scheme, and a Dim_to_heads scheme; The normal solution is applicable when the number of heads of Query, Key and Value exactly matches the specified tile_num; The Heads_to_dim solution is applicable when num_heads is greater than the specified tile_num. By splitting the num_heads dimension, the num_heads of the query matrix are decomposed so that one tile can contain data of multiple heads at the same time to meet the mapping requirements. The Dim_to_heads function is applicable when num_heads is less than the specified tile_num. By splitting the K' dimension of the Query matrix, the complete head_dim data of a single head is distributed to multiple tiles, thereby completing the mapping from num_heads to tile_num. Here, K'=head_dim / / K16.

9. The NPU layout device for Attention operator generalization according to claim 7, characterized in that: The second mapping scheme includes using the preset Heads_to_dim, dim_to_Heads and Pad_per_group_heads schemes for the Query matrix; and using the preset Heads_to_dim, Dim_to_heads and Repeat_kv_heads schemes for the Key-Value matrix; The Pad_per_group_heads scheme is used to increase the number of heads in a group by padding a certain number of heads in the group of the Query matrix, thereby achieving the mapping from num_heads to tile_num; The Repeat_kv_heads scheme is used to distribute the Key-Value matrix to more tiles by repeating the number of heads of each Key-Value matrix, thereby ensuring that more tiles contain valid data.

Images