Segmentation pre-filling method, system and equipment for reasoning video memory and medium

Abstract

The invention discloses a pre-filling segmentation method, system and device for an inference video memory, and a medium. The method comprises the following steps: segmenting to-be-input target segmentation data with the same length according to a sequence model architecture based on an attention mechanism in a pre-filling stage; after segmentation, multiple segments of input segmentation data are generated, and the multiple segments of generated segmentation data are input in sequence to complete input of target segmentation data; wherein the cache key value of the later segmentation data before the pre-filling is updated in sequence according to the previous segmentation data after the pre-filling, and the later segmentation data is pre-filled; and generating a first basic unit of the target segmentation data according to the cache key values of the previous segmentation data and the later segmentation data in sequence until the multiple segments of segmentation data generate all basic units input by the target segmentation data. According to the method, the pre-filling stage is optimized, the processing speed of the super-long text in the pre-filling stage is increased by segmenting the segmented data with the same length and a compiling technology, video memory occupation is reduced, and the processing cost of the super-long text is reduced.

Description

Technical Field

[0001] The present invention belongs to the technical field of artificial intelligence large models, and particularly relates to a pre-fill segmentation method, system, device and medium for inference video memory. Background Art

[0002] In the technical field of artificial intelligence large models, large models face two major challenges when processing long texts: slow inference speed and high video memory occupancy. Although traditional inference methods mainly focus on the optimization of the attention module (attention), such as the combination of flash attention, longformer, linformer and reformer. These architectures and technologies effectively reduce the video memory occupancy of the model during training and inference. However, the sequence model architecture (transformer) of the traditional attention mechanism generally has a linear relationship between the video memory occupancy during inference and the square of the length of the input text. The input of ultra-long texts makes the video memory demand far exceed its carrying capacity. In addition, when the input text becomes ultra-long, the inference speed slows down, seriously affecting the video memory efficiency of the large model when processing ultra-long texts.

[0003] For example, (1) The inference of the model is divided into a pre-fill stage and a generation stage. The combination of the sequence model architecture (transformer) of the traditional attention mechanism and the attention module (attention) is simple and easy to use. However, as the length of the text input increases, the video memory demand of the attention module (attention) increases in a square relationship, resulting in slow processing speed and intolerance to the peak consumption of video memory. (2) The transformer combined with the attention optimization model (flash attention) has improvements in both speed and video memory occupancy. However, during the pre-fill stage and the generation stage, the video memory consumption maintains a linear relationship with the length of the input text, and it is difficult to meet the processing requirements for ultra-long texts exceeding 200k.

[0004] (3) The sliding window works by sliding to fix the attention length of the attention module, and can fix the video memory in a certain range during the generation stage without being affected by the length of the input text, thereby accelerating the processing speed. However, it needs to sacrifice a certain amount of accuracy, and during the pre-fill stage, the video memory consumption still has a linear relationship with the length of the input text, and it is impossible to effectively reduce the peak video memory.

[0005] (4) The quantization model reduces memory requirements and improves processing speed by converting the weights and activation values of the model into a lower-precision data type (such as 8-bit or 4-bit). This quantization may result in a certain loss of accuracy and may not achieve the expected effect in some cases.

[0006] Therefore, it is particularly necessary to solve the problems of how to improve the processing speed of ultra-long texts in the pre-filling stage, reduce the video memory occupancy, and reduce the processing cost of ultra-long texts. Accordingly, the present application provides an optimized pre-filling segmentation method, system, device and medium for the inference video memory of large models for ultra-long texts. Summary of the Invention

[0007] The present invention provides a pre-filling segmentation method, system, device and medium for inference video memory, which is used to solve the problems of slow processing speed, large video memory occupancy and high processing cost of traditional ultra-long texts in the pre-filling stage. The present invention can perform corresponding optimizations according to the model architecture, and improve the processing speed of ultra-long texts in the pre-filling stage, reduce the video memory occupancy, and reduce the processing cost of ultra-long texts through segmentation and compilation techniques in the pre-filling stage of inference.

[0008] The object of the present invention and the technical problems to be solved are achieved by the following technical solutions.

[0009] The present invention provides a pre-filling segmentation method for inference video memory. The pre-filling segmentation method includes: based on the sequence model architecture of the attention mechanism, the target segmentation data to be input is sequentially segmented by length in the pre-filling stage; after segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture to complete the input of the target segmentation data; wherein, according to the previous segmentation data after pre-filling, the cache key values of the subsequent segmentation data before pre-filling are sequentially updated for pre-filling the subsequent segmentation data; based on the generation module of the sequence model architecture, the first basic unit of the target segmentation data is sequentially generated according to the cache key values of the previous segmentation data and the subsequent segmentation data.

[0010] As an optional implementation, after generating multiple segments of segmentation data by the segmentation, it further includes: initializing the cache key value of the first segment of the input segmentation data, and initializing the cache key value of the first segment of the segmentation data input into the sequence model architecture to an integer multiple of the length of each segment of the segmentation data input into the sequence model architecture.

[0011] As an optional implementation, before sequentially updating the cache key values of the subsequent segmentation data before pre-filling according to the previous segmentation data after pre-filling for pre-filling the subsequent segmentation data, it further includes: obtaining the cache key value of the first segment of the segmentation data based on the pre-filling module.

[0012] As an optional implementation, before obtaining the cache key value of the first segmented data, the method further includes: compiling the pre-filling module; and updating the cache key value of the first segmented data after initialization based on the compiled pre-filling module according to the first segmented data.

[0013] As an optional implementation, the length of each segmented data segment is the same as the length of the sliding window.

