Self-attention task dynamic splitting method and system

Abstract

The application discloses a self-attention task dynamic division method and system, relates to the technical field of computer science, and comprises determining a self-attention model structure and a hardware topology structure, wherein the self-attention model structure is determined by input self-attention model parameters, and the hardware topology structure is determined by input hardware parameters; based on input self-attention task parameters, a current reasoning stage of the self-attention model is determined; based on the current reasoning stage, preconfigured task division strategies, the self-attention model structure and the hardware topology structure, an optimal task division strategy corresponding to the current reasoning stage is determined. The application can flexibly determine the task division strategies of each stage, flexibly select the task division strategies of card granularity, core granularity and single-core granularity, reasonably divide and distribute the self-attention task to a multi-card multi-core system, and thus the system computing efficiency and storage utilization rate are improved.

Description

Technical Field

[0001] This invention relates to the field of computer science and technology, and in particular to a task segmentation algorithm for multi-GPU and multi-chip systems, specifically to a method and system for dynamic segmentation of self-attention tasks. Background Technology

[0002] In recent years, large language models have developed rapidly, such as the GPT series, LLaMa series, QWEN, and DeepSeek, which have demonstrated powerful capabilities in many fields of natural language processing, such as intelligent customer service, text generation, and machine translation. Self-attention, as the core of the Transformer architecture, is crucial for large language models to understand and process text. It can capture complex semantic dependencies in massive amounts of text data. However, as the model size and the length of the input text increase, the computational cost of self-attention tasks grows exponentially. Traditional single-card computing power can no longer meet the demands, urgently requiring the powerful computing capabilities of multi-card, multi-chip systems to accelerate processing.

[0003] Advances in hardware technology have led to the widespread adoption of multi-GPU, multi-core systems. GPU clusters and multi-core CPUs provide the hardware foundation for large-scale computing. However, resource allocation and task scheduling in these systems face challenges. How to rationally partition and distribute self-attention tasks across multi-GPU, multi-core systems to fully leverage their parallel computing advantages and improve overall computational efficiency has become an urgent problem to be solved. Summary of the Invention

[0004] To address at least one of the technical problems raised in the background section, this application provides a dynamic segmentation method and system for self-attention tasks, which can determine the optimal task segmentation strategy, rationally segment and allocate self-attention tasks to multi-card multi-core systems, thereby improving system computational efficiency and storage utilization to meet the needs of applications such as large language models.

[0005] In a first aspect, embodiments of the present invention provide a method for dynamic segmentation of a self-attention task, comprising:

[0006] The self-attention model structure and hardware topology are determined, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters;

[0007] Based on the input self-attention task parameters, determine the current inference stage of the self-attention model;

[0008] Based on the current inference stage, the pre-configured task segmentation strategies, the self-attention model structure, and the hardware topology, the optimal task segmentation strategy corresponding to the current inference stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy.

[0009] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the current inference stage based on the current inference stage, pre-configured task segmentation strategies, self-attention model structure, and hardware topology includes:

[0010] In response to the fact that the current inference stage is the pre-filling stage, the optimal task splitting strategy corresponding to the pre-filling stage is determined based on the performance mapping relationship between the pre-configured task splitting strategies and the pre-filling stage.

[0011] In response to the current inference stage being the decoding stage, the optimal task segmentation strategy corresponding to the decoding stage is determined based on the pre-configured task segmentation strategies, self-attention model structure, and hardware topology.

[0012] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the decoding stage based on pre-configured task segmentation strategies, self-attention model structures, and hardware topologies includes:

[0013] For each pre-configured task segmentation strategy, the single-core task computation volume is determined based on the hardware topology, the self-attention model structure, and the task mapping position corresponding to each pre-configured task segmentation strategy.

[0014] For each pre-configured task segmentation strategy, based on the self-attention model structure and the single-core task computational volume, the computational scale of the self-attention model structure under each task segmentation strategy is determined respectively.

[0015] Based on the scale to be computed, the computation delay of each task segmentation strategy is determined, and based on the computation delay of each task segmentation strategy, the optimal task segmentation strategy corresponding to the decoding stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy.

[0016] In some optional embodiments of this example, the self-attention model structure includes an RMSnorm layer, a Linear_In layer, a RoPE_Q layer, a RoPE_K layer, and a Q layer. K T Layer, softmax layer, V layer, Linear_Out layer, and Res_Add layer;

[0017] The self-attention model parameters include the total number of Q heads, the total number of KV heads, the batch size, the context length, the hidden layer dimension, and the head dimension.

[0018] The hardware parameters include the number of cards, the number of chips on each card, the number of cores on each chip, and the tensor parallelism.

[0019] In some optional embodiments of this example, determining the single-core task computation volume based on the hardware topology, the self-attention model structure, and the task mapping positions corresponding to the pre-configured task splitting strategies includes:

[0020] In response to the task partitioning strategy being the first task partitioning strategy, and the task mapping location being a single core, the single-core task computation amount allocated to the single core is determined to be all tasks of D1 heads, where D1 = ceil(Z / (K×N)); the KV Cache size stored on each single core is all data of E1 = ceil(W / (K×N)).

[0021] In response to the task partitioning strategy being the second task partitioning strategy, the task mapping location is a single core, wherein K inter-core tensors on the single core are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / K) of D2 heads, where D2 = ceil(Z / N); the KV Cache stored on each core is (1 / K) of E2 heads of data, where E2 = ceil(W / N) heads;

[0022] In response to the task partitioning strategy being the third task partitioning strategy, the task mapping location is the entire card, wherein (K×N) inter-core tensors on the entire card are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / (K×N)) of D3 heads, where D3=Z; the KV Cache amount stored on each single core is (1 / (K×N)) of E3 heads of data, where E3=W;

[0023] Where ceil() is the floor function; Z is the number of Q heads allocated to each card, Z=X / TP, where X is the total number of Q heads and TP is the tensor parallelism; W is the number of KV heads allocated to each card, where W=Y / TP, where Y is the total number of KV heads.

[0024] In some optional embodiments of this example, determining the computational scale of the self-attention model structure under each task segmentation strategy based on the self-attention model structure and the single-core task computational complexity includes:

[0025] Based on the single-core task computational load, the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. KT Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer;

[0026] Based on the aforementioned RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q... K T Layer, softmax layer, The uncomputable subscales of the V layer, Linear_Out layer, and Res_Add layer are used to determine the uncomputable scale of the self-attention model structure under each task segmentation strategy.

[0027] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the decoding stage based on the computational delay of each task segmentation strategy further includes:

[0028] Determine whether the layers of the self-attention model structure need to communicate;

[0029] For the layers that require communication, the communication delay is determined, and based on the communication delay and the computation delay, the optimal task splitting strategy corresponding to the decoding stage is determined.

