Large model memory management method and device, electronic device, and readable storage medium

Abstract

The present invention discloses a memory management method and device for a large model, an electronic device, and a readable storage medium. The method comprises inputting data into a trained text length estimation model, estimating the length of the text output by the large model and adjusting the length upward to an integer, calculating the number of cache blocks based on the memory page size and the KV cache dimension, and adjusting the number of cache blocks upward to an integer; and finally allocating KV cache memory blocks for large model decoding. The present invention effectively avoids delays caused by dynamic adjustments by allocating sufficient video memory or internal memory in advance; in batch inference scenarios, it can reasonably plan computing resources and improve throughput; in streaming output scenarios, under word-by-word generation scenarios, the present invention's estimated output length can optimize the user experience, such as displaying a progress bar or early truncation processing.

Description

Technical Field

[0001] The present invention relates to the field of large model data generation and reasoning, and in particular to a memory management method and device for a large model, an electronic device, and a readable storage medium. Background Art

[0002] With the rapid advancement of science and technology, large-model inference technology, as one of the core technologies in the field of artificial intelligence, is showing broad application prospects. In practical application scenarios, optimizing the inference performance of large models faces multiple challenges, among which the memory access efficiency of the compute unit has become a key bottleneck restricting overall system performance. Research has shown that the inference speed of large models is not only limited by the theoretical computing power of the compute unit, but is also closely related to the memory access characteristics of the computing architecture. This correlation is mainly manifested in the following aspects: First, the large number of parameters in large models leads to complex memory access patterns; second, the data reuse nature of neural network calculations places special requirements on the storage hierarchy; third, bandwidth limitations between compute units and storage units can create a serious "memory wall" problem. Therefore, existing memory management methods for large models suffer from low memory access efficiency, high data transfer overhead, and slow overall inference speed for large models. Summary of the Invention

[0003] In response to the shortcomings of the existing technology, the present invention provides a memory management method and device, electronic device, and readable storage medium for large models, which can avoid problems such as repeated memory reallocation and discontinuous memory access when decoding large models, and improve the inference speed of large models.

[0004] The purpose of the present invention is achieved through the following technical solutions:

[0005] A memory management method for a large model includes the following steps:

[0006] Step 1: Build and train a text length estimation model;

[0007] The input of the text length estimation model is the text data and the related information of the large model, and the output is the length of the text output by the large model; the text length estimation model includes an input layer, an embedding layer, three serially connected MQA-based decoder structures, a third root mean square normalization layer, a linear projection layer and an output layer; the MQA-based decoder structure includes a first root mean square normalization layer, a multi-query attention layer, a first splicing layer, a second root mean square normalization layer, a feedforward neural network and a second splicing layer; the multi-query attention layer shares some attention heads, that is, shares Key and Value, to capture contextual relationships and learn information from different subspaces; at the same time, different information is assigned different weights, and the learned information and the corresponding weights are calculated, and the calculation results are output as features; the first splicing layer is used to splice the input of the first root mean square normalization layer and the output of the multi-query attention layer; the feedforward neural network is used to perform a nonlinear mapping operation on the normalized result of the second root mean square normalization layer; the second splicing layer is used to perform the same splicing operation as the first splicing layer on the output of the feedforward neural network and the output of the first splicing layer;

[0008] Step 2: Adjust the predicted text length output by the text length estimation model in step 1 upward to an integer;

[0009] Step 3: Calculate the number of cache blocks based on the memory page size and the KV cache dimension;

[0010] Step 4: Adjust the number of cache blocks upward to an integer;

[0011] Step 5: Allocate kv cache memory blocks for large model decoding.

[0012] Furthermore, the predicted text length output by the output layer is an exact length or a quantized interval;

[0013] If the text length estimation model outputs an accurate length prediction, the upward adjustment in step 2 is specifically as follows: first, the length value is rounded up, and then an additional buffer value is added;

[0014] If the text length estimation model outputs a quantization interval, the upward adjustment in step 2 is specifically as follows: first, the upper limit value of the quantization interval is rounded up, and then an additional buffer value is added.

[0015] Furthermore, the rounding up is rounding to units, hundreds or thousands, which is selected according to business needs; the buffer value is a fixed value, or a percentage of the predicted value.

[0016] Furthermore, in step 4, the number of cache blocks is adjusted upward to an integer by one of the following methods:

[0017] Method 1: round up the unit digit;

[0018] Method 2: After rounding up the units digit, add an additional security block number;

[0019] Method 3: Memory alignment optimization, specifically:

[0020] N'=ceil(N / A)×A

[0021] Where N' is the adjusted number of cache blocks, N is the number of cache blocks calculated in step 3, A is the alignment coefficient, and ceil() is the rounding function.

[0022] Method 4: For very long text generation tasks, adopt a dynamic chunking strategy:

[0023] N' = min( ceil(N) + S, N_max )

[0024] Where S represents the safety margin and N_max represents the maximum number of blocks available on the device.

[0025] Furthermore, the step 5 is specifically as follows: using cudaMalloc (GPU) or posix_memalign (CPU) to perform aligned memory allocation, with each block corresponding to the continuous space of the memory page size as the allocation unit, and allocating a total of N×D bytes of continuous memory space; N' cache blocks are arranged in sequence in the N×D bytes of continuous memory space, and each memory block contains:

[0026] Key cache: occupies 50% of the space of a single block, that is, D / 2 bytes;

[0027] Value buffer: occupies 50% of the space of a single block, that is, D / 2 bytes;

[0028] Meta-information header: 8-byte block status identifier.

[0029] Furthermore, after completing the aligned memory allocation in step 5, the following memory protection strategy is immediately executed:

[0030] (1) Pre-clear processing, execute memset clearing immediately after allocation to avoid numerical anomalies caused by uninitialized memory;

[0031] (2) Backup mechanism: Automatically switch to shard allocation mode when continuous memory is insufficient.

[0032] Furthermore, the memory block allocation in step five is asynchronous allocation and is executed in parallel with the computing pipeline.

[0033] A memory management device for a large model, comprising:

[0034] An input data acquisition module, used to acquire and combine input data;

[0035] The text length estimation and adjustment module estimates the text length output by the large model through the text length estimation model and adjusts the estimated text length value upward to an integer;

[0036] The memory block calculation and adjustment module calculates the number of memory blocks to be allocated using the text length and memory page size after upward adjustment, and adjusts the number of memory blocks upward to an integer;

[0037] The kv cache memory block allocation module is used to allocate kv cache memory blocks for large model decoding.

[0038] An electronic device, comprising:

[0039] one or more processors;

[0040] The storage device is used to store one or more programs, and when the one or more programs are executed by the electronic device, the electronic device implements the memory management method of the large model.

[0041] A computer-readable storage medium stores a program, which, when executed by a processor, implements a memory management method for a large model.

[0042] The beneficial effects of the present invention are as follows:

[0043] This method uses a text length estimation model to estimate the output text length of a large model. It then calculates the number of memory blocks to be allocated based on the estimated output text length. This one-time allocation of memory required for the KV cache in large-model decoding avoids issues such as duplicate memory allocation and discontinuous memory access during inference, thereby improving the inference speed of the large model. The method can achieve different results in different application scenarios. By allocating sufficient video memory or internal memory in advance, it effectively avoids delays caused by dynamic adjustments. In batch inference scenarios, it rationally plans computing resources and improves throughput. In streaming output scenarios, such as token-by-token generation, estimating the output length can optimize the user experience, such as by displaying a progress bar or enabling early truncation. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] Figure 1 Schematic diagram of the flow of the memory management method of a large model according to an embodiment of the present invention.

[0045] Figure 2 Schematic diagram of the structure of a text length estimation model according to an embodiment of the present invention.

[0046] Figure 3 Schematic diagram of an MQA-based decoder structure according to an embodiment of the present invention.

[0047] Figure 4 Schematic diagram of a memory management device for a large model according to an embodiment of the present invention.

[0048] Figure 5 FIG. 4 is a schematic diagram of an electronic device according to an embodiment of the present invention. DETAILED DESCRIPTION

[0049] The present invention will be described in detail below with reference to the accompanying drawings and preferred embodiments, and the purpose and effects of the present invention will become more apparent. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0050] Exemplary embodiments will be described in detail herein, with examples illustrated in the accompanying drawings. In the following description, when referring to the drawings, identical numerals in different figures represent identical or similar elements, unless otherwise indicated. The embodiments described in the following exemplary embodiments are not intended to represent all embodiments consistent with the present application. Rather, they are merely examples of apparatus and methods consistent with certain aspects of the present application, as detailed in the appended claims.