[0014] As an optional implementation, the corresponding pre-filled module of each segmented data input is compiled based on the deep learning model library to fix the shape of each basic unit input variable.

[0015] As an optional implementation, each segmented data is optimized and calculated based on the attention optimization module; transfer calculation is performed based on each segmented data, and the value of the classification memory occupancy function of each basic unit is gradually calculated; and the memory peak value is reduced based on the value operation of the classification memory occupancy function of each basic unit.

[0016] As an optional implementation, the weights and activation values of the trained sliding window model are quantized and updated based on the quantization model.

[0017] The present invention also provides a pre-filling segmentation system for inference video memory, which includes: an input unit, which is used to sequentially split the target segmentation data to be input into lengths based on the sequence model architecture of the attention mechanism during the pre-filling stage; after the segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture to complete the input of the target segmentation data; wherein, the cache key value of the subsequent segmentation data before pre-filling is updated in sequence according to the previous segmentation data after pre-filling, so as to pre-fill the subsequent segmentation data; a generation unit, which is used to generate the first basic unit of the target segmentation data in sequence according to the cache key values of the previous segmentation data and the subsequent segmentation data based on the generation module of the sequence model architecture.

[0018] The present invention also provides an electronic device, which includes: a memory for storing non-temporary computer-readable instructions; and a processor for running the computer-readable instructions, so that when the computer-readable instructions are executed by the processor, the pre-filling segmentation method described above is implemented.

[0019] The present invention also provides a computer-readable storage medium, which includes computer instructions. When the computer instructions are executed on a device, the device executes the above-mentioned pre-filling segmentation method.

[0020] Compared with the prior art, the present invention has obvious advantages and beneficial effects. Based on the above technical solution, the present invention has at least one of the following advantages and effects:

[0021] 1. A pre - filling segmentation method for inference video memory provided by the present invention, the pre - filling segmentation method includes: the sequence model architecture based on the attention mechanism segments the length of the target segmentation data to be input in sequence during the pre - filling stage; after segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture in sequence to complete the input of the target segmentation data; wherein, according to the previous segmentation data after pre - filling, the cache key value of the next segmentation data before pre - filling is updated in sequence for pre - filling the next segmentation data; the generation module based on the sequence model architecture sequentially generates the first basic unit of the target segmentation data according to the cache key values of the previous segmentation data and the next segmentation data until all the basic units of the target segmentation data are generated from the multiple segments of segmentation data and input into the sequence model architecture. The pre - filling segmentation method of the present invention improves the processing speed of ultra - long text in the pre - filling stage, reduces video memory occupancy, and reduces the cost of ultra - long text processing through segmentation and compilation techniques by optimizing the model architecture accordingly during the pre - filling stage of inference.

[0022] 2. The present invention also provides a pre - filling segmentation system for inference video memory, the pre - filling segmentation system includes: an input unit, which segments the length of the target segmentation data to be input in sequence based on the sequence model architecture of the attention mechanism during the pre - filling stage; after segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture in sequence to complete the input of the target segmentation data; wherein, according to the previous segmentation data after pre - filling, the cache key value of the next segmentation data before pre - filling is updated in sequence for pre - filling the next segmentation data; a generation unit, which sequentially generates the first basic unit of the target segmentation data according to the cache key values of the previous segmentation data and the next segmentation data based on the generation module of the sequence model architecture until all the basic units of the target segmentation data are generated from the multiple segments of segmentation data and input into the sequence model architecture. The pre - filling segmentation system of the present invention improves the processing speed of ultra - long text in the pre - filling stage, reduces video memory occupancy, and reduces the cost of ultra - long text processing through segmentation and compilation techniques by optimizing the input unit of the model architecture accordingly during the pre - filling stage of inference.

[0023] 3. The present invention also provides an electronic device, the electronic device includes: a memory for storing non - transient computer - readable instructions; and a processor for running the computer - readable instructions, so that when the computer - readable instructions are executed by the processor, the above - mentioned pre - filling segmentation method is implemented. The electronic device of the present invention improves the processing speed of ultra - long text in the pre - filling stage, reduces video memory occupancy, and reduces the cost of ultra - long text processing through segmentation and compilation techniques by optimizing the input unit of the model architecture accordingly during the pre - filling stage of inference.

[0024] IV. The present invention also provides a computer-readable storage medium, which includes computer instructions. When the computer instructions run on a device, the device is caused to execute the pre-filling segmentation method described above. By correspondingly optimizing the input unit of the model architecture, the computer-readable storage medium of the present invention improves the processing speed of ultra-long texts in the pre-filling stage through segmentation and compilation techniques during the pre-filling stage of inference, reduces the video memory occupancy, and reduces the cost of processing ultra-long texts.

[0025] The above description is only an overview of the technical solution of the present invention. In order to be able to understand the technical means of the present invention more clearly, it can be implemented according to the content of the specification. And in order to make the above structure and other purposes, features and advantages of the present invention more obvious and understandable, the following specifically gives preferred embodiments and, in conjunction with the drawings, details are described as follows. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 It is a schematic flowchart of a pre-filling segmentation method for inference video memory in this embodiment.

[0027] Figure 2 It is a schematic flowchart of another pre-filling segmentation method for inference video memory in this embodiment.

[0028] Figure 3 It is a schematic flowchart of a pre-filling segmentation method for inference video memory of a transformer model in this embodiment.

[0029] Figure 4 It is a schematic flowchart of a pre-filling segmentation method for inference video memory of a flash attention model in this embodiment.

[0030] Figure 5 It is a schematic flowchart of a pre-filling segmentation method for inference video memory of a transformer + slding window model in this embodiment.