[0030] Secondly, embodiments of the present invention also provide a self-attention task dynamic segmentation system, the system comprising:

[0031] The structure determination module is configured to determine the self-attention model structure and the hardware topology, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters;

[0032] The current inference stage determination module is configured to determine the current inference stage of the self-attention model based on the input self-attention task parameters;

[0033] The optimal task segmentation strategy determination module is configured to determine the optimal task segmentation strategy corresponding to the current inference stage based on the current inference stage, pre-configured task segmentation strategies, self-attention model structure, and hardware topology, so as to segment the self-attention task according to the optimal task segmentation strategy.

[0034] Thirdly, embodiments of the present invention also provide a computer device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the above-mentioned self-attention task dynamic segmentation method.

[0035] Fourthly, embodiments of the present invention also provide a computer-readable storage medium storing a computer program, which, when executed by a processor, implements the above-described self-attention task dynamic segmentation method.

[0036] The self-attention task dynamic segmentation method and system provided in this invention can determine the current inference stage of the self-attention model by configuring self-attention task parameters. Furthermore, by using pre-configured task segmentation strategies, self-attention model structures, and hardware topologies, it can determine the optimal task segmentation strategy corresponding to the current inference stage. This allows for flexible determination of task partitioning strategies for each stage and flexible selection of task partitioning strategies at the card granularity, core granularity, and single-core granularity levels. This enables the reasonable segmentation and allocation of self-attention tasks to a multi-card, multi-core system, improving system computational efficiency and storage utilization to meet the needs of applications such as large language models. Attached Figure Description

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

[0038] Figure 1 This is a flowchart illustrating a self-attention task dynamic segmentation method according to an embodiment of the present invention;

[0039] Figure 2 This is a schematic diagram of the structure of the self-attention model in an embodiment of the present invention;

[0040] Figure 3 This is a schematic diagram of the hardware topology in an embodiment of the present invention;

[0041] Figure 4 This is a schematic diagram of the single-core task volume partitioning and mapping corresponding to Strategy 1 in this embodiment of the invention;

[0042] Figure 5 This is a schematic diagram of the single-core task volume partitioning and mapping corresponding to Strategy 2 in this embodiment of the invention;

[0043] Figure 6 This is a schematic diagram of the single-core task volume partitioning and mapping corresponding to Strategy 3 in this embodiment of the invention;

[0044] Figure 7 This is a schematic diagram of the KV Cache segmentation in an embodiment of the present invention;

[0045] Figure 8This is a schematic diagram of the structure of a self-attention task dynamic segmentation system according to an embodiment of the present invention;

[0046] Figure 9 This is a schematic diagram of the physical structure of a computer device provided in an embodiment of the present invention. Detailed Implementation

[0047] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the embodiments of the present invention will be further described in detail below with reference to the accompanying drawings. Here, the illustrative embodiments of the present invention and their descriptions are used to explain the present invention, but are not intended to limit the present invention.

[0048] The information collected in the technical solution of this application is information and data authorized by the user or fully authorized by all parties. The collection, storage, use, processing, transmission, provision, disclosure and application of the relevant data all comply with the relevant laws, regulations and standards of the relevant countries and regions, necessary confidentiality measures have been taken, and they do not violate public order and good morals. Corresponding operation portals are provided for users to choose to authorize or refuse.

[0049] The acquisition, transmission, storage, use, and processing of data in this application comply with relevant national laws and regulations. It should be noted that certain software, components, models, and other existing industry solutions may be mentioned in the embodiments of this application. These should be considered exemplary, intended only to illustrate the feasibility of implementing the technical solution of this application, and do not imply that the applicant has already used or necessarily used such solutions.

[0050] The key terms used in this application are explained below:

[0051] Self-attention: A mechanism in deep learning used to process sequential data. It allows each element of the input sequence to be compared with other elements in the sequence to compute a representation of the sequence, capturing complex dependencies within the sequence. By mapping the input vector to queries, keys, and values, the similarity between queries and keys is calculated to obtain attention weights, and then the values ​​are weighted and summed to obtain the output. Advantages include parallel computation, effective capture of long-distance dependencies, and applicability to various types of sequential data. It is an important component of models such as the Transformer. Advanced versions, such as multi-head self-attention, enhance the model's expressive power by using multiple self-attention mechanisms to focus on different parts of the input in parallel.

[0052] Chipset: A large chip is broken down into multiple smaller chip modules with specific functions; these modules are called chipsets. These chips are then integrated together using advanced packaging technologies to form a complete system-on-a-chip (SoC) or other chip component. This approach reduces the cost and complexity of chip manufacturing, improves design flexibility and scalability; for example, chips with different functions can be manufactured using different process nodes and combined together, and it also facilitates the integration of chips from different suppliers to achieve complex system functions.

[0053] Context, in fields such as natural language processing, refers to the surrounding information related to the currently processed text segment, task, or operation. For example, in language models, context can be other words or sentences before and after the current word or sentence, or even the entire document content. It helps the model understand the semantics and context of the current element, thereby generating a more appropriate response or performing accurate predictions and analyses. For large-scale language models, context length is an important metric, determining the size of the text range the model can process and consider simultaneously; for example, DeepSeek LLM supports a context length of 8192 tokens. In other fields such as image processing and video analysis, context can also refer to surrounding image information, time-series information, etc., related to the currently processed image region or video frame.

[0054] Reduce-Scatter: A many-to-many communication primitive, which can be understood as "Reduce + Scatter". In distributed computing, especially in the distributed training of large models, it reduces data based on shards. That is, it first performs reduction operations (such as summation, finding the maximum value, and other common reduction operations) on the data of several nodes, and then distributes the reduced results to multiple nodes. It is related to but different from operations such as AllReduce (which reduces the data and replicates the result to all nodes). It is one of the important operations for data communication and processing in distributed computing, and its communication mode and algorithm design have an important impact on improving the performance and efficiency of distributed systems.

[0055] Table 1 shows the complete English and Chinese expressions for the English abbreviations involved in this application, specifically:

[0056] Table 1

[0057]

[0058] In self-attention tasks, the storage of KV cache brings pressure from both memory access time and storage capacity overhead. Large model techniques use shared KV head methods such as QGA and MLA to reduce the storage size of KV cache, but in multi-GPU and multi-core systems, shared storage KV cache will bring high communication overhead, reduce parallel computing efficiency, and make it difficult to give full play to the synergy of multiple cores.

[0059] Some traditional task segmentation strategies are designed for specific task loads and context lengths, lacking versatility and adaptability. For example, task segmentation based on fixed rules does not consider real-time system load and dynamic task changes, which can easily lead to some chips or components being overloaded while others remain idle, resulting in wasted resources and failing to meet the complex computational requirements of large language model self-attention tasks.

[0060] In view of this, this application focuses on the efficient processing of self-attention tasks in multi-card multi-core systems, and proposes a dynamic segmentation method for self-attention tasks, such as... Figure 1 As shown, the method includes:

[0061] Step 10: Determine the self-attention model structure and hardware topology, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters.

[0062] In this embodiment, the self-attention model parameters include: batch size B, context length S, total number of Q heads X, total number of KV heads Y, hidden layer dimension H, and head dimension G.

[0063] The self-attention model structure is as follows: Figure 2 As shown, it includes the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer. K T Layer (i.e.) Figure 2 GEMM_1), softmax layer, V layer (i.e.) Figure 2 The GEMM_2, Linear_Out, and Res_Add layers are included.

[0064] Among them, Q K T The layer is a matrix multiplication of the output of RoPE_Q and the transpose of the output of RoPE_K layer; The V layer performs matrix multiplication between the output of the softmax layer and the Y V heads output by the Linear_In layer.

[0065] In this embodiment, the hardware parameters include the number of cards M, the number of chips on each card N, the number of cores on each chip K, the tensor parallelism TP, and the communication bandwidth.

[0066] Among them, the hardware topology is as follows Figure 3 As shown, there are M cards, where each card includes N cores, and each core includes K cores.

[0067] Step 20: Based on the input self-attention task parameters, determine the current inference stage of the self-attention model.

[0068] In this embodiment, the self-attention task parameters include a task identifier, wherein the task identifier is used to characterize the current inference stage of the self-attention model.

[0069] Step 30: Based on the current inference stage, the pre-configured task segmentation strategies, the self-attention model structure, and the hardware topology, determine the optimal task segmentation strategy corresponding to the current inference stage, and segment the self-attention task according to the optimal task segmentation strategy.

[0070] In this embodiment, the current inference stage includes a prefill stage and a decoding stage. Step 30 includes:

[0071] (1) In response to the current inference stage being the pre-filling stage, the optimal task splitting strategy corresponding to the pre-filling stage is determined based on the performance mapping relationship between the pre-configured task splitting strategies and the pre-filling stage.

[0072] Specifically, the performance of task segmentation strategies varies in different inference stages. It is necessary to determine the optimal task segmentation strategy for the pre-filling stage based on the performance mapping relationship between each task segmentation strategy and the pre-filling stage. For example, if strategy A performs better than strategy B in the pre-filling stage, then strategy A is directly determined as the optimal performance strategy for the pre-filling stage.

[0073] (2) In response to the current inference stage being the decoding stage, based on the pre-configured task segmentation strategies, self-attention model structure, and hardware topology, the optimal task segmentation strategy corresponding to the decoding stage is determined, which can be achieved through the following steps:

[0074] Step 301: For each pre-configured task segmentation strategy, determine the single-core task computation volume based on the hardware topology, the self-attention model structure, and the task mapping position corresponding to each pre-configured task segmentation strategy.

[0075] According to the self-attention model parameters and hardware parameters configured above, it can be seen that the self-attention task of the model has X Q heads and Y KV heads; the number of Q heads allocated to each card according to the inter-card tensor parallelism (TP) is Z=X / TP, and the number of KV heads is W=Y / TP.

[0076] In this embodiment, the pre-configured task splitting strategies include a first task splitting sub-strategy (i.e., Strategy 1 below), a second task splitting sub-strategy (i.e., Strategy 2 below) and a third task splitting sub-strategy (i.e., Strategy 3 below), specifically:

[0077] In Strategy 1, the computing task of each QKV head is mapped to a single core. As Figure 4 shown, each card has K×N cores, so the computing task allocated to a single core is Z / (K×N) heads. The amount of KV Cache stored on each core is W / (K×N) heads. The input and output linear layers, QKV matrix multiplication, and softmax of a single head are all completed locally on each computing core. All-Reduce communication between cores, between chiplets and between cards is performed at the end of the self-attention task.

[0078] It should be noted that the advantage of Strategy 1 is that each head is calculated on the same core, a single self-attention block only needs one more communication, and multi-layer operator calculation is completed locally within the core. Based on the task splitting method of Strategy 1, it has obvious advantages in terms of communication volume and communication times.

[0079] It also needs to be explained that Strategy 1 has the following two disadvantages:

[0080] 1. When the inter-card TP is large and the number of model Q heads is small, the number of Q heads allocated to each card (Z=Q / TP) is small. When it is less than the number of cores (K×N) of each card, that is, Z<K×N, there is a situation where only some computing cores are assigned the computing task of 1 head, and other computing cores are idle, resulting in waste of computing and bandwidth resources.

[0081] 2. Since the amount of KV Cache stored on each core is W / (K×N) heads, when the number of KV heads W of the model is small (for example, DeepSeek compresses KV to only one head), this results in multiple backups of KV Cache on a single card, reducing the utilization of local storage. For large-batch or context tasks, the capacity of KV Cache can easily exceed local storage, requiring memory access to global storage, which increases memory access overhead and affects model inference performance.

[0082] Therefore, Strategy 1 is suitable for the following scenarios: small inter-card TP; a large number of model QKV heads; the prefill stage; and the decode stage with small batches and short context.

[0083] In Strategy 2, the computation task of each QKV head (i.e., Q head and KV head) is mapped to a single chiplet, as Figure 5 shows that K cores on a single chiplet perform tensor parallelism, and each core is allocated 1 / K of the task of a head. Therefore, Z / N heads are allocated to a single chiplet, and 1 / K of the tasks of Z / N heads are allocated to a single core. The amount of KVCache stored on each chiplet is W / N heads. The amount of KV Cache stored on a single core is 1 / K of W / N heads. The input and output linear layers of a single head, the matrix multiplication of QKV, and softmax are all completed on a single chiplet, and inter-core communication within the chiplet exists during the computation. All-Reduce communication among cores, chiplets and cards is performed at the end of the self-attention task.

[0084] It should be noted that both the advantages and disadvantages of Strategy 2 are moderate. Specifically, the computation of each layer involves inter-core communication within the chiplet, with moderate communication volume and communication times. Task allocation is relatively even, as long as the number of Q heads allocated per card is greater than the number of chiplets Z<N in the card, all computing resources and memory access bandwidth can be utilized. KV Cache is stored repeatedly among chiplets, only one copy is stored inside a chiplet, and shared by K cores in the chiplet.

[0085] Strategy 2 is suitable for the following scenarios: the decode stage of most general models; when a single card is allocated a relatively large number of QKV heads; it is slightly inferior to Strategy 1 in the decode stage with small batches and short context; it is slightly inferior to Strategy 3 in the decode stage where the number of Q heads allocated to a single card is less than N and the batch and context are very large.

