Large language model speculation decoding acceleration method based on HBM-PIM near storage computing device

Abstract

The invention belongs to the related technical field of artificial intelligence acceleration, and particularly relates to a large language model speculation decoding acceleration method based on HBM-PIM near storage computing equipment, which comprises the following steps: updating the cumulative probability of a current draft sequence when a draft token is generated by a draft language model every time; and when the cumulative probability is lower than a threshold value obtained by pre-offline learning, stopping the generation of the draft of the request, and taking the current length as the self-adaptive draft length of the current round of the request. For each request, the draft language model puts generated draft tokens into a shared pool, and when the number of tokens in the shared pool reaches a preset value, the target large language model is triggered once for parallel verification; and temporarily storing draft tokens which are continuously generated in the verification process in the preparatory pool, and after the verification is finished, lifting the tokens corresponding to the requests in the preparatory pool to the shared pool in batches or discarding the tokens according to the verification result. According to the invention, the utilization efficiency of hardware resources can be improved to further accelerate speculative decoding.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence acceleration technology, and more specifically, relates to a method for accelerating large language model inference decoding based on HBM-PIM near-memory computing devices. Background Technology

[0002] Transformer-based generative large language models are widely used in tasks such as dialogue, code generation, and complex reasoning. Their reasoning typically includes a pre-filling stage and an autoregressive decoding stage. The pre-filling stage computes multiple requests in parallel; the decoding stage generates tokens autoregressively, with each token depending on the output of the previous token, thus throughput is limited by strict temporal dependencies. To improve the throughput of the decoding stage, a speculative decoding method introduces a lightweight draft language model. It first predicts a sequence of draft tokens of length d, and then the target large language model verifies these d tokens in parallel. If a draft token is rejected, the target large language model outputs the correct token and discards the subsequent autoregressively generated draft tokens. The draft language model's prediction speed is significantly faster than the target large language model's verification speed, and when the acceptance rate of draft tokens is high, the computational cost of the target model can be significantly reduced.

[0003] To adapt to the different computing modes of draft language models (memory-intensive) and target large language models (computation-intensive), and to further reduce data transfer overhead and improve bandwidth utilization, near-memory computing heterogeneous systems based on HBM-PIM have emerged. These systems place computationally intensive operators in high-computing-power units and push memory-bandwidth-limited operators down to in-memory computing units (HBM-PIM).

[0004] However, existing systems generally employ fixed draft sequence lengths, with draft language models and target large language models executed serially. Furthermore, they determine operator mapping strategies through offline analysis. This presents several challenges when handling large batches of requests: First, the fixed draft length leads to a significant amount of invalid computation. Draft acceptance rates vary considerably across different models, datasets, and task types. A single fixed draft length can generate a large number of draft tokens that are subsequently rejected by the target model in large-batch scenarios, making both their generation and verification computations redundant. Second, the serial execution of the draft language model and the target large language model creates pipeline cavitation: in a fixed serial process of first predicting all requests and then verifying the entire batch, the optimal draft length for different requests within a batch is inconsistent. Short-sequence requests must wait for long-sequence requests to complete their predictions before entering verification, resulting in severe idle time. Finally, static operator mapping is mismatched with dynamic workloads: the length of the draft token and the number of requests change continuously during the decoding process, causing the arithmetic strength (floating-point operations / data access volume) of operators such as attention and fully connected operations to switch between central processor bandwidth limitations and in-memory computing unit limitations. At this time, static mapping will cause operators to run on unsuitable accelerators, reducing overall efficiency.

[0005] Therefore, a novel speculative decoding acceleration architecture is needed that can adapt to draft length, decouple the draft model and target model pipeline, and dynamically sense arithmetic strength for operator scheduling when running in near-memory computing heterogeneous systems. Summary of the Invention

[0006] In view of the above-mentioned defects or improvement needs of the existing technology, the present invention provides a method for accelerating large language model speculation decoding based on HBM-PIM near-memory computing devices. The purpose is to improve the hardware resource utilization efficiency of HBM-PIM heterogeneous systems, so as to further accelerate speculation decoding.

[0007] To achieve the above objectives, according to one aspect of the present invention, a method for accelerating large language model speculative decoding based on an HBM-PIM near-memory computing device is provided, which employs an HBM-PIM near-memory computing device to perform pre-filling and speculative decoding, including:

[0008] Based on the embedding of each inference request, pre-fill the request and generate the first output token of the target large language model corresponding to the request; In the speculative decoding iteration phase, each round of speculative decoding iteration includes a draft prediction stage and a target model verification stage. The draft prediction stage is executed in parallel for all requests. The draft prediction stage for each request contains n iterations of prediction. The input of the first iteration of the draft prediction stage in the first round of speculative decoding for each request is the first output token corresponding to that request. The input of subsequent iterations of the draft prediction stage in each round of speculative decoding for each request is the output of the previous iteration. n is greater than or equal to 1, and the value of n is adaptively determined by the following method: the first output token generated in the draft prediction stage... t The next iteration of prediction determines the maximum probability from the probability distribution generated within the draft language model. p t The corresponding draft token, used as the token for the t-th iteration prediction, is the first... t The tokens predicted in the next iteration constitute the current draft token sequence for this request; update the cumulative probability H of the current draft token sequence for this request. t =H t-1 p t Initially, H0 = 1; set H t With preset threshold Comparison, when H t ≥ Continue with the next iteration of prediction to generate the next draft token, when H t < Draft token generation is terminated at this time, and n is adaptively determined at this point; In each round of speculation decoding, the draft tokens output by each iteration of the draft prediction stage for all requests are stored in a shared pool in real time. The number of tokens in the shared pool is checked in real time to see if it reaches a preset threshold. If so, a parallel verification of all draft tokens in the shared pool by the target large language model is triggered to determine if the output token of the target large language model matches the corresponding draft token in the shared pool. The shared pool is locked during the verification period. For requests that have not completed n iterations of prediction in this round of speculation decoding, iterative prediction continues during the verification period, and the predicted draft tokens are stored in a reserve pool. For each request, if there are inconsistent draft tokens in the shared pool, the draft tokens in the reserve pool for that request are discarded. All draft tokens are used, and the corrected output token of the target large language model is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. If all draft tokens in the shared pool of this request meet the consistency requirement in this round of verification, the draft tokens in the request's preparation pool are moved to the shared pool, and the last generated draft token in the preparation pool of this request is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. All the verified consistent draft tokens in each round of speculative decoding for each request, as well as the corrected output token of the target large language model for the verified inconsistent draft tokens, are used as the generation result of this request and output.

