A hierarchical tree attention routing method, device and equipment for super-long context

Abstract

The application provides a hierarchical tree attention routing method, device and equipment for an ultra-long context, cuts a KV cache according to token positions to each equipment and matches with a local Q-index sub-tree, so that each equipment searches from a root to a layer in the local sub-tree to obtain a local candidate Key position index set, simultaneously inputs a query vector into a Q-routing tree to select a path layer by layer to obtain an activated attention specialist set, a memory specialist set and an FFN specialist set, then the attention specialist and the memory specialist of each equipment perform calculation only in a range corresponding to the local candidate Key position index set and obtain a local output, finally a Q-communication tree aligned with a physical topology of the equipment performs tree-like reduction on the local output of each equipment stage by stage to obtain a layer output after being fused with a local FFN output; thus, three links of KV retrieval, specialist routing and cross-equipment reduction are converged to a logarithmic level by means of a tree structure, and collaborative dimension reduction is realized.

Description

Technical Field

[0001] This invention relates to the field of natural language processing, and in particular to a hierarchical tree-structured attention routing method, apparatus, and device for ultra-long contexts. Background Technology

[0002] The context window length of a large language model directly determines its information carrying capacity. The industry is evolving from 1M tokens to 1G tokens to support end-to-end understanding of enterprise-level full knowledge bases, full code repositories, and full document archives. However, when the context length L expands to 10... 9 At the scale of L, a single forward computation of the model needs to be completed sequentially in the following three steps: (a) retrieving the subset related to the query vector from the KV cache of length L; (b) selecting the target expert from the hybrid expert set containing N experts from the query vector; and (c) reducing the above retrieval and computation results among a distributed cluster composed of P GPUs to obtain the global output.

[0003] For the three stages mentioned above, existing technologies have proposed local improvements respectively: the scheme represented by Tree Attention utilizes the associativity of logsumexp to organize cross-device reduction into a tree structure, thereby reducing the number of communication steps in stage (c); the scheme represented by Fovea Transformer constructs a static multi-scale tree based on the token position to perform coarse-grained compression on stage (a); the scheme represented by Hierarchical Pruned Attention (HiP) relies on the attention locality assumption to perform tree search pruning on the key and perform top-k screening on stage (a); the scheme represented by the traditional hybrid expert model (MoE) uses a single-layer fully connected router to perform expert selection on stage (b).

[0004] While each of the above schemes achieves some optimization within its specific stage, stages (a), (b), and (c) still execute sequentially through their own independent data paths during a single forward computation, lacking structural alignment and intermediate result reuse. Specifically, the tree-structured reduction range of stage (c) is not aligned with the KV fragmentation range of stage (a), requiring a linear scan of all local keys before reduction; the key selection results from stage (a) are not provided to the expert computation in stage (b), necessitating independent re-searching of the entire sequence by the selected experts in stage (b); and the router in stage (b) continues to operate in O(N) linear matching mode as the expert scale N increases. The direct result is that the overall complexity of the forward computation cannot fall below the lower bound determined by the maximum complexity of the three stages, and the overall complexity still increases linearly with respect to at least one of the sequence length L, the number of experts N, and the number of devices P, making real-time inference at the 1G token scale impractical in engineering.

[0005] In view of the above, this application is hereby submitted. Summary of the Invention

[0006] This invention discloses a hierarchical tree-structured attention routing method, apparatus, and device for ultra-long contexts, aiming to solve the problems in existing ultra-long context processing schemes where the complexity of the three stages of KV retrieval, expert routing, and cross-device reduction cannot be simultaneously converged to the logarithmic level, and there is a lack of intermediate result reuse between stages, thus failing to form a collaborative dimensionality reduction.

[0007] The first embodiment of the present invention provides a hierarchical tree-based attention routing method for ultra-long contexts, including: A query vector is generated from the hidden state of the current token, and the query vector is broadcast to each device in the distributed device cluster that includes multiple devices; each device searches down from the root node to the leaf node in its local Q-index subtree to obtain the local candidate key position index set of the device. The query vectors are input into the Q-routing tree in parallel. The Q-routing tree is rooted at the query type and leafed at the experts. Paths that meet the preset conditions for scores are selected layer by layer from the root node to the leaf nodes to obtain the activated attention expert set, memory expert set, and FFN expert set. The attention expert set performs sparse attention computation within the key vector range corresponding to the local candidate key position index set of this device to obtain the local attention output of this device; the memory expert set performs memory retrieval within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of this device; the FFN expert set performs computation based on the hidden state to obtain the local FFN output of this device; The local attention output and local memory output of each device are respectively reduced in a hierarchical tree structure through a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output. The global attention output, the global memory output and the local FFN output are then merged to form the layer output of the current token.

