A method, system, device and medium for joint optimization of flash attention and paged attention in long sequence processing

Abstract

The application relates to the technical field of long sequence processing, and particularly relates to a joint optimization method, system, device and medium of Flash Attention and Paged Attention in long sequence processing, which comprises obtaining a paging KV cache and a corresponding block table of a target sequence; in a block iteration calculation process of the Flash Attention, a logical index of a key-value block required to be accessed by a current query block is converted into a physical address in real time according to the block table by a fusion attention kernel function, corresponding key tensors and value tensors are collected from the non-continuous physical pages according to the converted physical address, and the collected key tensors and value tensors are used to complete block attention score calculation with the query block. The application has the beneficial effect that through the collaborative design of the fusion attention operator and the paging memory management mechanism, the efficient access to the non-continuous physical pages is realized while the fusion performance advantage of the Flash Attention operator is maintained.

Description

Technical Field

[0001] This invention relates to the field of long sequence processing technology, and in particular to a joint optimization method, system, device and medium for Flash Attention and PagedAttention in long sequence processing. Background Technology

[0002] In long-sequence inference scenarios, Flash Attention requires physically contiguous storage of key-value tensors to maintain its operator fusion advantage, which fundamentally conflicts with the non-contiguous page storage method of Paged Attention. Direct combination requires explicit data collection before attention computation, copying data from scattered pages to a contiguous buffer. This not only introduces additional GPU memory access overhead but also undermines the operator fusion characteristics of Flash Attention, leading to a decline in inference performance. Therefore, achieving efficient access to paged, non-contiguous key-value caches while maintaining the operator fusion performance of Flash Attention is a pressing technical problem to be solved in long-sequence inference. Summary of the Invention

[0003] To solve the above-mentioned technical problems, the present invention provides the following technical solution: In a first aspect, the present invention provides a joint optimization method for Flash Attention and Paged Attention in long sequence processing, including obtaining the paged KV cache of the target sequence and the corresponding block table; wherein, the paged KV cache stores key tensors and value tensors in non-contiguous physical pages according to a fixed page size, and the block table records the mapping relationship between logical pages and physical pages. During the block-based iterative computation of Flash Attention, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time based on the block table. Based on the converted physical address, the corresponding key tensors and value tensors are collected from non-contiguous physical pages. The collected key tensors and value tensors are then used to complete the block-based attention score calculation with the query block, thereby achieving joint optimization of paginated KV caching and Flash Attention. The fusion attention kernel function has address conversion logic embedded through just-in-time compilation.

[0004] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing of the present invention, the method includes: obtaining the paged key-value cache and corresponding block table of the target sequence, including: Pre-allocate a global key-value cache buffer based on the target hardware parameters and construct a lock-free linked list of free physical pages; In response to receiving an input sequence, the required number of physical pages are atomically reserved from the free physical page list through atomic comparison and swap operations, and a block table recording the mapping relationship between logical pages and physical pages is constructed for the input sequence.

[0005] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing of the present invention, the address translation logic is provided in the form of a custom index mapping function and inlined into the fusion attention kernel function through just-in-time compilation mechanism; Before performing distributed collection, the method caches the block table of the target sequence in the shared memory of the fusion attention kernel function to accelerate the translation of logical indexes to physical addresses.

[0006] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing of the present invention, in response to the completion of sequence processing, the block table of the completed sequence is traversed, and the atomic reference counter of the physical page corresponding to each physical page identifier in the block table is decremented by 1. In response to the atomic reference counter being zeroed, the identifier of the physical page is returned to the list of free physical pages; In this system, in response to multiple sequences of block tables pointing to the same physical page, the atomic reference counter of the same physical page is incremented to achieve prefix sharing of the same physical page by multiple sequences.

[0007] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing according to the present invention, the method further includes: Maintain the global batch state, which records the current length, block table pointer, and estimated remaining steps for each sequence within the batch; The execution order of sequences within a batch is dynamically adjusted based on the global batch status and the business priority of each sequence. In response to the completion of processing of any sequence, reclaim the physical pages occupied by the completed sequence and accept new requests into the batch.

[0008] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing of the present invention, the method further includes: adaptively selecting the page size according to the length of the input sequence; wherein: If the length of the input sequence is less than a first threshold, the first page size is used. In response to the length of the input sequence being greater than or equal to a first threshold, a second page size greater than the first page size is adopted.

[0009] As a preferred embodiment of the joint optimization method of Flash Attention and Paged Attention in long sequence processing of the present invention, the page size is a power of 2.

[0010] Secondly, the present invention provides a joint optimization system of Flash Attention and Paged Attention in long sequence processing, comprising: an acquisition module for acquiring the paged KV cache and the corresponding block table of the target sequence; wherein, the paged KV cache stores key tensors and value tensors in non-contiguous physical pages according to a fixed page size, and the block table records the mapping relationship between logical pages and physical pages; The processing module is used to perform block-based iterative calculations for Flash Attention. During this process, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time based on the block table. Then, based on the converted physical address, it collects the corresponding key tensors and value tensors from non-contiguous physical pages and uses the collected key tensors and value tensors to complete the block-based attention score calculation with the query block. The fusion attention kernel function has address conversion logic embedded in it through just-in-time compilation.

[0011] Thirdly, the present invention provides a computer device, including a memory and a processor, wherein the memory stores a computer program, and the processor executes the computer program to implement the steps of the method described above.