[0051] The terms used in this application are for the purpose of describing specific embodiments only and are not intended to limit this application. As used in this application and the appended claims, the singular forms "a," "an," "the," and "the" are intended to include the plural forms, unless the context clearly indicates otherwise. It should also be understood that the term "and / or" as used herein refers to and encompasses any and all possible combinations of one or more of the associated listed items.

[0052] This application provides a memory management method for large models, please refer to Figure 1 ,like Figure 1 As shown, the method includes the following steps:

[0053] Step 1: Build and train a text length prediction model.

[0054] The text length estimation model takes as input text data and information about the large model, and outputs the length of the text output by the large model. The text data is the original text input by the user into the large model, such as questions, instructions, or natural language content to be processed. Information about the large model includes the model name (such as GPT-4, LLaMA-3), quantization bit count (such as FP16, INT8), model size (parameter size, such as 7B, 70B), and adjustable parameters (such as temperature coefficient, top-p sampling value).

[0055] like Figure 2 As shown in the figure, the text length estimation model includes an input layer (input), an embedding layer (embedding), three serially connected MQA-based decoder structures, a third root mean square normalization layer (Root MeanSquare, RMS), a linear projection layer (Linear) and an output layer (Output). Figure 3 It is an MQA-based decoder structure, including the first RMS normalization layer, the Multi Query Attention (MQA) layer, the first concatenation layer, the second RMS normalization layer, the Feed-Forward Network (FFN) and the second concatenation layer.

[0056] The input layer is used to receive input data, including the text data to be processed and related information of the large model (model name, quantization bit number, model size, adjustable parameters), and splice the text data to be processed and related information of the large model into a sequence.

[0057] The embedding layer is used to map the concatenated sequence of the input layer into a high-dimensional vector representation. In this embodiment, the embedding dimension of the embedding layer is 4096 and the vocabulary size is 32000.

[0058] In the three serially connected MQA-based decoder structures, the input to the first RMS normalization layer of the first MQA-based decoder structure is the high-dimensional vector representation after mapping by the embedding layer; the input to the first RMS normalization layer of the two subsequent MQA-based decoder structures is the output of the second concatenation layer of the previous MQA-based decoder structure. In this embodiment, the number of attention heads in each MQA-based decoder structure is 32, and the dimension of each head is 4096 / 32 = 128. The dimensions of the three matrices Q, K, and V are all (4096, 128), and the output dimension of each MQA-based decoder structure is (4096, 4096).

[0059] The first root mean square normalization layer is used to normalize the input to enhance the training stability of the model. The normalization calculation formula is as follows:

[0060]

[0061] Where d represents the dimension of the layer, which is 4096 in this case. i represents the i-th input data, ε is a parameter that prevents the RMS calculation result from being 0, ε=1e-7.

[0062] The multi-query attention layer is used to capture contextual relationships and learn information from different subspaces. It also assigns different weights to different information, calculates the learned information and the corresponding weights, and outputs the calculation results as features. The specific calculation formula is as follows:

[0063]

[0064]

[0065] Where Q represents Query, K represents key, V represents Value, d k Indicates the dimension of Query and Key, W O is the output projection matrix. i ,K,V) each query head Q i The attention weights are calculated independently, but all query heads share K and V. MQA(Q,K,V) means merging the output results of all query heads, and Concat() is a string concatenation function.

[0066] As a variant of Self-Attention, the multi-query attention layer reduces computational overhead by sharing some attention heads (Key and Value) while maintaining good expressiveness.

[0067] The first concatenation layer is used to concatenate the input of the first RMS normalization layer and the output of the multi-query attention layer. The concatenation operation is to add the first RMS normalization input and the output of the multi-query attention layer element-wise.

[0068] The second RMS normalization layer is responsible for normalizing the output of the first concatenation layer.

[0069] The feedforward neural network consists of a fully connected layer and an activation function, which is used to perform nonlinear mapping operations on the results normalized by the second root mean square normalization layer, thereby enhancing the model's nonlinear modeling capabilities. In this embodiment, the feedforward neural network includes three weight matrices, of which the dimensions of the two dimensionality-increasing matrices are (4096, 16384) and the dimension of the dimensionality-reducing matrix is ​​(16384, 4096).

[0070] The second concatenation layer is used to perform the same concatenation operation on the output of the feedforward neural network and the output of the first concatenation layer as the first concatenation layer. The output of the second concatenation layer serves as the output of the entire MQA-based decoder structure.