[0008] A second embodiment of the present invention provides a hierarchical tree-structured attention routing device for ultra-long contexts, comprising: The query and retrieval unit is used to generate a query vector from the hidden state of the current token and broadcast the query vector to each device in the distributed device cluster that includes multiple devices; each device searches down from the root node to the leaf node in its local Q-index subtree to obtain the local candidate key position index set of the device. The expert routing unit is used to input the query vector into the Q-routing tree in parallel. The Q-routing tree is rooted at the query type and leafed at the experts. It selects paths from the root node to the leaf nodes based on the scores that meet the preset conditions, thereby obtaining the activated attention expert set, memory expert set, and FFN expert set. The local expert computing unit is used to perform sparse attention computation on the attention expert set within the key vector range corresponding to the local candidate key position index set of the device to obtain the local attention output of the device; to perform memory reading on the memory expert set within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of the device; and to perform computation on the FFN expert set based on the hidden state to obtain the local FFN output of the device. The cross-device reduction and fusion unit is used to reduce the local attention output and local memory output of each device through a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output respectively; and fuse the global attention output, the global memory output and the local FFN output as the layer output of the current token.

[0009] The third embodiment of the present invention provides a hierarchical tree-shaped attention routing device for ultra-long contexts, including a memory and a processor. The memory stores a computer program, which can be executed by the processor to implement a hierarchical tree-shaped attention routing method for ultra-long contexts as described in any of the above embodiments.

[0010] The fourth embodiment of the present invention provides a computer-readable storage medium storing a computer program, which can be executed by the processor of the device in which the computer-readable storage medium is located, to implement a hierarchical tree-shaped attention routing method for ultra-long contexts as described in any of the above embodiments.

[0011] Based on the hierarchical tree-structured attention routing method, apparatus, and device for ultra-long contexts provided by this invention, the KV cache is partitioned into devices according to token positions and equipped with local Q-index subtrees. Each device searches its local subtree from the root to the layer to obtain a local candidate key position index set. Simultaneously, the query vector is input into the Q-routing tree to select paths layer by layer to obtain the activated attention expert set, memory expert set, and FFN expert set. Then, the attention experts and memory experts of each device only perform calculations within the range corresponding to the local candidate key position index set and obtain local outputs. Finally, the local outputs of each device are reduced step by step by a Q-communication tree aligned with the physical topology of the devices and fused with the local FFN outputs to obtain the layer output. Thus, the three stages of KV retrieval, expert routing, and cross-device reduction all converge to the logarithmic level through the tree structure, and the output local candidate key position index set is reused as the range of subsequent expert calculations. The leaf positions of the Q-communication tree and the Q-index subtree fragments correspond to each other in the token position interval, thereby achieving collaborative dimensionality reduction among the three stages. Attached Figure Description

[0012] Figure 1 This is a flowchart illustrating a hierarchical tree-based attention routing method for ultra-long contexts provided in the first embodiment of the present invention; Figure 2 This is a schematic diagram of a hierarchical tree-shaped attention routing device for ultra-long contexts provided in the second embodiment of the present invention. Detailed Implementation

[0013] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0014] This invention discloses a hierarchical tree-structured attention routing method, apparatus, and device for ultra-long contexts, aiming to solve the problems in existing ultra-long context processing schemes where the complexity of the three stages of KV retrieval, expert routing, and cross-device reduction cannot be simultaneously converged to the logarithmic level, and there is a lack of intermediate result reuse between stages, thus failing to form a collaborative dimensionality reduction.

[0015] Please see Figure 1The first embodiment of this invention provides a hierarchical tree-structured attention routing method for ultra-long contexts, which can be executed by a routing device (hereinafter referred to as the system), specifically by one or more processors within the routing device. It is applied to a large language model comprising a key-value cache, an expert set, and a distributed device cluster. The experts are sub-neural network modules with independent parameters, including FFN experts, attention experts, and memory experts. The key-value cache is divided into multiple shards according to token positions and deployed on each device in the device cluster. Each device locally maintains a Q-index subtree for the corresponding shard. S101, Generate a query vector from the hidden state of the current token, and broadcast the query vector to each device in the distributed device cluster that includes multiple devices; Each device searches down from the root node to the leaf node in its local Q-index subtree to obtain the local candidate key position index set of the device; In this embodiment, the KV cache to be processed is first divided into consecutive segments according to the token position. Each shard is deployed to a separate device cluster. On a GPU device, this embodiment takes Each GPU device pre-constructs a local Q-index subtree for the KV shards it holds. The Q-index subtree uses the cluster centers of the Key vectors in the KV shards held by the device as internal nodes and the Key vectors as leaf nodes. Each leaf node stores both the Key vector and its token position in the KV cache. Specifically, this embodiment uses a hybrid approach of token position and semantics to construct the Q-index subtree, that is, it constructs the set of Key vectors covered by the current node. , according to the formula

[0016] Calculate its semantic diversity, where The set of key vectors covered by the current node. The number of Key vectors in this set. For the set of A key vector, This is the mean vector of all key vectors in the set. In this embodiment, a preset semantic diversity threshold is used. Set the value to 0.3 and set the preset minimum sub-block size. Take 64 tokens; when the calculated and At that time Perform further sub-node partitioning, with even-depth partitioning performed by binary search based on token position to preserve positional locality, and odd-depth partitioning performed using the K-means++ algorithm. Semantic clustering is performed on the inner key vectors to improve semantic hit rate; otherwise, it will... Retained as leaf nodes; where the number of cluster centers used by the K-means++ algorithm is... Instead of being fixed in advance, it is done according to the following formula.