[0031] Figure 6 It is a schematic flowchart of a pre-filling segmentation method for inference video memory of a flash attention + slding window model in this embodiment.

[0032] Figure 7 It is a schematic flowchart of a pre-filling segmentation method for inference video memory of a quantification + sliding window model in this embodiment.

[0033] Figure 8Schematic diagram of the process structure of a prefill segmentation method for the inference video memory of the prefill chunk + sliding window model in this embodiment.

[0034] Figure 9 Schematic diagram for comparing the video memory occupancy of the prefill segmentation method for each model inference in the above embodiments of the present invention.

[0035] Figure 10 Schematic diagram for comparing the inference speed of the prefill segmentation method for each model inference in the above embodiments of the present invention.

[0036] Figure 11 Schematic diagram of the structure of a prefill segmentation system for inference video memory in this embodiment.

[0037] Figure 12 Schematic diagram of the structure of an electronic device in this embodiment.

[0038] Explanation of the reference numerals in the drawings:

[0039] 200: Prefill segmentation system 210: Input unit

[0040] 220: Generation unit 300: Electronic device

[0041] 310: Memory 320: Processor

[0042] 330: Computer-readable instructions Detailed implementation manners

[0043] To further elaborate on the technical means and effects adopted by the present invention to achieve the intended invention purpose, the following describes in detail the specific implementation manners, structures, features, and effects thereof according to the present invention in combination with the accompanying drawings and preferred embodiments.

[0044] The present invention provides a prefill segmentation method for inference video memory. As Figure 1 shown, the prefill segmentation method includes the following processes:

[0045] S1: In the prefill stage, the sequence model architecture based on the attention mechanism sequentially segments the length of the target segmentation data to be input; after segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture in order to complete the input of the target segmentation data; among them, according to the previous segmentation data after prefill, the cache key value of the next segmentation data before prefill is sequentially updated for pre-filling the next segmentation data.

[0046] S2: The generation module based on the sequence model architecture sequentially generates the first basic unit of the target segmentation data according to the cache key values of the previous segmentation data and the next segmentation data until all the basic units of the multi-segment segmentation data are generated as the input sequence model architecture of the target segmentation data.

[0047] Specifically, for example, the sequence model architecture of the attention mechanism can select the transformer model. The target objects for segmentation of the target segmentation data in the pre-filling stage are not limited to including any one of text data, image data, audio data, and video data (target segmentation data of a single data type), or target segmentation data of a mixed data type composed of multiple types among text data, image data, audio data, and video data. For example, when the target segmentation data is of a mixed data type, it can be segmented based on the segmentation length in the target segmentation data. For example, when segmenting text data, the storage space value occupied by the text data is used as the segmentation length for segmentation; when segmenting image data, the pixel value of the image in the target segmentation data is used as the segmentation length for segmentation; when segmenting audio data or video data, the duration of the audio data or video data in the target segmentation data can be used as the segmentation length for segmentation; or when segmenting video data, a certain number of frames can also be used as the segmentation length for segmentation. The segmentation of the target segmentation data is based on the application of the actual data type. The above segmentation lengths can be in an equal-length manner or a non-equal-length manner, which will not be elaborated here.

[0048] In the present invention, as Figure 2As shown, the input target segmentation data T needs to be explained. The above-mentioned target segmentation data T can be an input text data. Based on the type of the target segmentation data (text data), the segmentation length type (by the storage space value occupied by the text data) is selected, and the target segmentation data T is segmented into n segments, denoted as C1 to Cn respectively. The length of each segmented data after segmentation is t, and the cache key value (kvcache) before segmentation in the transformer model is defined as an integer multiple of t. After the segmentation operation is completed, an initialization operation is performed on the cache key value Cn kv cache before segmentation. The length of the cache key value of C1 is initialized to N (N is divisible by t) so that the cache key value of the first segmented data can be updated smoothly. Execute 1. Input C1 and the cache key value of C1 after initialization into the prefill module of the transformer model in sequence, generate and update C1, and update the generated C1 after prefill into the cache key value (kv cache) of C2; Execute 2. Input C2 and the cache key value C2 kv cache of C2 into the prefill module of the transformer model in sequence, generate and update C2, and update the generated C2 after prefill into the cache key value (kv cache) of C3 for prefill of C3; Execute 3. Input C3 and the cache key value C3 kv cache of C3 into the prefill module of the transformer model in sequence, generate and update C3, and update the generated C3 after prefill into the cache key value (kv cache) of the next segmented data for prefill of the next segmented data,... Execute n. Input the nth Cn and the cache key value Cn kv cache of Cn into the prefill module of the transformer model in sequence, generate and update Cn, and update the generated Cn after prefill until the prefill module of the transformer model completes the prefill of each segmented data in the target segmentation data T in sequence according to the segmentation order, obtaining the first basic unit (token) of the target segmentation data T and the cache key value (kv cache) of each segmented data Cn kv cache.The first basic unit (token) obtained according to the segmentation order and the subsequent segmentation data Cn kv cache corresponding to the first basic unit (token) are input into the generation module of the Transformer model, and the subsequent basic unit (token) is returned. Then, the above subsequent basic unit (token) and the subsequent segmentation data Cn kv cache corresponding to the subsequent basic unit (token) are input into the generation module of the transformer model, and the next basic unit (token) is sequentially returned in order, …, until the generation module sequentially returns all L basic units (tokens) of the input target segmentation data T and sequentially outputs the generation results R of the L basic units (tokens). The result R includes the outputs of the L basic units (tokens) output. It should be further noted that the token is the smallest unit for the transformer model to encode the target segmentation data, and the kv cache is a certain intermediate vector generated by the previous segmentation data of the target segmentation data during the inference of the above model.