[0071] The third root mean square normalization layer is used to normalize the output of the third MQA-based decoder structure. The specific normalization operation is:

[0072]

[0073]

[0074] Among them, RMSNorm is the final calculation result, where x represents the input vector and γ represents the scaling factor, which is a trainable parameter.

[0075] The linear projection layer is used to map the output of the third RMS normalization layer to the target dimension in preparation for the final prediction.

[0076] The output layer is used to output the predicted text length, which can be used as a basis for subsequent reasoning optimization (such as memory pre-allocation, batch scheduling, etc.). In this embodiment, the dimension of the output layer is 32,000, and 32,000 also corresponds to the size of the vocabulary. The predicted text length here can be expressed in different ways, including precise length prediction (accurate to the single digit, such as outputting "256"), and quantitative interval prediction (such as taking 1000 as the order of magnitude, outputting "≈2k" to indicate that the predicted length is around 2000). The specific method to be adopted depends on the requirements of the application scenario. For example, precise prediction is suitable for scenarios that require strict control of memory usage (such as embedded devices), and quantitative prediction is suitable for resource scheduling optimization (such as batch tasks in cloud computing environments).

[0077] The MQA-based decoder structure in the text length prediction model of this invention is a variant of the traditional transformer architecture. Compared to the standard Transformer's Self-Attention mechanism, this structure has advantages in computational efficiency and is particularly suitable for lightweight prediction tasks. The text length prediction model of this invention uses three MQA-based decoder structures as core data processing units to gradually extract deep features of the input data and improve prediction accuracy.

[0078] During training, the text length prediction model of the present invention uses the actual output text length (such as the number of tokens or characters) as the supervisory signal during training, that is, the expected output of the text length prediction model during training. The input data during training includes the input text of the large model (such as user queries, prompts, etc.) and the relevant information of the large model (model name, quantization bit number, model size, adjustable parameters). The input data during online prediction needs to be consistent with the structure of the input data format during training to ensure that the effects during training and online reasoning are consistent. During the training process, the text length prediction model learns the mapping relationship between input data and output length to optimize prediction accuracy. After training is completed, the model can be used to dynamically predict the text output length of the large model in the reasoning stage, and its prediction results can be used as the key input for subsequent optimization strategies (such as batch scheduling or video memory allocation), thereby improving the overall efficiency of large model reasoning.

[0079] Step 2: Adjust the predicted text length output by the text length estimation model in step 1 upward to an integer.

[0080] The predicted values ​​output by the text length estimation model may fluctuate or deviate to a certain extent. To ensure the stability and reliability of subsequent memory allocation, the predicted text length needs to be adjusted upward.

[0081] The specific methods of upward adjustment are as follows:

[0082] If the text length estimation model outputs a length value accurate to the single digit (e.g., a predicted value of 1256), first perform a ceiling adjustment (e.g., rounding to 1300) and then add a buffer (e.g., +500, ultimately adjusting to 1,800). The buffer can be adjusted based on the actual application scenario to reserve additional memory space to prevent insufficient memory due to prediction bias. Adjustments to the quantized interval predictions.

[0083] If the predicted text length output by the text length estimation model is a quantized interval, such as "≈2k", which means approximately 2000, or 1500–2500, then round up to an upper limit of 2500 and increase the buffer value by 500 to avoid insufficient memory in critical situations; the adjusted value is 3000.

[0084] Adjusting your strategy can effectively avoid the following problems:

[0085] (1) Memory Reallocation: When inferring a large model, if the initially allocated memory is insufficient, the system may be forced to dynamically expand the memory, resulting in additional computational overhead and latency. By adjusting the predicted value upward, sufficient memory can be allocated at one time, reducing the performance loss caused by runtime adjustments.

[0086] (2) Preventing Out-of-Memory (OOM): If the predicted value is too low, it may cause memory / memory overflow, which in turn causes reasoning failure. Increasing the buffer value can improve fault tolerance and ensure that large models can still run stably when generating long texts.

[0087] The buffer size can be flexibly selected based on actual conditions: a fixed value (such as +500) is suitable for most common scenarios and is computationally simple; a dynamic adjustment (such as increasing it by 20% of the predicted value) is suitable for tasks with large variations in output length (such as dialogue generation vs. summary generation).