[0017] Using the silhouette coefficient in the range The adaptive determination is as follows, where For the first The average Euclidean distance (cohesion) from each key vector to other key vectors in its cluster. For the first The minimum of the average Euclidean distances (separation) from each key vector to all key vectors in all other clusters. In this embodiment, it is taken as Furthermore, to overcome the local optima problem of single-step clustering, for each candidate... The value is evaluated by performing K-means++ clustering 10 times and taking the mean of the silhouette coefficients; after the above Q-index subtree is constructed, for each current token entering S101, this embodiment uses its hidden state. After linear transformation Generate query vectors and through collection communication Broadcast to Each device; on each device, starting from the root node of the Q-index subtree, at each current level, the following formula is applied.

[0018] Calculate the normalized similarity between the query vector and the cluster centers of each candidate child node in this layer, where For the current layer The cluster center vectors corresponding to each candidate child node, and siblings is the set of all candidate child nodes in the current layer. The temperature parameter is used; then, all candidate child nodes in the current layer are calculated according to the normalized similarity. Sort by size from largest to smallest, select The largest value Each child node serves as the starting point for the search of the next layer (beam search width in this embodiment). (Set to 3), and for the selected Each child node repeats the above calculation and selection process at the next level until the selected child node is a leaf node. The token positions of the Key vectors stored in all the reached leaf nodes are then aggregated to obtain the local candidate Key position index set of this device. In this embodiment To enable the Q-index subtree to be jointly trained end-to-end with the large language model, this embodiment uses the Gumbel-Softmax method based on the normalized similarity during the forward computation. The sampling beam search branch path involves injecting Gumbel noise into the similarity distribution of each layer and then performing argmax to obtain discrete branch selections. In the backpropagation, the gradient of the language modeling loss of the large language model is backpropagated along the sampled branch path to the cluster center vector of each internal node on the sampled path. This allows the internal node parameters of the Q-index subtree to be jointly updated with the parameters of the large language model, thereby enabling the retrieval behavior of the Q-index subtree to adaptively optimize according to the distribution of training data and the needs of downstream tasks.

[0019] S102, the query vector is input into a Q-routing tree in parallel, the Q-routing tree is rooted at the query type and leafed at experts; from the root node, paths that meet preset conditions are selected layer by layer to the leaf nodes, resulting in the activated attention expert set, memory expert set, and FFN expert set (the memory expert in this invention is a dedicated neural network module designed for ultra-long context scenarios, used for storing and quickly retrieving high-frequency long-range information, complementing the attention expert: the attention expert is responsible for dynamically calculating context relevance, and the memory expert is responsible for pre-storing high-frequency semantic features to reduce redundant computation overhead). In this embodiment, the query vector generated from the current token hidden state is started in parallel with the Q-index subtree retrieval of each device. Parallel input is used to construct a pre-built Q-routing tree to complete expert selection. This embodiment's Q-routing tree has three layers: a query type layer, a semantic dimension layer, and an expert leaf layer, from top to bottom. The query type layer predefines 4 to 8 coarse-grained types as nodes based on prior knowledge of human query intent (including fact queries, sentiment queries, instruction queries, reasoning queries, multi-turn queries, and meta-queries). Nodes in the semantic dimension layer are automatically generated using unsupervised clustering (hierarchical clustering is performed on query vectors belonging to the same coarse-grained type to obtain fine-grained subclasses). Each expert leaf node in the expert leaf layer corresponds to a specific FFN expert, attention expert, or memory expert, and is pre-associated with an expert type identifier indicating its category. Furthermore, each node in both the query type layer and the semantic dimension layer is associated with a trainable weight vector. With a trainable bias term Starting from the root node of the Q-routing tree, in the current layer of the query type layer and semantic dimension layer, the following formula is used:

[0020] Calculate the scores of the query vector and each candidate child node in the current layer, where The current layer number (value is 1 or 2, corresponding to the query type layer and semantic dimension layer respectively). The index of the candidate child node in the current layer. For the query vector at the th Layer for the first The scores calculated for each candidate child node; then all candidate child nodes in the current layer are processed according to the scores. Sort by size from largest to smallest, retain only those scores greater than the preset route retention threshold. child nodes (in this embodiment) (Take a value of 0.5), and when the number of child nodes retained exceeds the preset maximum number of child nodes to retain in a single-level path. Only the first one in the sorted list is taken. Each child node serves as the starting point for the search in the next layer (in this embodiment) (Set to 3), through this combined retention mechanism of threshold truncation and upper limit truncation, a query involving multiple semantic categories can activate multiple paths with different intensities, thereby avoiding the hard truncation of multi-label queries when traditional MoE only takes the Top-k. For each retained child node, the above score calculation and retention process is repeated in the next layer until the retained child node is an expert leaf node in the expert leaf layer. The experts corresponding to all the reached expert leaf nodes are respectively assigned to the activated attention expert set according to their associated expert type identifier. A collection of memory experts Gathering with FFN experts To further ensure that each activated expert in the three sets can participate in the calculation according to its path strength in subsequent fusion, this embodiment assigns a normalized weight to each activated expert, firstly according to the following formula.