[0049] The optimization of the present invention for the prefill stage is mainly based on the operation of the input data T. The input data of length T is divided into segmentation data with each segment of length t. By sharing the cache key-value storage area of each segment of the segmented data and adopting the update and reuse method, on the one hand, the problem of generating the cache key-value storage in the prefill stage is overcome, and on the other hand, the storage of the cache key-value of each segment of the segmented data in the prefill segment during the inference process is saved by optimizing the cache key-value storage method, reducing the situation of over-occupying the peak video memory in the prefill stage, and improving the filling efficiency of the target segmentation data in the prefill stage. Further, by reducing the video memory requirement in the prefill stage, the model inference process can run more efficiently, especially in an environment with limited video memory resources.

[0050] As an optional implementation manner, after generating multiple segments of segmentation data after the segmentation, it further includes: initializing the cache key-value of the first segment of the input segmentation data, and initializing the cache key-value of the first segment of the input segmentation data of the sequence model architecture to an integer multiple of the length of each segment of the input segmentation data of the sequence model architecture.

[0051] It should be noted that in the sequence model architecture transformer, after the input data is segmented into n segments according to the preset length t, multiple segments of segmented data C1 to Cn are obtained. After the segmentation is completed, generally, the cache key-value of the first segment of the segmented data needs to be initialized. The cache key-value Cn kv cache of the number of segments n is initialized to N (N is divisible by t). When t is a positive integer after the segmentation, the length N of the initialized cache key-value Cn kv cache is also a positive integer, effectively ensuring that each segment of the segmented data can be accurately recorded and stored by the initialized cache key-value Cn kv cache without loss of precision.

[0052] As an alternative implementation, before updating the cache key value (kv cache) of the subsequent segmentation data before pre-filling in sequence according to the pre-filled previous segmentation data, and for pre-filling the subsequent segmentation data, it further includes: based on the pre-filling module, obtaining the cache key value (kv cache) of the first-segment segmentation data from the first-segment segmentation data.

[0053] It should be noted that in the sequence model architecture transformer, the cache key value (kvcache) before splitting the input data is defined as C0. The initial value of the length of C0 before splitting the input data is generally 0, and the length of the cache key value corresponding to C0 is also 0. After the input data is split into n segments according to the preset length t, the segmented multi-segment segmentation data C1 to Cn are obtained. After the splitting is completed, generally, the cache key value of the first-segment segmentation data needs to be initialized. During the initialization process, the cache key value Cn kv cache of the number of segments n is initialized to N (N is divisible by t). Specifically, based on C0, the cache key value of C1 is updated, and the length of the C1 cache key value (C1 kv cache) is updated to N (N is divisible by t). In the present invention, the cache key value (kv cache) before splitting the input data is defined as C0, and the cache key value of the subsequent segmentation data before pre-filling is updated in sequence according to the pre-filled previous segmentation data for pre-filling the subsequent segmentation data; thus, the cache key value C1 kvcache of the first-segment segmentation data C1 is obtained, based on C1 kv cache and C1, the pre-filled C1 is obtained by inputting into the pre-filling module, and based on the pre-filled C1, the cache key value C2 kv cache of C2 before pre-filling is updated for the pre-filling operation of the subsequent non-pre-filled segmentation data C2, so as to ensure that the cache key value of the subsequent segmentation data before pre-filling is updated in sequence according to the pre-filled previous segmentation data for pre-filling the subsequent segmentation data until the pre-filling operations of all the segmentation data with the number of segments n are completed.

[0054] As an alternative implementation, before obtaining the cache key value (C1kv cache) of the first-segment segmentation data, it further includes: compiling the pre-filling module; updating the cache key value after initialization of the first-segment segmentation data based on the compiled pre-filling module according to the first-segment segmentation data.

[0055] It should be noted that before pre-filling the n segmented data in the pre-fill module, the pre-fill module is further compiled. For example, but not limited to, in the compiled Transformer, the pre-fill chunk module is compiled based on the n segmented data for feature extraction of the input n segmented data, providing effective semantic information for the decoding process. Overall, the compilation of the pre-fill chunk module can be stacked by n encoding operations. The following is a simple example showing how to define an encoding module and stack and configure the encoding structure of the encoding module.

[0056]

[0057] The above xi is the output of the i-th segmented data layer, and feature extraction is performed based on the input of the segmented data layer. The above feature extraction can be based on positional encoding or mask tokens. The information of positional encoding is fused into the features of the segmented data layer to solve the accurate synchronous operation of the parallel inference of all n segmented data layers in the Transformer when input simultaneously. The above positional encoding can have various ways, not limited to fixed positional encoding and encoding set as learnable parameters. Mask tokens can make the features fused in the features of each segmented data layer be marked with a unique encoding, and the difference between any two basic units (tokens) remains consistent, making the value of each basic unit (token) have a clear boundary. Therefore, the input of the final model includes a comprehensive feature information of the current i-th segmented data layer, that is, it contains its own semantic information, and also contains the position information and masked token information of the current i-th segmented data layer in the whole sentence. It should be noted that when building the compilation operation in the attention module, it is not limited to using multi-head attention, and multiple attention modules can be combined together to form a stacked configuration encoding module.

[0058] As an optional implementation manner, the length of each segmented data is the same as the length of the sliding window.