[0009] Furthermore, in each round of speculative decoding, based on the number of requests that will continue to execute the next iteration of prediction and the number of current draft tokens corresponding to all requests, the floating-point operations and memory access bytes required by various computational operators of the prediction model to execute the next iteration of prediction are periodically calculated to correspond to the arithmetic strength of the computational operators. According to the arithmetic strength of each operator and the known peak computing power and bandwidth of the device, the device mapping of the operator is updated, and the operator is assigned to the high-computing-power central processor or high-bandwidth in-memory computing unit in the HBM-PIM near-memory computing device for execution, thereby completing one round of speculative decoding iteration.

[0010] Furthermore, the model computation operators include attention operators and fully connected operators; the device mapping update method for each operator is as follows: Calculate the arithmetic strength of the fully connected operator in the prediction phase. , represented as In the formula, Represents the output dimension of the known draft language model; calculates the arithmetic strength of the attention operator during the validation phase, denoted as... In the formula, L represents the total length consisting of the request length and the draft token that the request has been verified to be valid; Determine the actual computing power of the central processor , represented as ,in, and These represent the floating-point arithmetic capability and memory access bandwidth of the central processor, respectively; determining the actual computing power of the in-memory computing unit. , represented as ,in, and This indicates the floating-point arithmetic capability and memory access bandwidth of the in-memory computing unit; Will Substitute respectively and In the formula, the central processor or in-memory computing unit corresponding to the maximum actual computing power is used as the execution unit of the fully connected operator in the next prediction stage; Will Substitute respectively and In the formula, the central processor or in-memory computing unit corresponding to the maximum actual computing power is used as the execution unit of the attention operator in the next verification stage.

[0011] Furthermore, preset thresholds Determined offline, the determination method is as follows: Based on the sample dataset, a complete speculative decoding process is run for each sample data point, with varying numbers of draft tokens. j The following records the corresponding cumulative probability H j In addition to the target large language model validation results, the condition pass rate is determined based on the target large language model validation results, and the condition pass rate corresponding to the sample data is statistically analyzed. H j The curve of change; In the transformation curve of each sample, determine the interval where the condition pass rate is not lower than a preset lower limit, and find the cumulative probability that maximizes the average number of draft tokens among all samples. H j The interval is used to perform a discrete grid search within the interval to determine the threshold for deployment, which serves as the preset threshold. .

[0012] Furthermore, regarding cumulative probability H t = H t-1 p t The calculations are performed in the logarithmic field, that is, in the first stage of the draft prediction phase. t Accumulated predictions in each iteration logp t get ; Then, when determining whether to perform iterative prediction, With log When comparing, ≥log Continue the prediction in the next iteration when <log The generation of the draft token for the current round of speculative decoding is terminated at this time.

[0013] Furthermore, the threshold is dynamically adjusted based on the system load during runtime. This includes: (1) reducing the number of requests when the equipment utilization rate is lower than the first threshold and the number of requests is lower than the second threshold. (1) Allow more draft tokens to be generated; (2) Increase the efficiency when the device utilization rate is higher than the third threshold or the number of requests is higher than the fourth threshold. Reduce the number of draft tokens.

[0014] Furthermore, before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests within each batch are encoded in parallel to obtain the embedding of each request. The inference decoding of the target large language model is accomplished by executing the draft language model and the target large language model in parallel on multiple pipeline stages. The deployment of the model on the HBM-PIM near-memory computing device includes: The near-memory computing devices are divided into S groups, each language model is divided into S pipelines, and all requests are divided into more than S batches. The near-memory computing devices in each group process the corresponding pipelines in sequence, and each batch is processed by the near-memory computing devices in each group in sequence. For each group of near-memory computing devices, based on the weight of each pipeline segment and its corresponding request key-value cache dimension, a combination of row and column sharding is used to map the weight matrix and cache matrix to different library groups of each near-memory computing device in order to balance the load and maximize internal bandwidth utilization. Each near-memory computing device includes a buffer layer chip and multiple vertically stacked DRAM chips. Each DRAM chip contains multiple library groups, and each library group is equipped with a global buffer and an accumulator. Each library in each library group integrates a computing unit, which includes several multipliers, an adder tree, and an output buffer. The computing units corresponding to the libraries in the same pseudo-channel of each DRAM chip work in parallel to make full use of the internal bandwidth of the pseudo-channel. The buffer chip integrates a dedicated function unit, including a comparator, adder, exponentiation unit, adder tree, divider, and multiplier, for performing non-matrix operations.