[0021] Calculate the path weight of each activated expert before normalization, where path is the path from the root node to the leaf node corresponding to that activated expert. For the path in the th The score corresponding to the selected node in the layer. The set of sibling nodes at the same level as the selected node is obtained by multiplying the softmax probabilities of the selected nodes at each level along the path to get the path probability of the leaf expert. , , The path weights of all activated experts in the three sets are globally normalized to obtain the final weights used for calculation and output fusion by each expert. , , .

[0022] S103, perform sparse attention calculation on the attention expert set within the key vector range corresponding to the local candidate key position index set of this device to obtain the local attention output of this device; perform memory reading on the memory expert set within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of this device; perform calculation on the FFN expert set based on the hidden state to obtain the local FFN output of this device (the memory slot is the basic storage unit of the global key-value storage matrix, which corresponds one-to-one with the token position in the KV cache); In this embodiment, the local candidate key location index set obtained by each device is adopted. With activated attention expert group A collection of memory experts Gathering with FFN experts Local computations of the three types of experts are performed in parallel on each device; for the attention expert set Every attention expert This embodiment only applies to the local candidate key location index set of this device. The importance scores of the query vector and each key vector are calculated internally, i.e., according to the following formula.

[0023] Calculate the importance score, where for The token position number in the middle, Location in the KV cache The key vector at that location, For the first An attention expert in position Trainable biases at the location (used to enable different attention experts to develop differentiated attention preferences for the same location). For the first An attention expert analyzed the current token's position. The importance score calculated for the location; then... All token positions are ranked according to their importance score. Sort by size from largest to smallest, and select the first element in the sorted list. The token positions constitute the sparse attention mask of this attention expert (the sparse attention window size in this embodiment). (Set to 2048), and then based on the sparse attention mask, select... The key vector and value vector corresponding to the i-th token position are subjected to standard softmax attention computation to obtain the i-th token position. The output of each attention expert, and normalized according to the assigned weights. After weighting, the local attention output of this device is obtained by summing the results from each expert. Therefore, the scope of attention computation extends from the length of the entire sequence. First compressed using a Q-index subtree to Then, it is sorted and compressed by importance score to... , achieved from arrive Then The dual dimensionality reduction reduces the computational complexity of a single attention expert from the original... Down to The constant order; for the set of memory experts Each memory expert in the local candidate key location index set... Memory retrieval is performed within the corresponding memory slot range (that is, the memory matrix inside the memory expert is accessed according to...). After filtering out the corresponding slots, key-value matching is performed using the query vector (according to the normalized weights assigned in S102). The local memory output of this device is obtained by weighted summation. For the FFN expert set Each FFN expert in the game directly uses the hidden state of the current token. The expert performs its own two-layer fully connected feedforward computation as input, according to the normalized weights assigned in S102. The local FFN output of this device is obtained by weighted summation. The above three types of local output , , That is, it serves as input for subsequent cross-device reduction and final fusion.

[0024] S104, the local attention output and local memory output of each device are respectively reduced by a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output respectively; the global attention output, the global memory output and the local FFN output are fused together as the layer output of the current token.

[0025] In this embodiment, the local attention output obtained by each device With local memory output First, a pre-constructed Q-communication tree reduces the local output of each device to a global output, and then merges it with the local FFN output to obtain the layer output of the current token. The Q-communication tree is constructed hierarchically according to the physical topology of the device cluster. Devices within the same NVLink domain (in this embodiment, every 4 GPUs share one NVLink domain) are organized into the bottom-level subtree of the Q-communication tree to utilize their high bandwidth for fast reduction; devices within the same InfiniBand domain (in this embodiment, every 16 GPUs share one InfiniBand domain) are organized into the middle-level subtree of the Q-communication tree for medium-bandwidth reduction; and devices across... Devices in the InfiniBand domain only undergo reduction via RDMA at higher levels of the Q-communication tree to minimize the use of low-speed cross-domain links. Furthermore, the position of each device's corresponding leaf node in the Q-communication tree corresponds to the token position range of the KV shard held by that device in the KV cache. This alignment ensures that the token position range covered by the Q-index subtree maintained locally by each device precisely corresponds to the leaf position of that device in the Q-communication tree. This allows cross-device communication to only transmit the reduction result without transmitting the candidate key itself, thus decoupling communication volume from the number of devices. For each non-leaf node of the Q-communication tree, the following formula is used...

[0026]