[0086] In Strategy 3, the computation task of each QKV head is mapped to the entire card, and K×N cores on the entire card perform tensor parallelism, as Figure 6 shows, each core is allocated 1 / (K×N) of the task of a head. Therefore, 1 / (K×N) of the tasks of Z heads are allocated to each core. The amount of KV Cache stored on a single core is 1 / (K×N) of W heads. During the computation of the input and output linear layers of a single head, the matrix multiplication of QKV, and softmax, there are inter-core communication within chiplets and inter-chiplet communication. All-Reduce communication among cores, chiplets and cards is performed at the end of the self-attention task.

[0087] It should be noted that Strategy 3 has the following two advantages:

[0088] 1. Even if only one head is allocated to each card, there are no idle computing cores. Computing and bandwidth resources are fully utilized. 2. The KV Cache exists only once on the entire card, and its capacity is unlikely to exceed local storage, thus improving storage utilization.

[0089] It should also be noted that computation of multi-level operators involves communication between cores and between kernels, and Strategy 3 has a significant disadvantage in terms of communication volume and number of communication sessions.

[0090] Therefore, Strategy 3 is suitable for scenarios where: the TP between cards is large; the number of QKV heads in the model is small; and the decoding stage has a very large batch size and context.

[0091] Therefore, based on the pre-configured task segmentation strategies mentioned above, and considering the self-attention model structure and hardware topology, the specific task volume mapped onto a single core is as follows:

[0092] In response to the task partitioning strategy being the first task partitioning strategy, and the task mapping location being a single core, the single-core task computation amount allocated to the single core is determined to be all tasks of D1 heads, where D1 = ceil(Z / (K×N)); the KV Cache size stored on each single core is all data of E1 = ceil(W / (K×N)).

[0093] In response to the task partitioning strategy being the second task partitioning strategy, the task mapping location is a single core, wherein K inter-core tensors on the single core are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / K) of D2 heads, where D2 = ceil(Z / N); the KV Cache stored on each core is (1 / K) of E2 heads of data, where E2 = ceil(W / N) heads;

[0094] In response to the task partitioning strategy being the third task partitioning strategy, the task mapping location is the entire card, wherein (K×N) inter-core tensors on the entire card are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / (K×N)) of D3 heads, where D3=Z; the KV Cache amount stored on each single core is (1 / (K×N)) of E3 heads of data, where E3=W;

[0095] Where ceil() is the floor function; Z is the number of Q heads allocated to each card, Z=X / TP, where X is the total number of Q heads and TP is the tensor parallelism; W is the number of KV heads allocated to each card, where W=Y / TP, where Y is the total number of KV heads.

[0096] Step 302: Based on the self-attention model structure and the single-core task computational complexity, determine the computational scale of the self-attention model structure under each task segmentation strategy.

[0097] In this embodiment, after determining the amount of tasks mapped to a single core by each task partitioning strategy, it is necessary to further determine the workload under each task partitioning strategy. Figure 2 The computational size of the self-attention model structure shown.

[0098] In some optional embodiments of this example, the computational scale of the self-attention model structure under each task segmentation strategy is determined based on the self-attention model structure and the amount of task mapped to a single core, including:

[0099] Step 3021: Based on the single-core task computational load, determine the RMSnorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer respectively. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer;

[0100] Step 3022: Based on the RMSnorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q... K T Layer, softmax layer, The uncomputable subscales of the V layer, Linear_Out layer, and Res_Add layer are used to determine the uncomputable scale of the self-attention model structure under each task segmentation strategy.

[0101] Specifically, the computational scale of the self-attention model structure is determined by determining the computational scale of each layer. The sum of the computational scales of all layers is the computational scale of the self-attention model structure.

[0102] When using different task segmentation strategies in this application, the computational scale of the self-attention model structure varies, where:

[0103] (1) In response to the task segmentation strategy being the first task segmentation strategy, the computational scale of the self-attention model structure is as follows:

[0104] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0105] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of task mapped to a single core corresponding to the first task segmentation strategy;

[0106] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The subscale to be computed in layer V is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the first task segmentation strategy.

[0107] The sub-scale to be computed in the softmax layer is determined based on the batch size, context length, and the amount of task mapped to a single core corresponding to the first task segmentation strategy.

[0108] Wherein: the sub-scale to be computed in the RMSNorm layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Res_Add layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Linear_In layer is [B×S, H]×[H, (D1+2×E1)×G], the sub-scale to be computed in the Linear_Out layer is [B×S, D1×G]×[D1×G, H]; the sub-scale to be computed in the RoPE_Q layer is [B×S, D1×G], that is, the operation of B×S vectors with dimension D1×G; the sub-scale to be computed in the RoPE_K layer is [B×S, E1×G], that is, the operation of B×S vectors with dimension E1×G; the Q K T The subscale to be computed for the layer is (B×D1) times [S, G]×[G, S]; the The subscale to be computed in layer V is (B×D1) times [S, S]×[S, G]; the subscale to be computed in layer softmax is [B×D1×S, S], which is the operation of B×D1×S vectors with dimension S.

[0109] (2) In response to the task segmentation strategy being the second task segmentation strategy, the computational scale of the self-attention model structure is as follows:

[0110] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0111] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of task mapped to a single core corresponding to the second task segmentation strategy;

[0112] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The subscale to be computed in layer V is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the second task segmentation strategy.

[0113] The sub-size to be computed in the softmax layer is determined based on the batch size, context length, and the amount of task mapped to a single core corresponding to the second task segmentation strategy.

[0114] Wherein: the sub-scale to be computed in the RMSNorm layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Res_Add layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Linear_In layer is [B×S, H]×[H, align((D2+2×E2)×ceil(G / K))]; the sub-scale to be computed in the Linear_Out layer is [B×S, D2×G]×[D2×G, align(H / K)], where align() is a hardware affinity function, which is used to determine the computational amount corresponding to the delayed computation based on the sub-scale to be computed and the hardware processing capability; the sub-scale to be computed in the RoPE_Q layer is [B×S, align(D2×ceil(G / K))]; the sub-scale to be computed in the RoPE_K layer is [B×S, align(E2×ceil(G / K))]; the Q K T The subscale to be computed for the layer is (B×D²) times [S, G]×[G, align(S / K)]; The subscale to be computed in layer V is (B×D2) times [S, S]×[S, align(S / K)]; the subscale to be computed in layer softmax is [B×D2×S, S], which is the operation of B×D2×S vectors with dimension S.

[0115] (3) In response to the task segmentation strategy being the third task segmentation strategy, the computational scale of the self-attention model structure is as follows:

[0116] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0117] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy;

[0118] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The size of the V-layer to be computed is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy.

[0119] The sub-size to be computed in the softmax layer is determined based on the batch size, context length, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy.