[0015] Furthermore, before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests within each batch are encoded in parallel to obtain the embedding of each request. The pre-configured content-addressable memory serves as the hardware implementation of a shared pool shared by all batches; the pre-configured on-chip SRAM cache serves as the hardware implementation of a reserve pool for each batch.

[0016] According to another aspect of the present invention, a large language model speculation decoding acceleration system based on HBM-PIM near-memory computing devices is provided for executing the large language model speculation decoding acceleration method as described above, comprising: a host, a hardware manager, multiple high-bandwidth near-memory computing devices, and a target large language model and a draft language model running on the computing devices. The host receives inference requests and encodes each request to obtain an embedding for each request. Multiple high-bandwidth near-memory computing devices, along with a target large language model and a draft language model running on the computing devices, are used to perform pre-filling and speculative decoding processes based on the request embeddings. During this process, the hardware manager is used to calculate the draft token probability and cumulative probability for each request to control the number of draft tokens for that request; to cache draft tokens to a shared pool or a pre-set pool; and to schedule operators.

[0017] In summary, compared with the prior art, the technical solutions conceived by this invention have the following main advantages: This invention proposes a method to accelerate large language model inference decoding based on HBM-PIM near-memory computing devices, introducing an adaptive draft sequence length adjustment mechanism: each time the draft language model generates a draft token, the prediction probability p of that token is calculated. t And update the cumulative probability of the current draft sequence. When H t Below the threshold obtained through pre-learning offline When the draft generation for a request is terminated, the current length is used as the adaptive draft length for this round of the request. A decoupled prediction-verification asynchronous pipeline shared pool and reserve pool mechanism is also introduced: For each batch of requests, the draft language model first places the generated draft tokens into the shared pool. When the number of tokens in the shared pool reaches the maximum value of the verification parallel capability, a parallel verification of the target large language model is triggered. To avoid the shared pool becoming too full prematurely, forcing requests that should continue prediction into verification, this invention temporarily stores the draft tokens generated during the verification process in the reserve pool. After the current round of verification is completed, the tokens of the corresponding requests in the reserve pool are either promoted to the shared pool or discarded based on the verification results. Through the above mechanism, the draft language model and the target large language model are decoupled in time, allowing prediction to continue during the verification phase, significantly reducing pipeline bubbling. Therefore, this invention can reduce the calculation of invalid draft tokens and avoid generating a large number of low-confidence tokens that will inevitably be rejected during the verification phase through adaptive draft length control based on cumulative probability, reducing redundant calculations from the source. In addition, it can alleviate the serial bottleneck between the draft language model and the target large language model, improve pipeline utilization, and achieve decoupled asynchronous pipeline for prediction and verification by using a shared pool and a preparatory pool. This allows the draft language model to continue predicting draft tokens during the verification of the target large language model, effectively reducing pipeline cavitation and improving the average resource utilization of the central processor and in-memory computing units. Attached Figure Description

[0018] Figure 1 The flowchart of a method for accelerating large language model inference decoding based on HBM-PIM near-memory computing device is provided for an embodiment of the present invention.

[0019] Figure 2 This is a schematic diagram of the time axis of the asynchronous pipeline for prediction and verification decoupling provided in an embodiment of the present invention.

[0020] Figure 3 This is a schematic diagram of data partitioning provided in an embodiment of the present invention; wherein, (a) is the key-value partitioning at the header level corresponding to the in-memory computing unit; and (b) is the key-value partitioning within a single header corresponding to a library group and a library.

[0021] Figure 4 The diagram below is a schematic diagram of the overall architecture of the accelerator hardware system provided in the embodiment of the present invention; wherein, (a) is the overall architecture diagram; and (b) is the near-memory computing device architecture diagram.

[0022] Figure 5 This is a schematic diagram of the hardware manager architecture provided in an embodiment of the present invention.

[0023] Figure 6A schematic diagram of the internal processing unit structure of the in-memory computing unit provided in an embodiment of the present invention; wherein, (a) is a schematic diagram of a partial pseudo-channel architecture; (b) is a schematic diagram of the computing unit; and (c) is a schematic diagram of the dedicated function unit. Detailed Implementation

[0024] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention. Furthermore, the technical features involved in the various embodiments of this invention described below can be combined with each other as long as they do not conflict with each other.