[0027] The attention outputs and corresponding logsumexp values ​​of its left and right child nodes are merged into the attention output and logsumexp value of this non-leaf node, where the logsumexp value is... The logarithm of the local softmax denominator, synchronously saved by each device during sparse attention computation, is obtained by recursively reducing the attention output from the leaf nodes of the Q-communication tree upwards to the root node. The attention output obtained at the root node is then used as the global attention output. Local memory output for each device Similarly, the global memory output is obtained by reducing the Q-communication tree according to the above formula. The logsumexp reduction structure in this embodiment utilizes the associativity of the logsumexp operation, making it possible to reduce the total number of operations that would otherwise need to be performed on the entire system. Execution between devices Step-by-step linear reduction is transformed into a step-by-step linear reduction along the tree depth. Step-by-step parallel reduction, in The required reduction steps are reduced from 1024 steps to 10 steps; then, the following formula is used.

[0028] The global attention output The global memory output Local FFN output The hidden state of the current token (As a residual term, it participates in the fusion to preserve the original information flow and stabilize deep network training.) After summation, it is normalized by the LayerNorm layer to obtain the layer output of the current token. The final output of the current layer of the large language model is sent to the next layer. The parameters of the large language model, the Q-index subtree, and the Q-routing tree on which the above reasoning process depends are pre-trained according to the following phased dynamic learning strategy. In this embodiment, firstly, on training samples with sequence lengths of 4k to 32k tokens, only the cluster center vectors of each internal node of the Q-index subtree and the trainable weight vectors of each layer node of the Q-routing tree are trained. The loss function adopts contrastive loss so that queries with the same semantics can be mapped to the path of similar nodes in the Q-index subtree and the Q-routing tree. Then, in the second stage, the sequence length of the training samples is gradually doubled and expanded to 1G tokens, and the language modeling loss is jointly optimized (used to drive the...). The parameters of the large language model are updated together with the parameters of the Q-index subtree and the Q-routing tree, including dynamic language generation quality, tree path supervision loss (used to make the path sampled by the beam search close to the true optimal path), and expert load balancing loss (used to avoid a few experts being overactivated while the rest are idle for a long time). In the second stage, whenever the sequence length of the training samples doubles, the Q-index subtree is reconstructed and the cluster center vector of its internal nodes is inherited from the previous stage (that is, the cluster center vector obtained by training in the previous stage is used as the initial value after reconstruction in this stage and then training continues), so that the structure of the Q-index subtree can adapt to the expanded KV cache space and maintain the continuity of the parameters accumulated during training.

[0029] Experimental data and effect verification To verify the technical effectiveness of this invention, we conducted a comparative experiment on a distributed cluster consisting of 20 A100 GPUs, testing the inference performance and long document question answering accuracy of this embodiment compared to existing mainstream technical solutions under different context lengths. The experiment used a 7B-parameter large language model as the benchmark, and the test dataset was the LongBench long text understanding benchmark. The experimental results are shown in Table 1.

[0030] Table 1 Performance Comparison of Different Technical Solutions

[0031] It should be noted that the 1G context scenario adopts a three-level KV cache architecture consisting of GPU memory and high-speed SSD. Hot Key-Value Layer: The most recent 2048 tokens + key-value pairs of the current routing path are stored in GPU memory (58GB). Warm KV layer: The most recent 1M tokens are stored in the node's local Optane persistent memory (access latency <1μs). Cold KV layer: The remaining 99.9% of tokens are stored in a distributed high-speed SSD cluster (access latency <10μs). The tree-structured routing mechanism of this invention ensures that 99.9% of queries only need to access the hot KV layer, so the inference latency is not significantly increased.

[0032] Key parameter selection criteria: Beam search width b=3: While ensuring a retrieval hit rate of over 99%, the search complexity of the Q-index subtree is controlled within... ; Sparse attention window size : Reduce the computational complexity of attention by two orders of magnitude, provided that the accuracy of long document question answering decreases by no more than 2%; Q-Route Tree Single-Level Path Retention Limit = 3: It can support multiple intent queries simultaneously, and the number of expert activations is controlled within a reasonable range.

[0033] Experimental results show that the present invention achieves an inference latency of less than 100ms / token in a 1G token ultra-long context through a three-tree collaborative dimensionality reduction architecture, with a video memory usage of only 1 / 4 of the existing technology, and a reduction of cross-device communication volume by an order of magnitude, while maintaining long text understanding capabilities comparable to the existing technology.

[0034] Please see Figure 2 A second embodiment of the present invention provides a hierarchical tree-structured attention routing device for ultra-long contexts, comprising: The query retrieval unit 201 is used to generate a query vector from the hidden state of the current token and broadcast the query vector to each device in the distributed device cluster including multiple devices; each device searches down from the root node to the leaf node in its local Q-index subtree to obtain the local candidate key position index set of the device. The expert routing unit 202 is used to input the query vector into the Q-routing tree in parallel. The Q-routing tree has the query type as the root and the experts as the leaf nodes. It selects the path from the root node to the leaf node with the score meeting the preset condition, and obtains the activated attention expert set, memory expert set and FFN expert set. The local expert computing unit 203 is used to perform sparse attention calculation on the attention expert set within the key vector range corresponding to the local candidate key position index set of the device to obtain the local attention output of the device; to perform memory reading on the memory expert set within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of the device; and to perform calculation on the FFN expert set based on the hidden state to obtain the local FFN output of the device. The cross-device reduction and fusion unit 204 is used to reduce the local attention output and local memory output of each device through a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output respectively; and fuse the global attention output, the global memory output and the local FFN output as the layer output of the current token.