[0120] Wherein: the sub-scale to be computed in the RMSNorm layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Res_Add layer is [B×S, H], that is, the operation of B×S vectors with dimension H; the sub-scale to be computed in the Linear_In layer is [B×S, H]×[H, align((D3+2×E3)×ceil(G / (K×N)))]; the sub-scale to be computed in the Linear_Out layer is [B×S, D3×G]×[D2×G, align(H / (K×N))]; the sub-scale to be computed in the RoPE_Q layer is [B×S, align(D3×ceil(G / (K×N)))]; the sub-scale to be computed in the RoPE_K layer is [B×S, align(E3×ceil(G / (K×N)))]; the Q K T The subscale to be computed for the layer is (B×D3) times [S, G]×[G, align(S / (K×N))]; The subscale to be computed in layer V is (B×D3) times [S, S]×[S, align(S / (K×N))]; the subscale to be computed in layer softmax is [B×D3×S, S], which is the operation of B×D3×S vectors with dimension S.

[0121] Step 303: Based on the scale to be computed, determine the computation delay of each task segmentation strategy, and based on the computation delay of each task segmentation strategy, determine the optimal task segmentation strategy corresponding to the decoding stage, so as to segment the self-attention task according to the optimal task segmentation strategy.

[0122] Specifically, after determining the computational scale, the computational latency of each task partitioning strategy can be determined through model simulation, which will not be elaborated here. Furthermore, by comparing the computational latency of each task partitioning strategy, the task partitioning strategy with the lowest computational latency is selected from the card-level, core-level, and single-core-level task partitioning strategies, and this strategy is determined as the optimal task partitioning strategy. After determining the optimal task partitioning strategy, the self-attention tasks are partitioned according to the optimal task partitioning strategy and allocated to the multi-card, multi-core system.

[0123] So far, compared with the traditional fixed segmentation strategy that lacks the adaptive ability to task scale (batch size, context length) and hardware configuration (number of cards, number of cores), which leads to rigid resource scheduling. When the number of QKV heads of the model does not match the number of hardware cores (for example, when the inter-card TP is large and the number of Q heads is small), some computing cores are idle (for example, idle cores when Z < K×N in Strategy 1), resulting in low utilization of computing resources. The dynamic self-attention task segmentation method proposed in the present application calculates the inference time consumption of three different strategies under the current self-attention computing task scale and hardware parameters by configuring computing task scale parameters such as inference batch size, context length and the number of QKV heads of self-attention, as well as configuring hardware system-related parameters such as number of cards, inter-card parallelism, number of chiplets, number of cores and computing power bandwidth, then determines the adaptive optimal strategy scheme under the current configuration, so as to achieve the effect of flexibly selecting task division strategies at card granularity, chiplet granularity and single-core granularity.

[0124] Among them, Strategy 1: a single core independently calculates a complete head, and the communication traffic in the prefill stage is reduced by 70% (only one All-Reduce is required), which is suitable for high QKV head scenarios; Strategy 2: chiplet-level calculation takes into account load balancing, suitable for general scenarios; Strategy 3: a whole card cooperatively calculates a single head, which can completely eliminate the problem of idle cores, suitable for low QKV head scenarios.

[0125] It should also be noted that in Strategy 2 and Strategy 3, each core only calculates part of the tasks of the self-attention head. If KV Cache is repeatedly stored on each core, it is easy to exceed the local storage capacity in the case of a long context in the decode stage. In this case, the data exceeding the capacity can only be stored in the global storage, resulting in performance degradation. If KV Cache only stores a part of the data on each core, the utilization of local storage will be significantly improved, but there will be a problem of data communication, and the impact of communication on performance needs to be minimized.

[0126] The traditional repeated storage strategy (for example, each core stores W / (K×N) heads in Strategy 1) causes storage redundancy and capacity overflow risk when the number of KV heads is small, and long context tasks require frequent access to global storage.

[0127] Based on this, as Figure 7 shown, since the data required on the algorithm core needs to be stored and excessive storage will increase computing overhead, the present application splits KV Cache into L blocks along the S dimension of the cache. Wherein, Figure 7In this context, `head_dim` represents the head dimension, `Cache_len` represents the cache length, and `L` represents the number of chunks. In strategy 1, `L` is 1, storing all chunks. In strategy 2, `L` is the number of cores `K` within a single core, and in strategy 3, `L` is the number of cores `K×N` within the entire card. This way, no communication is needed during the multiplication of `Q` by `K`. When calculating the softmax, each token needs to communicate to find the maximum and sum of the scores within its own core's interval, but the communication volume is small. Finally, after multiplying by `V`, Reduce-Scatter communication is required. This KV cache local storage strategy effectively balances local storage utilization with the performance overhead of communication.

[0128] This application can design storage levels according to the partitioning strategy, where strategy 1 is core-level repeated storage; strategy 2 is intra-core shared storage (only 1 copy / K cores are stored); strategy 3 is single-copy storage for the entire card (storage volume compressed by K×N times); this application can also optimize the .KV Cache block partitioning, partitioning the Cache into L blocks (L=K or K×N) in the S dimension; and reduce memory access pressure through local softmax calculation + Reduce-Scatter communication.

[0129] In some optional embodiments of this example, it is also necessary to determine whether the layers in the self-attention model structure need to communicate; for the layers that need to communicate, the communication latency is determined by simulation, and the optimal task splitting strategy is determined based on the communication latency and the computation latency. The communication latency, for example, is the All-Reduce communication latency between cores, between individual cores, and between cards at the end of the self-attention task.

[0130] As an example rather than a limitation, the task partitioning strategy that minimizes the sum of communication latency and computation latency is selected as the optimal task partitioning strategy.

[0131] This application is mainly applied to the field of large-scale artificial intelligence model inference computation, and is especially suitable for large language model (LLM) inference scenarios that require efficient processing of long sequences and high-concurrency self-attention tasks, such as intelligent customer service dialogue systems, long text generation, machine translation, code generation and other AI applications that require low latency and high throughput.

[0132] Based on the same inventive concept, this application also provides a self-attention task dynamic segmentation system, which can be used to implement the method described in the above embodiments, as described in the following embodiments. Since the principle by which this self-attention task dynamic segmentation system solves the problem is similar to that of a self-attention task dynamic segmentation method, the implementation of a self-attention task dynamic segmentation device can refer to the implementation of a self-attention task dynamic segmentation method, and repeated details will not be elaborated further. As used below, the terms "unit" or "module" can refer to a combination of software and / or hardware that implements a predetermined function. Although the system described in the following embodiments is preferably implemented in software, hardware implementation, or a combination of software and hardware, is also possible and contemplated.

[0133] like Figure 8 As shown, the self-attention task dynamic segmentation system includes:

[0134] The structure determination module 801 is configured to determine the self-attention model structure and the hardware topology, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters.

[0135] The current inference stage determination module 802 is configured to determine the current inference stage of the self-attention model based on the input self-attention task parameters.

[0136] The optimal task segmentation strategy determination module 803 is configured to determine the optimal task segmentation strategy corresponding to the current inference stage based on the current inference stage, pre-configured task segmentation strategies, self-attention model structure, and hardware topology.