[0025] Example 1 A method for accelerating large language model speculative decoding based on HBM-PIM near-memory computing devices is proposed, which employs HBM-PIM near-memory computing devices to perform pre-filling and speculative decoding, such as... Figure 1 As shown, it includes: Based on the embedding of each inference request, pre-fill the request and generate the first output token of the target large language model corresponding to the request; In the speculative decoding iteration phase, each round of speculative decoding iteration includes a draft prediction stage and a target model verification stage. The draft prediction stage is executed in parallel for all requests. The draft prediction stage for each request contains n iterations of prediction. The input of the first iteration of the draft prediction stage in the first round of speculative decoding for each request is the first output token corresponding to that request. The input of subsequent iterations of the draft prediction stage in each round of speculative decoding for each request is the output of the previous iteration. n is greater than or equal to 1, and the value of n is determined adaptively. The determination method is: in the t-th iteration prediction generated in the draft prediction stage, the maximum probability p is determined from the probability distribution generated internally by the draft language model. t The corresponding draft token is used as the token for the t-th iteration prediction. The tokens output from the previous t iterations constitute the current draft token sequence for this request; the cumulative probability H of the current draft token sequence for this request is updated. t =H t-1 p t Initially, H0 = 1; set H t With preset threshold Comparison, when H t ≥ Continue with the next iteration of prediction to generate the next draft token, when H t < Draft token generation is terminated at this time, and n is adaptively determined at this point; In each round of speculation decoding, the draft tokens output by each iteration of the draft prediction stage for all requests are stored in a shared pool in real time. The number of tokens in the shared pool is checked in real time to see if it reaches a preset threshold. If so, a parallel verification of all draft tokens in the shared pool by the target large language model is triggered to determine if the output token of the target large language model matches the corresponding draft token in the shared pool. The shared pool is locked during the verification period. For requests that have not completed n iterations of prediction in this round of speculation decoding, iterative prediction continues during the verification period, and the predicted draft tokens are stored in a reserve pool. For each request, if there are inconsistent draft tokens in the shared pool, the draft tokens in the reserve pool for that request are discarded. All draft tokens are used, and the corrected output token of the target large language model is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. If all draft tokens in the shared pool of this request meet the consistency requirement in this round of verification, the draft tokens in the request's preparation pool are moved to the shared pool, and the last generated draft token in the preparation pool of this request is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. All the verified consistent draft tokens in each round of speculative decoding for each request, as well as the corrected output token of the target large language model for the verified inconsistent draft tokens, are used as the generation result of this request and output.

[0026] This embodiment proposes an adaptive draft sequence length adjustment mechanism: Each time the draft language model generates a draft token, it calculates the predicted probability p of that token. t And update the cumulative probability of the current draft sequence. When H t Below the threshold obtained through pre-learning offline When the time is reached, terminate the draft generation for this request and use the current length as the adaptive draft length for this round of the request; threshold The selection can be based on the validation set statistics, provided that a certain acceptance rate is guaranteed.

[0027] This embodiment also proposes a prediction-verification decoupling mechanism for the shared pool and reserve pool of asynchronous pipelines: like Figure 2As shown, for each batch of requests, the draft language model first places the generated draft tokens into a shared pool. When the number of tokens in the shared pool reaches the maximum value of the verification parallel capability or the GPU is idle, a parallel verification of the target large language model is triggered. To avoid the shared pool from being "filled" too early, forcing requests that should continue prediction into verification, this embodiment temporarily stores the draft tokens generated during the verification process in a reserve pool. After the current round of verification is completed, the tokens of the corresponding requests in the reserve pool are either promoted to the shared pool or discarded in batches according to the verification results. Through the above mechanism, the draft language model and the target large language model are decoupled in time, and prediction can continue during the verification phase, significantly reducing pipeline bubbling.

[0028] Therefore, the method in this embodiment can reduce the computation of invalid draft tokens. By using adaptive draft length control based on cumulative probability, it avoids generating a large number of low-confidence tokens that will inevitably be rejected during the verification stage, thus reducing redundant computation from the source. In addition, it can alleviate the serial bottleneck between the draft language model and the target large language model, improve pipeline utilization, and achieve decoupled asynchronous pipeline for prediction and verification by using a shared pool and a reserve pool. This allows the draft language model to continue predicting draft tokens during the verification of the target large language model, effectively reducing pipeline cavitation and improving the average resource utilization of the central processor and in-memory computing units.

[0029] As a preferred implementation, in each round of speculative decoding, the number of floating-point operations and memory access bytes required by various computational operators of the prediction model to execute the next iteration prediction is periodically calculated based on the number of requests that will continue to execute the next iteration prediction and the number of current draft tokens corresponding to all requests. This is used to correspond to the arithmetic strength of the computational operators. According to the arithmetic strength of each operator and the known peak computing power and bandwidth of the device, the device mapping of the operator is updated, and the operator is assigned to the high-computing-power central processor or high-bandwidth in-memory computing unit in the HBM-PIM near-memory computing device for execution, thereby completing one round of speculative decoding iteration.

[0030] This implementation proposes dynamic operator mapping to fully leverage the performance of heterogeneous systems. Based on online estimation of arithmetic strength and threshold comparison, it dynamically selects the central processor or in-memory computing unit to execute different operators, ensuring that operators such as attention and fully connected always run on units more suitable for their computational / bandwidth characteristics, thereby improving the overall system throughput.

[0031] The floating-point arithmetic capability (F) and memory access bandwidth (BW) of the computation unit and the in-memory computation unit are known; among them, the computation unit has a higher floating-point arithmetic capability (denoted as F). ) and relatively limited memory access bandwidth (denoted as The in-memory computing unit has relatively limited floating-point arithmetic capabilities (denoted as ). ) and higher memory access bandwidth (denoted as The arithmetic strength of an operator is defined as the number of floating-point operations (...). ) and the number of bytes accessed ( The ratio of ) Operators should be scheduled to units more suitable for their execution. Assuming computation and memory access can overlap, the performance of each unit in processing operators is as follows: Computational unit:

[0032] In-memory computing unit:

[0033] This represents the computing power required for the operator to execute on the central processor, through taking... and The minimum computing power among them is taken as the actual computing power of the central processor. Similarly, This represents the computational power required for the operator to execute on the in-memory computing unit, by fetching... and The minimum computing power among them is taken as the actual computing power of the in-memory computing unit. By comparison and Each operator is assigned to a unit with higher prediction power.