[0012] Fourthly, the present invention provides a computer-readable storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements the steps of the method described above.

[0013] Compared with existing technologies, the beneficial effects of this invention are as follows: By paging the KV cache into fixed-size pages and dynamically allocating physical pages, combined with lock-free free lists and atomic reservation operations, the huge waste caused by pre-allocating contiguous GPU memory for each sequence in traditional solutions is avoided; furthermore, by utilizing the Flex Attention API and just-in-time compilation technology, the address translation logic is directly embedded into the CUDA kernel function of Flash Attention, and the address is translated in real time according to the block table during block iteration, and KV data is collected from non-contiguous pages, eliminating the tedious step of copying and then calculating; even further, by combining three levels of optimization—inter-head parallelism, inter-sequence dynamic load balancing, and intra-block warp-level parallelism—GPU computing resources are fully utilized. Attached Figure Description

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

[0015] Figure 1 This is a flowchart illustrating a joint optimization method of Flash Attention and Paged Attention in long sequence processing, provided as an embodiment.

[0016] Figure 2 This is a graph comparing inference latency and throughput performance.

[0017] Figure 3 A graph comparing video memory usage and evaluating paging optimization efficiency. Detailed Implementation

[0018] To make the above-mentioned objects, features, and advantages of the present invention more apparent and understandable, specific embodiments of the present invention will be described in detail below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of the present invention, and not all of them. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the protection scope of the present invention.

[0019] Example 1, referring to Figure 1 This is the first embodiment of the present invention, which provides a joint optimization method for FlashAttention and PagedAttention in long sequence processing, including: This application addresses the problem of how to achieve efficient access to non-contiguous physical pages in a paginated key-value cache without sacrificing the performance of Flash Attention operator fusion. A method that effectively solves the aforementioned problem will be provided, and the following will elaborate on how to implement this joint optimization method of Flash Attention and PagedAttention in long sequence processing using several embodiments. Figure 1 A flowchart illustrating a joint optimization method of Flash Attention and Paged Attention in long sequence processing is shown, including: S100: Obtain the paginated KV cache and corresponding block table of the target sequence.

[0020] It should be noted that the paginated KV cache for the target sequence stores the key tensors and value tensors in non-contiguous physical pages at a fixed page size, while the block table records the mapping relationship between logical pages and physical pages.

[0021] It is understood that the target sequence refers to the specific input sequence for which attention computation is currently required. Obtaining the paged key-value cache and corresponding block table of the target sequence is done before attention computation. In this embodiment, "obtaining" encompasses various possible implementations of these two data structures: building them from scratch, reading them from memory, or receiving them from an external source.

[0022] In one specific implementation, the paginated key-value cache and corresponding block table of the target sequence are obtained in the following way: Pre-allocate a global key-value buffer based on the target hardware parameters and construct a lock-free list of free physical pages; and in response to receiving an input sequence, atomically reserve the required number of physical pages from the list of free physical pages through atomic comparison and swap operations, and construct a block table for the input sequence that records the mapping relationship between logical pages and physical pages.

[0023] It should be noted that the cache buffer is divided into a separate key cache and a value cache, and the total capacity is calculated based on the maximum number of pages supported.

[0024] The formula for calculating the total capacity of each buffer is as follows: In the formula: This represents the maximum number of pages supported by the system. The hidden layer dimension of the model; This represents the number of attention layers.

[0025] Furthermore, a lock-free linked list of free physical pages is constructed, and all physical page identifiers are added to the linked list in sequence at the beginning.

[0026] When the system receives an input sequence, it first calculates the required number of physical pages n based on the length N of the current sequence.

[0027] For a sequence of length N, the required number of pages n is calculated as follows: Furthermore, by atomically comparing and swapping operations, n consecutive physical page identifiers are atomically reserved from the head of the free list, with a time complexity of O(n). (1). Next, a block table is created for the sequence. It should be noted that the block table is a 32-bit integer array and is stored in the GPU global memory. The i-th entry records the physical page identifier corresponding to the i-th logical page of the sequence.

[0028] In a preferred embodiment, to optimize memory utilization and page table overhead, the page size is adaptively selected based on the length of the input sequence. Experiments show that: when When: Page table overhead is relatively large, but internal fragmentation is minimal (average waste of 15.5 tokens), therefore it is suitable for short-sequence intensive scenarios; when At this time: it is at a balance point, with moderate page table overhead and an average internal fragmentation of 31 tokens, therefore it is suitable for mixed-length batch inference; when When: page table overhead is minimal, but internal fragmentation increases to an average of 63.5 tokens, therefore it is suitable for long-sequence single-task scenarios; while when While page tables are more compact, they suffer from excessive internal fragmentation, and the 64MB granularity per page makes allocation inflexible, so their use is generally not recommended.

[0029] Preferably, for a sequence of length L, if L < 512 (the first threshold), use (First page size); If 512 L<2048, use (Second page size); If L For 2048, it is more recommended to use .

[0030] Preferably, by adopting the above-mentioned preferred scheme, the average memory utilization rate can be increased to 97.3%.

[0031] In another specific implementation, the page size P is fixed as a power of 2.

[0032] Preferably, by using a power of 2, the calculation of logical block index and page offset can be performed by shift and bitwise AND operations instead of division and modulo operations during subsequent address translation, thereby reducing computational overhead.