[0088] Furthermore, rounding can be tailored to business needs, with different strategies available, such as rounding to the hundreds place (1234 → 1300) or rounding to the thousands place (1234 → 2000). This optimization step allows the system to more reliably manage memory resources, improving the stability and efficiency of inference for large models. The adjusted length value serves as a key parameter for subsequent memory pre-allocation, thereby optimizing the overall inference process.

[0089] Step 3: Calculate the number of cache blocks based on the memory page size and KV cache dimensions.

[0090] During inference on large autoregressive models, efficient management of the KV Cache (Key-Value Cache) is crucial to inference speed. To optimize graphics memory utilization and reduce dynamic memory allocation overhead, the number of KV Cache blocks needs to be appropriately calculated based on the estimated text length. The specific calculation method is as follows:

[0091] N=W*L / D

[0092] Where N represents the number of KV Cache blocks to allocate; W represents the KV Cache dimension, which is determined by the model architecture, such as 128 bytes per token; L represents the adjusted estimated output text length; and D represents the system memory page size, such as 4KB or 16KB per page, depending on the hardware architecture. W × L represents the total memory size required to store the entire sequence in the KV Cache. Dividing this by D represents the number of physical storage units required for partitioning by memory page size.

[0093] Example calculation: Assuming the adjusted output length L = 1,800 tokens, the KV Cache dimension W = 128 bytes / token, and the memory page size D = 4KB (4,096 bytes), then: N = 128 × 1800 / 4096 ≈ 56.25.

[0094] Step 4: Adjust the number of cache blocks calculated in step 3 upward to an integer.

[0095] After completing the initial calculation N = W × L / D, the resulting value is usually a floating point number (such as 56.25). Because memory allocation must be in integer blocks, the number of blocks needs to be adjusted. There are several ways to adjust the number of blocks, which can be flexibly selected according to different application scenarios:

[0096] (1) Forced rounding up, using ceil(N) operation to ensure sufficient memory space is allocated. For example, if N=56.25, 57 blocks are taken. This operation can completely avoid the risk of insufficient memory, but the cost is that it may cause about 1.5% memory waste (taking 56.25 as an example).

[0097] (2) Minimum safety margin strategy: additional safety blocks are added on the basis of rounding, and the formula is adjusted: N' = ceil(N) + S (S is the number of safety blocks set based on experience, usually 1-3 blocks). In this way, adding safety blocks can provide memory safety guarantees for scenarios with zero tolerance for memory overflows.

[0098] (3) Memory alignment optimization: consider the hardware memory access characteristics to adjust the alignment: calculate the number of aligned blocks: N'=ceil(N / A)×A (A is the architecture-specific alignment coefficient, such as a multiple of 4), which can effectively improve memory access efficiency.

[0099] (4) A dynamic block strategy is adopted for the task of generating very long text: a segmented allocation mechanism is adopted, the initial number of basic blocks allocated is set (e.g. 32 blocks), and a monitoring threshold is set to trigger dynamic expansion, and CUDA's cudaMallocAsync is used to achieve non-stop expansion. The adjusted complete calculation formula is:

[0100] N'= min( ceil(N) + S, N_max )

[0101] Where S represents the safety margin, which defaults to 1; N_max represents the maximum number of blocks available on the device.

[0102] Step 5: Allocate kv cache memory blocks for large model decoding.

[0103] Based on the integer number N of KV Cache blocks calculated in step 4, the system performs the video memory / memory allocation operation. The specific allocation method is as follows:

[0104] Use cudaMalloc (GPU) or posix_memalign (CPU) for aligned memory allocation; each block strictly corresponds to D bytes (memory page size) of continuous space as the allocation unit, and a total of N'×D bytes of continuous memory space is allocated.

[0105] Memory layout optimization uses the optimal memory arrangement scheme: [Block 0][Block 1]...[Block N-1]. Each block contains:

[0106] Key cache: occupies 50% of the space of a single block (D / 2 bytes);

[0107] Value buffer: occupies 50% of the space of a single block (D / 2 bytes);

[0108] Meta-information header: 8-byte block status identifier (including valid bits, etc.).

[0109] At the same time, some memory protection strategies are adopted according to different scenarios, such as:

[0110] Pre-clear processing, execute memset clearing immediately after allocation to avoid numerical anomalies caused by uninitialized memory;