[0137] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the current inference stage based on the current inference stage, pre-configured task segmentation strategies, self-attention model structure, and hardware topology includes:

[0138] In response to the fact that the current inference stage is the pre-filling stage, the optimal task splitting strategy corresponding to the pre-filling stage is determined based on the performance mapping relationship between the pre-configured task splitting strategies and the pre-filling stage.

[0139] In response to the current inference stage being the decoding stage, the optimal task segmentation strategy corresponding to the decoding stage is determined based on the pre-configured task segmentation strategies, self-attention model structure, and hardware topology.

[0140] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the decoding stage based on pre-configured task segmentation strategies, self-attention model structures, and hardware topologies includes:

[0141] Based on the hardware topology and the task mapping positions corresponding to the pre-configured task splitting strategies, the computational load of a single-core task is determined.

[0142] Based on the self-attention model structure and the computational complexity of the single-core task, the computational scale of the self-attention model structure is determined.

[0143] Based on the scale to be computed, the computational delay of each task segmentation strategy is determined, and based on the computational delay of each task segmentation strategy, the optimal task segmentation strategy corresponding to the decoding stage is determined.

[0144] In some optional embodiments of this example, the self-attention model structure includes an RMSnorm layer, a Linear_In layer, a RoPE_Q layer, a RoPE_K layer, and a Q layer. K T Layer, softmax layer, V layer, Linear_Out layer, and Res_Add layer;

[0145] The self-attention model parameters include the total number of Q heads, the total number of KV heads, the batch size, the context length, the hidden layer dimension, and the head dimension.

[0146] The hardware parameters include the number of cards, the number of chips on each card, the number of cores on each chip, and the tensor parallelism.

[0147] In some optional embodiments of this example, determining the single-core task computation volume based on the hardware topology and the task mapping positions corresponding to each pre-configured task splitting strategy includes:

[0148] In response to the task partitioning strategy being the first task partitioning strategy, and the task mapping location being a single core, the single-core task computation amount allocated to the single core is determined to be all tasks of D1 heads, where D1 = ceil(Z / (K×N)); the KV Cache size stored on each single core is all data of E1 = ceil(W / (K×N)).

[0149] In response to the task partitioning strategy being the second task partitioning strategy, the task mapping location is a single core, wherein K inter-core tensors on the single core are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / K) of D2 heads, where D2 = ceil(Z / N); the KV Cache stored on each core is (1 / K) of E2 heads of data, where E2 = ceil(W / N) heads;

[0150] In response to the task partitioning strategy being the third task partitioning strategy, the task mapping location is the entire card, wherein (K×N) inter-core tensors on the entire card are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / (K×N)) of D3 heads, where D3=Z; the KV Cache amount stored on each single core is (1 / (K×N)) of E3 heads of data, where E3=W;

[0151] Where ceil() is the floor function; Z is the number of Q heads allocated to each card, Z=X / TP, where X is the total number of Q heads and TP is the tensor parallelism; W is the number of KV heads allocated to each card, where W=Y / TP, where Y is the total number of KV heads.

[0152] In some optional embodiments of this example, determining the computational scale of the self-attention model structure based on the self-attention model structure and the computational complexity of the single-core task includes:

[0153] Based on the single-core task computational load, the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer;

[0154] Based on the aforementioned RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q... K T Layer, softmax layer, The computational subscales of the V layer, Linear_Out layer, and Res_Add layer are used to determine the computational scale of the self-attention model structure.

[0155] In some optional embodiments of this example, in response to the task segmentation strategy being a first task segmentation strategy, the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer are respectively determined. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer include:

[0156] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0157] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of task mapped to a single core corresponding to the first task segmentation strategy;

[0158] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The subscale to be computed in layer V is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the first task segmentation strategy.

[0159] The sub-scale to be computed in the softmax layer is determined based on the batch size, context length, and the amount of task mapped to a single core corresponding to the first task segmentation strategy.

[0160] Wherein: the sub-scale to be computed in the RMSNorm layer is [B×S, H]; the sub-scale to be computed in the Res_Add layer is [B×S, H]; the sub-scale to be computed in the Linear_In layer is [B×S, H]×[H, (D1+2×E1)×G]; the sub-scale to be computed in the Linear_Out layer is [B×S, D1×G]×[D1×G, H]; the sub-scale to be computed in the RoPE_Q layer is [B×S, D1×G]; the sub-scale to be computed in the RoPE_K layer is [B×S, E1×G]; and the Q... K T The subscale to be computed for the layer is (B×D1) times [S, G]×[G, S]; the The subscale to be computed in layer V is (B×D1) times [S, S]×[S, G]; the subscale to be computed in layer softmax is [B×D1×S, S].

[0161] In some optional embodiments of this example, in response to the task segmentation strategy being the second task segmentation strategy, the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer include:

[0162] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0163] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of task mapped to a single core corresponding to the second task segmentation strategy;

[0164] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The subscale to be computed in layer V is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the second task segmentation strategy.

[0165] The sub-size to be computed in the softmax layer is determined based on the batch size, context length, and the amount of task mapped to a single core corresponding to the second task segmentation strategy.

[0166] Wherein: the sub-scale to be computed in the RMSNorm layer is [B×S, H]; the sub-scale to be computed in the Res_Add layer is [B×S, H]; the sub-scale to be computed in the Linear_In layer is [B×S, H]×[H, align((D2+2×E2)×ceil(G / K))]; the sub-scale to be computed in the Linear_Out layer is [B×S, D2×G]×[D2×G, align(H / K)], align() is a hardware affinity function, which is used to determine the computational amount corresponding to the delayed computation based on the sub-scale to be computed and the hardware processing capability; the sub-scale to be computed in the RoPE_Q layer is [B×S, align(D2×ceil(G / K))]; the sub-scale to be computed in the RoPE_K layer is [B×S, align(E2×ceil(G / K))]; the sub-scale to be computed in the Q ... K T The subscale to be computed for the layer is (B×D²) times [S, G]×[G, align(S / K)]; The subscale to be computed in layer V is (B×D2) times [S, S]×[S, align(S / K)]; the subscale to be computed in layer softmax is [B×D2×S, S].

[0167] In some optional embodiments of this example, in response to the task segmentation strategy being a third task segmentation strategy, the determination of the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer is performed respectively. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer include:

[0168] The subscale to be computed for the RMSNorm layer and the Res_Add layer is determined, wherein the subscale to be computed for the RMSNorm layer and the Res_Add layer is determined based on the batch size, context length, and hidden layer dimension;

[0169] The sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined respectively, wherein the sub-scales to be computed for the Linear_In layer and the Linear_Out layer are determined based on the batch size, context length, hidden layer dimension, head dimension, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy;