[0035] The third embodiment of the present invention provides a hierarchical tree-shaped attention routing device for ultra-long contexts, including a memory and a processor. The memory stores a computer program, which can be executed by the processor to implement a hierarchical tree-shaped attention routing method for ultra-long contexts as described in any of the above embodiments.

[0036] The fourth embodiment of the present invention provides a computer-readable storage medium storing a computer program, which can be executed by the processor of the device in which the computer-readable storage medium is located, to implement a hierarchical tree-shaped attention routing method for ultra-long contexts as described in any of the above embodiments.

[0037] Based on the hierarchical tree-structured attention routing method, apparatus, and device for ultra-long contexts provided by this invention, the KV cache is partitioned into devices according to token positions and equipped with local Q-index subtrees. Each device searches its local subtree from the root to the layer to obtain a local candidate key position index set. Simultaneously, the query vector is input into the Q-routing tree to select paths layer by layer to obtain the activated attention expert set, memory expert set, and FFN expert set. Then, the attention experts and memory experts of each device only perform calculations within the range corresponding to the local candidate key position index set and obtain local outputs. Finally, the local outputs of each device are reduced step by step by a Q-communication tree aligned with the physical topology of the devices and fused with the local FFN outputs to obtain the layer output. Thus, the three stages of KV retrieval, expert routing, and cross-device reduction all converge to the logarithmic level through the tree structure, and the output local candidate key position index set is reused as the range of subsequent expert calculations. The leaf positions of the Q-communication tree and the Q-index subtree fragments correspond to each other in the token position interval, thereby achieving collaborative dimensionality reduction among the three stages.

[0038] Exemplary examples show that the computer program described in the third and fourth embodiments of the present invention can be divided into one or more modules, which are stored in the memory and executed by the processor to complete the present invention. The one or more modules can be a series of computer program instruction segments capable of performing specific functions, which describe the execution process of the computer program in implementing a hierarchical tree-shaped attention routing device oriented towards ultra-long contexts. For example, the apparatus described in the second embodiment of the present invention.

[0039] The processor referred to can be a Central Processing Unit (CPU), or other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor can be a microprocessor or any conventional processor. This processor is the control center of the hierarchical tree-shaped attention routing method for ultra-long contexts, connecting the various parts of the implementation of this method using various interfaces and lines.

[0040] The memory can be used to store the computer program and / or modules. The processor, by running or executing the computer program and / or modules stored in the memory, and by calling data stored in the memory, implements various functions of a hierarchical tree-shaped attention routing method for ultra-long contexts. The memory may mainly include a program storage area and a data storage area. The program storage area may store the operating system, at least one application program required for a function (such as sound playback function, text conversion function, etc.), etc.; the data storage area may store data created according to the use of the mobile phone (such as audio data, text message data, etc.). In addition, the memory may include high-speed random access memory, and may also include non-volatile memory, such as hard disk, memory, plug-in hard disk, smart media card (SMC), secure digital (SD) card, flash card, at least one disk storage device, flash memory device, or other volatile solid-state storage device.

[0041] If the implemented module is implemented as a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, all or part of the processes in the above embodiments of the present invention can also be implemented by a computer program instructing related hardware. The computer program can be stored in a computer-readable storage medium, and when executed by a processor, it can implement the steps of the various method embodiments described above. The computer program includes computer program code, which can be in the form of source code, object code, executable files, or certain intermediate forms. The computer-readable medium can include: any entity or device capable of carrying the computer program code, recording media, USB flash drives, portable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signals, telecommunication signals, and software distribution media, etc. It should be noted that the content included in the computer-readable medium can be appropriately added or removed according to the requirements of legislation and patent practice in the jurisdiction. For example, in some jurisdictions, according to legislation and patent practice, computer-readable media do not include electrical carrier signals and telecommunication signals.

[0042] It should be noted that the device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, 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 can be selected to achieve the purpose of this embodiment according to actual needs. Furthermore, in the accompanying drawings of the device embodiments provided by this invention, the connection relationships between modules indicate that they have communication connections, which can be specifically implemented as one or more communication buses or signal lines. Those skilled in the art can understand and implement this without any creative effort.

[0043] The above description is merely a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the technical scope disclosed in the present invention should be included within the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be determined by the scope of the claims.

Claims