[0111] Backup mechanism: Automatically switch to shard allocation mode when continuous memory is insufficient.

[0112] The following performance optimization measures are also applied to memory allocation:

[0113] Asynchronous allocation and parallel execution with the computing pipeline improve efficiency;

[0114] Unified memory management, supporting the Unified Virtual Memory (UVM) access mode;

[0115] Intelligent caching automatically reclaims idle blocks based on the LRU (Least Recently Used) strategy.

[0116] The allocated memory area will be used exclusively to store the Key-Value matrix generated for each token during the decoding process.

[0117] This embodiment uses a text length estimation model to estimate the length of the text output by a large model. It then uses the memory page size and the KV cache dimension to calculate the number of cache blocks required. This avoids issues such as repeated memory reallocation and discontinuous memory access during decoding for large models. This embodiment can be effective in different application scenarios. By allocating sufficient video memory or memory in advance, it effectively avoids delays caused by dynamic adjustments. In batch inference scenarios, it rationally plans computing resources and improves throughput. In streaming output scenarios, estimating the output length during token-by-token generation can optimize the user experience, such as by displaying a progress bar or pre-emptively truncating the output.

[0118] On the other hand, Figure 4 As shown, the present application also provides a memory management device for a large model, comprising:

[0119] An input data acquisition module, used to acquire and combine input data;

[0120] The text length estimation and adjustment module estimates the text length output by the large model through the text length estimation model and adjusts the estimated text length value upward to an integer;

[0121] The memory block calculation and adjustment module calculates the number of memory blocks to be allocated using the text length and memory page size after upward adjustment, and adjusts the number of memory blocks upward.

[0122] The kv cache memory block allocation module is used to allocate kv cache memory blocks for large model decoding.

[0123] This application also provides an electronic device, see Figure 5 , Figure 5 This is a structural diagram of an embodiment of the electronic device of the present application, except Figure 5 In addition to the processor, memory, network interface, and non-volatile memory shown, any device with data processing capabilities in which the apparatus in the embodiment is located may also include other hardware, generally based on the actual functions of the device with data processing capabilities, which will not be described in detail.

[0124] The implementation process of the functions and effects of each unit in the above-mentioned device is specifically described in the implementation process of the corresponding steps in the above-mentioned method, and will not be repeated here.

[0125] For the device embodiments, since they basically correspond to the method embodiments, the relevant parts can be referred to the partial description of the method embodiments. The device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the present invention. A person of ordinary skill in the art can understand and implement the present invention without inventive work.

[0126] An embodiment of the present invention further provides a computer-readable storage medium having a program stored thereon. When the program is executed by a processor, the memory management method of the large model in the above embodiment is implemented.

[0127] The computer-readable storage medium may be an internal storage unit of any device with data processing capabilities described in any of the aforementioned embodiments, such as a hard disk or memory. The computer-readable storage medium may also be an external storage device, such as a plug-in hard disk, a smart memory card (SmartMedia card, SMC), an SD card, a flash card, etc. equipped on the device. Furthermore, the computer-readable storage medium may also include both an internal storage unit of any device with data processing capabilities and an external storage device. The computer-readable storage medium is used to store the computer program and other programs and data required by any device with data processing capabilities, and may also be used to temporarily store data that has been output or is to be output.

[0128] Those skilled in the art will readily appreciate other embodiments of the present application after considering the specification and practicing the contents disclosed herein. This application is intended to cover any variations, uses, or adaptations of the present application that follow the general principles of the present application and include common knowledge or customary techniques in the art not disclosed herein. The description and examples are to be considered merely as exemplary, and the true scope and spirit of the present application are indicated by the claims.

[0129] It should be understood that the present application is not limited to the exact structures described above and shown in the drawings, and that various modifications and changes may be made without departing from the scope thereof. The scope of the present application is limited only by the appended claims.

Claims