[0034] For attention operators in the prediction phase, due to their low arithmetic strength, they are always allocated to in-memory computation units. For fully connected operators in the verification phase, due to their high arithmetic strength, they are always allocated to computation units. However, in a preferred implementation, the fully connected operators in the prediction phase and the attention operators in the verification phase need to be dynamically scheduled to higher-performance units at runtime. The arithmetic strength of the fully connected operators in the prediction phase is expressed as... The arithmetic strength of the attention operator in the verification phase is expressed as: Among them, B and D in D out N, L, D head These represent the number of requests, model input dimension, model output dimension, number of draft tokens per request, request sequence length, and key-value cache dimension, respectively.

[0035] That is, the method in this embodiment also proposes an arithmetic strength-aware operator scheduling strategy: Before each prediction and verification round, the arithmetic strength of the draft language model fully connected operator and the target large language model attention operator is estimated based on the current effective batch size and draft token length. The arithmetic strength is compared with the pre-defined thresholds of "limited in-memory computing unit" and "limited central processor bandwidth" to determine whether the operator should be mapped to the central processor or the in-memory computing unit in the current round. For operators with stable arithmetic strength, such as the draft language model attention operator and the target large language model fully connected operator, a fixed mapping is used, while the remaining operators are dynamically remapped according to the above rules to improve the overall throughput and energy efficiency.

[0036] As a preferred implementation method, a preset threshold can be used. Determined offline, the determination method is as follows: Based on the sample dataset, a complete speculative decoding process is run for each sample data point, with varying numbers of draft tokens. j The following records the corresponding cumulative probability H j In addition to the target large language model validation results, the condition pass rate is determined based on the target large language model validation results, and the condition pass rate corresponding to the sample data is statistically analyzed. H j The curve of change; In the transformation curve of each sample, determine the interval where the condition pass rate is not lower than a preset lower limit, and find the cumulative probability that maximizes the average number of draft tokens among all samples. H j The interval is used to perform a discrete grid search within the interval to determine the threshold for deployment, which serves as the preset threshold. .

[0037] To reduce multiplication operations, a preferred implementation method is to accumulate probabilities. H t = H t-1 p t The calculations are performed in the logarithmic field, that is, in the first stage of the draft prediction phase. t Accumulated predictions in each iteration logp t get ; Then, when determining whether to perform iterative prediction, With log When comparing, ≥log Continue the prediction in the next iteration when <log The generation of the draft token for the current round of speculative decoding is terminated at this time.

[0038] As a further preferred implementation, the threshold can be dynamically adjusted based on the system load during operation. This includes: (1) reducing the number of requests when the equipment utilization rate is lower than the first threshold and the number of requests is lower than the second threshold. (1) Allowing the generation of longer draft sequences to improve the parallel validation of the target model and the overall system throughput; (2) When the device utilization rate is higher than the third threshold or the number of requests is higher than the fourth threshold, increase Shorten the draft sequence length to reduce the number of tokens to be verified and control the average latency.

[0039] As a preferred implementation, before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests in each batch are encoded in parallel to obtain the embedding of each request. The inference decoding of the target large language model is accomplished by executing the draft language model and the target large language model in parallel on multiple pipeline stages. The deployment of the model on the HBM-PIM near-memory computing device includes: The near-memory computing devices are divided into S groups, each language model is divided into S pipelines, and all requests are divided into more than S batches. The near-memory computing devices in each group process the corresponding pipelines in sequence, and each batch is processed by the near-memory computing devices in each group in sequence. For each group of near-memory computing devices, based on the weight of each pipeline segment and its corresponding request key-value cache dimension, a combination of row and column sharding is used to map the weight matrix and cache matrix to different library groups of each near-memory computing device in order to balance the load and maximize internal bandwidth utilization. Each near-memory computing device includes a buffer layer chip and multiple vertically stacked DRAM chips. Each DRAM chip contains multiple library groups, and each library group is equipped with a global buffer and an accumulator. Each library in each library group integrates a computing unit, which includes several multipliers, an adder tree, and an output buffer. The computing units corresponding to the libraries in the same pseudo-channel of each DRAM chip work in parallel to make full use of the internal bandwidth of the pseudo-channel. The buffer chip integrates a dedicated function unit, including a comparator, adder, exponentiation unit, adder tree, divider, and multiplier, for performing non-matrix operations.

[0040] The weight matrix includes: query, key and value generation matrix, attention output projection matrix, and two-layer feedforward network matrix. An alternating row and column partitioning strategy is used for the weight matrix. 1) The query, key, and value generation layers perform row-wise sharding based on the attention header dimension, so that the computation of each header is completed locally as much as possible; 2) The attention output projection layer is divided into column-oriented slices according to the output dimension, and the results are calculated by each in-memory computing device and then converged through reduction.

[0041] 3) For the weight matrix of the first layer of the feedforward network, column-oriented partitioning is adopted to improve the arithmetic strength of the matrix multiplication operator in the in-memory computing device; 4) For the second layer of the feedforward network, row-wise slicing is used to naturally balance the output dimension across devices.

[0042] In this implementation, the model layers of the target large language model and the draft language model are mapped to multiple near-memory computing devices in a pipelined parallel manner. The model layer is divided into S consecutive pipeline segments, each pipeline segment is mapped to a set of near-memory devices, and the batch size is divided into more than S to ensure that each pipeline segment is kept as fully loaded as possible.

[0043] Figure 3 This explains that the key-value cache mapping (i.e., cache matrix mapping) uses a header-based heap, dimension-based grouping, and library-based striping approach: Figure 3 (a) shows that each attention head in each language model is mapped to a specific in-memory computing device, and multiple attention heads corresponding to different requests can share the device but do not interfere with each other; Figure 3 (b) shows that the key values ​​are dimensionally sharded at the pseudo-channel and library group levels, striped at the library level, and local multiplication and accumulation are achieved through pseudo-channel-level broadcasting during querying.