[0170] The RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layers and The subscale to be computed in layer V, wherein the RoPE_Q layer, RoPE_K layer, and Q layer are... K T Layers and The size of the V-layer to be computed is determined based on the batch size, context length, head dimension, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy.

[0171] The sub-size to be computed in the softmax layer is determined based on the batch size, context length, and the amount of tasks mapped to a single core corresponding to the third task segmentation strategy.

[0172] Wherein: the subscale to be computed in the RMSnorm layer is [B×S, H]; the subscale to be computed in the Res_Add layer is [B×S, H]; the subscale to be computed in the Linear_In layer is [B×S, H]×[H, align((D3+2×E3)×ceil(G / (K×N)))]; the subscale to be computed in the Linear_Out layer is [B×S, D3×G]×[D2×G, align(H / (K×N))]; the subscale to be computed in the RoPE_Q layer is [B×S, align(D3×ceil(G / (K×N)))]; the subscale to be computed in the RoPE_K layer is [B×S, align(E3×ceil(G / (K×N)))]; the subscale to be computed in the Q ... K T The subscale to be computed for the layer is (B×D3) times [S, G]×[G, align(S / (K×N))]; The subscale to be computed in layer V is (B×D3) times [S, S]×[S, align(S / (K×N))]; the subscale to be computed in layer softmax is [B×D3×S, S].

[0173] In some optional embodiments of this example, determining the optimal task segmentation strategy corresponding to the decoding stage based on the computational delay of each task segmentation strategy further includes:

[0174] Determine whether the layers of the self-attention model structure need to communicate;

[0175] For the layers that require communication, the communication delay is determined, and based on the communication delay and the computation delay, the optimal task splitting strategy corresponding to the decoding stage is determined.

[0176] According to embodiments of this disclosure, this disclosure also provides a computer device, including: at least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, the instructions being executed by the at least one processor to enable the at least one processor to perform the steps of a self-attention task dynamic segmentation method of the foregoing embodiments.

[0177] According to embodiments of this disclosure, this disclosure also provides a computer-readable storage medium storing a computer program that, when executed by a processor, implements a self-attention task dynamic segmentation method of the foregoing embodiments.

[0178] Figure 9 A schematic block diagram of an example computer device 900 that can be used to implement embodiments of the present disclosure is shown. The computer device is intended to represent various forms of digital computers, such as laptop computers, desktop computers, workstations, personal digital assistants, servers, blade servers, mainframe computers, and other suitable computers. The computer device may also represent various forms of mobile devices, such as personal digital processors, cellular phones, smartphones, wearable devices, and other similar computing devices. The components shown herein, their connections and relationships, and their functions are merely illustrative and are not intended to limit the implementation of the present disclosure described and / or claimed herein.

[0179] like Figure 9 As shown, device 900 includes a computing unit 901, which can perform various appropriate actions and processes based on a computer program stored in read-only memory (ROM) 902 or a computer program loaded from storage unit 908 into random access memory (RAM) 903. RAM 903 may also store various programs and data required for the operation of device 900. The computing unit 901, ROM 902, and RAM 903 are interconnected via bus 904. Input / output (I / O) interface 905 is also connected to bus 904.

[0180] Multiple components in device 900 are connected to I / O interface 905, including: input unit 906, such as keyboard, mouse, etc.; output unit 907, such as various types of monitors, speakers, etc.; storage unit 908, such as disk, optical disk, etc.; and communication unit 909, such as network card, modem, wireless transceiver, etc. Communication unit 909 allows device 900 to exchange information / data with other devices through computer networks such as the Internet and / or various telecommunications networks.

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

[0182] For example, in some embodiments, a self-attention task dynamic segmentation method can be implemented as a computer software program tangibly contained in a machine-readable medium, such as storage unit 908. In some embodiments, part or all of the computer program can be loaded and / or installed on device 900 via ROM 902 and / or communication unit 909. When the computer program is loaded into RAM 903 and executed by computing unit 901, one or more steps of the self-attention task dynamic segmentation method described above can be performed. Alternatively, in other embodiments, computing unit 901 can be configured to perform a self-attention task dynamic segmentation method by any other suitable means (e.g., by means of firmware).

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

[0184] The program code used to implement the methods of this disclosure may be written in any combination of one or more programming languages. This program code may be provided to a processor or controller of a general-purpose computer, special-purpose computer, or other programmable data processing apparatus, such that when executed by the processor or controller, the program code causes the functions / operations specified in the flowcharts and / or block diagrams to be implemented. The program code may be executed entirely on a machine, partially on a machine, as a standalone software package partially on a machine and partially on a remote machine, or entirely on a remote machine or server.

[0185] In the context of this disclosure, a machine-readable medium can be a tangible medium that may contain or store a program for use by or in conjunction with an instruction execution system, apparatus, or device. A machine-readable medium can be a machine-readable signal medium or a machine-readable storage medium. A machine-readable medium can be, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, apparatus, or devices, or any suitable combination of the foregoing. More specific examples of machine-readable storage media include electrical connections based on one or more wires, portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fiber, portable compact disk read-only memory (CD-ROM), optical storage devices, magnetic storage devices, or any suitable combination of the foregoing.

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

[0187] The systems and technologies described herein can be implemented in computing systems that include backend components (e.g., as a data server), or computing systems that include middleware components (e.g., an application server), or computing systems that include frontend components (e.g., a user computer with a graphical user interface or web browser through which a user can interact with implementations of the systems and technologies described herein), or any combination of such backend, middleware, or frontend components. The components of the system can be interconnected via digital data communication of any form or medium (e.g., a communication network). Examples of communication networks include local area networks (LANs), wide area networks (WANs), and the Internet.

[0188] Computer systems can include clients and servers. Clients and servers are generally located far apart and typically interact via communication networks. Client-server relationships are created by computer programs running on the respective computers and having a client-server relationship with each other. Servers can be cloud servers, servers in distributed systems, or servers incorporating blockchain technology.

[0189] It should be noted that in the description of this application, relational terms such as "first" and "second" are used only to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitations, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element.

[0190] In the embodiments of this application, the singular forms "a," "the," etc., including the plural forms, should be broadly understood as "a kind" or "a class" rather than limited to the meaning of "an." Furthermore, the term "the" should be understood to include both the singular and plural forms, unless the context explicitly indicates otherwise. Additionally, the term "according to" should be understood as "at least partially based on…," and the term "based on" should be understood as "at least partially based on…," unless the context explicitly indicates otherwise.

[0191] It should be understood that the various forms of processes shown above can be used to rearrange, add, or delete steps. For example, the steps described in this disclosure can be executed in parallel, sequentially, or in different orders, as long as the desired result of the technical solution of this disclosure can be achieved, and this is not limited herein.