[0059] It should be noted that the length of each segmented data in the above real-time manner can be selected differently (specifically, it can be determined according to the actual length of each segmented data information), or it can be selected that the length of each segmented data is the same (such as the above embodiment). In the embodiment where the length of each segmented data is the same, in the pre-filling stage, the sliding window attention is used to relatively slide in each segmented data with the same length of the sliding window attention in sequence from the first segmented data of each segmented data, and one sliding covers one segmented data until all features of the entire input data T are covered. This sliding window attention operation can not only avoid the high computing cost of global attention during the pre-filling stage from generating an excessive video memory peak and occupying the video memory resources, but also ensure that the local information of each segmented data of the input data T is fully captured by the sliding window attention operation, avoiding the phenomenon of screen freezing caused by the over-occupation of the video memory peak during the capture of the local information of the input data T, and even possibly causing the occurrence of picture errors or program crashes, thus improving the efficiency of capturing the local information of the input data T.

[0060] As an alternative embodiment, based on the deep learning model library (Torch), the corresponding pre-fill module (prefill) for each segmented data input after segmentation is compiled to fix the shape of the input variables of each basic unit (token).

[0061] It should be noted that the deep learning model library Torch is the abbreviation of the open-source deep learning framework (PyTorch) for machine learning and deep learning, and it is an open-source machine learning library mainly used for the construction, training and application of deep learning models. For example, Torch models are constructed and trained in fields such as images, speech, video, natural language processing, and autonomous driving, and it has good flexibility and scalability. In this embodiment, an image is selected as an example for illustration.

[0062] In PyTorch, torch.compile is an experimental API for inference, which can compile the pre-filled (prefill) model into optimized running code with fixed input variable shapes to improve the efficient performance of inference. In the present invention, by creating a specific input object (dummy_input) whose shape is consistent with the expected input, and then passing it to the interface (API) of torch.compile, an example of compiling and fixing the shape of the input variables in the pre-filling stage is as follows:

[0063]

[0064] Based on the above, inference is performed using the corresponding prefill model obtained after compilation. This corresponding prefill model can optimize the variable shapes input for subsequent prefill operations based on the shape information prefilled by the input object (dummy_input). In the above example, the shape of dummy_input is preset to the shape of a grayscale image (number of channels, height, width). When torch.compile is called, the compiler uses the shape of the above grayscale image to optimize the execution process of the model, so that when subsequent inference is performed through the compiled_model, it is not necessary to specify the input shape each time, thereby improving the prefill efficiency in the prefill stage during the inference process.

[0065] As an alternative implementation, each split data of the computational split is optimized based on the attention optimization module (flash attention); transfer calculations are performed based on each split data of the split, and the values of the classification video memory occupancy function (softmax) of each basic unit (token) are gradually calculated, and the video memory peak is reduced based on the values of the classification video memory occupancy function (softmax) of each basic unit (token).

[0066] It should be noted that in the present invention, by segmenting the above input text in the prefill stage and updating the results of each segment to the kv cache, by designing the segmentation method in the prefill stage, the improvement in the prefill stage can greatly reduce the video memory peak and save the video memory of the transformer model in the prefill stage. Generally, the video memory required by the transformer model mainly depends on the peak video memory, and the peak video memory is mainly concentrated in the classification video memory occupancy function (softmax) of the attention module attention. The video memory occupancy evaluation value Score of the softmax calculation is as shown in Equation (1):

[0067]

[0068] i and j are the i-th and j-th basic units (tokens) in the input data; the length of the input target split data is T, the number of heads of the transformer model is c, and the dimension of each head of the model is H. Generally, the model is fixed, and its video memory occupancy before splitting and segmenting is 3 * c * T 2 *H.

[0069] The optimization of the present invention for the pre-filling stage is mainly based on the operation of the input data T. The input data with a length of T is divided into segmented data with each segment having a length of t, so T = n * t. If the sliding window module is not used, after the segmentation based on the present invention, the video memory occupancy of the input data with a length of T is 3 * c * T * t * H, which is reduced by n times. If used in combination with the sliding window module, after the segmentation based on the present invention, the sliding window length is N (N < T), and the video memory occupancy is 3 * c * t * N * H. It can be seen that t, N are not related to the length of the input data T. That is to say, the above-mentioned inference model of the present invention can accept input data T of infinite length and efficiently complete the inference.

[0070] As an optional implementation manner, based on the quantization model, the weights and activation values of the trained sliding window model are quantized and updated.

[0071] It should be noted that in the neural network of the quantization model based on the sliding window, the floating-point weights and activations are converted into a lower precision (such as 8 bits or less) through the quantization model to reduce the overall size of the model, so as to improve the inference efficiency in the pre-filling stage. Further, by combining the sliding window, during the pre-filling stage, the convolution operation is performed by reducing the attention length of the fixed attention module, and the weights of the sliding window model are quantized, so as to fix the peak video memory in a preset range during the generation stage, which is not affected by the input length, thereby avoiding the over-occupation of the peak video memory and improving the processing efficiency. The Python-based example is as follows:

[0072]

[0073]

[0074] In the above embodiments, the present invention defines a QuantConv2d class that quantifies weights but not biases. By inheriting from nn.Conv2d, it adds quantization and dequantization operations. When actually quantizing and dequantizing QuantConv2d, weights and biases are set in combination with the sliding window. During the forward pass, first, the first segment of the input target split data T is quantized, then the step-by-step sliding quantization convolution operation is performed, and finally, the result of the quantized convolution operation is dequantized to obtain a non-quantized result for output. It should be noted that the above example is carried out when the inference experimental API exists, and PyTorch support for the corresponding API is required during actual use. During actual deployment, it may be necessary to use specialized hardware or quantization tools provided by a specific framework to implement the operations described above.

[0075] In the above embodiments of the present invention, an exemplary introduction is given below based on the implementation processes, environments, hyperparameters, and experimental results under different technologies:

[0076] Implementation environment: NVIDIA 3090 24G graphics card, the Volcano Engine of ByteDance's cloud service platform under the linux operating system (Linux ubuntu 18.04), computer program programming language and version (python = 3.10.10), deep learning model library and library class number (torch = 2.30), general parallel computing platform and programming model developed by NVIDIA and version (CUDA = 11.8).

[0077] When the sequence model architecture based on the attention mechanism, such as the transformer, is implemented alone as Figure 3 , the largest computational overhead of the transformer lies in the score n = softmax(q * k) in the attention module, where q represents the effective information of the split data sequence elements of the input target split data T, k represents the elements of the split data sequence of the input target split data T, v represents the feature information values related to each element in the split data sequence, and o is the output of the attention operation in the attention module. The lengths of q and k are the same as the length of the input target split data T. At this time, the video memory occupancy of the output o has a quadratic relationship with the input target split data T.

[0078] When the sequence model architecture based on the attention mechanism is implemented in combination with the attention optimization model (flash attention) as Figure 4As shown, the largest cost of flash attention calculation lies in the evaluation value score n = softmax(qi * [k1~kn]) of the segmented data operated in sequence in the attention module attention (it should be noted that the segmentation of this data only occurs during the qk matrix multiplication, which is different from the segmentation in the input stage of the present invention). qi ∈ [q1~qn], the length of each segmented data is a preset constant t, and the length of k1~kn is the length of the input target segmented data T. At this time, the video memory occupancy of o = score * v has a linear relationship with the input target segmented data T.

[0079] When implementing based on the sequence model architecture with attention mechanism combined with the attention sliding window model (transformer + slding window) as Figure 5 , the largest cost of the attention sliding window model transformer + slding window calculation lies in the softmax(q * k of the segmented data operated in sequence in the attention module attention c ), where k c represents the segmented data sequence of the input target segmented data T sliding to the corresponding segmented data element. The length of k c is the same as the length of the sliding window, and the length of q is the length of the input target segmented data T. At this time, the video memory occupancy of the output o = score * v has a linear relationship with the input T.

[0080] When implementing based on the sequence model architecture with attention mechanism combined with the attention optimized sliding window model (flash + slding window) as Figure 6 , generally speaking, the largest cost of flash + slding window calculation is in the softmax(qi * k in the attention module attention c ), qi ∈ [q1~qn], and qi is the effective information of the i-th input target segmented data sequence element of the target segmented data T. The length of qi is the same as the preset length t of each segment. Thus, the classification video memory occupancy function softmax(qi * k c ) is a fixed value and does not change with the change of the input length. However, the total video memory occupancy evaluation value score = cat[score 1, score 2,.., score n]. It can be seen that the length of score is the same as the length of the target segmented data T, that is, the output o = score * v will increase linearly with the change of the length of the input target segmented data T.

[0081] When implementing a sliding window model (quantification + sliding window) that combines an attention mechanism-based sequence model architecture with attention quantization processing, such as Figure 7 , generally speaking, the maximum overhead of quantification + sliding window calculation is in the softmax(q quant *k c_quant ) in the attention. Here, q quant represents the data of q after corresponding quantization processing, and k c_quant represents the data of k c after corresponding quantization processing, and v quant represents the data of v after corresponding quantization processing. In the original attention mechanism-based sequence model architecture transformer, the original q is of the float16 type. At this time, the lengths of q and k are the same as the length of the input target segmented data T. After quantization processing (taking float16 quantization processing to int8 as an example), the video memory required for each operation can be reduced by half. For example, the output of the prefill operation after quantization processing of softmax(q quant *k c_quant ) is reduced by half compared to the output of the prefill operation before quantization processing, but it still has a quadratic relationship with the target segmented data T.

[0082] When implementing a sliding window model (prefill chunk + sliding window) that combines an attention mechanism-based sequence model architecture with attention prefill segmentation, as Figure 8 shown, generally speaking, the maximum overhead of the prefill chunk + sliding window model calculation is in the softmax(qi*k c ), where qi ∈ [q1~qn]. The length of qi is the same as the length t of each segment preset for segmentation, and the length of k c is a preset fixed value. At this time, softmax(qi*k c ) is also a fixed value and does not change with the length of the input target segmented data T. Only the input qi of the i-th segment of the segmented data in each segment is calculated during each sliding window. At this time, the obtained score i is also segmented, but finally, no cat[] connection calculation processing is performed. At this time, the output oi = scorei*v does not change with the change of the input length. In the present invention, by adopting the sliding window model with attention prefill segmentation, the video memory occupancy ratio in the prefill stage does not change with the change of the length of the input target segmented data T, effectively reducing the peak video memory and the video memory resource occupancy ratio when inputting an extremely long target segmented data.

[0083] In an embodiment of the present invention, the implementation of the environmental parameters in the hyperparameter state is as follows. Implementation environment: NVIDIA 3090 24G graphics card, Linux ubuntu 18.04, python = 3.10.10, torch = 2.30, CUDA version is 11.8.

[0084] Hyperparameter implementation example data:

[0085] The model has approximately 300 million parameters (0.3B). The inference is performed using the inference mode of the attention prefill chunk + sliding window model of the present invention (prefill chunk + sliding window). The preset number of generated tokens is 1, and the input length range is from 256 to 10 5 any value in. The batch size is 1, the size of each segment of the prefill chunk is 256, the chunk size of q and k in the flash attention model is 256, and it is quantized with int8. The sliding window size is 4096 for inference. The comparison of the video memory (HBM) occupancy represented by the output results of the above model experiments is as follows Figure 9 The comparison of the video memory (HBM) occupancy shown.