[0044] As a preferred implementation, before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests in each batch are encoded in parallel to obtain the embedding of each request. The pre-configured content-addressable memory serves as the hardware implementation of a shared pool shared by all batches; the pre-configured on-chip SRAM cache serves as the hardware implementation of a reserve pool for each batch.

[0045] This implementation utilizes a shared pool and a reserve pool to achieve decoupled asynchronous pipelines for prediction and verification, while also enabling cross-batch aggregation when batching requests.

[0046] In other words, the shared pool is implemented using content-addressable storage or on-chip memory with an index table. Its capacity is configured according to the parallel verification capability of the target large language model and the maximum number of draft tokens. Draft tokens are stored and retrieved according to the request identifier and location index. The preparation pool is an on-chip SRAM cache configured separately for each batch, which can accommodate all new draft tokens generated in each batch within a verification cycle. It is used to store new draft tokens generated before the previous verification cycle is completed.

[0047] This embodiment proposes an in-memory computation speculation decoding acceleration method that supports adaptive draft sequence length. While ensuring generation quality, it reduces redundant generation and verification calculations caused by rejected draft tokens; reduces pipeline cavitation caused by serial execution of draft and target models, and improves parallelism; and dynamically switches operator mapping between the central processor and PIM based on the runtime operator arithmetic strength, fully leveraging the hardware potential of heterogeneous systems.

[0048] Example 2 A large language model speculation decoding acceleration system based on HBM-PIM near-memory computing devices is used to execute the large language model speculation decoding acceleration method as described above, including: a host, a hardware manager, multiple high-bandwidth near-memory computing devices, and a target large language model and a draft language model running on the computing devices. The host receives inference requests and encodes each request to obtain an embedding for each request. The host and multiple near-memory computing devices are connected via a high-speed interconnect unit. Multiple high-bandwidth near-memory computing devices, along with a target large language model and a draft language model running on the computing devices, are used to perform pre-filling and speculative decoding processes based on the embedding of the request. During this process, the hardware manager is used to calculate the draft token probability and cumulative probability for each request to control the number of draft tokens for that request; to cache draft tokens to a shared pool or a pre-set pool; and to schedule operators.

[0049] In practical implementation, the system is as follows: Figure 4 As shown, the host is used to receive user requests and issue batch tasks, run the inference service framework, and output decoding results; the hardware manager is used to adaptively control the draft token length, maintain the shared pool and the pre-pool, coordinate the prediction and verification pipeline, and perform arithmetic strength-aware operator scheduling; each near-memory computing device includes: a central processor for performing the pre-filling stage, executing computationally intensive operators, and some control logic; a router for managing data transmission between the central processor in this near-memory computing device and other near-memory computing devices; and multiple in-memory computing units based on HBM-PIM, each in-memory computing unit including a buffer layer chip and multiple vertically stacked DRAM chips, in each DRAM The chip integrates a library-side computing unit and a dedicated function unit on the buffer layer. Model weights and key-value caches are distributed and stored in various in-memory computing units. Among them, a row / column alternating partitioning strategy is adopted for multiple model weight matrices, and they are evenly distributed among different in-memory computing units, pseudo-channels, and library groups to improve bandwidth utilization and data reuse. The key-value cache is mapped to different in-memory computing units by header. At the same time, the key-value matrix is ​​partitioned in column-first and row-first order within the library group to make full use of the row buffer and improve internal bandwidth utilization.

[0050] The hardware manager can be deployed on the host or a standalone control chip, such as Figure 5 As shown, it includes: (1) Draft generator, used to control the draft sequence of each batch and calculate the current token probability and cumulative acceptance probability at each decoding step; in implementation, by configuring exponent unit, division unit, fixed-point multiplier, comparator and small-capacity on-chip temporary storage, it supports making draft termination decisions quickly, while controlling area and power consumption.

[0051] (2) The controller module, located in the draft generator, is used to collect the utilization rate of the central processor and in-memory computing units, the request queue length, and the operator execution status, providing a basis for threshold adjustment and operator scheduling. (3) Shared pool module, located in the draft generator, is used to cache draft tokens to be verified by the target model across multiple batches; (4) The preparation pool module is used to temporarily store newly added draft tokens when the previous round of verification is not completed, and to migrate the tokens in the pool to the shared pool or discard them according to the verification results after the verification is completed; (5) Operator scheduler module, used to estimate the arithmetic strength of each operator based on the current batch size and draft length and select the corresponding execution device; Each DRAM chip in the in-memory computing unit has multiple pseudo-channels, each pseudo-channel contains multiple library groups, each library group contains multiple libraries and is equipped with a certain number of accumulators and global buffers, and each library is configured with a computing unit, which includes multipliers, adders and buffers, for parallel processing of two data streams from the local row buffer and the pseudo-channel global buffer within one clock cycle.

[0052] like Figure 6 As shown, each DRAM chip is divided into multiple pseudo-channels. Figure 6 (a) illustrates that each pseudo-channel comprises multiple library groups, each containing several libraries, and each library group has a global buffer and an accumulator. In this embodiment, a computing unit is integrated next to each library. Figure 6 (b) shows a computation unit comprising several multipliers, an adder tree, and an output buffer, capable of processing two streams of data from the library interface and the global buffer in parallel within one clock cycle. Multiple computation units operate in parallel within the same pseudochannel, fully utilizing the pseudochannel's internal bandwidth. Figure 6 (c) shows the dedicated function units integrated on the buffer chip, including comparators, adders, exponentiation units, addition trees, dividers, and multipliers, for performing non-matrix operations such as normalization and activation functions.