[0192] The specific embodiments described above further illustrate the purpose, technical solution, and beneficial effects of the present invention. It should be understood that the above descriptions are merely specific embodiments of the present invention and are not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A dynamic segmentation method for self-attention tasks, characterized in that, include: The self-attention model structure and hardware topology are determined, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters; Based on the input self-attention task parameters, determine the current inference stage of the self-attention model; In response to the fact that the current inference stage is the pre-filling stage, based on the performance mapping relationship between the pre-configured task segmentation strategies and the pre-filling stage, the optimal task segmentation strategy corresponding to the pre-filling stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy. In response to the current inference stage being the decoding stage, for each pre-configured task segmentation strategy, the single-core task computation volume is determined based on the hardware topology, the self-attention model structure, and the task mapping position corresponding to each pre-configured task segmentation strategy, including: In response to the task partitioning strategy being the first task partitioning strategy, and the task mapping location being a single core, the single-core task computation amount allocated to the single core is determined to be all tasks of D1 heads, where D1 = ceil(Z / (K×N)); the KV Cache size stored on each single core is all data of E1 = ceil(W / (K×N)). In response to the task partitioning strategy being the second task partitioning strategy, the task mapping location is a single core, wherein K inter-core tensors on the single core are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / K) of D2 heads, where D2 = ceil(Z / N); the KV Cache stored on each core is (1 / K) of E2 heads of data, where E2 = ceil(W / N) heads; In response to the task partitioning strategy being the third task partitioning strategy, the task mapping location is the entire card, wherein (K×N) inter-core tensors on the entire card are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / (K×N)) of D3 heads, where D3=Z; the KV Cache size stored on each single core is (1 / (K×N)) of E3 heads of data, where E3=W; Where ceil() is the floor function; Z is the number of Q heads allocated to each card, Z=X / TP, where X is the total number of Q heads and TP is the tensor parallelism; W is the number of KV heads allocated to each card, where W=Y / TP, where Y is the total number of KV heads and N is the number of cores on each card; For each pre-configured task segmentation strategy, based on the self-attention model structure and the single-core task computational volume, the computational scale of the self-attention model structure under each task segmentation strategy is determined respectively. Based on the scale to be computed, the computation delay of each task segmentation strategy is determined, and based on the computation delay of each task segmentation strategy, the optimal task segmentation strategy corresponding to the decoding stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy.

2. The task dynamic segmentation method according to claim 1, characterized in that, The self-attention model parameters include the total number of Q heads, the total number of KV heads, the batch size, the context length, the hidden layer dimension, and the head dimension. The hardware parameters include the number of cards, the number of chips on each card, the number of cores on each chip, and the tensor parallelism.

3. The task dynamic segmentation method according to claim 1, characterized in that, The self-attention model structure includes an RMSnorm layer, a Linear_In layer, a RoPE_Q layer, a RoPE_K layer, and a Q layer. K T Layer, softmax layer, The V layer, Linear_Out layer, and Res_Add layer, wherein the computational scale of the self-attention model structure under each task segmentation strategy is determined based on the self-attention model structure and the computational complexity of the single-core task, including: Based on the computational load of the single-core task, the RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q layer are determined respectively. K T Layer, softmax layer, The subscales to be computed in the V layer, Linear_Out layer, and Res_Add layer; Based on the aforementioned RMSNorm layer, Linear_In layer, RoPE_Q layer, RoPE_K layer, and Q... K T Layer, softmax layer, The uncomputable subscales of the V layer, Linear_Out layer, and Res_Add layer are used to determine the uncomputable scale of the self-attention model structure under each task segmentation strategy.

4. The task dynamic segmentation method according to claim 1, characterized in that, The determination of the optimal task segmentation strategy corresponding to the decoding stage based on the computational delay of each task segmentation strategy further includes: Determine whether the layers of the self-attention model structure need to communicate; For the layers that require communication, the communication delay is determined, and based on the communication delay and the computation delay, the optimal task splitting strategy corresponding to the decoding stage is determined.

5. A self-attention task dynamic segmentation system, characterized in that, include: The structure determination module is configured to determine the self-attention model structure and the hardware topology, wherein the self-attention model structure is determined by the input self-attention model parameters, and the hardware topology is determined by the input hardware parameters; The current inference stage determination module is configured to determine the current inference stage of the self-attention model based on the input self-attention task parameters; The optimal task segmentation strategy determination module is configured to determine the optimal task segmentation strategy corresponding to the current inference stage based on the current inference stage, pre-configured task segmentation strategies, self-attention model structure, and hardware topology, so as to segment the self-attention task according to the optimal task segmentation strategy. The optimal task splitting strategy determination module is further configured as follows: In response to the fact that the current inference stage is the pre-filling stage, based on the performance mapping relationship between the pre-configured task segmentation strategies and the pre-filling stage, the optimal task segmentation strategy corresponding to the pre-filling stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy. In response to the current inference stage being the decoding stage, for each pre-configured task segmentation strategy, the single-core task computation volume is determined based on the hardware topology, the self-attention model structure, and the task mapping position corresponding to each pre-configured task segmentation strategy, including: In response to the task partitioning strategy being the first task partitioning strategy, and the task mapping location being a single core, the single-core task computation amount allocated to the single core is determined to be all tasks of D1 heads, where D1 = ceil(Z / (K×N)); the KV Cache size stored on each single core is all data of E1 = ceil(W / (K×N)). In response to the task partitioning strategy being the second task partitioning strategy, the task mapping location is a single core, wherein K inter-core tensors on the single core are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / K) of D2 heads, where D2 = ceil(Z / N); the KV Cache stored on each core is (1 / K) of E2 heads of data, where E2 = ceil(W / N) heads; In response to the task partitioning strategy being the third task partitioning strategy, the task mapping location is the entire card, wherein (K×N) inter-core tensors on the entire card are parallelized, and the single-core task computation amount allocated to the single core is determined to be (1 / (K×N)) of D3 heads, where D3=Z; the KV Cache size stored on each single core is (1 / (K×N)) of E3 heads of data, where E3=W; Where ceil() is the floor function; Z is the number of Q heads allocated to each card, Z=X / TP, where X is the total number of Q heads and TP is the tensor parallelism; W is the number of KV heads allocated to each card, where W=Y / TP, where Y is the total number of KV heads and N is the number of cores on each card; For each pre-configured task segmentation strategy, based on the self-attention model structure and the single-core task computational volume, the computational scale of the self-attention model structure under each task segmentation strategy is determined respectively. Based on the scale to be computed, the computation delay of each task segmentation strategy is determined, and based on the computation delay of each task segmentation strategy, the optimal task segmentation strategy corresponding to the decoding stage is determined, so as to segment the self-attention task according to the optimal task segmentation strategy.

6. A computer device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the computer program, it implements a self-attention task dynamic segmentation method according to any one of claims 1 to 4.

7. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed by a processor, implements a self-attention task dynamic segmentation method as described in any one of claims 1 to 4.

Images