[0086] (1) From Figure 9 It can be seen that the prefill chunk + sliding window of the present invention saves the most video memory. As the length changes, the output always maintains the same video memory occupancy ratio. That is, the attention prefill chunk + sliding window model of the present invention can be applied when the input target segmented data T is of infinite length, and the processing conditions depend on the total amount of video memory resources, and have no correlation with the length of the input target segmented data T.

[0087] (2) From Figure 9It can be seen that in the present invention, the compilation of the prefill module (torch.compile) is adopted from a deep learning model library to improve the computational efficiency of the operations in the prefill stage. The present invention overcomes the problem that the traditional prefill stage cannot fix the shape of the input variables during compilation, resulting in the inability to compile when the traditional prefill operation is combined with torch.compile, and thus it cannot be applied to the scenario with variable input lengths in combination with torch.compile. In the above implementation of the present invention, the prefill segmentation exactly meets the requirements of fixed-length input, so we can use torch.compile to accelerate the calculation process. By applying torch.compile in this way to the ultra-long target segmentation data with variable input lengths, the present invention can not only improve the processing speed of inference in the prefill stage, but also maintain the flexibility and adaptability of the model combination to adapt to the ultra-long target segmentation data with the length of the input target segmentation data T being ultra-long.

[0088] In the present invention, the speed comparison of the output results of the above model experiments is as Figure 10 shown. From Figure 10 it can be seen that although the prefill segmentation compilation, flash attention processing, and quantization processing are slightly slower in terms of the speed of result output, it is found by comparison that the difference is not significant, which well maintains the speed of the above inference processing and keeps the high-efficiency inference performance of the model of the present invention.

[0089] The present invention also provides a prefill segmentation system for inference video memory, as Figure 11 shown. The prefill segmentation system 200 includes: an input unit 210, which is used to sequentially segment the length of the target segmentation data to be input during the prefill stage based on the sequence model architecture (transformer) of the attention mechanism; after segmentation, multiple segments of segmentation data are generated, and the generated multiple segments of segmentation data are sequentially input into the sequence model architecture (transformer) in sequence to complete the input of the target segmentation data; wherein, the cache key-value (kvcache) of the subsequent segmentation data before prefill is sequentially updated according to the previous segmentation data after prefill for pre-filling the subsequent segmentation data; a generation unit 220, which is used to sequentially generate the first basic unit (token) of the target segmentation data based on the generation module of the sequence model architecture (transformer) according to the cache key-values of the previous segmentation data and the subsequent segmentation data until all the basic units (tokens) of the target segmentation data are generated for input into the sequence model architecture from the multiple segments of segmentation data.

[0090] It should be noted that when the input unit 210 and the generation unit 220 of the pre-filling segmentation system 200 of the present invention are running, by performing the above corresponding optimizations on the input unit of the model architecture, the processing speed of ultra-long texts in the pre-filling stage is improved through the above segmentation and compilation technologies during the pre-filling stage of inference, the video memory occupancy is reduced, and the cost of processing ultra-long texts is reduced.

[0091] The present invention also provides an electronic device, as Figure 12 shown. The electronic device 300 includes: a memory 310 for storing non-transitory computer-readable instructions 330; and a processor 320 for running the computer-readable instructions 330, so that when the computer-readable instructions 330 are executed by the processor 320, the pre-filling segmentation method described above is implemented. When the electronic device of the present invention is running, by performing the above corresponding optimizations on the input unit of the model architecture, the processing speed of ultra-long texts in the pre-filling stage is improved through the above segmentation and compilation technologies during the pre-filling stage of inference, the video memory occupancy is reduced, and the cost of processing ultra-long texts is reduced.

[0092] The present invention also provides a computer-readable storage medium, which includes computer instructions. When the computer instructions are run on a device, the device is caused to execute the pre-filling segmentation method described above. When the computer-readable storage medium of the present invention is running, by performing the above corresponding optimizations on the input unit of the model architecture, the processing speed of ultra-long texts in the pre-filling stage is improved through the above segmentation and compilation technologies during the pre-filling stage of inference, the video memory occupancy is reduced, and the cost of processing ultra-long texts is reduced.

[0093] It should be noted that any process or method description shown in a flowchart or described in other ways herein can be understood to represent a module, segment, or part of code including one or more executable instructions for implementing a specific logical function or process. And the scope of the preferred embodiments of the present invention includes additional implementations, where the functions can be executed in a way that is not shown or discussed, including in a substantially simultaneous manner or in a reverse order according to the involved functions, which should be understood by those skilled in the technical field to which the embodiments of the present invention belong.

[0094] The logic and / or steps represented in the flowchart or otherwise described herein can, for example, be considered a definitional sequence list of executable instructions for implementing logical functions, and can be embodied specifically in any computer-readable medium for use by an instruction execution system, apparatus, or device (such as a computer-based system, a system including a processor, or other systems that can fetch and execute instructions from the instruction execution system, apparatus, or device), or in connection with these instruction execution systems, apparatus, or devices. For the purposes of this specification, a "computer-readable medium" can be any device that can contain, store, communicate, propagate, or transport a program for use by or in connection with an instruction execution system, apparatus, or device. More specific examples (non-exhaustive list) of computer-readable media include the following: an electrical connection portion with one or more wirings (electronic device), a portable computer diskette (magnetic device), a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), an optical fiber device, and a portable compact disc read-only memory (CDROM). Additionally, a computer-readable medium can even be paper or other suitable media on which a program can be printed, as the program can be obtained electronically, for example, by optically scanning the paper or other media, followed by editing, interpretation, or other appropriate processing as necessary, and then stored in a computer memory.