[0053] In this specific implementation, by integrating computing units and accumulators under DRAM technology, up to several terabytes per second of internal memory bandwidth can be provided with approximately 13% DRAM chip area overhead and a limited increase in power consumption, achieving efficient acceleration of matrix-vector and matrix-matrix multiplication in attention and fully connected operators.

[0054] The model layers of the target large language model and the draft language model are mapped to multiple near-memory computing devices in a pipelined parallel manner. The model layer is divided into S consecutive pipeline segments, each pipeline segment is mapped to a set of near-memory devices, and the batch size is divided into more than S to ensure that each pipeline segment is kept as fully loaded as possible.

[0055] The relevant technical solutions are the same as above, and will not be repeated here.

[0056] In summary, this invention discloses a heterogeneous in-memory computation speculative decoding acceleration system and method supporting adaptive draft sequence length. The system includes a host, a hardware manager, and multiple heterogeneous in-memory computing devices based on HBM-PIM. This invention adaptively adjusts the draft sequence length of each request at runtime, dynamically truncating draft generation based on cumulative acceptance probability to reduce invalid computation. Simultaneously, it introduces a prediction-verification decoupled asynchronous pipeline, utilizing shared pool and reserve pool mechanisms to eliminate pipeline bubbles caused by draft length differences. Furthermore, this invention proposes an arithmetic strength-aware operator scheduling strategy, dynamically mapping operators to in-memory computing units or the central processor based on real-time changes in effective batch size and draft length to maximize hardware utilization. This invention can significantly improve the inference throughput of large language models on heterogeneous computing systems and reduce energy consumption. Therefore, it can improve the hardware resource utilization efficiency of HBM-PIM heterogeneous systems, further accelerating speculative decoding.

[0057] Those skilled in the art will readily understand that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A method for accelerating large language model inference decoding based on HBM-PIM near-memory computing devices, characterized in that, Pre-filling and speculative decoding are performed using the HBM-PIM near-memory computing device, including: Based on the embedding of each inference request, pre-fill the request and generate the first output token of the target large language model corresponding to the request; In the speculative decoding iteration phase, each round of speculative decoding iteration includes a draft prediction stage and a target model verification stage. The draft prediction stage is executed in parallel for all requests. The draft prediction stage for each request contains n iterations of prediction. The input of the first iteration of the draft prediction stage in the first round of speculative decoding for each request is the first output token corresponding to that request. The input of subsequent iterations of the draft prediction stage in each round of speculative decoding for each request is the output of the previous iteration. n is greater than or equal to 1, and the value of n is determined adaptively. The determination method is: in the t-th iteration prediction generated in the draft prediction stage, the maximum probability p is determined from the probability distribution generated internally by the draft language model. t The corresponding draft token is used as the token for the t-th iteration prediction. The tokens output from the previous t iterations constitute the current draft token sequence for this request; the cumulative probability H of the current draft token sequence for this request is updated. t =H t-1 p t Initially, H0 = 1; set H t With preset threshold Comparison, when H t ≥ Continue with the next iteration of prediction to generate the next draft token, when H t < Draft token generation is terminated at this time, and n is adaptively determined at this point; In each round of speculation decoding, the draft tokens output by each iteration of the draft prediction stage for all requests are stored in a shared pool in real time. The number of tokens in the shared pool is checked in real time to see if it reaches a preset threshold. If so, a parallel verification of all draft tokens in the shared pool by the target large language model is triggered to determine if the output token of the target large language model matches the corresponding draft token in the shared pool. The shared pool is locked during the verification period. For requests that have not completed n iterations of prediction in this round of speculation decoding, iterative prediction continues during the verification period, and the predicted draft tokens are stored in a reserve pool. For each request, if there are inconsistent draft tokens in the shared pool, the draft tokens in the reserve pool for that request are discarded. All draft tokens are used, and the corrected output token of the target large language model is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. If all draft tokens in the shared pool of this request meet the consistency requirement in this round of verification, the draft tokens in the request's preparation pool are moved to the shared pool, and the last generated draft token in the preparation pool of this request is used as the input for the first iteration of the draft prediction stage in the next round of speculative decoding. All the verified consistent draft tokens in each round of speculative decoding for each request, as well as the corrected output token of the target large language model for the verified inconsistent draft tokens, are used as the generation result of this request and output.

2. The method for accelerating large language model inference decoding as described in claim 1, characterized in that, In each round of speculative decoding, the number of requests to continue with the next iteration of prediction is periodically determined. The number of floating-point operations and memory access bytes required by the fully connected operator in the prediction phase of the draft language model to perform the next iteration of prediction is used to calculate the arithmetic strength of the fully connected operator in the prediction phase. This is synchronously and periodically based on the number of draft tokens in the shared pool corresponding to each request that will continue to the next verification. The length of the request and the draft token that has been verified are used to calculate the floating-point operations and memory access bytes required by the attention operator in the verification phase of the target large language model to perform the next verification. This is used to calculate the arithmetic strength of the attention operator in the verification phase. Based on the arithmetic strength of each operator and the known floating-point operation capabilities and memory access bandwidth of the central processor and in-memory computing units in the HBM-PIM near-memory computing device, the device mapping of the operator is updated, and the operator is assigned to the high-performance central processor or high-bandwidth in-memory computing unit in the HBM-PIM near-memory computing device for execution, thereby completing one round of speculative decoding iteration.