1. A memory management method for a large model, characterized in that: The steps include: Step 1: Build and train a text length estimation model; The input of the text length estimation model is the text data and the related information of the large model, and the output is the length of the output text of the large model; the text length estimation model includes an input layer, an embedding layer, three serially connected MQA-based decoder structures, a third root mean square normalization layer, a linear projection layer and an output layer; the MQA-based decoder structure includes a first root mean square normalization layer, a multi-query attention layer, a first splicing layer, a second root mean square normalization layer, a feedforward neural network and a second splicing layer; the multi-query attention layer shares part of the attention head, that is, shares the key and value, for capturing contextual relationships and learning information from different subspaces; at the same time, different weights are assigned to different information, and the learned information and the corresponding weights are calculated, and the calculation results are output as features; the first splicing layer is used to splice the input of the first root mean square normalization layer and the output of the multi-query attention layer; The feedforward neural network is used to perform a nonlinear mapping operation on the normalized result of the second root mean square normalization layer; the second splicing layer is used to perform the same splicing operation as the first splicing layer on the output of the feedforward neural network and the output of the first splicing layer; Step 2: Adjust the predicted text length output by the text length estimation model in step 1 upward to an integer; Step 3: Calculate the number of cache blocks based on the adjusted text length, memory page size, and KV cache dimensions; Step 4: Adjust the number of cache blocks upward to an integer; Step 5: Allocate kv cache memory blocks for large model decoding based on the number of cache blocks adjusted upward.

2. The memory management method of a large model according to claim 1, characterized in that: The predicted text length output by the output layer is an exact length or a quantized interval; If the text length estimation model outputs an accurate length prediction, the upward adjustment in step 2 is specifically as follows: first, the length value is rounded up, and then an additional buffer value is added; If the text length estimation model outputs a quantization interval, the upward adjustment in step 2 is specifically as follows: first, the upper limit value of the quantization interval is rounded up, and then an additional buffer value is added.

3. The memory management method of a large model according to claim 2, characterized in that: The rounding up is rounding to the unit place, the hundred place, or the thousand place, which is selected according to business needs; the buffer value is a fixed value, or a percentage of the predicted value.

4. The memory management method of a large model according to claim 3, characterized in that: In step 4, the number of cache blocks is adjusted upward to an integer by one of the following methods: Method 1: round up the unit digit; Method 2: After rounding up the units digit, add an additional security block number; Method 3: Memory alignment optimization, specifically: N'=ceil(N / A)×A Where N' is the adjusted number of cache blocks, N is the number of cache blocks calculated in step 3, A is the alignment coefficient, and ceil() is the rounding function. Method 4: For very long text generation tasks, adopt a dynamic chunking strategy: N' = min( ceil(N) + S, N_max ) Where S represents the safety margin and N_max represents the maximum number of blocks available on the device.

5. The memory management method of a large model according to claim 4, characterized in that: The step 5 is specifically as follows: using cudaMalloc or posix_memalign to perform aligned memory allocation, with each block corresponding to the continuous space of the memory page size as the allocation unit, and a total of N×D bytes of continuous memory space allocated; N' cache blocks are arranged in sequence in the N×D bytes of continuous memory space, and each memory block contains: Key cache: occupies 50% of the space of a single block, that is, D / 2 bytes; Value buffer: occupies 50% of the space of a single block, that is, D / 2 bytes; Meta information header: 8 bytes of block status identification; Where D represents the system memory page size.

6. The memory management method of a large model according to claim 1 or 5, characterized in that: After completing the aligned memory allocation in step 5, immediately execute the following memory protection strategy: (1) Pre-clear processing, execute memset clearing immediately after allocation to avoid numerical anomalies caused by uninitialized memory; (2) Backup mechanism: Automatically switch to shard allocation mode when continuous memory is insufficient.

7. The memory management method of a large model according to claim 1, characterized in that: The memory block allocation in step five is asynchronous allocation and is executed in parallel with the computing pipeline.

8. A memory management device for a large model, characterized in that: include: An input data acquisition module, used to acquire and combine input data; The text length estimation and adjustment module estimates the text length output by the large model through the text length estimation model and adjusts the estimated text length value upward to an integer; The memory block calculation and adjustment module calculates the number of memory blocks to be allocated using the adjusted text length, memory page size, and KV cache dimensions, and adjusts the number of memory blocks upward to an integer. The kv cache memory block allocation module is used to allocate kvcache memory blocks for large model decoding based on the number of memory blocks adjusted upward.

9. An electronic device, characterized in that: include: one or more processors; A storage device for storing one or more programs, which, when executed by the electronic device, enables the electronic device to implement the memory management method of the large model as described in any one of claims 1 to 7.

10. A computer-readable storage medium, characterized in that A program is stored thereon, and when the program is executed by a processor, the memory management method of the large model as claimed in any one of claims 1 to 7 is implemented.

Images