[0033] In another preferred embodiment, to optimize memory utilization and page table overhead, the system supports adaptive selection of page size based on the length of the input sequence. Specifically, the system maintains three page pools: P=32, P=128, and P=512. For the first 512 tokens of the sequence, P=128 is used by default to balance index overhead and internal fragmentation. When the sequence length continues to grow beyond 512 tokens, the subsequent portion automatically switches to P=512, significantly reducing the number of page table entries and the frequency of index calculations with a larger page granularity. When the sequence reaches its end and the remaining length is less than 512 tokens, a smaller P=32 is used for padding to minimize memory waste caused by free space within the page.

[0034] Preferably, the above-mentioned preferred scheme can reduce the overall size of the page table by about 40%, while keeping the internal fragmentation rate below 5%.

[0035] In another preferred implementation, for very long contexts, such as 1 million tokens, when a single GPU memory is insufficient to accommodate the entire KV cache, it supports expansion to a three-tier storage architecture. The specific three-tier storage architecture is as follows: L1 layer stores recently accessed hot pages with a capacity of approximately 10GB, which can accommodate approximately 160 pages; L2 layer stores warm pages with a capacity of approximately 100GB, which can accommodate approximately 1600 pages, transmitted via PCIe 4.0 with a bandwidth of approximately 32GB / s and a single-page transfer latency of approximately 2ms; and L3 layer stores cold pages with a capacity of approximately 1TB, which can accommodate approximately 16,000 pages, transmitted via NVMe with a bandwidth of approximately 7GB / s and a single-page transfer latency of approximately 9ms.

[0036] Furthermore, the page migration strategy adopts the Least Recently Used algorithm, which maintains the access timestamp of each page. When the L1 layer is full, the oldest page is expelled to the L2 layer, and when the L2 layer is full, it is expelled to the L3 layer.

[0037] Furthermore, it also includes predicting the page to be accessed based on attention patterns, that is, loading it from the L2 / L3 layer to the L1 layer in advance. Ideally, for sequential access patterns, the hit rate of predicting the page to be accessed based on attention patterns can reach 92%, which can effectively hide transmission delays.

[0038] In an optional implementation, obtaining the paged key-value cache and corresponding block table can also be achieved by loading a pre-serialized paged key-value cache and block table from external persistent storage. This involves serializing the previously calculated paged key-value cache and corresponding block table according to a specific format and writing them to disk or non-volatile memory. When it is necessary to continue processing the same sequence, the loading module reads this data from external storage and reloads the key-value pages into free physical pages in the GPU memory, while simultaneously restoring the block table, thereby obtaining the complete paged key-value cache and block table.

[0039] In another specific implementation, after obtaining the paginated KV cache and corresponding block table of the target sequence, the model begins layer-by-layer forward propagation. For the current layer, the query Q, key K, and value V projection are first calculated based on the input of the current layer. Then, the ASSIGN allocation operation is performed, writing the newly generated key and value tensors into the reserved physical pages.

[0040] Specifically, for the token at position t in the sequence, its physical address is calculated using a two-level indexing method: In the formula: b is the logical page index, o is the page offset, and p is the final physical address.

[0041] Furthermore, the system writes K[t] and V[t] into the corresponding positions in the global KV cache buffer according to the calculated physical address.

[0042] It should be noted that the writing process is performed incrementally on a token-by-token basis, and when the sequence length increases to the point that the current page is full, the system will trigger a new reservation operation to allocate additional physical pages and update the block table.

[0043] Ideally, the allocation operation allows logically continuous key-value tensors to be distributed and stored across multiple non-contiguous physical pages, laying the foundation for subsequent Paged Attention management.

[0044] S200: During the block-based iterative calculation of Flash Attention, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time according to the block table. Based on the converted physical address, the corresponding key tensors and value tensors are collected from non-contiguous physical pages. The collected key tensors and value tensors are used to complete the block-based attention score calculation with the query block, thereby realizing the joint optimization of pagination KV caching and Flash Attention.

[0045] In one specific implementation, the attention kernel function is built based on the PyTorch 2.x Flex Attention API and incorporates address translation logic through just-in-time compilation. Specifically, two custom functions are defined: Used to implement virtual column isolation within a batch. Used to implement address translation.

[0046] in, The specific manifestations are as follows: In the formula: For batch indexing; For attention head index; For querying the index; Key-value index; The sequence identifier to which the currently queried token belongs; This is the sequence identifier to which the current key-value token belongs; This represents the total number of tokens in the current query sequence. Table logic AND operations.

[0047] Furthermore, The specific manifestations are as follows: In the formula: For block tables; The ID of the sequence to which the key value belongs; Round down the logical page index; This is the offset within the page.

[0048] Furthermore, by inlining these two functions into the block loop of Flash Attention using the Torch Inductor compiler, the generated CUDA kernel automatically performs address translation when accessing KV tensors.

[0049] Ideally, the compiler can identify the regularity of address calculations through the above steps. By elevating page table lookups to the outer loop and using warp-level shuffle instructions to share page table entries among threads, global memory access can be reduced.