1. A hierarchical tree-based attention routing method for ultra-long contexts, characterized in that, include: A query vector is generated from the hidden state of the current token, and the query vector is broadcast to each device in a distributed device cluster that includes multiple devices; Each device searches its local Q-index subtree from the root node down to the leaf node to obtain its local candidate key location index set. The Q-index subtree is constructed using a hybrid approach of token location and semantics. For the set B of key vectors covered by the current node in the Q-index subtree, its semantic diversity is calculated using the following formula. When the semantic diversity exceeds a preset threshold and... When the size of a sub-block exceeds the preset minimum size, sub-node partitioning is performed on block B. During partitioning, even-numbered depths are bisected by token position, and odd-numbered depths are subjected to semantic clustering using the K-means++ algorithm. ; in, is the number of Key vectors in set B; i is the index of the Key vector in set B; Let i be the i-th Key vector in set B; It is the mean vector of all Key vectors in set B; Let L2 be the norm of the vector; This is a measure of semantic diversity; The odd-numbered depths are subjected to semantic clustering using the K-means++ algorithm, specifically as follows: ; Where n is the number of candidate cluster centers to be evaluated; This is the preset upper limit for the number of cluster centers; Let be the average Euclidean distance from the i-th Key vector to other Key vectors in its cluster; It is the minimum value of the average Euclidean distance from the i-th Key vector to all Key vectors in other clusters; for and The larger of the two; The final number of cluster centers; The query vectors are input into the Q-routing tree in parallel. The Q-routing tree is rooted at the query type and leafed at the experts. Paths that meet the preset conditions for scores are selected layer by layer from the root node to the leaf nodes to obtain the activated attention expert set, memory expert set, and FFN expert set. The attention expert set performs sparse attention computation within the key vector range corresponding to the local candidate key position index set of this device to obtain the local attention output of this device; the memory expert set performs memory retrieval within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of this device; the FFN expert set performs computation based on the hidden state to obtain the local FFN output of this device; The local attention output and local memory output of each device are respectively reduced in a hierarchical tree structure through a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output. The global attention output, the global memory output and the local FFN output are then merged to form the layer output of the current token.

2. The hierarchical tree-based attention routing method for ultra-long contexts according to claim 1, characterized in that, Each device searches its local Q-index subtree from the root node down to the leaf node to obtain its local candidate key location index set, specifically: Each device's Q-index subtree uses the cluster center of the Key vector in the KV shard held by this device as the internal node and the Key vector as the leaf node. Each leaf node stores the Key vector and its token position in the KV cache. Starting from the root node of the Q-index subtree, each device calculates the normalized similarity between the query vector and the cluster centers of each candidate child node in that layer using the following formula: in, is the query vector; i is the index of the candidate child node in the current layer; This is the cluster center vector corresponding to the i-th candidate child node in the current layer; is the set of all candidate child nodes in the current layer; j is the index of the child node in this set; The preset temperature parameters; Select the normalized similarity of the i-th candidate child node for the query vector; All candidate child nodes in the current layer are compared according to the normalized similarity. Sort by size from largest to smallest, and select The b child nodes with the largest values ​​are used as the starting point for the search in the next layer; The normalized similarity calculation and Top-b selection process is repeated at the next level for the selected b child nodes until the selected child nodes are leaf nodes; the token positions of the key vectors stored in all the leaf nodes are collected to obtain the local candidate key position index set of this device.

3. The hierarchical tree-based attention routing method for ultra-long contexts according to claim 2, characterized in that, Also includes: End-to-end joint training of the Q-index subtree: In the forward computation, the branch path of the sampled beam is searched according to the normalized similarity in the Gumbel-Softmax manner. In the backward computation, the gradient of the language modeling loss of the large language model is backpropagated along the sampled branch path to the cluster center vector of each internal node on the sampled path, so that the cluster center vector of the internal node of the Q-index subtree is jointly updated with the parameters of the large language model.

4. The hierarchical tree-based attention routing method for ultra-long contexts according to claim 1, characterized in that, The process of selecting paths from the root node to the leaf nodes based on scores that meet preset conditions, specifically, is as follows: The Q-routing tree has three layers, from top to bottom: the query type layer, the semantic dimension layer, and the expert leaf layer. Each node in the query type layer and the semantic dimension layer is associated with a trainable weight vector and a trainable bias term. Each expert leaf node in the expert leaf layer corresponds to an expert and is associated with an expert type identifier. The expert type identifier is used to indicate whether the expert corresponding to the expert leaf node belongs to the attention expert, memory expert, or FFN expert. Starting from the root node of the Q-routing tree, at the current layer in the query type layer and the semantic dimension layer, the score of the query vector and each candidate child node in the current layer is calculated using the following formula: in, Let l be the query vector; l is the current layer number in the Q-routing tree, and j is the index of the candidate child node in the current layer; The trainable weight vector associated with the j-th candidate child node in layer l; The trainable bias term associated with the j-th candidate child node in layer l; The score calculated for the query vector at the l-th layer for the j-th candidate child node; Sort all candidate child nodes in the current layer by score. Sort by size from largest to smallest, retain only child nodes whose scores are greater than the preset route retention threshold, and when the number of retained child nodes exceeds the preset single-layer path retention limit, only the top child nodes in the sorted list are taken as the search starting point for the next layer. For each retained child node, repeat the above score calculation and retention process in the next layer until the retained child node is an expert leaf node in the expert leaf layer; then, according to the expert type identifier associated with all the reached expert leaf nodes, classify the experts into the activated attention expert set, memory expert set, or FFN expert set respectively.