3. The method for accelerating large language model inference decoding as described in claim 2, characterized in that, The model computation operators include attention operators and fully connected operators; the device mapping update method for each operator is as follows: Calculate the arithmetic strength of the fully connected operator in the prediction phase. , represented as In the formula, This represents the output dimension of a known draft language model; The arithmetic strength of the attention operator during the verification phase is denoted as: In the formula, L represents the total length consisting of the request length and the draft token that the request has been verified to be valid; Determine the actual computing power of the central processor , represented as ,in, and These represent the floating-point arithmetic capability and memory access bandwidth of the central processor, respectively; determining the actual computing power of the in-memory computing unit. , represented as ,in, and This indicates the floating-point arithmetic capability and memory access bandwidth of the in-memory computing unit; Will Substitute them separately and In the formula, the central processor or in-memory computing unit corresponding to the maximum actual computing power is used as the execution unit of the fully connected operator in the next prediction stage; Will Substitute them separately and In the formula, the central processor or in-memory computing unit corresponding to the maximum actual computing power is used as the execution unit of the attention operator in the next verification stage.

4. The method for accelerating large language model inference decoding as described in claim 1, characterized in that, Preset threshold Determined offline, the determination method is as follows: Based on the sample dataset, a complete speculative decoding process is run for each sample data point, with varying numbers of draft tokens. j The following records the corresponding cumulative probability H j In addition to the target large language model validation results, the condition pass rate is determined based on the target large language model validation results, and the condition pass rate corresponding to the sample data is statistically analyzed. H j The curve of change; In the variation curve of each sample, determine the interval where the condition pass rate is not lower than a preset lower limit, and find the cumulative probability that maximizes the average number of draft tokens among all samples. H j The interval is used to perform a discrete grid search within the interval to determine the threshold for deployment, which serves as the preset threshold. .

5. The method for accelerating large language model inference decoding as described in claim 4, characterized in that, For cumulative probability H t = H t-1 p t The calculations are performed in the logarithmic field, that is, in the first stage of the draft prediction phase. t Accumulated predictions in each iteration logp t get ; Then, when determining whether to perform iterative prediction, With log When comparing, ≥log Continue the prediction in the next iteration when <log The generation of the draft token for the current round of speculative decoding is terminated at this time.

6. The method for accelerating large language model inference decoding as described in claim 4, characterized in that, The threshold is dynamically adjusted based on the system load during runtime. This includes: (1) reducing the number of requests when the equipment utilization rate is lower than the first threshold and the number of requests is lower than the second threshold. (1) Allow more draft tokens to be generated; (2) Increase the efficiency when the device utilization rate is higher than the third threshold or the number of requests is higher than the fourth threshold. Reduce the number of draft tokens.

7. The method for accelerating large language model inference decoding as described in claim 1, characterized in that, Before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests in each batch are encoded in parallel to obtain the embedding of each request. The inference decoding of the target large language model is accomplished by executing the draft language model and the target large language model in parallel on multiple pipeline stages. The pre-deployment of the model on the HBM-PIM near-memory computing device includes: The near-memory computing devices are divided into S groups, each language model is divided into S pipelines, and all requests are divided into more than S batches. The near-memory computing devices in each group process the corresponding pipelines in sequence, and each batch is processed by the near-memory computing devices in each group in sequence. For each group of near-memory computing devices, based on the weight of each pipeline segment and its corresponding request key-value cache dimension, a combination of row and column sharding is used to map the weight matrix and cache matrix to different library groups of each near-memory computing device in order to balance the load and maximize internal bandwidth utilization. Each near-memory computing device includes a buffer layer chip and multiple vertically stacked DRAM chips. Each DRAM chip contains multiple library groups, and each library group is equipped with a global buffer and an accumulator. Each library in each library group integrates a computing unit, which includes several multipliers, an adder tree, and an output buffer. The computing units corresponding to the libraries in the same pseudo-channel of each DRAM chip work in parallel to make full use of the internal bandwidth of the pseudo-channel. The buffer chip integrates a dedicated function unit, including a comparator, adder, exponentiation unit, adder tree, divider, and multiplier, for performing non-matrix operations.

8. The method for accelerating large language model inference decoding as described in claim 1, characterized in that, Before performing pre-filling using the HBM-PIM near-memory computing device, the received user requests are divided into multiple batches, and the requests in each batch are encoded in parallel to obtain the embedding of each request. The pre-configured content-addressable memory serves as the hardware implementation of a shared pool shared by all batches; The on-chip SRAM cache is pre-configured as a hardware implementation of the reserve pool for each batch.

9. A large language model inference decoding acceleration system based on HBM-PIM near-memory computing device, characterized in that, The method for performing the large language model speculation decoding acceleration method as described in any one of claims 1 to 8 includes: a host, a hardware manager, multiple high-bandwidth near-memory computing devices, and a target large language model and a draft language model running on the computing devices. The host receives inference requests and encodes each request to obtain an embedding for each request. Multiple high-bandwidth near-memory computing devices, along with a target large language model and a draft language model running on the computing devices, are used to perform pre-filling and speculative decoding processes based on the request embeddings. During this process, the hardware manager is used to calculate the draft token probability and cumulative probability for each request to control the number of draft tokens for that request; to cache draft tokens to a shared pool or a pre-set pool; and to schedule operators.

Images