[0050] Specifically, during the inference phase, for the l-th layer attention calculation of the target sequence, the system calls the aforementioned fusion attention kernel function to execute the block-based strategy of Flash Attention: First, calculate the row block size based on the GPU's SRAM capacity M and attention head dimension d. and column block size The specific calculation formula is as follows: Furthermore, the query matrix Q is divided into multiple query blocks. The key matrix K and value matrix V are divided into multiple key-value blocks. Specifically, it manifests as follows: For each query block The kernel function iterates through all the key-value blocks that need to be accessed. During the iteration, the kernel function performs the following operations: First, based on the logical index j of the current key-value block, the built-in score_mod function is used to query the block table of the target sequence, converting the logical index j into a physical address, and obtaining... The actual location of the non-contiguous physical page is determined by directly reading the corresponding address from the scattered non-contiguous physical pages based on the converted physical address. Regarding tensors, it's important to note that this process doesn't involve pre-copying data from scattered pages into a contiguous buffer. Instead, it uses index mapping within the kernel function for on-demand access. (It's crucial to emphasize that the key and value tensors currently being accessed are written into the non-contiguous physical pages via the aforementioned ASSIGN operation; ASSIGN handles writing, and GATHER handles reading, both sharing the same block table and address translation logic.) Finally, the read data is used... With the current query block Calculate local attention score The calculation formula is as follows: Furthermore, to achieve numerical stability, the maximum value of each row is maintained. and exponential sum And it iterates according to the FlashAttention update formula, specifically as follows: In the formula: rowmax(·) is to take the maximum value of each row; rowsum(·) is to sum the values ​​of each row.

[0051] It is important to emphasize that when processing the next key-value block, the global maximum value and cumulative sum need to be updated, specifically as follows: In the formula: This is the updated global row maximum value; For the updated global row sum; This is the correction factor for the old maximum value; This is the correction coefficient for the new local maximum.

[0052] Compared to traditional softmax, which requires all data to compute, our Flash Attention method can achieve block-based incremental computation using just these two formulas. Furthermore, it doesn't require storing the complete N×N attention matrix; it only needs to maintain the values ​​for each row. and That's all.

[0053] The update formula for the output block is: In the formula: This is the final normalization factor; Scale the old output by the exponential difference and the old sum; This signifies the contributions of Xin Kuai.

[0054] It is understandable that after all key-value blocks have been traversed, the output block will be... That is, the final attention result of this query block.

[0055] Understandably, the block size of Flash Attention directly affects SRAM utilization and HBM access count. Therefore, for L4 GPUs, with each SM having 128KB of SRAM, the recommended configuration includes: when the attention head dimension... hour: =64, =64, each block occupies 64KB, leaving 64KB for intermediate variables and instruction cache; when focusing on the head dimension hour: =128, =128, single block occupies 112KB, SRAM utilization rate is 87.5%.

[0056] For the H100 GPU, the optimal block size can be automatically calculated using the following formula: In the formula: 0.8 is the SRAM capacity; 0.8 is the safety factor; 4d is the QKV block overhead. The cost is for the attention score matrix.

[0057] In a preferred embodiment, a three-level parallel strategy is employed to optimize page access. The specific three-level parallel strategy includes inter-head parallelism, inter-sequence parallelism, and intra-block parallelism. Inter-head parallelism involves distributing 32 attention heads across 32 thread blocks, with each thread block independently calculating the attention output of one head. This means there is no data communication between thread blocks, achieving parallelism. However, the number of thread blocks is adaptively adjusted for models with different numbers of heads. Inter-sequence parallelism, in batch inference scenarios, involves distributing B sequences across B thread block groups, each containing 32 blocks. Dynamic load balancing is used to complete the blocks of shorter sequences ahead of time, assisting in processing the remaining blocks of longer sequences. Intra-block parallelism refers to the parallel computation of a single attention head at the warp level, meaning each warp processes... One row in a row, 32 warps process 32 rows in parallel, reducing global memory access by sharing page table entries between threads using the warpshuffle instruction.

[0058] In another specific implementation, when sequence generation is complete or interrupted, the system performs a page reclamation operation. Specifically, the reclamation process involves batch returning all physical page identifiers occupied by the sequence to the free list, using atomic operations to ensure thread safety. It should be noted that for scenarios where multiple sequences share the same prefix, such as batch inference or dialogue history, this method supports shared reference counts for prefix pages. That is, an atomic reference counter is maintained for each physical page. When the block tables of multiple sequences point to the same physical page, the counter increments. It is crucial to emphasize that the page is only truly reclaimed when the counter drops to zero.

[0059] In an alternative implementation, for newly allocated pages, whose lifespan is typically very short, such as intermediate results of temporary calculations, reference counting can be used, i.e., the page is immediately reclaimed when the count drops to zero.

[0060] In another alternative implementation, for pages that have survived more than 10 generation steps, such as the key-value cache of dialogue history, a mark-and-sweep method can be used. This involves periodically scanning all old generation pages, marking pages that are still referenced, and clearing unmarked pages. The scanning cycle can be triggered when there are 100 generation steps or when the video memory usage exceeds 80%.

[0061] In another alternative implementation, for cases where a single sequence occupies more than 50 pages, a separate linked list of large objects can be maintained and prioritized for reclamation to quickly release large blocks of video memory.

[0062] In another alternative implementation, for predictive recycling, sequences that are about to be completed can be marked in advance based on the sequence generation progress and estimated remaining length, and these sequences can be evicted first when memory is scarce, reducing interference with long-running tasks.