[0095] It should be understood that various parts of the present invention can be implemented by hardware, software, firmware, or a combination thereof. In the above-described embodiments, multiple steps or methods can be implemented by software or firmware stored in a memory and executed by a suitable instruction execution system. For example, if implemented in hardware, as in another embodiment, any one or a combination of the following techniques well known in the art can be used: discrete logic circuits having logic gate circuits for implementing logical functions on data signals, application-specific integrated circuits having appropriate combinational logic gate circuits, programmable gate arrays (PGAs), field-programmable gate arrays (FPGAs), etc.

[0096] Those of ordinary skill in the art of this technology can understand that all or part of the steps included in the methods for implementing the above embodiments can be completed by instructing relevant hardware through a program, and the said program can be stored in a computer-readable storage medium. When the program is executed, it includes one or a combination of the steps of the method embodiments.

[0097] In addition, in each embodiment of the present invention, each functional unit may be integrated into a processing module, may exist physically alone for each unit, or two or more units may be integrated into one module. The above-mentioned integrated module may be implemented in part in the form of hardware or in the form of a software functional module. When the integrated module is implemented in the form of a software functional module and sold or used as an independent product, it may also be stored in a computer-readable storage medium.

[0098] The above-mentioned storage medium may be a read-only memory, a magnetic disk, an optical disc, etc. Although the embodiments of the present invention have been shown and described above, it can be understood that the above embodiments are exemplary and should not be construed as limiting the present invention. Those of ordinary skill in the art can make changes, modifications, substitutions, and variations to the above embodiments within the scope of the present invention.

[0099] The above are only the preferred embodiments of the present application, and do not impose any formal limitations on the present application. Although the present application has been disclosed above with the preferred embodiments, it is not intended to limit the present application. Any person skilled in the art can make some changes or modifications to the above-disclosed technical content to form equivalent embodiments with equivalent changes within the scope of the technical solution of the present application. However, as long as it does not depart from the technical content of the technical solution of the present application, any simple modification, equivalent change, and modification made to the above embodiments based on the technical essence of the present application still fall within the scope of the technical solution of the present application.

Claims

1. A pre-filling segmentation method for inference video memory, characterized in that, Including: In the pre-filling stage, the sequence model architecture based on the attention mechanism sequentially segments the length of the target segmentation data to be input; After segmentation, multiple segments of segmented data are generated, and the generated multiple segments of segmented data are sequentially input into the sequence model architecture to complete the input of the target segmentation data; among them, according to the previous segmented data after pre-filling, the cache key values of the subsequent segmented data before pre-filling are sequentially updated for pre-filling the subsequent segmented data; The generation module based on the sequence model architecture sequentially generates the first basic unit of the target segmentation data according to the cache key values of the previous segmented data and the subsequent segmented data.

2. The pre-filling segmented method according to claim 1, wherein After generating multiple segments of segmented data after the segmentation, it further includes: Initializing the cache key value of the first segment of segmented data input, and initializing the cache key value of the first segment of segmented data input into the sequence model architecture to an integer multiple of the length of each segment of segmented data input into the sequence model architecture.

3. The pre-filled segmentation method according to claim 2, wherein Before updating the cache key values of the subsequent segmented data before pre-filling according to the previous segmented data after pre-filling for pre-filling the subsequent segmented data, it further includes: Based on the pre-filling module, obtaining the cache key value of the first segment of segmented data.

4. The pre-filling segmented method according to claim 3, characterized in that, Before obtaining the cache key value of the first segment of segmented data, it further includes: Compiling the pre-filling module; Based on the first segment of segmented data, updating the initialized cache key value of the first segment of segmented data based on the compiled pre-filling module.

5. The pre-filling segmentation method according to claim 1, wherein: The length of each segment of segmented data is the same as the length of the sliding window.

6. The pre-filling segmentation method according to claim 5, wherein: Based on the deep learning model library, compiling the corresponding pre-filling module for each segmented data input after segmentation to fix the shape of the input variables of each basic unit.

7. The pre-filling segmentation method according to claim 5, wherein: Optimizing and calculating each segmented data after segmentation based on the attention optimization module; Performing transfer calculation based on each segmented data after segmentation, gradually calculating the value of the classification video memory occupancy function of each basic unit, and operating to reduce the video memory peak value based on the value of the classification video memory occupancy function of each basic unit.

8. The pre-filling segmentation method according to claim 5, wherein: Based on the quantization model, performing quantization processing and updating on the weights and activation values of the trained sliding window model.

9. A pre-fill segmented system for inference video memory, characterized in that, Including: An input unit, configured to sequentially segment the length of the target segmentation data to be input in the pre-filling stage based on the sequence model architecture of the attention mechanism; After segmentation, multiple segments of segmented data are generated, and the generated multiple segments of segmented data are sequentially input into the sequence model architecture to complete the input of the target segmentation data; among them, according to the previous segmented data after pre-filling, the cache key values of the subsequent segmented data before pre-filling are sequentially updated for pre-filling the subsequent segmented data; A generation unit, configured to sequentially generate the first basic unit of the target segmentation data according to the cache key values of the previous segmented data and the subsequent segmented data based on the generation module of the sequence model architecture.

10. An electronic device, characterized in that, Including: A memory, configured to store non-temporary computer-readable instructions; And A processor for running computer-readable instructions, such that when the computer-readable instructions are executed by the processor, the pre-filled segmentation method according to any one of claims 1 to 8 is implemented.

11. A computer-readable storage medium, characterized in that, Comprising computer instructions, such that when the computer instructions are run on a device, the device executes the pre-filled segmentation method according to any one of claims 1 to 8.

Images