5. A hierarchical tree-based attention routing method for ultra-long contexts according to claim 1, characterized in that, The sparse attention computation performed on the attention expert set within the key vector range corresponding to the local candidate key position index set of this device is specifically as follows: For each attention expert with index m in the set of attention experts, the following formula is applied only to the local candidate key location index set of this device. The importance scores of the query vector and each key vector are calculated internally, and in... The importance score is ranked first. The token positions form a sparse attention mask, and then the sparse attention calculation is performed based on the sparse attention mask: Where m is the index of the attention expert in the attention expert set, and j is... The token's position number, where t is the current time index of the token. For query vector, Let j be the key vector at position j in the KV cache. Let m be the trainable bias term for the m-th attention expert at position j. Let j be the importance score calculated by the m-th attention expert for the current token at position j. The default sparse attention window size.

6. The hierarchical tree-based attention routing method for ultra-long contexts according to claim 1, characterized in that, The stepwise tree-like reduction, achieved through a Q-communication tree aligned with the device's physical topology, yields the global attention output and the global memory output, specifically: For each non-leaf node of the Q-communication tree, the output to be reduced corresponding to its left and right child nodes. The corresponding logsumexp value is combined with the output of the non-leaf node and the logsumexp value according to the following formula, and then recursively reduced from the leaf node of the Q-communication tree from bottom to top until the root node. The output obtained at the root node is used as the global attention output or the global memory output: Specifically, for the reduction of attention output, the output corresponding to each leaf node of the Q-communication tree... For the local attention output of each device; for the reduction of the memory output, the output corresponding to each leaf node of the Q-communication tree. Output local memory for each device; and These are the logsumexp values ​​corresponding to the left and right child nodes of the current non-leaf node, respectively. and These are the outputs corresponding to the left and right child nodes of the current non-leaf node, respectively. This is the logsumexp value obtained after reducing the current non-leaf nodes; This is the output obtained after reducing the current non-leaf nodes.

7. A hierarchical tree-based attention routing device for ultra-long contexts, characterized in that, include: A query retrieval unit is used to generate a query vector from the hidden state of the current token and broadcast the query vector to each device in a distributed device cluster including multiple devices. Each device searches down from the root node to the leaf node in its local Q-index subtree to obtain its local candidate key position index set. The Q-index subtree is constructed using a hybrid token position and semantic approach. For the set B of key vectors covered by the current node in the Q-index subtree, its semantic diversity is calculated using the following formula. When the semantic diversity is greater than a preset threshold and When the size of a sub-block exceeds the preset minimum size, sub-node partitioning is performed on block B. During partitioning, even-numbered depths are bisected by token position, and odd-numbered depths are subjected to semantic clustering using the K-means++ algorithm. ; in, is the number of Key vectors in set B; i is the index of the Key vector in set B; Let i be the i-th Key vector in set B; It is the mean vector of all Key vectors in set B; Let L2 be the norm of the vector; This is a measure of semantic diversity; The odd-numbered depths are subjected to semantic clustering using the K-means++ algorithm, specifically as follows: ; Where n is the number of candidate cluster centers to be evaluated; This is the preset upper limit for the number of cluster centers; Let be the average Euclidean distance from the i-th Key vector to other Key vectors in its cluster; It is the minimum value of the average Euclidean distance from the i-th Key vector to all Key vectors in other clusters; for and The larger of the two; The final number of cluster centers; The expert routing unit is used to input the query vector into the Q-routing tree in parallel. The Q-routing tree is rooted at the query type and leafed at the experts. It selects paths from the root node to the leaf nodes based on the scores that meet the preset conditions, thereby obtaining the activated attention expert set, memory expert set, and FFN expert set. The local expert computing unit is used to perform sparse attention computation on the attention expert set within the key vector range corresponding to the local candidate key position index set of the device to obtain the local attention output of the device; to perform memory reading on the memory expert set within the memory slot range corresponding to the local candidate key position index set to obtain the local memory output of the device; and to perform computation on the FFN expert set based on the hidden state to obtain the local FFN output of the device. The cross-device reduction and fusion unit is used to reduce the local attention output and local memory output of each device through a Q-communication tree aligned with the physical topology of the device to obtain the global attention output and global memory output respectively; and fuse the global attention output, the global memory output and the local FFN output as the layer output of the current token.

8. A hierarchical tree-based attention routing device for ultra-long contexts, characterized in that, The system includes a memory and a processor, wherein the memory stores a computer program that can be executed by the processor to implement a hierarchical tree-based attention routing method for ultra-long contexts as described in any one of claims 1 to 6.

Images