[0063] In summary, the beneficial effects of this joint optimization method of Flash Attention and Paged Attention in long sequence processing are as follows: By paging the KV cache into fixed-size pages and dynamically allocating physical pages, combined with lock-free free lists and atomic reservation operations, the huge waste caused by pre-allocating contiguous GPU memory for each sequence in traditional solutions can be avoided; furthermore, by utilizing the Flex Attention API and just-in-time compilation technology, the address translation logic is directly embedded into the CUDA kernel function of Flash Attention, and the address is translated in real time according to the block table and KV data is collected from non-contiguous pages during block iteration, which can save the tedious step of copying and then calculating; even further, by combining three levels of optimization—inter-head parallelism, dynamic load balancing between sequences, and intra-block warp-level parallelism—GPU computing resources can be fully utilized.

[0064] Example 2: This example will take the generation of a summary from a legal document with 100,000 tokens as an example to illustrate the complete implementation process of this method in a single-sequence long text reasoning scenario.

[0065] Scenario description: A user inputs a legal document with 100,000 tokens and requests the model to generate a summary.

[0066] Parameter configuration: Model is LLaMA-7B (32 layers, 32 attention heads, 4096 hidden dimensions), page size The GPU is an NVIDIA L4 (24GB VRAM), and the inference precision is FP16.

[0067] Execution process: 1. During the initialization phase, 150 physical pages are allocated, each page is 64MB, for a total of 9.6GBKV cache space.

[0068] 2. The input document is processed by a tokenizer to obtain 100,000 tokens. Calculate the required number of pages: =782. Due to exceeding the available page count, the system has implemented a segmentation strategy, dividing the document into six segments of approximately 16,667 tokens each, with each segment requiring approximately 131 pages.

[0069] 3. Processing the first segment: The RESERVE operation allocates 131 pages and constructs a block table. Propagate forward layer by layer, calculating the QKV projection at each layer, and then writing the KV to the paging cache via the ASSIGN operation. The attention fusion kernel dynamically accesses non-contiguous pages based on the block table and calculates the attention output.

[0070] 4. After processing the first segment, extract key information (approximately 2000 tokens) as context and reclaim the remaining pages. When processing the second segment, mark the context key-value page as shared (reference count +1), and assign new content to a new page.

[0071] 5. Repeat the above process to process all fragments, eventually generating a digest of approximately 500 tokens. The total inference time is approximately 45 seconds, with an average latency of 0.45ms per token and a peak GPU memory usage of 23.1GB (13.4GB for model weights + 1GB for activation values ​​+ 8.7GB for KV cache).

[0072] Performance comparison: Traditional continuous KV caching schemes require pre-allocation of 51.2GB of video memory, far exceeding the L4 capacity, and must use CPU offload, resulting in inference time exceeding 8 minutes. In contrast, this method completes the task on a single GPU through paging management and segmented processing, improving speed by about 10 times.

[0073] Example 3 illustrates an optimized implementation of this method in a batch inference scenario, using the example of an online service simultaneously processing 16 requests of different lengths.

[0074] The scenario and parameter configuration are as follows: 16 requests are processed simultaneously, with the sequence length distribution as follows: {512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 512, 1024, 1536, 2048, 2560, 3072, 3584, 4096} token. The model configuration is the same as in Example 2 (LLaMA-7B, FP16). To optimize the efficiency of short sequences, this example sets the page size P=64 and the batch size B=16.

[0075] First, calculate the total page requirements: =8+16+24+32+40+48+56+64+8+16+24+32+40+48+56+64=576 pages, occupying approximately 36.9GB. Since this exceeds the single GPU capacity of 24GB, dynamic batch management is enabled.

[0076] The first 8 requests in the initial batch are selected, and the RESERVE operation is performed to allocate pages. To reduce kernel function branch divergence, a padding strategy is adopted, that is, the sequences within the batch are grouped by length, and the block table lengths of sequences in the same group are aligned. For example, sequences with lengths of 512 and 1024 are aligned to 1024.

[0077] Furthermore, batch inference is performed using the fusion attention kernel function. The `mask_mod` function within the kernel function implements sequence isolation. It is important to emphasize that this mask is used to ensure that each query only focuses on valid key-value pairs within the same sequence, and that different sequences do not interfere with each other.

[0078] When certain sequences in a batch are generated, such as two sequences of 512 tokens (approximately 0.5 seconds later), the physical pages they occupied are immediately reclaimed, and the two subsequent waiting 512 token requests are added to the batch. Ideally, with this pipelined approach, all 16 requests are completed within approximately 3.2 seconds.

[0079] In this embodiment, the weighted shortest job first (WSJF) algorithm is used to dynamically adjust the execution order: In the formula: value is the business priority weight, estimated_pages is the estimated page requirement, and remaining_steps is the remaining generation steps.

[0080] Furthermore, sequences with the same prefix are detected within the batch. An atomic reference counter is added to the shared key-value pages to avoid duplicate storage. Prefix identification is implemented using a hash table, where the key is the hash value of the prefix token sequence and the value is the corresponding list of physical pages. When multiple sequences' block tables point to the same physical page, the reference counter is incremented; the page is only truly reclaimed when the counter reaches zero.

[0081] Ultimately, this embodiment achieves the following performance metrics on an L4 GPU: average latency of 0.38ms per token, throughput of approximately 2630 tokens / s, and peak memory usage of 22.8GB. Compared to the traditional approach of pre-allocating a contiguous cache of maximum length (4096), this method saves approximately 67% of the KV cache memory (the traditional approach requires 16...). 4096 4096 32 2 2 bytes = 34.4 GB, this method actually uses about 11.5 GB).

[0082] Example 4: This example uses a text summarization task to illustrate the application of this method in an encoder-decoder architecture and further demonstrates how it works in conjunction with a sparse attention mechanism.

[0083] For key-value cache management in encoder-decoder architectures, specifically: For encoder-decoder models such as T5 and BART, this method needs to be adapted to incorporate a cross-attention mechanism. Key modifications: Encoder KV Cache: During the encoding phase, the KV of all tokens is calculated at once and stored in a separate encoder page pool. Since the encoder output is fixed, these pages remain unchanged throughout the decoding process and can be shared by multiple decoding sequences.

[0084] Decoder self-attention key-value cache: Managed according to the standard process of this method, and updated incrementally for each token.

[0085] Cross-attention: The decoder's query Q accesses the encoder's key-value pairs. Cross-pooling access is achieved by adding an encoder sequence ID check to the mask_mod. This design makes the memory usage of the encoder KV independent of the decoding length, making it particularly suitable for long input-short output scenarios.

[0086] Specifically, for working in conjunction with sparse attention mechanisms: For models employing sparse attention patterns (such as Longformer's sliding window + global token), this method supports this by extending mask_mod: Local window: mask = (|q_idx - kv_idx| ≤ window_size) Global token: mask=(kv_idxingglobal_token_ids) Composite pattern: mask = local_mask ∨ global_mask In sparse mode, most key-value pairs do not participate in computation, allowing for further optimization of page allocation: only pages are allocated within the window and for the global token, with placeholders used for other positions. Experiments show that for a Longformer with a window size of 512, this optimization can save approximately 75% of the key-value cache memory.

[0087] Example 5: This example uses a lip-reading speech generation task to illustrate the extended application of this method in a visual-language multimodal scenario.

[0088] Scenario description: Input a 128-frame silent lip-reading video (approximately 5 seconds) and generate the corresponding Chinese speech mel spectrogram.

[0089] Parameter configuration: The visual encoder uses a Win Transformer (4 layers, window size 7). 7) The audio generator has a 3-layer iterative structure, with Flash Attention embedded in the 2nd layer. Page size (High audio frame rate)

[0090] Execution process: 1. Visual Feature Extraction. Input Video Encoded as local features by Swin Transformer and global features Local features capture lip shape details in each frame, while global features model temporal dependencies.

[0091] 2. First-layer generator. It converts local features... With noise The images are stitched together and then coarsely generated using 3D convolution. (80 represents the Mel frequency band number, and 512 represents the time frame number, corresponding to a 4x upsampling of the video frame). At this point, the spectrogram lacks global consistency and exhibits phoneme misalignment.

[0092] 3. FlashAttention fusion. As a query Global visual features Key obtained through projection Sum Since the number of audio frames (512) is much larger than the number of visual frames (128), a paginated key-value cache is used to manage visual features: allocation Page storage KV projection. Fusion attention calculation: This operation allows the audio representation to acquire global visual context and corrects the temporal alignment of local phonemes. The block computation of FlashAttention will reduce 512... The attention matrix of 128 is decomposed into 32. 8 small pieces (piece size 16) 16) SRAM usage increased from 512 128 4 bytes = 256KB reduced to 16 16 4 bytes = 1KB, significantly reducing HBM access.

[0093] 4. Second-level generator. Concatenation. Upsampled to transposed convolution Further refinement of the spectrum details.

[0094] 5. Third-level generator. Final output. The waveform is converted using the Griffin-Lim algorithm.

[0095] 6. Training losses include: reconstruction loss (L1 distance), adversarial loss (multi-scale discriminator), and synchronization loss (InfoNCE). Synchronization loss ensures accurate lip-sync between the generated audio and video. in , This is the balance coefficient.

[0096] In summary, the experimental results show that the word error rate using this method is reduced to 43.19%, which is an improvement of about 5.5% compared to the baseline VCA-GAN's 48.7%, and the inference latency is reduced from 23.67ms to 18.24ms, which is an acceleration of about 23%.

[0097] Example 6, refer to Figures 2-3 This embodiment tests the performance of the method on an NVIDIA L4 GPU and compares it with a traditional continuous key-value caching scheme. Specifically: The test model was LLaMA-7B (32 layers, 32 attention heads, 4096 hidden dimensions), with an inference accuracy of FP16 and a page size of P=128. The test data covered the first token latency, average latency of subsequent tokens, throughput, memory usage, and numerical accuracy under different sequence lengths.

[0098] ① In the latency and throughput tests, the test sequence lengths were 128, 512, 2048, and 8192 tokens, respectively. For each length, 100 inference runs were performed and the average value was taken. The results are as follows (e.g.) Figure 2 (as shown) For a sequence length of 128 tokens: the first token latency (TTFT) is 12.3ms, the average latency of subsequent tokens is 0.28ms, and the throughput is 3571 tokens / s. Compared with the traditional continuous KV cache (TTFT 11.8ms, latency 0.26ms), the overhead increases by about 8%.

[0099] For a sequence length of 512 tokens: the TTFT is 15.7ms, the average latency is 0.32ms, and the throughput is 3125 tokens / s. The traditional solution has a TTFT of 14.9ms and a latency of 0.30ms, narrowing the performance gap to approximately 7%.

[0100] For a sequence length of 2048 tokens: the TTFT is 28.4ms, the average latency is 0.38ms, and the throughput is 2632 tokens / s. The traditional solution has a TTFT of 52.1ms (due to insufficient video memory triggering partial CPU offload) and a latency of 0.71ms. This method achieves an acceleration of approximately 1.9 times.

[0101] For a sequence length of 8192 tokens: the TTFT is 89.2ms, the average latency is 0.51ms, and the throughput is 1961 tokens / s. Traditional solutions cannot run due to memory overflow and must use CPU offload, resulting in a TTFT exceeding 800ms and a latency of approximately 6.8ms. This method achieves approximately 13 times the speedup.

[0102] In summary, when the sequence length exceeds 2048 tokens, the advantages of this method compared to traditional schemes gradually become apparent, and the delay increases approximately linearly with the sequence length, rather than exponentially.

[0103] ② In the memory efficiency test, the memory usage of the LLaMA-7B model under different sequence lengths and batch configurations was recorded. It should be noted that the model weights were fixed at 13.4GB, and the activation value buffer was approximately 0.8~1.2GB, depending on the batch size.

[0104] like Figure 3 As shown in the figure, the left figure shows that the memory usage of both schemes increases linearly with the number of tokens, and the curves almost overlap, indicating that the paging mechanism itself does not increase the additional memory overhead. The right figure reveals that the memory saving rate of the paging scheme exhibits obvious periodic fluctuation characteristics, that is, a significant peak (up to 50%) appears at specific token count points (such as page boundaries such as 128, 256, 512, etc.), while the saving rate drops to 0% or even negative values ​​in the internal position of the page. This sawtooth fluctuation is due to the internal fragmentation effect of the paging mechanism: when the sequence length exactly fills an integer number of physical pages, there is no fragmentation loss, but when the end of the sequence only occupies part of the page, free space in the page is generated, resulting in a short-term decrease in memory utilization. However, as the sequence length increases, the relative impact of internal fragmentation gradually dilutes, the fluctuation range of the saving rate narrows and approaches 0%, and the overall average saving rate stabilizes at a low overhead level of 2.8%-4.7%, verifying the effectiveness of the paging scheme in controlling memory waste.

[0105] Specifically, for single-sequence KV caching: for a 2048-token sequence, the traditional solution requires pre-allocation of 1.05GB of contiguous video memory, while this method actually occupies 1.02GB, saving about 3%. For an 8192-token sequence, the traditional solution requires 4.19GB, while this method occupies 4.10GB, saving about 2%.

[0106] For mixed-length batch inference: the traditional approach requires pre-allocation of 16 × 1.05 GB = 16.8 GB, while this method, through dynamic allocation and prefix sharing, actually occupies about 11.3 GB, saving about 33%.

[0107] Regarding the ratio of memory overhead to the theoretical minimum: the theoretical minimum is the actual number of tokens × d × L × 2 × 2 bytes. The additional overhead of this method includes page tables, internal fragmentation, and free list metadata. Actual testing shows that the overhead ratio is between 2.8% and 4.7%, still meeting the design goal of less than 5%.

[0108] To verify numerical accuracy, this embodiment calculates perplexity on the WikiText-103 dataset to confirm that pagination key-value caching and address translation do not affect the model output quality. Baseline (continuous KV buffer): PPL=7.32 This method (paged key-value cache, P=64): PPL=7.31 This method (paged key-value cache, P=128): PPL=7.31 Furthermore, 100 samples were randomly selected, and the L2 distance between the output logits of this method and the baseline was calculated. The average value was 1.3 × 10^(-5), and the maximum value was 8.7 × 10^(-5), both within the FP16 accuracy error range (approximately 6 × 10^(-5)). This demonstrates that the paging mechanism and address translation do not introduce numerical errors.

[0109] In summary, this method demonstrates significant advantages in latency and throughput for long sequence inference, while maintaining near-theoretical minimum memory usage and consistent numerical accuracy with the baseline. Particularly noteworthy is its ability to efficiently complete inference in scenarios with 8192 tokens, where traditional methods are insufficient for single-GPU operation.

[0110] Example 7: This example also provides a joint optimization system for Flash Attention and Paged Attention in long sequence processing, including: The acquisition module is used to acquire the paginated KV cache and the corresponding block table of the target sequence. The paginated KV cache stores key tensors and value tensors in non-contiguous physical pages with a fixed page size, and the block table records the mapping relationship between logical pages and physical pages. The processing module is used to perform block-based iterative calculations for Flash Attention. During this process, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time based on the block table. Then, based on the converted physical address, it collects the corresponding key tensors and value tensors from non-contiguous physical pages and uses the collected key tensors and value tensors to complete the block-based attention score calculation with the query block. The fusion attention kernel function has address conversion logic embedded in it through just-in-time compilation.

[0111] This embodiment also provides an electronic device suitable for the joint optimization of Flash Attention and Paged Attention in long sequence processing, including: a memory and a processor; the memory is used to store computer-executable instructions, and the processor is used to execute the computer-executable instructions to implement the joint optimization method of Flash Attention and Paged Attention in long sequence processing as proposed in the above embodiment.

[0112] This embodiment also provides a storage medium storing a computer program that, when executed by a processor, implements the joint optimization method for Flash Attention and Paged Attention in long sequence processing as proposed in the above embodiments.

[0113] The storage medium proposed in this embodiment and the method for jointly optimizing Flash Attention and PagedAttention in long sequence processing proposed in the above embodiments belong to the same inventive concept. Technical details not described in detail in this embodiment can be found in the above embodiments, and this embodiment has the same beneficial effects as the above embodiments.

[0114] Based on the above description of the implementation methods, those skilled in the art can clearly understand that the present invention can be implemented using software and necessary general-purpose hardware, and of course, it can also be implemented using hardware. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product can be stored in a computer-readable storage medium, such as a computer floppy disk, read-only memory (ROM), random access memory (RAM), flash memory, hard disk, or optical disk, etc., including several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods of the various embodiments of the present invention.

[0115] It should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and not to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the technical solutions of the present invention, and all such modifications or substitutions should be covered within the scope of the claims of the present invention.

Claims

1. A joint optimization method of Flash Attention and Paged Attention in long sequence processing, characterized in that... include, Obtain the paginated KV cache and corresponding block table of the target sequence; wherein, the paginated KV cache stores key tensors and value tensors in non-contiguous physical pages according to a fixed page size, and the block table records the mapping relationship between logical pages and physical pages; During the block-based iterative computation of Flash Attention, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time according to the block table. Based on the converted physical address, the corresponding key tensors and value tensors are collected from the non-contiguous physical pages. The collected key tensors and value tensors are then used with the query block to complete the block-based attention score calculation, thereby achieving joint optimization of paginated KV caching and Flash Attention. The fusion attention kernel function has address conversion logic embedded through just-in-time compilation.

2. The joint optimization method of Flash Attention and Paged Attention in long sequence processing as described in claim 1, characterized in that: The step of obtaining the paginated key-value cache and corresponding block table of the target sequence includes: Pre-allocate a global key-value cache buffer based on the target hardware parameters and construct a lock-free linked list of free physical pages; In response to receiving an input sequence, the required number of physical pages are atomically reserved from the free physical page list through an atomic comparison swap operation, and a block table recording the mapping relationship between logical pages and physical pages is constructed for the input sequence.

3. A joint optimization method for Flash Attention and Paged Attention in long sequence processing as described in claim 1 or 2, characterized in that: The address translation logic is provided in the form of a custom index mapping function and is inlined into the fusion attention kernel function through just-in-time compilation. Before performing the distributed collection, the method caches the block table of the target sequence in the shared memory of the fusion attention kernel function to accelerate the translation from logical index to physical address.

4. The joint optimization method of Flash Attention and Paged Attention in long sequence processing as described in claim 3, characterized in that: The method further includes: In response to the completion of sequence processing, the block table of the completed sequence is traversed, and the atomic reference counter of the physical page corresponding to each physical page identifier in the block table is decremented by 1; In response to the atomic reference counter being zeroed, the identifier of the physical page is returned to the list of free physical pages; In response to multiple sequences of block tables pointing to the same physical page, the atomic reference counter of the same physical page is incremented to achieve prefix sharing of the same physical page by the multiple sequences.

5. The joint optimization method of Flash Attention and Paged Attention in long sequence processing as described in claim 4, characterized in that: The method further includes: Maintain a global batch state, which records the current length, block table pointer, and estimated remaining steps for each sequence within the batch; Based on the global batch status and the business priority of each sequence, the execution order of sequences within the batch is dynamically adjusted. In response to the completion of processing of any sequence, reclaim the physical pages occupied by the completed sequence and accept new requests into the batch.

6. A joint optimization method for Flash Attention and Paged Attention in long sequence processing as described in any one of claims 4 or 5, characterized in that: The method further includes: adaptively selecting the page size based on the length of the input sequence; wherein: In response to the length of the input sequence being less than a first threshold, a first page size is adopted; In response to the length of the input sequence being greater than or equal to the first threshold, a second page size greater than the first page size is adopted.

7. The joint optimization method of Flash Attention and Paged Attention in long sequence processing as described in claim 6, characterized in that: The page size is a power of 2.

8. A joint optimization system for Flash Attention and Paged Attention in long sequence processing, employing the method described in any one of claims 1-7, characterized in that, include: The acquisition module is used to acquire the paginated KV cache and the corresponding block table of the target sequence; wherein, the paginated KV cache stores key tensors and value tensors in non-contiguous physical pages according to a fixed page size, and the block table records the mapping relationship between logical pages and physical pages; The processing module is used to perform block-based iterative calculations for Flash Attention. During this process, the fusion attention kernel function converts the logical index of the key-value block to be accessed by the current query block into a physical address in real time based on the block table. Then, based on the converted physical address, it collects the corresponding key tensors and value tensors from the non-contiguous physical pages and uses the collected key tensors and value tensors to complete the block-based attention score calculation with the query block. The fusion attention kernel function has address conversion logic embedded in it through just-in-time compilation.

9. A computer device comprising a memory and a processor, wherein the memory stores a computer program, characterized in that, When the processor executes the computer program, it implements the steps of the method according to any one of claims 1 to 7.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the steps of the method according to any one of claims 1 to 7.

Images