Method and apparatus for optimising quantisation of an ML model

Abstract

A method for optimising quantisation of an autoregressive machine learning (ML) model, particularly a large language model, LLM wherein the model is to be deployed on a user device. The quantisation method uses a prefilling phase and a decoding phase, wherein the decoding phase uses multiple iterations and wherein the quantisation method: uses performance metrics that include a set of weight quantisation bitwidths supported by the user device; selects a pre-filling quantisation bitwidth according to the quantisation bitwidth selected; uses the performance metrics to select a decoding quantisation bitwidth and deploys the model on the device based on a quantisation schedule. The schedule may selectively define decoding quantisation bitwidths to be used at the decoding phase during different phases of inference of the model. The method may use a dynamic scheduler and include: running combinations of decoding precision, constructing a dataset with precision-performance pairs for training, defining the model architecture of the scheduler and training the scheduler to optimise a precision switch policy. The switching policy may define when to switch between different decoding quantisation bitwidths. In this way, an optimal bitwidth value most suitable for the user device may be selected.

Description

Field [001 ] The present application generally relates to a method for optimising quantisation of an autoregressive ML model, particularly a large language model, LLM. Background

[002] Large language models (LLMs) have recently emerged as a powerful machine learning model that drives new use cases across various settings, including mobile phones, where features such as real-time voice translation, chat completion and style transfer, and voice assistance, can serve a single user for daily activities and data centers, where various requests are raised from millions of users for chat bot, machine translation and code generation applications.

[003] LLMs present significant computational challenges due to their massive number of parameters, which can range from a few billion to hundreds of billions. This vast scale imposes a substantial burden on both computational resources and memory. As such, there is a growing demand for more efficient LLM inference methods to meet the latency and throughput requirements of LLM-driven applications.

[004] Modern LLMs have demonstrated unprecedented capabilities across various natural language understanding and generation tasks. For example, LLMs have been able to generate human-like text and engage into meaningful dialog with humans, due to their unprecedented language modeling capabilities, including natural language understanding (NLU) and generation (NLG). This advancement came partly as a result of LLM architectures utilizing billions of parameters and thousand-of-tokens long contexts. Nonetheless, the scale of today’s LLMs also leads to significant deployment challenges, particularly in the landscape of embedded and mobile devices.

[005] For example, as shown in Figure 1a, LLM inference is broken down into two stages: i) the prefill stage, which runs once and comprises the processing of the input prompt by the LLM up until the prediction of the first new token; and ii) the decoding stage, which follows prefilling and consists of multiple forward passes predicting a single token each time to eventually generate the output response. In the prefilling stage, prompt tokens are executed in parallel whereas in the decoding phase, output tokens are generated one-by-one in an autoregressive manner. Thus more generally, the machine learning model may be termed an autoregressive ML model with a prefill phase and a decoding phase. The prefilling phase may be termed an initial phase of inference or an in-filling phase and is typically used for decoder-only models. The prefill phase is typically used to generate key-value (KV) pairs which are stored in a KV cache for use in the decoding phase. The decoding phase then generates the output from the model using the KV cache generated in the prefill stage. One output is generated for each cycle of inference of the decoding phase.

[006] Computationally, the prefill stage is typically compute-bounded, as all the tokens of the input prompt can be processed in parallel sharing the same LLM weights, thus increasing the arithmetic intensity of the computation. In contrast, the decoding stage is memory-bound due to its autoregressive nature that severely restricts weight re-use in each iteration, resulting in a notably lower computation-to-communication ratio. As a result, the processing speed of the decoding stage largely depends on the memory bandwidth of the system and often leads to severe underutilization of the available computational resources. With its multiple iterations and the hardware undersutilization in each iteration, the decoding stage commonly dominates the LLM inference latency, with more severe impact on resource-constrained mobile devices.

[007] Figure 1 b is an alternative representation of Figure 1a. During the prefilling phase, all input tokens are processed in parallel using a full attention mask, making it a compute-bound workload due to its highly parallel nature. In this phase, different keys and values are stored in the KV cache as context information of the input prompts. In the decoding phase, output tokens are generated sequentially in an autoregressive manner. This phase is memory-bound, as the primary computation involves vector-matrix multiplication. The keys and values generated during the prefilling phase are used to retrieve the necessary information for decoding and new keys and values are appended to the KV cache as more tokens are generated. Both phases utilize the same set of weights.

[008] Several effective quantisation solutions have been developed, primarily for Graphic Processing Units (GPUs). These units are optimized for floating point computation but are hindered by memory access bottlenecks. Among these solutions, weight-only quantisation converts model weights into low-bitwidth integers (e.g., 4-bit) while maintaining activations in the original 16-bit floating point format. On-the-fly weight dequantisation is used at inference time to perform computations in floating-point. Weight-only quantisation can preserve accuracy and reduce inference latency by avoiding memory access overheads. Unfortunately, peak-memory use, energy consumption, and latency can still be high. To address this, weightactivation quantisation methods also quantise the activations and use efficient fixed-point operators, eliminating the need for on-the-fly dequantisation. Yet, quantising activations beyond 16 bits typically degrades accuracy, especially with static per-tensor statistics. Previous strategies have included quantising activations for only certain expensive operations, e.g. matrix multiplication, or employing dynamic per-token quantisation, which tends to be slow on GPUs and lacks robust hardware support on edge devices such as mobile phones.

[009] Hardware support is another crucial practical aspect when considering on-device deployment of LLMs. For example, edge devices often have a dedicated compute unit, which is optimized for operations in a specific bitwidth (e.g. int8-int8 operations) and rely on private libraries, hence preventing the use of custom kernels. In contrast, GPUs (even edge device ones) are more flexible at the expense of lower efficiency. Thus, quantisation approaches compatible with tensor compute units are particularly attractive in practice.

[010] Notably, none of these methods provides lossless int8 activation quantisation, thus failing to concurrently achieve compute efficiency, memory efficiency, and energy efficiency in on-device settings.

[011] The applicant has therefore identified the need for a way to implement LLMs on constrained-resource devices. Summary

[012] In a first approach of the present techniques, there is provided a computer-implemented method for optimising an autoregressive machine learning, ML, model before deployment to a user device, wherein the autoregressive ML model comprises a plurality of model weights and comprises a prefilling phase in which an input to the model is pre-processed and a decoding phase in which the pre-processed input is processed through multiple iterations to generate an output. The method comprises receiving performance metrics for the ML model, the performance metrics comprising at least one of an output metric for the output from the ML model and a set of weight quantisation bitwidths which are supported by the user device; selecting, based on the received performance metrics, a prefilling quantisation bitwidth for the prefilling phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the prefilling phase is converted into an integer having the selected prefilling quantisation bitwidth; selecting, based on the received performance metrics, at least one decoding quantisation bitwidth for the decoding phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the decoding phase is converted into an integer having the selected at least one decoding quantisation bitwidth; and outputting, to a user device, the ML model with a quantisation schedule which specifies the selected prefilling quantisation and the at least one decoding quantisation bitwidth whereby the ML model is optimised when used on the user device. The selected prefilling quantisation bitwidth is greater than or equal to the at least one decoding quantisation bitwidths whereby the precision applied in the prefill stage is greater or the same as the precision applied in at least the start of the decoding phase.

[013] Inference using the ML model is broken down into two stages: i) the prefill stage, which runs once and comprises the pre-processing of the input (e.g. an input prompt for an LLM) up until the prediction of the first input to the decoding phase (e.g. the first new token); and ii) the decoding stage, which follows prefilling and consists of multiple forward passes predicting a single internmediate output (e.g. token) each time to eventually generate the output response. Each forward pass may be termed an iteration of the decoding phase. By selecting a prefilling quantisation bitwidth which is greater than or equal to the at least one decoding quantisation bitwidths, the method may be termed a phase-aware method that selectively allocates precision during different phases of inference. The performance metrics may be received from a user or may be determined automatically based on the capabilities of the user device. The method may thus be considered to provide customisation (or optimisation) of the ML model for the hardware constraints of the user device (e.g. the quantisation bitwidths which are available) so that an efficient and effective result will be achieved.

[014] Selecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth may comprise obtaining a prefilling quantisation bitwidth and a decoding quantisation bitwidth from the set of weight quantisation bitwidths; quantising the plurality of weights of the ML model using the obtained prefilling quantisation bitwidth and decoding quantisation bitwidth to generate a variant ML model; generating a set of variant ML models by repeating the obtaining and quantising steps; evaluating, using a calibration dataset, the generated set of variant ML models; identifying the variant ML models which meet the output metric; and selecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth from the identified variant ML model having lowest values for the prefilling quantisation bitwidth and decoding quantisation bitwidth. The set of weight quantisation bitwidths may be termed a set of precisions and the terms may be used interchangeably. Similarly, the prefilling quantisation bitwidth may be termed a prefilling precision and the or each decoding quantisation bitwidth may be termed a decoding precision.

[015] The output metric may be an accuracy metric or may be a generation quality metric. The latter may be used for generative tasks from the ML model which cannot be evaluated using accuracy metrics.

[016] The selection of the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth may be done before deployment, e.g. in a calibration phase. The objective of the calibration phase may be defined as finding the pair of smallest weight precisions for the prefill and decoding phase, respectively, that achieves equal or greater algorithmic performance. In other words, given a pretrained ML model m (e.g. an LLM), a set of variably quantised models [mp) containing one quantised model variant per arithmetic precision p in precision set y using quantiser Q(-) may be obtained. Next, the set of variably quantised models may be evaluated on a calibration dataset and the lowest pair of precisions (i.e. lowest values for prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth) that meet the quality target is determined.

[017] As described above, the decoding phase may comprise a plurality of iterations or passes, each of which generates a token and which may be combined or grouped to form distinct sequences which are time-ordered as the decoding phase is carried out. Selecting at least one decoding quantisation bitwidth may comprise: selecting a first decoding quantisation bitwidth for a first sequence of the decoding phase; and selecting a second decoding quantisation bitwidth for a second (later or subsequent) sequence of the decoding phase. The second decoding quantisation bitwidth is lower than the first decoding quantisation bitwidth. The first sequence may comprise a first number of passes or iterations of the decoding phase and the second sequence may comprise the remaining number of passes which are required to generate the output. In this way, there is a gradual lowering of the precision applied during the decoding phase. The method may further comprise setting a switching point which defines an end of the first sequence whereby when the ML model is being used on the user device, the switching point defines when to switch from the first decoding quantisation bitwidth to the second decoding bitwidth.

[018] In the example above, there are two decoding quantisation bitwidths but it will be appreciated that there may be multiple bitwidths. In other words, there may be a plurality of switching points, wherein each switching points defines an end of a sequence of the decoding phase. In other words, the decoding phase may comprise a plurality of time-ordered sequences each of which generates an intermediate output which is processed by the subsequent time-ordered sequence until the last sequence in which the output is generated. Each intermediate output may be a token which is generated by the last iteration in the sequence. Selecting at least one decoding quantisation bitwidth may comprise selecting a different decoding quantisation bitwidth for each one of the time-ordered sequences to obtain a plurality of decoding quantisation bitwidths which decrease in size in time-order across the decoding phase.

[019] This method may be termed Progressive Mixed-Precision Decoding (PMPD) because it is a technique that enables the gradual lowering of precision deeper into the generated sequence during the decoding phase. The technique of using multiple decoding quantisation bitwidths may be combined with the phase aware precision allocation described above. Alternatively, the techniques may be used independently.

[020] According to a second aspect of the present techniques, there is a method for optimising a large language model, LLM for deployment to a user device, the method comprising: receiving user input comprising at least one of a quantisation requirement for the LLM and accuracy requirements for the LLM; selecting, based on the received user input, a prefilling quantisation bitwidth for a prefilling phase of the LLM; selecting, based on the received user input, a quantisation schedule for a decoding phase of the LLM, wherein the quantisation schedule defines at least a first decoding quantisation bitwidth which is applied when processing a first set of tokens in the decoding phase and a second quantisation bitwidth which is applied when processing a second set of tokens in the decoding phase, wherein the second decoding quantisation bitwidth is lower than the first decoding quantisation bidwidth whereby in the decoding phase the first set of tokens are processed with greater precision than the second set of tokens, and outputting the LLM with the quantisation schedule to the user device.

[021] The features described above with respect to the first approach apply equally to the second approach and therefore, are not repeated. Similarly, the subsequent features apply equally to both approaches and are not repeated separately for both approaches.

[022] The quantisation schedule may define a plurality of decoding quantisation bitwidths, each having a progressively smaller value and thus the quantisation schedule may comprise a decoding quantisation schedule specifying the precisions or bitwidths which are used in the decoding phase. In this way, the first set of tokens which are generated earlier in the output sequence are processed with higher precision than the second set of tokens which are generated deeper ln the output sequence. The selected prefilling quantisation bitwidth may be greater than a largest of the decoding quantisation bitwidths. Thus, the method may be termed Progressive Mixed-Precision Decoding (PMPD), a technique that enables the gradual lowering of precision deeper into the generated sequence. The lowest bitwidths for prefilling and decoding may be selected. By observing distinct redundancy and sensitivity across the two LLM inference phases (i.e. prefilling and decoding), the proposed phase-aware precision allocator applies different arithmetic precisions and scheduling strategies in the prefill and decoding phases.

[023] Setting the or each switching point and selecting the at least one decoding quantisation bitwidth may be done prior to outputting the quantisation schedule. Outputting the ML model with the quantisation schedule may comprise outputting a quantisation schedule which comprises a static decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point. A static decoding quantisation schedule is a schedule which does not change when the ML model is used on the user device to which it has been deployed.

[024] The method may comprise selecting a static scheduler to implement the selected quantisation schedule and outputting the static scheduler with the LLM (ML model) and the quantisation schedule. A static scheduler is fixed and performs prompt-agnostic precision switching and so it is the same across input prompts, i.e. the same across all inputs to the ML model.

[025] Setting the or each switching point and selecting the at least one decoding quantisation bitwidth may comprise: obtaining an initial decoding quantisation bitwidth from the set of weight quantisation bitwidths, wherein when the ML model is being used on the user device, each of the model weights in the initial sequence of the decoding phase which follows the prefilling phase is converted into an integer having the initial decoding quantisation bitwidth; obtaining at least one switching point from a set of possible switching points; obtaining, from the set of weight quantisation bitwidths, an associated decoding quantisation bitwidth for each obtained at least one switching point, wherein when the ML model is being used on the user device, each of the model weights in a sequence of the decoding phase which follows a switching point is converted into an integer having the associated decoding quantisation bitwidth; quantising the plurality of weights of the ML model using the obtained at least one switching point and obtained initial and subsequent decoding quantisation bitwidths to generate a variant ML model; generating a set of variant ML models by repeating the obtaining and quantising steps; evaluating, using a calibration dataset, the generated set of variant ML models; and selecting the at least one switching point and obtained initial and subsequent decoding quantisation bitwidths from the identified variant ML model having a highest performance,

[026] The method may comprise selecting a dynamic scheduler to implement the selected quantisation schedule (specifically a decoding quantisation schedule having multiple decoding quantisation bitwidths). The method may comprise configuring the dynamic scheduler with a runtime quantisation schedule which is based on the selected quantisation schedule and which selects a switching point for changing the decoding quantisation bitwidth in response to a received input; and outputting the LLM with the dynamic scheduler and the runtime quantisation schedule. The dynamic scheduler may be configured by training the dynamic scheduler using an initial dataset which comprises a plurality of different inputs. In other words, the scheduler may be configured with a set of available precisions and a runtime quantisation schedule in the form of a dynamic switching policy ready for the deployment stage. In this way, the dynamic scheduler thus has prompt-specific precision switching and so it is different across Input prompts.

[027] In other words, the method may comprise outputting a dynamic scheduler when outputting the ML model with the quantisation schedule, wherein, when using the ML model on the user device, the dynamic scheduler determines, for each input to the ML model, a decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point. The method may comprise training the dynamic scheduler using a training dataset prior to outputting the dynamic scheduler, wherein the training dataset comprises a plurality of prefilling stage outputs each of which is labelled with a decoding quantisation schedule. The dynamic scheduler may also be termed a learned scheduler because it is trainable. The dynamic scheduler may be a lightweight ML model.

[028] Configuring the dynamic scheduler may comprise creating a training dataset comprising a plurality of data points by generating, using the initial dataset, a key value cache (KV cache) comprising a plurality of KV cache tensors each of which are generated by running an input from the received initial dataset through the prefill stage of the LLM; and pairing each KV cache tensor with a label representing a token index for a switch point in the quantisation schedule. The dynamic scheduler may then be trained to generate a runtime quantisation schedule having optimised pairs of labels and KV cache tensor which optimises performance of the LLM for different inputs.

[029] In other words, the method may comprise creating the training dataset by receiving a plurality of inputs for the ML model. For each received input, the method may comprise generating a prefilling phase output by running the input through the prefill phase of the ML model; obtaining a set of decoding quantisation schedules, wherein each decoding quantisation schedule comprises at least one switching point and an associated decoding quantisation bitwidth for each at least one switching point; quantising the plurality of weights of the ML model using each decoding quantisation schedule in the set of decoding quantisation schedules to generate a set of variant ML models; evaluating, using the prefilling phase output, the generated set of variant ML models; selecting the decoding quantisation schedule from the identified variant ML model having a highest performance; labelling the prefilling phase output with the selected decoding quantisation schedule; and outputting the prefilling phase outputs and associated labels as the training dataset.

[030] The LLM may comprise a plurality of separate models each having separate model weights and each having a different quantisation. Alternatively, the LLM may be a quantised model comprising a unified set of model weights and the selected prefilling quantisation bitwidth and the selected at least one decoding quantisation bitwidth are extracted from the unified set of model weights when quantisation is implemented on the user device. The quantisation may be at 4 bits, 3 bits or 2 bits. Each LLM comprises a plurality of layers and a different level of quantisation may be provided to sets of layers within the LLM. In other words, the method may comprise outputting the ML model as a plurality of separate models each having separate model weights and each having a different quantisation which is defined by the quantisation schedule. Alternatively, the method may comprise outputting the ML model as a unified set of model weights.

[031] The techniques described above may be carried out on an electronic device, e.g. a server, to optimise a ML model before deployment to a user device which typically has lower resources and thus would not be able to optimise the model locally. Once the ML model is deployed to the user device, the user device may quantise the model as defined in the quantisation schedule, e.g. to include phase aware precision allocation and / or PMPD and then use the quantised model. More specifically, the method at the user device may comprise receiving the ML model and the quantisation schedule as described above; receiving an input to be processed through the ML model; quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule; quantising the decoding phase of the ML model using the at least one decoding quantisation bitwidth defined in the quantisation schedule; and processing, using the quantised prefill and decoding phases, the input to generate an output from the ML model.

[032] When a dynamic scheduler is used, the user device may receive at the user device, the ML model, the quantisation schedule and the dynamic scheduler as described above. In this case, the decoding quantisation schedule needs to be determined before inference can be carried out. Thus, the method may further comprise receiving an input to be processed through the ML model; quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule; generating a prefilling phase output by running the input through the quantised prefill phase; inputting the prefilling phase output to the dynamic scheduler to generate a decoding quantisation schedule; quantising the decoding phase of the ML model using the decoding quantisation schedule; and processing, using the quantised decoding phase, the prefilling phase output to generate an output from the ML model. The method comprises generating the decoding quantisation schedule before all the steps of inference are carried out.

[033] In a third approach of the present techniques, there is provided an electronic device for optimising an autoregressive machine learning, ML, model, particularly a large language model, LLM, the electronic device comprising: at least one processor coupled to memory, wherein the at least one processor is configured individually or collectively to carry out the method described above.

[034] The features described above with respect to the first and second approaches apply equally to the third and fourth approaches and therefore, are not repeated.

[035] In a fifth approach of the present techniques, there is provided an electronic device for receiving an LLM which has been optimised using the methods described herein and using the received LLM.

[036] In a sixth approach of the present techniques, there is provided a system comprising the electronic device for optimising the LLM, and a user device which is configured to receive, from the electronic device, the optimised LLM and the quantisation schedule; receive an input; and process the received input using the received optimised LLM and quantisation schedule to generate an output.

[001] The electronic device, in particular the user device which receives the LLM, may be a constrained-resource device, but which has the minimum hardware capabilities to quantise an LLM. The device may be any one of: a smartphone, tablet, laptop, computer or computing device, virtual assistant device, a vehicle, an autonomous vehicle, a robot or robotic device, a robotic assistant, image capture system or device, an augmented reality system or device, a virtual reality system or device, a gaming system, an Internet of Things device, or a smart consumer device (such as a smart fridge, smart vacuum cleaner, smart oven, or smart lawn mower). It will be understood that this is a non-exhaustive and non-limiting list of example devices.

[002] In a related approach of the present techniques, there is provided a computer-readable storage medium comprising instructions which, when executed by a processor, causes the processor to carry out any of the methods described herein.

[003] As will be appreciated by one skilled in the art, the present techniques may be embodied as a system, method or computer program product. Accordingly, present techniques may take the form of an entirely hardware embodiment, an entirely software embodiment, or an embodiment combining software and hardware aspects.

[004] Furthermore, the present techniques may take the form of a computer program product embodied in a computer readable medium having computer readable program code embodied thereon. The computer readable medium may be a computer readable signal medium or a computer readable storage medium. A computer readable medium may be, for example, but is not limited to, an electronic, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any suitable combination of the foregoing.

[005] Computer program code for carrying out operations of the present techniques may be written in any combination of one or more programming languages, including object oriented programming languages and conventional procedural programming languages. Code components may be embodied as procedures, methods or the like, and may comprise subcomponents which may take the form of instructions or sequences of instructions at any of the levels of abstraction, from the direct machine instructions of a native instruction set to high-level compiled or interpreted language constructs.

[006] Embodiments of the present techniques also provide a non-transitory data carrier carrying code which, when implemented on a processor, causes the processor to carry out any of the methods described herein.

[007] The techniques further provide processor control code to implement the abovedescribed methods, for example on a general purpose computer system or on a digital signal processor (DSP). The techniques also provide a carrier carrying processor control code to, when running, implement any of the above methods, in particular on a non-transitory data carrier. The code may be provided on a carrier such as a disk, a microprocessor, CD- or DVD-ROM, programmed memory such as non-volatile memory (e.g. Flash) or read-only memory (firmware), or on a data carrier such as an optical or electrical signal carrier. Code (and / or data) to implement embodiments of the techniques described herein may comprise source, object or executable code in a conventional programming language (interpreted or compiled) such as Python, C, or assembly code, code for setting up or controlling an ASIC (Application Specific Integrated Circuit) or FPGA (Field Programmable Gate Array), or code for a hardware description language such as Verilog (RTM) or VHDL (Very high speed integrated circuit Hardware Description Language). As the skilled person will appreciate, such code and / ordata may be distributed between a plurality of coupled components in communication with one another. The techniques may comprise a controller which includes a microprocessor, working memory and program memory coupled to one or more of the components of the system.

[008] It will also be clear to one of skill in the art that all or part of a logical method according to embodiments of the present techniques may suitably be embodied in a logic apparatus comprising logic elements to perform the steps of the above-described methods, and that such logic elements may comprise components such as logic gates in, for example a programmable logic array or application-specific integrated circuit. Such a logic arrangement may further be embodied in enabling elements for temporarily or permanently establishing logic structures in such an array or circuit using, for example, a virtual hardware descriptor language, which may be stored and transmitted using fixed or transmittable carrier media.

[009] In an embodiment, the present techniques may be realised in the form of a data carrier having functional data thereon, said functional data comprising functional computer data structures to, when loaded into a computer system or network and operated upon thereby, enable said computer system to perform all the steps of the above-described method.

[010] The method described above may be wholly or partly performed on an apparatus, i.e. an electronic device, using a machine learning or artificial intelligence model. The model may be processed by an artificial intelligence-dedicated processor designed in a hardware structure specified for artificial intelligence model processing. The artificial intelligence model may be obtained by training. Here, "obtained by training" means that a predefined operation rule or artificial intelligence model configured to perform a desired feature (or purpose) is obtained by training a basic artificial intelligence model with multiple pieces of training data by a training algorithm. The artificial intelligence model may include a plurality of neural network layers. Each of the plurality of neural network layers includes a plurality of weight values and performs neural network computation by computation between a result of computation by a previous layer and the plurality of weight values.

[011] As mentioned above, the present techniques may be implemented using an Al model. A function associated with Al may be performed through the non-volatile memory, the volatile memory, and the processor. The processor may include one or a plurality of processors. At this time, one or a plurality of processors may be a general purpose processor, such as a central processing unit (CPU), an application processor (AP), or the like, a graphics-only processing unit such as a graphics processing unit (GPU), a visual processing unit (VPU), and / or an Al-dedicated processor such as a neural processing unit (NPU). The one or a plurality of processors control the processing of the input data in accordance with a predefined operating rule or artificial intelligence (Al) model stored in the non-volatile memory and the volatile memory. The predefined operating rule or artificial intelligence model is provided through training or learning. Here, being provided through learning means that, by applying a learning algorithm to a plurality of learning data, a predefined operating rule or Al model of a desired characteristic is made. The learning may be performed in a device itself in which Al according to an embodiment is performed, and / o may be implemented through a separate server / system.

[012] The Al model may consist of a plurality of neural network layers. Each layer has a plurality of weight values, and performs a layer operation through calculation of a previous layer and an operation of a plurality of weights. Examples of neural networks include, but are not limited to, convolutional neural network (CNN), deep neural network (DNN), recurrent neural network (RNN), restricted Boltzmann Machine (RBM), deep belief network (DBN), bidirectional recurrent deep neural network (BRDNN), generative adversarial networks (GAN), and deep Q-networks.

[013] The learning algorithm is a method for training a predetermined target device (for example, a robot) using a plurality of learning data to cause, allow, or control the target device to make a determination or prediction. Examples of learning algorithms include, but are not limited to, supervised learning, unsupervised learning, semi-supervised learning, or reinforcement learning. Brief description of the drawings

[014] Implementations of the present techniques will now be described, by way of example only, with reference to the accompanying drawings, in which:

[015] Figures 1a and 1b are two alternative Illustrations of the conventional paradigm of quantised LLM inference;

[016] Figure 2 is an illustration of LLM inference using phase-aware precision allocation and progressive mixed-precision decoding as described in the present techniques;

[017] Figures 3a and 3b are algorithms for phase aware precision allocation and progressive mixed-precision decoding as described in the present techniques;

[018] Figure 4a is a schematic illustration of the offline calibration process of the algorithms in Figures 3a and 3b;

[019] Figure 4b is a flowchart showing the steps in the offline calibration process to achieve phase aware precision allocation;

[020] Figures 5a to 5c are examples showing the effect of Phase-Aware Precision Allocation;

[021] Figure 6a plots the model precision (at three different bit-levels) against the ID of the token in the decoding phase for three different prompts as controlled by a static scheduler;

[022] Figure 6b is a flowchart showing the steps in the offline calibration process to achieve progressive mixed-precision decoding using a static scheduler;

[023] Figure 7 plots the model precision (at three different bit-levels) against the ID of the token in the decoding phase for three different prompts as controlled by a dynamic scheduler;

[024] Figure 8 plots shows the different steps in the offline calibration of Figure 4a depending on whether a static or a dynamic scheduler is selected;

[025] Figure 9a illustrates an example architectural design for a dynamic scheduler;

[026] Figure 9b is a flowchart showing the steps in the offline calibration process to train a dynamic scheduler;

[027] Figure 9c is a flowchart showing the steps in the offline calibration process for generating a data set to train the dynamic scheduler;

[028] Figures 10a and 10b are schematic illustrations showing the interaction between the precision switching scheduler and the LLM during runtime deployment as described in the present techniques;

[029] Figure 10c is a flowchart showing the steps at inference when using a dynamic scheduler;

[030] Figures 11a and 11b illustrate two implementations of the present techniques which may be termed Variably Quantised Weights and Single-Copy Any-Precision Model respectively;

[031] Figures 12a and 12b show two embodiments which are running an LLM inference based on the techniques describe above;

[032] Figure 13 is a block diagram of a system that may be used to perform the methods described above and / or to use an LLM with progressive mixed precision decoding as described above;

[033] Figures 14a and 14b illustrate how the decoding stage of LLM inference is memorybound, with the weights dominating the memory transfers;

[034] Figures 15a to 18b are pairs of plots showing BERTscore or Rouge-L score against average bits when using a static scheduler with different types of LLM and different datasets as described above;

[035] Figures 19a to 20b are pairs of plots showing BERTscore or Rouge-L score against average bits when using a dynamic scheduler with different types of LLM and different datasets which have been trained as described above;

[036] Figure 21a shows different schedules of mixing 3-bit and 4-bit model in the decoding phase running an LLM inference based on the techniques describe above;

[037] Figure 21b shows Rouge-L score against average bits for an LLM using the four different schedules shown in Figure 21a;

[038] Figures 22a to 22c are bar charts showing the relative speeds of different models when compared to the baseline fp16 model;

[039] Figures 23a to 23c are bar charts presenting an ablation study of the impact of the phase aware precision allocation on performance; and

[040] Figure 23d is a graph plotting the latency overhead of the prefill stage as a function of the bitwidth used during prefilling; and

[041] Figure 24 plots the normalized scores of different Learned Scheduler Design across different datasets. Detailed description of the drawings

[042] Broadly speaking, embodiments of the present techniques provide a method for optimising a large language model, LLM for deployment, e.g. to a resource-constrained device, such as a smartphone. LLM inference is broken down into two stages: i) the prefill stage, which runs once and comprises the processing of the input prompt by the LLM up until the prediction of the first new token; and ii) the decoding stage, which follows prefilling and consists of multiple forward passes predicting a single token each time to eventually generate the output response. One aspect of the method is a phase-aware method that selectively allocates precision during different phases of LLM inference. To further optimise the decoding phase, a technique that enables the gradual lowering of precision deeper in the generated sequence is termed Progressive Mixed-Precision Decoding (PMPD). Precision-switching schedulers that dynamically drive the precision-lowering decisions in either a task- and prompt-adaptive manner are also proposed. More generally, the proposed techniques may be used with any autoregressive ML model with a prefill phase and an iterative decoding phase, with an LLM being a specific example of such a model. Thus, throughout the subsequent description, it will be appreciated that the techniques being described in relation to an LLM apply to any autoregressive ML model with a prefill phase and an iterative decoding phase.

[043] To alleviate the computational burden of LLMs, several approaches have been proposed, including model compression and adaptive inference techniques. Particularly, model compression approaches, such as quantisation, have been widely studied, adopting low-precision representations for the weights, KV Cache entries and activations of the LLM. To facilitate more effective and lossless compression, hybrid dense-and-sparse representations are becoming increasingly popular, relying on the support of specialised accelerators. Additionally, recent literature broadly focuses on the large (server-grade) end of the model spectrum, where representation redundancy is more profound. As a result, the naive adoption of such methods on smaller footprint (mobile grade) LLMs, often induces significant quantisation errors.

[044] In parallel, adaptive inference methodologies, such as speculative decoding, have been proposed aiming to dynamically reduce the amount of computation required for the prediction of “easier" tokens in the response sequence. Such methods however, rely on the availability of plentiful compute resources used that can be exploited in parallel, offering an improved speed-accuracy trade-off at the expense of an increased workload involving redundant computations. Thus, such solutions are targeted for multi-node systems, such as GPU clusters, but are not well-suited for the resource-constrained landscape.

[045] In spite of the great potential of large language models (LLMs) across various tasks, their deployment on resource-constrained devices remains challenging due to to their excessive computational and memory demands. Quantisation has emerged as an effective solution by storing weights in reduced precision. However, utilizing precisions that would be low enough ( / .e. 2 / 3-bit) to substantially alleviate the memory-boundedness of LLM decoding, still suffers from prohibitive performance drop. Existing approaches fail to explore the diversity in computational patterns, redundancy, and sensitivity to approximations of the different phases of LLM inference, resorting to a uniform quantisation policy throughout. As detailed below an alternative phase-aware method that selectively allocates precision during different phases of LLM inference is proposed. The proposed method achievesboth strong context extraction during prefill and efficient memory bandwidth utilization during decoding. To further address the memory-boundedness of the decoding phase, Progressive Mixed-Precision Decoding (PMPD) is proposed. This is a technique that enables the gradual lowering of precision deeper in the generated sequence, together with a spectrum of precision-switching schedulers that dynamically drives the precision-lowering decisions in either a task- and prompt-adaptive manner. Extensive evaluation across diverse language tasks shows that PMPD-enhanced LLMs achieve significant gains in speedup over both conventional LLM inference methods and over uniform quantisation approaches for LLMs. Notably, when targeting Nvidia GPUs, PMPD achieves 1.4-12.2x speedup in LLM linear layers overfp16 models, while when targeting an LLM-optimized NPU. The proposed approach delivers a throughput gain of 3.8-8.Ox over fp16 models and up to 1.54x over uniform quantisation approaches while preserving the output quality.

[046] In other words, the aim of the present techniques is to focus on embedded and / or mobile deployment settings, where the scarcity of computational resources restricts the applicability of Inference schemes that incur redundant computations or require heterogeneous processing units, as well as limits the parameter scale of the models that can be deployed on-device. Most existing approaches apply a uniform quantisation scheme across both prefill and decoding stages and as recognised in the proposed techniques this presents two key limitations: i) They overlook the distinct fault tolerances against quantisation errors in prefill and decoding stages, leading to a significant accuracy drop while reducing the bit-widty for mobile-scale LLMs. li) Since the decoding phase is memory-bound, reducing the bit-width in this stage is more effective in improving hardware performance. As such, previous uniform quantisation approaches can only achieve limited speedups, underscoring the need for more targeted quantisation strategies specifically tailored to LLM decoding.

[047] By recognising the key limitations in prior work, the proposed techniques identify a novel insight for LLM weight quantisation. The prefill phase, as well as early parts of the decoding phase, are more sensitive to approximation errors from quantisation, than later parts of the autoregression generation process. The present techniques thus introduce a novel LLM inference method that counteracts the limitations of existing approaches by means of a phase-aware and progressive reduced-precision approach. First, the proposed method departs from the existing deployment approaches that employ uniform quantisation precision throughout the LLM inference process and introduces phase-aware precision allocation. By considering the distinct redundancy, sensitivity to approximations, and arithmetic intensity of the prefill and decoding stage, the proposed scheme tailors the arithmetic precision to each phase’s characteristics. This leads to both maintained generation quality through high-precision prompt encoding during prefill and improved sustained throughput through reduced precision during decoding.

[048] To further address the memory-boundedness of decoding, the proposed techniques introduce progressive mixed-precision decoding (PMPD), a novel decoding scheme that gradually reduces precision throughout the generation of the output response. Based on the observation that tokens generated later in the output sequence are more resilient to approximations while earlier tokens are more sensitive, the proposed method employs quantisation as an approximation mechanism and progressively reduces the arithmetic precision during the decoding phase. To balance generation quality with decoding throughput, precision scheduling is formulated as a constrained optimisation problem with the objective of minimising the average bitwidth while preserving output quality. Two complementary precision switching schedulers are proposed. A high performance task-specific static scheduler and a flexible dynamic precision-switching scheduler that strategically determines the most highest-performing precision-reduction time for each Input prompt and which is task-agnostic.

[049] As demonstrated below, the method can be applied on Vicuna-7B, MobileLLaMa-14.B, Stable LM Zephyr-3B and Phi-1.5 across diverse language tasks to achieve decoding throughput gains of 3.8-8.0 times on NPU platform and 1.40-12.20 times speedup on LLM linear layer computations on GPU. As noted in the background section, although the detail is described in relation to LLM, the proposed method has general applicable to any autoregressive machine learning (ML) model with a prefill stage and an autoregressive decoding phase. Thus, the terms LLM and autoregressive ML model may be used interchangeably through.

[050] Figure 2 illustrates the proposed method that comprises phase-aware precision allocation and progressive mixed-precision decoding. The proposed approach considers the respective characteristics of the different LLM inference phases. As such, high precision is assigned during prefill and reduced precision during decoding, leading to more efficient utilization of the memory bandwidth. In other words, the phase-aware precision allocation strategy optimises precisions differently for the prefill and decoding phases, leveraging the distinct error resilience of each stage to achieve an extremely low average bitwidth for mobile-grade LLMS.

[051] As shown in Figure 2, the LLM inference process consists of two phases: prefilling and decoding. By observing distinct redundancy and sensitivity across the two phases, the proposed phase-aware precision allocator applies different arithmetic precisions and scheduling strategies in the prefill and decoding phases.

[052] In the prefill phase, where the context information of input prompts is crucial for generating accurate output tokens, a high-precision quantisation scheme is employed. To further enhance performance and support efficient parallel processing, a uniform precision scheme is utilized, maximizing hardware efficiency during this phase,

[053] In the decoding phase, the empirical experiments reveal that the initial output tokens are particularly sensitive to quantisation errors. Based on this observation, a progressive lowering quantisation scheme is introduced in the decoding phase.

[054] As shown in the right-hand part of Figure 2, the precision during the decoding phase is gradually reduced, which significantly eases the memory bandwidth overhead. In other words, there is a progressive mixed-precision decoding scheme that progressively reduces the precision as decoding progresses to longer token sequences. Given the memory-bound nature of LLM decoding, this dynamic reduction scheme can effectively enhance the hardware performance of LLM inference. Not shown in Figure 2 but explained in more detail below, the proposed method may also include a prompt-agnostic scheduler and / or a task-agnostic scheduler for precision switching, enabling the flexible deployment of the proposed approach across diverse scenarios, accommodating varying data availability constraints and generation quality requirements.

[055] Method

[056] The algorithm in Figure 3a presents the currently proposed method for Progressive Mixed-Precision Decoding and Figure 3b shows an alternative version of the approach. In Figure 3a, lines 1 to 8 represent the offline stage and lines 9 to 22 represent the deployment stage. In summary, the offline stage is responsible for producing variably quantised variants of a given LLM (line 1), conducting the phase-aware precision allocation (line 2), and configuring the selected type of precision-switching scheduler with optimized PMPD parameters for the target use case (lines 3-8). Upon deployment, the scheduler assigns the allocated precisions to the prefill and decoding stages, and triggers precision-switching events for each encountered prompt (line 20).

[057] In Figure 3b, the approach begins with an offline calibration process (corresponding to lines 1 and 2 of the algorithm in Figure 3b) which takes user requirements as inputs to optimize precision-switching schedulers for LLM inference. A precision scheduler adopts a phase-aware precision allocation strategy. During the decoding stage, two types of precision schedulers are introduced: prompt-agnostic static schedulers and prompt-adaptive dynamic schedulers.

[058] Offline Calibration Process. A first part of the offline calibration process of the algorithm in Figures 3a and 3b is illustrated in Figure 4. To offer users the flexibility to explore the trade-off between algorithmic and hardware performance, the proposed method may adopt an offline calibration process that allows for the customized optimization of LLMs for downstream tasks. The process starts with receiving user requirements as inputs, such as accuracy requirements and target LLM models and quantisation approach. quantisA calibration dataset may be leveraged to evaluate for better algorithmic performance on downstream tasks.

[059] Phase-Aware Precision Allocation

[060] More specifically considering the detail of the algorithm shown in Figure 3a, the following input parameters are considered: • a quantiser Q(-), which is typically a post-training quantisation (PTQ) method, • a set of precisions P, indicating the supported weight bitwidths, • a reference value qref in a task-specific quality metric, and • a quality drop tolerance e. The objective of the calibration phase is to find the pair of smallest weight precisions for the prefill and decoding phase, respectively, that achieves equal or greater algorithmic performance than qref- e. This objective is solved by means of a calibration step during the offline stage of the proposed methodology. First, given a pretrained LLM m, a set of variably quantised models {mp} containing one quantised model variant per arithmetic precision p in precision set J5 using quantiser Q() is obtained (see line 1 of Figure 3a). Next, the phase- aware precision allocation step is performed (see line 2 of Figure 3a). Concretely, the pretrained LLM is evaluated on a calibration dataset, by assessing its algorithmic performance using different combinations of prefill and decoding bitwidths. Through this process, the lowest pair of precisions that meet the quality target is determined.

[061] As shown in Figure 4, the model calibration comprises evaluating models with different bitwidths and selecting the lowest bitwidth for prefilling and decoding. Recently, it has been observed that allocating larger model capacity to the prefill phase for natural language understanding (NLU) allows for reduced capacity in the decoding phase for natural language generation (NLG) in order to produce responses with both high quality and efficiency [Aishwarya et al. (2024)]. The proposed method Phase-Aware Precision Allocation typically uses a high-precision model during the prefill phase, enabling the use of a low-precision model in the decoding phase. Since the prefill phase is compute-bound and the LLM inference latency is typically dominated by the decoding phase, using a high-precision model in the prefill phase introduces negligible overhead.

[062] Figure 4b is a flowchart showing the steps which may be carried out in the phase-aware precision allocation phase. In a first step S400, inputs in the form of performance metrics which are specific to the user device may be obtained. These inputs may be obtained from the user or may be obtained automatically by analysing the user device and its intended use. Next at step S402 a prefilling quantisation bitwidth and a decoding quantisation bitwidth may be obtained or selected from the set of weight quantisation bitwidths. These selected bitwidths are then used to quantise the plurality of weights of the ML model to generate a variant ML model at step S404. The quantisation is done using any suitable technique, for example, the SqueezeLLM [Kim et al.(2024)] or the GPTQ quantisation method [Frantar et al.(2023)].

[063] As shown at step S406, there is a decision as to whether to generate more variant ML models. If not, the process loops back to step S402. If so, at step S408, the generated set of variant ML models are evaluated using a calibration dataset. At step S410, any or all variant ML models which meet the output metric (e.g. the accuracy metric) are identified and those that do not are disregarded or discarded. Finally, the identified variant ML model having the lowest values for the prefilling quantisation bitwidth and decoding quantisation bitwidth is identified at step S412; and these identified values are output at step S414.

[064] Figures 5a to 5c shows the effect of Phase-Aware Precision Allocation. Three diverse tasks are targeted and responses from two variants of the same LLM Vicuna-7B (as described in Chiang et al, 2023), are compared. In each of Figures 5a to 5c, the top part of the Figure illustrates the results which are achieved using a uniformly low-precision variant with 2-bit quantised weights in both prefill and decoding. Tthe bottom part of each Figure illustrates the results with a two-precision variant with higher-precision 3-bit weights during prefill and lower-precision 2-bit weight during decoding.

[065] Various improvements. Reasoning abilities are enhanced, with the model demonstrating higher accuracy in dealing with numerical tasks, such as correctly calculating Fibonacci numbers (Figure 5a). Instruction-following improves, as the model better understands and responds to user instructions (Figure 5b). Third, an emerging ability like creativity is noticeably strengthened, with the model being able to handle open-ended tasks such as poem writing (Figure 5c).

[066] In summary, as shown in Figures 5a to 5c, empirically it can be observed that using the 3-bit model for prefilling significantly enhances the generative (e.g. NLP) capabilities when compared to the uniform 2-bit model. This result may be atributable to the significant steering role of the prefill phase in the output generation process. Specifically, a high-quality KV cache extracted from the input prompt during prefilling provides good understanding of the task and also relevant information-rich context to enhance the generation capability of the low-precision model.

[067] In summary, by observing distinct redundancy and sensitivity across the two LLM inference phases (i.e. prefilling and decoding), the proposed phase-aware precision allocator applies different arithmetic precisions and scheduling strategies in the prefill and decoding phases. In the prefill phase, where the context information of input prompts is crucial for generating accurate output tokens, a high-precision quantisation scheme is used. To further enhance performance and support efficient parallel processing, a uniform precision scheme is utilized, maximizing hardware efficiency during this phase.

[068] Progressive Mixed-Precision Decoding.

[069] As shown in Figure 2, the precision during the decoding phase is gradually reduced, significantly easing the memory bandwidth overhead. Given the memory-bound nature of the LLM decoding phase, this dynamic precision reduction scheme can effectively enhance the hardware performance of LLM inference, leading to higher decoding throughput in tokens per second. In other words, higher precision for earlier tokens, which often play a more decisive role in the response quality, is maintained while generating deeper tokens more rapidly through the use of narrower-bitwidth weights.

[070] Precision Switching Scheduler

[071] Returning to Figure 4a, scheduler preparation to enable the phase-aware precision allocation and progressive mixed-precision decoding is also carried out in the offline calibration. This involves selecting an appropriate scheduler approach which may be a static scheduler or a dynamic scheduler. The scheduler is configured with the PMPD parameters <P.S >and issues precision-switching actions to the system processor following schedule S for each input prompt. Under this setup, two types of schedulers are proposed: i) a prompt-agnostic static scheduler (static scheduler and the terms are used interchangeably), where S is determined prior to deployment, and ii) a task-agnostic learned scheduler (dynamic scheduler and the terms are used interchangeably), where S is derived dynamically for each prompt.

[072] Prompt-Agnostic Static Scheduler

[073] A static scheduler is fixed and has prompt-agnostic precision switching and so it is the same across Input prompts. This is illustrated for example In Figure 6a which plots the precision against decoded token ID to show the switching points for three different prompts. In each case, the initial precision is 8 bit which switches to 4 bit and then to 2 bit and there is the same switching point for each reduction per prompt. The switching to lower precision allows more efficient BW utilization which leads to higher throughput.

[074] Figure 6b is a flowchart showing how the schedule S for the static schedule is selected through an offline optimization process (line 4 of Figure 3a) that utilizes a calibration set. In a first step S600, inputs in the form of performance metrics which are specific to the user device may be obtained. These inputs may be obtained from the user or may be obtained automatically by analysing the user device and its intended use and may be the same as those identified above. Next at step S602, at least one switching point may be obtained or selected from a set of switching points and obtained initial and subsequent decoding quantisation bitwidths_may be obtained or selected from the set of weight quantisation bitwidths. In other words, PMPD may be parametrised with i) the precision set P and ii) the precision-switching schedule S = {st(p) \p EP], which consists of the possible switching points of each precision defined as st(p) e [1,CL), where CL is the maximum context length of a given LLM.

[075] These selected bitwidths and switching point(s) are then used to quantise the plurality of weights of the ML model to generate a variant ML model at step S604. As explained above, the quantisation is done using any suitable technique, for example, As shown at step S606, there is a decision as to whether to generate more variant ML models. If not, the process loops back to step S602. If so, at step S608, the generated set of variant ML models are evaluated using a calibration dataset.

[076] It will be appreciated that there may be overlap between the set of variably quantised models {mp} that was generated in the phase-aware precision allocation step (line 1 of Figure 3a) and the flowchart for PMPD. In this part of the method, the starting point st(p) (also termed switching point and the terms are used interchangeably) may be defined as the token index of the output sequence where a switch from the current precision (i.e bitwidth) to the lower precision p is performed. Under this setup, the evalulation to obtain the highest-performing PMPD configuration may be done using the following optimization problem: min st(p), Vp £ J’MPmin} (1) s. t. qref — e <q(S) &0 <st(p) <CL, p E P (2) p >qst(p) <st(q), p.q EP (3) where q(S) is the achieved quality of schedule S. The objective function aims to minimize the number of tokens to be processed at each precision in the precision set subject to the following additional constraints. The lowest precision pmin will be used last and until the end of the sequence for higher throughput. The precision used must meet the specified quality target (first constraint in Eq. (2)) and not exceed the LLM’s maximum context length (second constraint in Eq. (2)). There is also a progressive precision lowering approach (precedence constraint in Eq. (3)).

[077] When deciding at step S606 whether to generate more variant models, it is noted that the objective function of Eq. (1) can be evaluated for all combinations of precision-switching schedules. Ideally, the highest-performing schedule considers all p e P per-prompt per-step when deciding fora precision switch without the precedence constraint shown in Eq. (3). This schedule can be obtained through an exhaustive enumeration over \Pf- schedules. However, for realistic values of 1^1 and CL, this computation quickly becomes intractable. Besides, even if such a schedule can be efficiently found, it is, in most cases, impractical as it would require excessively frequent switching, e.g. per decoding step, between variably quantised weights, aggravating the memory bandwidth demand.

[078] Instead, the proposed progressive precision lowering during decoding means that the total number of schedules to be examined is reduced to: ZI^I-1 CL\ (1^1 - 1)! r=o (Ci — H (1^1 — 1 — r)! where r is the number of times the precision switches. Nonetheless, the search time for the optimal schedule increases with CL and IJ5], which is still prohibitive at run time. Hence, to further reduce the search time for the optimal S to 0(1), the number of candidate precisionswitching points is parameterised as denoted by N. A range constraint to the second constraint in Eq. (2) is also added as: a rni . . r CL 2CL (N-2)CL stwemut—.....

[079] While the precedence constraint in Eq. (3) and the optimization objective in Eq. (1) still hold, the achieved quality in Eq. (2) is approximated as I H 4—ixev where V is the validation set and <?X(S) is the measured quality score of the schedule on a sample x.

[080] As the validation set is defined per task, the static scheduler remains prompt-agnostic (i.e. uses the same precision-switching schedule for all prompts) and task-specific. As such, this type of scheduler is tailored to the resilience of the task at hand and has the advantage of inducing close to no runtime overhead. However, by adopting a uniform schedule within a given task, it requires the availability of a task-representative validation set, which is often not a realistic assumption, and it does not exploit the variability in difficulty and resilience to approximation across different prompts.

[081] Task-Agnostic Learned Scheduler

[082] To eliminate the need for a per-task validation set, a trainable scheduler (also termed a learned scheduler) that adjusts the schedule S to each input prompt may be used. For a given LLM, the scheduler is trained on a generic, task-agnostic dataset and then applied across various downstream tasks, amortizing in this way its training cost. Specifically, a schedule is devised given the features of each input prompt. Training the scheduler to make very fine-grained per-step decisions, however, would require both substantial model capacity and excessively frequent scheduler invocation, which induce significant runtime overhead and counteract PMPD’s speedup benefits.

[083] In other words, a dynamic scheduler has prompt-specific precision switching and so it is different across input prompts. This is illustrated for example in Figure 7 which plots the precision against decoded token ID to show the switching points for three different prompts (they are the same prompts as in Figure 6). In each case in Figure 7, there is a different switching point per prompt. For example, for the first prompt, there is only one switch from 8 bit to 4 bit. For the second and third prompts there are two swtiches from 8 bit to 4 bit and finally to 2 bit. However, the timing of the switches is different. In other words, the dynamic scheduler can have a different number of switches and / or different timing of the switching depending on the input.

[084] Figure 8 schematically shows the different steps in the offline calibration depending on whether a static or a dynamic scheduler is selected. If a static scheduler is selected, it is not possible to have different switching points for different prompts and thus the phase-aware precision alllocation and progressive mixed-precision decoding is set. Thus, the phase of designing and training the scheduler is skipped and the calibration process is ended.

[085] The optimal precision lowering locations are likely to be highly dependant on the specific task and input prompt. This motivates the design of a runtime scheduler that dynamically decides the precision-reduction positions by considering the relevant context information through the KV cache. The dynamic scheduler allows the implementation of this second aspect of the present techniques which can be used together with or separate from the first aspect.

[086] As shown in Figure 8, the training of the dynamic scheduler includes running different combinations of decoding precision, constructing a dataset with precision-performance pairs for training, defining the model architecture of the scheduler and training the scheduler to optimise the precision switch policy.

[087] Figure 9a illustrates an example architectural design for a dynamic scheduler which can be trained to enhance the performance of LLM inference. An example dynamic scheduler is a lightweight trainable multi-layer perceptron which at interference minimizes the runtime overhead of generating precision-lowering decisions. As shown in the example, it is an MLP classifier with mainly two linear layers and a lightweight scheduler with low runtime overhead only invoked once per inference instance. The Input: KV cache is generated from the prefill stage and the Output is the Token index of precision switch. The scheduler leverages the KV cache generated by the LLM to effectively capture and utilize context information. Concretely, the scheduler processes as input the KV cache generated from the prefill stage and generates the output token index where a precision switch should occur. The scheduler leverages the KV cache generated by the LLM to effectively capture and utilize context information.

[088] More specifically, as an example of a suitable architecture and to avoid runtime overhead, a lightweight attention module may be used for the Learned scheduler. Given an input Keys cache K e WT><Dfc and Values cache V e nV*0”, the output of Learned scheduler may be computed as: 0 = Softmax K (4) Logits = Softmax^MLP (0))(5) where MLP is a feedforward neural network with one hidden layer and a ReLU as its activation, and qe e RDt> is a learned query vector. In other words, the trainable parameters of the network include the query vector and the parameters of the MLP.

[089] Figure 9b shows a flowchart setting out the steps of one option for training the scheduler which as shown in step S900 is received by the user device. The scheduler is trained using a dataset which is received at step S902 and where each data point is represented as a pair in the form: < KVCache, Label > The scheduler may be trained using any standard process, e.g. at step S904 an input KVCache is selected from the dataset. The scheduler makes a prediction of the label at step S906. An error may be determined between the prediction and label at step S908 and adjustments may be made to the weights of the dynamic scheduler at step S910. As noted above, the trainable parameters may be the parameters (e.g. weights) of the MLP. If there is convergence in the adaptations to the model, as determined at step S912, the scheduler may be output at step S914. If not, the process loops back to select another datapoint.

[090] For KVCache, the KV cache tensor from the final Transformer block of the given LLM is used. For example, the shape [2, 128, 55, 32] represents 2 (key and value matrices), 128 (attention heads), 55 (sequence length), and 32 (hidden dimension per head). Each Label is a one-hot vector of dimension P, representing the token index for precision switching. For example, when P = 4, the four different token indices can be [0], [1 / 3 • n], [2 / 3 • n] and [n], with n being the maximum supported response length of the given LLM.

[091] Figure 9c is an example flowchart showing one method for generating the training dataset. At step S920, an ML model (e.g. the LLM and the terms are used interchangeably) is received. At step S922, an input to be processed by the ML model is received. The input (e.g. a prompt) is processed at step S924 through the prefill stage of the received ML model to generate an output from this stage, e.g. the KV cache described above. The result is output at step S926. A set of decoding quantisation schedules is then obtained at step S928. These quantisation schedules may be in the form of switching points (e.g. the token indices described above) together with the quantisation bitwidths which are applied before and after each switching point. The received ML model is quantised using each of the obtained schedules at step S930. In this way, the LLM can be evaluated on the validation dataset using various precision combinations (e.g. a mix of 4-bit, 3-bit, and 2-bit). In other words, at step S923, the set of variant ML models is evaluated. At step S934, the precision combination that yields the best algorithmic performance is identified. The “best algorithmic performance” may be determined as described above using the accuracy metric and other constraints. The schedule of this best variant is assigned as the label for the corresponding KV cache at step S936. There is a determination at step S938 as to whether more data points are required for the dataset and if so, the process loops back to consider another input and if not, the dataset is output at step S940.

[092] In other words, the features extracted in the prefill phase are used as inputs to the scheduler. With this approach, a lightweight architecture is maintained for the scheduler. Minimal overhead during inference is also introduced because the scheduler generates schedule S once before decoding (line 13 of Figure 3a). As described above, different input features, i.e. KV caches and activations, from various layers of the model are investigated. The results show that using the prefilled KV cache is efficacious in most cases, which indicates that the prefilled KV cache provides good task-specific context. Even without a task-specific validation set and labeled ground truths, the proposed scheduler performs close to the static and, in some cases, even outperforms it in both algorithmic and hardware performance.

[093] As an example of the training, to prepare the training dataset, the C4 dataset is used as the seed dataset. The C4 dataset is a collection of about 750GB of English-language text sourced from the public Common Crawl web scrape. Each sample is randomly truncated to create a completion task. For each sample, generate N sequences, each corresponding to a high-precision step from: (0 CL 2CL (N-Z)CL , ' N-l ’ The ground-truth label is set to the lowest high-precision step such that the Rouge-L score of the sequence matches or exceeds that of the fp16 model output. During training, a crossentropy loss function may be used as the objective.

[094] Runtime Deployment

[095] The detailed steps of the deployment stage are set out in lines 9 to 22 of the algorithm of Figure 3a and lines 3 to 13 of the algorithm of Figure 3b. These steps of runtime deployment are schematically illustrated in Figures 10a and 10b which shows the interaction between the precision switching scheduler and the LLM. Figure 10a shows the runtime deployment which may be used with or without the dynamic scheduler. The precision-switching scheduler receives an input which is the runtime information, for example the decoded token index and optionally the KV cache for the dynamic scheduler. The precision-switching scheduler outputs the runtime precision policy which as shown by the dotted arrows controls the number of bits used in each layer of the LLM. Thus, unlike existing quantisation approaches for LLMs which employ a single arithmetic precision that remains uniform across both the prefill and decoding phases. Instead, the proposed method enables the usage of different precision in each phase. Importantly, it tailors the arithmetic precision to each phase’s characteristics. In the computebound prefill phase, higher precision is used to maintain accuracy without impacting processing speed. In the memory-bound decoding phase, lower precision is used; significantly reducing the memory bandwidth requirements and improving decoding throughput (tokens / s).

[096] The proposed approach improves decoding throughput further by exploiting the following property: tokens generated earlier in the output sequence are more sensitive and would lead to accuracy degradation if aggressive quantisation is used, while tokens generated deeper in the output sequence are more resilient. For example, the prefill stage has a fixed precision, e.g. 4-bit in this example and the decoding stage uses lower precision (e.g. 3-bit and 2-bit). The proposed method exploits this observation by generating the first few sensitive tokens using high precision - so that accuracy is maintained - and then switching to lower precision for the later more resilient tokens - so that they are generated more rapidly and achieving higher sustained throughput (tokens / s). The exact precision-switching time is decided by the precision-switching scheduler at run time.

[097] Figure 10b illustrates more detail of the runtime deployment forthe dynamic scheduler. Once trained, the scheduler is configured with i) the set of available precisions and ii) the dynamic switching policy ready forthe deployment stage. The LLM generates the KV cache. At runtime, the scheduler receives and processes the KV cache and evaluates the policy to determine when a precision switching should be triggered by generating lowering decisions. In this way, upon deployment, the dynamic scheduler determines the highest-performing precision-switching time in a prompt-adaptive manner, leading to a PMPD configuration tailored to the prompt and task at hand. In other words, there is decoding with mixed precision. The scheduler leverages the KV cache generated by the LLM to effectively capture and utilize context information. In other words, when using a dynamic trainable scheduling approach, the scheduler selects the most suitable time to reduce precision based on the characteristics of the current prompt, leading to a prompt-optimized precision-switching schedule that maximizes throughput.

[098] Figure 10c illustrates the runtime deployment of Figure 10b as a sequence of steps in a flowchart with steps on the left carried out by the ML model and steps on the right by the dynamic scheduler (after the initial receiving step). Initially at step S1000, the user device receives an ML model, a quantisation schedule (specifically a prefilling quantisation bitwidth) and a dynamic scheduler model. At step S1002, an input which is to be processed by the ML model is also received, e.g. a prompt. The prefilling phase of the ML model is quantised at step S1004 using the prefilling quantisation bitwidth specified above. At step S1006, a prefilling phase output is then generated using the quantised prefilling phase of the ML model. This prefilling phase output is then input at step S1008 to the dynamic scheduler to generate a decoding quantisation schedule which is output at step S1010. At step S1012, the decoding phase of the ML model is then quantised using the decoding quantisation schedule and at step S1014, the prefilling phase output is processed through the quantised decoding phase of the ML model to generate the final output which is then output from the user device at step S1016.

[099] Figures 11a and 11 b illustrate two deployment implementations which may be termed Variably Quantised Weights and Single-Copy Any-Precision Model. Each comprises a Neural Processing Unit (NPU) and a precision switching scheduler coupled to the main memory via a multiplexer (MUX). In the first implementation shown in Figure 11a, different precisions (e.g. 8-bit, 4-bit and 2-bit quantisation) are stored as separate model weights and distinct sets of weights are loaded for mixed precision. This is suitable for large, distributed systems with large batch size and will tend to achieve higher algorithmic performance than the implementation of Figure 11 b.

[100] In Figure 11b, multiple precisions (e.g. 8-bit, 4-bit and 2-bit quantisation) are stored as unified model weights, and the same set of weights is loaded for mixed precision. This is described for example in [Park et al.(2024)]. This is more suited to resource-constrained edge and mobile devices with limited memory and computational capacity. There may be some accuracy loss but this can be mitigated by an optimized scheduler.

[101] Figures 12a and 12b show two embodiments which are running an LLM inference based on the techniques describe above. Figure 12a schematically shows the deployment of LLMs on mobile apps to provide new and improved user experiences. Key LLM-enabled use cases include real-time speech translation, LLM-based voice assistant (e.g. Bixby) with advanced multi-turn conversation capabilities, and on-device subtitle translation on TVs. For these user-facing applications to run on-device, they need to meet stringent latency constraints in order to guarantee high responsiveness and maintain the user experience. This is particularly challenging on such resource-constrained devices, as the memory bandwidth is limited. The proposed PMPD method accelerates the on-device execution of LLM inference, by selectively lowering the memory bandwidth requirements of a given LLM and delivering improved efficiency and performance. In this case, the PMPD implementation which is used is the arrangement of Figure 11b. It is also especially suitable for NPU-equipped platforms, such as mobile phones and TVs, which support multiple arithmetic precisions, such as FP16, INT8 and INT4.

[102] Figure 12b schematically shows the deployment of LLMs in cloud-based systems. LLMs are increasingly hosted on servers to support services for thousands of users. Such services are exposed either externally to users as a service or internally to employees for higher productivity. Applications include translation services (e.g. translate.samsung.com), generic assistants (e.g. ChatGPT, Gemini, Gauss) and code generation models (e.g. CodeLlama, CodeQwenl .5). In both real-time and offline deployments, the proposed method improves the latency of LLM inference, leading to faster response to the users and to higher throughput (sustained served users / s) and more efficient system utilisation. In this example, the PMPD implementation which is used in the arrangement of Figure 11a. The proposed design is parametrized with respect to the arithmetic precisions it utilizes. As such, it can be seamlessly deployed to different types of server platforms, and it can be adapted to the traffic of the server system. For instance, under higher traffic load, PMPD can switch to more aggressive precision during decoding (2-bit) to decongest the system, while under lowertraffic, the proposed method can utilize higher precision to provide higher-quality responses.

[103] Figure 13 is a block diagram of a system that may be used to perform the methods described above. The system may comprise an electronic device in the form of a user device 150 to use an LLM with progressive mixed precision decoding as described above. The user device 150 comprises: at least one processor 152 coupled to a memory 154. The at least one processor 152 may comprise one or more of: a microprocessor, a microcontroller, and an integrated circuit. The memory 154 may comprise volatile memory, such as random access memory (RAM), for use as temporary memory, and / or non-volatile memory such as Flash, read only memory (ROM), or electrically erasable programmable ROM (EEPROM), for storing data, programs, or instructions, for example.

[104] The electronic device 150 comprises at least one LLM 160 with at least one set of weights 168. There is also a scheduler 170 which is configured with the set of available precisions for the LLM and (optionally) a dynamic switching policy. The LLM and the scheduler may be received from a server as described above.

[105] The system may thus further comprise an electronic device 100, in the form of a server. The electronic device 100 comprises: at least one processor 102 coupled to a memory 104. The at least one processor 102 may comprise one or more of: a microprocessor, a microcontroller, and an integrated circuit. The memory 104 may comprise volatile memory, such as random access memory (RAM), for use as temporary memory, and / or non-volatile memory such as Flash, read only memory (ROM), or electrically erasable programmable ROM (EEPROM), for storing data, programs, or instructions, for example.

[106] The electronic device 100 comprises at least one LLM 106 with at least one set of weights 108. There is also a scheduler 112 which is configured with the set of available precisions 112 for the LLM and (optionally) a dynamic switching policy 114. The LLM 106 and / or the scheduler 112 may have been trained as described above using a dataset 120. The dataset may be stored remotely from both the electronic device and the user device as illustrated. It will also be appreciated that a whole or part of the dataset may be stored on one or both of the electronic device and the user device.

[107] Figures 14a and 14b illustrate how the decoding stage of LLM inference is memorybound, with the weights dominating the memory transfers. Figure 14a shows the breakdown of R / W transactions in DRAM using a prefill heavy and a decoding heavy LLM respectively. Figure 14a plots Rouge-L against timestep of perturbation as described in [Schuster et al 2022], Figure 14b shows that approximations earlier in the decoded sequence have more severe impact on the output quality, while deeper tokens demonstrate higher resilience. To remedy memory boundedness during LLM inference, Progressive Mixed-Precision Decoding (PMPD) is introduced. This is a novel method that gradually reduces the arithmetic precision deeper in the decoded sequence while matching the precision to the characteristics of the prefill and decoding phase, respectively. PMPD co-optimizes both the time-to-first-token (TTFT) and the sustained throughput (tokens / s) while maintaining a target generation quality.

[108] Evaluation

[109] Experiments were conducted on edge-deployable models, including Vicuna-7B [Chiang et al, 2023], MobileLLaMA-1.4B [Chu et al, 2023], Stable LM Zephyr-3B [https: / / stability.ai / news / stablelm-zephyr-3b-stability-llm], and Phi-1.5 [“Textbooks Are All You Need II: phi-1.5 technical report,” Microsoft Research, 2023], The models were evaluated on various tasks including zero-shot generative performance on news summarization, dialogue summarization, and translation tasks using the CNN / DM [Hermann et al, 2015], Dialogsum [Chen et al, 2021], and IWSLT French-English [Cettolo et al, 2017] datasets, respectively. These tasks encompass context understanding and text generation, providing a comprehensive evaluation of PMPD. Open-ended question answering was tested on MT-Bench [Chiang et al, 2023], This used a special case where Static scheduler is infeasible due to the absence of a validation set.

[110] Single-precision quantised models and quantised models enhanced with Dense-and-Sparse decomposition (DNS) [Kim et al, 2024] were compared. The DNS ratios were adjusted to match PMPD’s average bitwidth. PMPD’s average bitwidth is calculated as the weighted average of bitwidths over decoding steps, weighted by the number of tokens generated at each precision.

[111] The nested quantisation method of Any-Precision LLM [Parket al.(2024)] was used for the post-training quantisation PTQ method. This ensured that using multiple weight precisions incurred no memory footprint overhead. The Static scheduler finds a schedule that minimizes high-precision steps on each benchmark’s validation set while maintaining lossless 5 performance. The Learned scheduler was trained using the first 256 samples from the C4 test dataset as the seed dataset. The high-precision variant of each model is quantised to the lowest lossless precision determined by perplexity on the C4 dataset, while low precision is defined to be one bit lower than the high precision. 10

[112] Table 1 presents the algorithmic performance of PMPD in comparison with both singleprecision baselines and DNS. Both Static and Learned schedulers achieve significant bitwidth reduction while maintaining competitive performance. Pairwise winners between the scheduler and DNS are underlined whilst the highest overall scores are in bold. Dataset CNN / DM Dialogsum IWLST Model Method Bit Rouge-L or BERTScore Bit Rouge-L or BERTScore Bit BLEU or SacreBLEU Vicuna-7B Baseline-I 2 8.30 / 78.4 2 10.2 / 75.5 2 1.2 / 1.2 Baseline-h 3 24.2 / 86.9 3 24.4 / 88.2 3 31.6 / 31.6 DNS 2.39 24.2 / 86.8 2.0 - 2.68 27.6 / 27.6 PMPD-Static 2.39 24.3 / 87.0 2.0 25.0 / 88.2 2.68 31.0 / 31.1 PMPD-Learned 2.43 24.0 / 86.7 2.74 24.5 / 88.2 2.37 29.9 / 29.9 Mobile LLaMA Baseline-I 3 16.3 / 83.3 3 15.8 / 84.1 3 9.8 / 9.83 Baseline-h 4 17.2 / 83.5 4 16.8 / 84.9 4 12.7 / 12.7 DNS 3.37 17.4 / 83.5 3.21 14.7 / 84.4 3.65 12.0 / 12.0 PMPD-Static 3.37 17.6 / 83.7 3.0 17.0 / 85.0 3.65 12.6 / 12.6 pmpd-Learned 3.19 16.6 / 83.2 3.21 17.1 / 85.0 3.48 11.8 / 11.8 Phi-1.5 OR Zephyr-3B For Baseline-I 3 13.4 / 82.4 3 15.3 / 85.1 3 21.1 / 21.1 Baseline-h 4 16.2 / 84.0 4 18.0 / 86.1 4 30.4 / 30.4 DNS 3.71 12.4 / 81.8 3.30 16.1 / 85.7 3.34 28.2 / 28.2 PMPD-Static 3.71 16.2 / 84.0 3.30 18.1 / 86.2 3.0 29.7 / 29.7 IWLST PMPD- Learned 3.09 15.5 / 83.4 3.52 17.9 / 86.1 3.34 29.8 / 29.8

[113] Both Static and Learned schedulers deliver lossless performance on the CNN / DM and Dialogsum datasets with up to 33% reduction in bitwidth, highlighting their ability to effectively capture context and generate high-quality outputs. In some cases, PMPD even outperforms the high-precision baseline by up to 2.4%. On the IWSLT dataset, Static experiences a moderate performance drop of up to 2.3%, while still reducing the average bitwidth by 9-33%. Notably, the results show that using high-precision models in the decoding phase is not always necessary. Also, Static is able to accurately identify such scenarios, achieving the maximum bitwidth reduction possible.

[114] DNS is a robust baseline that maintains low-bit precision by selectively handling outlier values. Nonetheless, the propsed Static scheduler consistently outperforms DNS at the same bitwidth across all tasks, achieving up to 29% higher BLEU on the IWSLT dataset. Additionally, the performance gain increases with smaller models: DNS performs comparably to Static on CNN / DM for Vicuna-7B but lags behind by 23% in Rouge-L on MobileLLaMA. This indicates that PMPD is particularly well-suited for edge deployment scenarios.

[115] On average, the Static scheduler achieves a slightly higher bitwidth reduction (0.66 vs. 0.63 bits) and frequently outperforms Learned, showing more stable performance with less variance. This is likely because Static leverages task-specific data, while Learned may face challenges from data distribution shifts between its training set and each test set.

[116] Table 2 plots the performance of the proposed method with fp16 output as a reference and using MT-Bench as a challenging benchmark covering a wide range of open-ended tasks. Model Method Bit Rouqe-L / BERTScore Vicuna-7B 2bit 2 16.7 / 80.4 3bit 3 41.0 / 89.9 PMPD-Learned 2.68 39.6 / 88.9 DNS 2.68 36.7 / 88.2 Zephyr-3B 3bit 3 28.6 / 87.2 4bit 4 40.8 / 89.6 PMPD-Learned 3.48 40.7 / 89.7 DNS 3.48 35.3 / 88.6 MobileLLaMa 3bit 3 22.3 / 82.4 4bit 4 28.9 / 85.8 PMPD-Learned 3.39 29.1 / 85.4 DNS 3.39 25.4 / 84.2 Phi-1.5 3bit 3 27.9 / 85.7 4bit 4 35.4 / 87.8 PMPD-Learned 3.56 35.7 / 87.9 DNS 3.56 29.9 / 86.7

[117] In particular, the usefulness of Learned in this scenario is demonstrated where Static is not feasible due to the absence of task-specific validation set. The faithfulness of the quantised models to the fp16 model was measured by evaluating the similarity of their outputs. As shown in Table 2, Learned achieves an average bitwidth reduction of 0.47 with no drop in BERTScore and only up to 1.4 points drop in Rouge-L. Additionally, it consistently outperforms DNS with 13.9-18.4% higher Rouge-L, demonstrating that PMPD better preserves model fidelity in complex tasks. Moreover, Learned’s relative performance improves with smaller models, making it particularly promising for reducing quantisation error in mobile-friendly models.

[118] Figures 15a to 22b plot similar results to those presented above. Figures 15a and 15b plot scores against average bits when using a static scheduler as described above. The score is BERTscore in Figure 15a and Rouge-L score in Figure 15b. For both tests, the LLM is MobileLlama-1.4B, the dataset is CNN / DailyMail and the hardware includes an NPU with 4K Multiply-Accumulate (MAC) units and 64 GB / s off-chip memory bandwidth. The decoding starts with 3-bit and reduces to 2-bit. The switching at 20, 40 and 60 tokens (and so on) is shown. Each point on the plot has a static schedule. Speedup measurements obtained from in-house mobile NPU simulator. There is a speedup of approximately 1.55 times compared to using 4-bit throughout. Due to the memory-boundness of the decoding stage, using lower-precision weights leads to significant speedups on conventional NPUs.

[119] Figures 16a and 16b plot the same scores against average bits as Figures 15a and 15b. All the details are the same as for the previous results except that the dataset is MT bench. In this case, the speedup is approximately 1.25 times.

[120] Figures 17a and 17b plot the same scores against average bits as Figures 16a and 16b. All the details are the same as for Figures 16a and 16b.except that the LLM is Phi-1.5 (1,3B) [3], Similarly, Figures 18a and 18b plot the same scores against average bits as Figures 16a and 16b. All the details are the same as for Figures 16a and 16b.except that the LLM is Qwen2-1.5B. In these cases, there is still a speed-up but it is smaller than the other examples.

[121] Figures 19a and 19b plot scores against average bits when using a dynamic scheduler as described above. The score is BERTscore in Figure 19a and Rouge-L score in Figure 19b. For both tests, the LLM is MobileLlama-1.4B, the dataset is MT Bench and the hardware includes an NPU with 4K Multiply-Accumulate (MAC) units and 64 GB / s off-chip memory bandwidth. The decoding starts with 4-bit and reduces to 32-bit. Speedup measurements obtained from in-house mobile NPU simulator. There is a speedup of approximately 1.16 times compared to using 4-bit throughout. Due to the memory-boundness of the decoding stage, using lower-precision weights leads to significant speedups on conventional NPUs.

[122] Figures 20a and 20b plot the same scores against average bits as Figures 19a and 19b. All the details are the same as Figures 19a and 19b except that the dataset is CNN / Daily Mail. In this case, the speedup is approximately 1.17 times.

[123] Figure 21a shows different schedules of mixing 3-bit and 4-bit model in the decoding phase. Figure 21b shows the results which are obtained on summarization task on CNN / DM dataset with Phi-1.5. To ensure fair comparison, the same number of tokens are generated by the 4-bit model baseline and 4-bit model is used in the prefill phase for all schedules.

[124] Figures 21a and 21b show an exploration of the optimization space of mixing generation of high-precision and low-precision model in the decoding phase to further increase the output quality of low-precision model. Different ways to schedule the two models are explored. As shown in Figure 21a, the high-precision model generation is arranged in the first half, middle half, and last half of the generated tokens, respectively. Figure 21 b which plots the Rouge-L score shows that scheduling the high-precision model generation for the first half of the decoding stage leads to the best performance, matching the performance of using a high-precision model throughout the inference as described above in the Progressive Mixed-Precision Decoding techniques. One intuitive explanation is that using high-precision model to generate the first few tokens minimizes the error accumulation that affects the later inference process. This observation also aligns with the “attention sink" phenomenon [Xiao et al.(2023)], where the attention scores tilt heavily towards the initial tokens. As a reference, the performance of alternating the two models at each token was also tested, which also proves to work well as shown in Figure 21b. However, frequent precision switching introduces hardware overhead such as additional memory traffic so this schedule is not practical in practice.

[125] Figures 22a to 22c compare the end-to-end NPU throughput. To estimate the processing speed of PMPD when deployed on an NPU, an analytical performance model of the hardware architecture of FlightLLM [Zeng et al, 2024] was developed. This is an LLM-optimized accelerator, adapted to support multi-precision weight loading. FlightLLM consists of: i) a unified Matrix Processing Unit that can perform multiple types of multiplications between matrices and / or vectors through a hierarchical structure of multiply-accumulate (MAC) units, ii) a Special Function Unit where LLM-specific operations such as softmax, layer normalisation etc are mapped and Hi) an optimized memory hierarchy for the computational pattern of LLMs, while keeping all activation on-chip during decoding. The underlying design space of FlightLLM’s NPU architecture is traversed through a design space exploration method that takes into consideration the need to support different precisions of LLM weights. For the experiments, there are two NPU configurations consisting of: 4K and 16K MAC units with 1 GHz clock frequency (i.e. 8 and 16 teraops / sec (TOPS) peak throughput, respectively) for the deployment of smaller- (MobileLLaMa-1.4B, Phi-1.5) and larger-scale LLMs (Vicuna-7B), respectively, and with 32 GB / s off-chip memory bandwidth.

[126] As shown in Figures 22a to 22c, the proposed methods achieves a significant speedup of between 3.8-8.0 times over a baseline model termed fp16 across various datasets. Figures 22a to 22c show a comparison in terms of end-to-end NPU throughput. The value over each bar is the speedup ratio obtained in comparison to an fp16 model (i.e. using 16-bit halfprecision floating-point as the data type). Figure 22a is for the CNN / DM datasets, Figure 22b for Dialogsum and Figure 22c for IWSLT. There are more pronounced speedups for the larger Vicuna-7B, due to its higher resilience to quantisation errors, enabling more aggressive 2-bit precision lowering. Overall, the proposed method achieves more than 50 tokens / s and 15 tokens / s for mobile-grade LLMs and for the larger Vicuna-7B, respectively. This is in spite of the limited bandwidth (32 GB / s for NPU vs. 1008 GB / s for RTX4090 GPU). These results showcase the methods effectiveness and suitability for LLM deployment on mobile devices.

[127] As mentioned above, the present techniques can also be deployed on a GPU. Following Any-Precision LLM, the latencies of linear layers in the LLMs across different Nvidia GPUs, including RTX 4090 and A40 are also evaluated. The implementation was based on PyTorch [Paszke et al, 2019], and its built-in Profiler tool was used. The average self CUDA time metric [https: / / pytorch.org / tutorials / recipes / recipes / profiler_recipe.html] over 100 forward passes was used to ensure reliable results. Table 3 compares the latency of fp16 and quantised linear layers, reporting CUDA kernel runtimes with batch size of 1.

[128] Table 3 GPU speedup ratios Vicuna-7B MobileLlama Phi-1.5 Zephyr-3B Attn Proj. MLP Proj. Attn Proj. MLP Proj. Attn Proj. MLP Proj. Attn Proj. MLP Proj. RTX 4090 Baseline-h 6.25* 11.32* 1.25* 2.50* 1.32* 1.70* 3.33* 2.54* PMPD 6.25* 12.20* 1.40* 2.81* 1.51* 1.84* 2.73* 2.82* A40 Baseline-h 5.81* 3.77* 3.00* 3.96* 2.57* 2.83* 3.50* 3.02* PMPD 6.58* 4.60* 3.37* 4.74* 2.77* 3.51* 3.81* 4.27*

[129] As shown above, PMPD achieves speedups of 1.40-6.58x for smaller attention projection layers and 1.84-12.20x for larger feedforward projection layers. It is hypothesized that the lower speedup on smaller layers is primarily due to the memory-bound nature of matrix-vector multiplication. Despite the accelerated GPU execution, CPU-side processes for GPU kernel launching can be observed to become a bottleneck, resulting in suboptimal GPU utilization. Potential solutions include leveraging CUDA Graph [https: / / pytorch.org / docs / stable / generated / torch.cuda.graph.html] to launch multiple GPU operations simultaneously.

[130] Tables 4 and 5 report the GPU kernel latency using the self CUDA time metric on RTX 4090 and A40 GPUs. For PMPD, the latency is calculated as a weighted average of the fixed-bitwidth latencies, where the weights correspond to the average number of decoding steps performed at each bitwidth across datasets. To select between the Static and Learned schedulers, the one with the highest performance is chosed. For MobileLLaMA and Vicuna-7B, bitwidths are averaged over the CNN / DM, Dialogsum, and IWSLT datasets. For Phi-1.5, bitwidths are averaged over the CNN / DM and Dialogsum datasets, while for Zephyr-3B, the bitwidth is from the IWSLT dataset.

[131] Table 4: GPU Matrix kernel latency (jas) on RTX 4090: Model q proj k proj v proj o proj | up proj down proj | gate proj VlobileL LaMA-1,4B-Chat fp16 5 5 5 5 13.3 15 13 w2 3 3 3 3 4 5 4 w3 3 3 3 3 4.3 5 4.57 w4 4 4 4 4 5.1 6 5.97 phi-1 5 fp16 5 5 5 5 10.6 11.8 - w2 2.96 2.82 2.93 2.82 4.96 4.94 - w3 3 2.81 2.89 2.91 5.36 5.92 - w4 3.98 3.79 3.86 3.9 6.92 6.93 - stablelm-zep iyr-3b fp16 10 10 10 10 11 18 11 w2 2 2 2 4 6 4 6 w3 4 4 4 4 6 6 6 w4 5 3 5 3 8 7.1 6 vicuna-7b-v1.5 fp16 25 25 25 25 97.1 101.93 97 w2 4 4 3.8 4 7 8 7 w3 4 4 4 4 8.1 9 8 w4 5 5 5 5 12 12 12

[132] Table 5: GPU Matrix kernel latency (jas) on A40: Model | q proj | k proj | v proj | o proj | up proj | down proj | gate proj MobileLLaMA-1,4B-Chat fp16 18 18 18 18 42.4 47.5 42 w2 4.9 4.9 5 4.7 7.4 8 7.5 w3 5 5 5 5 8 9 8 w4 6 6 6 6 11 12 10.2 phi-1 5 fp16 18 18 18 18 62.2 60.3 - w2 5.9 5.7 5.9 5.9 10 11 - w3 6 6 6 6 12 13 - w4 7 7 7 7 20.3 21.3 - stablelm-zephyr-3b fp16 28 28 28 28 66.1 70.1 66.4 w2 6 6 6 6 11 10 11 w3 7 7 7 7 13 12.9 13 w4 8 8 8 8 22.1 23.2 22 vicuna-7b-v1.5 fp16 64 64 64 64 170 162 170 w2 9 9 9 9 24.3 30.9 24.2 w3 11.2 11 11 11 33.9 43 33.4 w4 21.2 21.5 21.3 21.2 46.6 58.9 45.9

[133] Figures 23a to 23d analyze the impact of phase-aware precision allocation on performance. Rouge-L on CNN / DM improves by an average of 4.3 points, BERTScore on Dialogsum increases by 3.6, and BLEU improves by 9.45. The 2-bit variant of Vicuna-7B benefits significantly from using a higher-precision model during the prefill phase, as it struggles with instruction following and often outputs incomplete tokens or repeated prompts. Using a 3-bit variant during prefilling improves language comprehension and generation. Latency overhead from using higher precision is minimal, ranging from 0.07% to 1.05%, confirming that the prefill phase is compute-bound. This supports the efficiency of phase-aware precision allocation for system optimization.

[134] Figure 24 plots the normalized scores of different Learned Scheduler Design across different datasets (CNN / DM, IWSLT and MT-bench). Different choices of inputs are used, including activations and KV caches from the first, middle, and last Transformer attention blocks. Specifically, the input activations to the final projection layer of the attention block are studied. Two random schedulers are tested: i) using a uniform distribution (random uniform) and ii) matching the label distribution ofthe training dataset (randomprior). As shown in Figure 24, using the KV cache from the last attention block achieves the highest normalized score with significant bitwidth reduction, highlighting its effectiveness in capturing contextual information for precise switch point prediction. This scheduler also outperforms the two random schedulers, indicating that it is effectively learning patterns in the KV cache.

[135] The proposed approach offers a plug-and-play framework compatible with any PTQ method. To showcase its generalizability, the performance of both Static and Learned schedulers using GPTQ is analysed using a uniform quantisation technique [Frantar et al.(2023)]. As shown in the table below, both schedulers deliver performance that closely matches the high-precision baseline. Method Bit Rouge-L BertScore 3bit 3 9.71 82.7 4bit 4 16.0 84.2 PMPD- static 3.63 16.0 84.6 PMPD- Learned 3.02 14.2 84.2

[136] In summary, the core of these techniques are 1) a phase-aware precision allocation method and 2) a progressive mixed-precision decoding scheme for LLM inference workloads, abbreviated PMPD. A novel LLM inference method is introduced that has a phase-aware approach which matches the properties of each phase, together with a progressive mixed-precision scheme that alleviates the memory-bounded of the decoding phase.

[137] In phase-aware precision allocation, given the observation that the prefill stage output is critical for the context of the generated response, the proposed method: i) exploits the compute-boundedness of the prefill phase to allocate higher arithmetic precision while maintaining high hardware utilization and ii) alleviates the memory-boundedness of the decoding phase by reducing the arithmetic precision to lower the memory bandwidth requirements. Overall, the phase-aware precision allocation happens statically prior to deployment.

[138] In progressive mixed-precision decoding, a precision-switching scheduler that adaptively reduces the arithmetic precision used during the decoding phase is introduced, by monitoring runtime information, such as the token index, the token-specific prediction confidence and / or the current state of the model (KV cache), for further optimization. Key components of the scheduler are i) the set of available precisions and ii) the precisionswitching policy that determines when the LLM should utilize lower bitwidth without severely impacting the output quality. The precision-switching policy can be either static, i.e. switch precisions at predefined token indices, or dynamic, i.e. tune the switching points based on the runtime information.

[139] A comprehensive evaluation across diverse LLM tasks shows that the proposed method achieves significant gains in speedup over conventional LLM inference methods and over state-of-the-art single-precision quantisation approaches for LLMs.

[140] Related Work

[141] LLM Quantisation

[142] As noted above, the prefilling stage of LLMs is generally compute-bound, while the decoding stage tends to be memory-bound [Patel et al, 2024], To enhance hardware performance, particularly due to the memory-bound nature of the decoding stage, applying aggressive quantisation techniques during this phase proves to be more effective. However, prior quantisation approaches have two main limitations. Existing quantisation approaches [Kim et al.(2024), Frantar et al. (2023), Tan et al. (2024)] typically apply a uniform quantisation policy and precision to both the prefilling and decoding stages. These approaches overlook the varying arithmetic intensity and redundancy characteristics that occur at different stages throughout the LLM inference process.

[143] Additionally, previous quantisation methods experience significant accuracy degradation in smaller LLMs ( / .e. models with fewerthan 7 billion parameters). While state-of-the-art quantisation techniques deliver near-lossless performance on larger models, they continue to struggle with performance drops in smaller models, which are critical for mobile and edge deployment. Instead, the proposed PMPD approach addresses these limitations by means of a phase-aware and adaptive precision-switching method, which effectively improves the decoding throughput of quantised mobile-friendly LLMs.

[144] Orthogonally to the above, any-precision LLM [Park et al.(2024)] enables the posttraining quantisation of a single set of weights, from which multiple variably quantised weights can be extracted. As demonstrated in the experiments detailed above, quantisation approach of any-precision LLM can be used in combination with PMPD to enable an efficient deployment with single-copy of the model.

[145] Precision-Adaptive Approaches

[146] Besides uniform quantisation methods that adopt a single precision throughout the computation, another line of work has proposed methods that dynamically adapt arithmetic at run time.

[147] Precision-Adaptive Training. MuPPET [Rajagopal et al.(2020)], CPT [Fu et al.(2021)] and AdaPT [Kummer et al.(2023)] aim to improve the efficiency of the training stage by dynamically adjusting the arithmetic precision throughout the training process. This family of methods focuses primarily on classification tasks. In contrast, the proposed PMPD method is optimized for the characteristics of LLMs and focuses on improving the efficiency of the inference stage.

[148] Mixed-Precision Inference. HAQ [Wang et al.(2019)] and HAWQ [Dong et al.(2019), Dong et al.(2020)] allow for different precision across the layers of a given DNN. Nonetheless, the selected precision remains fixed across all processed input samples. Bit-mixer [Bulat Tzimiropoulos(2021)] allows a similar utilization of mixed precision across layers, but can adjust the selected bitwidths in a per-sample manner.

[149] All the above works either focus on a different setting, such as the training stage, or are not applicable for the iterative, autoregressive decoding process of LLM inference.

[150] In summary, prior quantisation approaches have two main limitations: The prefilling stage of LLMs is generally compute-bound, while the decoding stage tends to be memorybound [4], To enhance hardware performance, particularly due to the memory-bound nature of the decoding stage, applying aggressive quantisation techniques during this phase proves to be more effective. However, existing quantisation approaches [Kim et al.(2024), Frantar et al.(2023), Tan et al.(2024)] typically apply a uniform quantisation policy and precision to both the prefilling and decoding stages. These approaches overlook the varying redundancy characteristics that occur at different stages throughout the LLM inference process. Previous quantisation methods also experience significant accuracy degradation in smaller LLMs (i.e. models with fewer than 7 billion parameters). While state-of-the-art quantisation techniques deliver near-lossless performance on larger models, they continue to struggle with performance drops in smaller models, which are critical for mobile and edge deployment. The techniques described above address these limitations by means of a phase-aware and adaptive precision-switching method, which effectively improves the decoding throughput of quantised mobile-friendly LLMs.

[001] Any-precision LLM [Park et al. (2024)] enables the post-training quantisation of a single set of weights, from which multiple variably quantised weights can be extracted. The quantisation approach of Any-precision LLM can be used in combination with the present techniques to enable an efficient deployment with single-copy of the model.

[002] Prior precision-adaptive approaches have focused on different settings. MuPPET [Rajagopal et al.(2020)], CPT [Fu et al.(2021)] and AdaPT [Kummer et al.(2023)] aim to improve the efficiency of the training stage by dynamically adjusting the arithmetic precision throughout the training process. Both methods are primarily applicable to convolutional neural networks (CNNs). In contrast, the techniques described above target LLMs and focuses on improving the efficiency of the inference stage. HAQ [Wang et al.(2019)] and HAWQ [Dong et al.(2020)] allow for different precision across layers, which remains fixed across all processed input samples. Bit-mixer [Bulat Tzimiropoulos (2021)] and any-precision DNNs [Yu et al, 2021] allow the utilization of mixed precision across layers, which can be adjusted in a per sample manner. CascadeCNN [Kouris et al (2018)] adopts different precision variants of the same DNN, organized in a low-to-high precision cascade of classifiers. Samples that fail to meet hand-tuned confidence-based criteria on the low-precision prediction stage are propagated for re-computation with higher precision. In contrast, the proposed techniques described above adapt the precision internally throughout the decoding stage of each input prompt (equivalent to input sample in the case of LLMs). All the above works either focus on a different setting (training stage) or are not applicable for the iterative, autoregressive decoding process of LLM inference.

[003] Those skilled in the art will appreciate that while the foregoing has described what is considered to be the best mode and where appropriate other modes of performing present techniques, the present techniques should not be limited to the specific configurations and methods disclosed in this description of the preferred embodiment. Those skilled in the art will recognise that present techniques have a broad range of applications, and that the embodiments may take a wide range of modifications without departing from any inventive concept as defined in the appended claims. References [Aishwarya et al.(2024)] PS Aishwarya, PranavAjit Nair, Yashas Samaga BL, Toby James Boyd, Sanjiv Kumar, Prateek Jain, and Praneeth Netrapalli. Tandem transformers for inference efficient Urns. In Forty-first International Conference on Machine Learning, 2024. [Bulat Tzimiropoulos (2021)] Adrian Bulat and Georgios Tzimiropoulos. “Bit-mixer: Mixed-precision networks with runtime bit-width selection.” In International Conference on Computer Vision (ICCV), 2021. [Cettolo et al, 2017] Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stuker, Katsuhito Sudoh, Koichiro Yoshino, and Christian Federmann. “Overview of the IWSLT 2017 evaluation campaign.” In Proceedings of the 14th International Conference on Spoken Language Translation, pp. 2-14, Tokyo, Japan, December 14-15 2017. International Workshop on Spoken Language Translation. [Chen et al, 2021] Yulong Chen, Yang Liu, Liang Chen, and Yue Zhang. “DialogSum: A real-life scenario dialogue summarization dataset.” In Findings of the Association for Computational Linguistics: ACLIJCNLP 2021, 2021. [Chiang et al, 2023] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric R Xing. “Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90%* ChatGPT Quality”, March 2023. URL https: / / lmsys.org / blog / 2023-03-30-vicuna / . [Chu et al, 2023] Xiangxiang Chu, Limeng Qiao, Xinyang Lin, Shuang Xu, Yang Yang, Yiming Hu, Fei Wei, Xinyu Zhang, Bo Zhang, Xiaolin Wei, and Chunhua Shen. “MobileVLM : A Fast, Strong and Open Vision Language Assistant for Mobile Devices.” arXiv preprint arXiv:2312.16886, 2023. [Dong et al.(2019)] Zhen Dong, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. “HAWQ: Hessian Aware Quantisation of Neural Networks with Mixed-Precision.” In Proceedings of the IEEE / CVF International Conference on Computer Vision (ICCV), pp. 293-302, 2019. [Dong et al.(2020)] Zhen Dong, Zhewei Yao, Yaohui Cai, Daiyaan Arfeen, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. “HAWQ-v2: Hessian Aware Trace-Weighted Quantisation of Neural Networks.” In Advances in Neural Information Processing Systems (NeurlPS), 2020. [Frantar et al. (2023)] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. “GPTQ: Accurate post-training compression for generative pretrained transformers.” In International Conference on Learning Representations, 2023. [Fu et al.(2021)] Yonggan Fu, Han Guo, Meng Li, Xin Yang, Yining Ding, Vikas Chandra, and Yingyan Lin. “CPT: Efficient Deep Neural Network Training via Cyclic Precision.” In International Conference on Learning Representations (ICLR), 2021. [Hermann et al, 2015] Karl Moritz Hermann, Tomas Kocisk, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. “Teaching Machines to Read and Comprehend.” In Advances on Neural Information Processing Systems, 2015. [Kim et al.(2024)] Sehoon Kim, Coleman Hooper, Amir Gholami, Zhen Dong, Xiuyu Li, Sheng Shen, Michael Mahoney, and Kurt Keutzer. “SqueezeLLM: Dense-and-Sparse Quantisation.” In International Conference on Machine Learning (ICML), 2024. [Kouris et al (2018)] A. Kouris, S.l. Venieris, and C.S. Bouganis, “CascadeCNN: Pushing the Performance Limits of Quantisation in Convolutional Neural Networks,” FPL, 2018 [Kummer et al. (2023)] Lorenz Kummer, Kevin Sidak, Tabea Reichmann, and Wilfried Gansterer. “Adaptive Precision Training (AdaPT): A dynamic quantised training approach for DNNs.” In SIAM International Conference on Data Mining (SDM), 2023. [Park et al. (2024)] Yeonhong Park, Jake Hyun, SangLyul Cho, Bonggeun Sim, and Jae W. Lee. “Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs.” In International Conference on Machine Learning (ICML), 2024. [Paszke et al. (2019)] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” In Advances on Neural Information Processing Systems, 2019. [Patel et al (2024)] R Patel et al., “Splitwise: Efficient Generative LLM Inference Using Phase Splitting,” ACM / IEEE 51st Annual International Symposium on Computer Architecture (ISCA), 2024. [Rajagopal et al.(2020)] Aditya Rajagopal, Diederik Vink, Stylianos Venieris, and Christos-Savvas Bouganis. “Multi-Precision Policy Enforced Training (MuPPET): A Precision-Switching strategy for Quantised Fixed-Point Training of CNNs.” In International Conference on Machine Learning (ICML), 2020. [Schuster et al. (2022)] Tai Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler, “Confident Adaptive Language Modeling,” NeurlPS, 2022. [Tan et al. (2024)] Fuwen Tan, Royson Lee, Lukasz Dudziak, Shell Xu Hu, Sourav Bhattacharya, Timothy Hospedales, Georgios Tzimiropoulos, and Brais Martinezs. “MobileQuant: Mobile-friendly Quantisation for On-device Language Models.” In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2024. [Wang et al. (2019)] Kuan Wang, Zhijian Liu, Yujun Lin, Ji Lin, and Song Han. “HAQ: Hardware-Aware Automated Quantisation with Mixed Precision.” In Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8612-8620, 2019. 5 [Xiao et al.(2023) Xiao, Tian, Chen, Han, and Lewis] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. “Efficient streaming language models with attention sinks.” arXiv preprint arXiv:2309.17453, 2023. 10 [Yu et al. (2021)] Haichao Yu, et al. “Any-Precision Deep Neural Networks,” AAAI, 2021. [Zeng et al, 2024] Shulin Zeng, Jun Liu, Guohao Dai, Xinhao Yang, Tianyu Fu, Hongyi Wang, Wenheng Ma, Hanbo Sun, Shiyao Li, Zixiao Huang, et al. “FlightLLM: Efficient Large Language Model Inference with a Complete Mapping Flow on FPGAs.” In ACM / SIGDA 15 International Symposium on Field Programmable Gate Arrays (ISFPGA), 2024.

Claims

1. A computer-implemented method for optimising an autoregressive machine learning, ML, model before deployment to a user device, wherein the autoregressive ML model comprises a plurality of model weights which are arranged in a prefilling phase in which an input to the model is pre-processed and a decoding phase in which the pre-processed input is processed through multiple iterations to generate an output, the method comprising:receiving performance metrics for the ML model, the performance metrics comprising at least one of an output metric for the output from the ML model and a set of weight quantisation bitwidths which are supported by the user device;selecting, based on the received performance metrics, a prefilling quantisation bitwidth for the prefilling phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the prefilling phase is converted into an integer having the selected prefilling quantisation bitwidth;selecting, based on the received performance metrics, at least one decoding quantisation bitwidth for the decoding phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the decoding phase is converted into an integer having the selected at least one decoding quantisation bitwidth; andoutputting, to a user device, the ML model with a quantisation schedule which specifies the selected prefilling quantisation and the at least one decoding quantisation bitwidth whereby the ML model is optimised when used on the user device;wherein the selected prefilling quantisation bitwidth is greater than or equal to the at least one decoding quantisation bitwidths.

2. The method of claim 1, wherein selecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth comprisesselecting, from the set of weight quantisation bitwidths, multiple pairs of bitwidths with each pair comprising a prefilling quantisation bitwidth and a decoding quantisation bitwidth;generating multiple variant ML models, wherein each variant ML model is generated by quantising the plurality of weights of the ML model using one of the selected pairs of bitwidths;evaluating, using a calibration dataset, the generated set of variant ML models;identifying the variant ML models which meet the output metric; andselecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth from the identified variant ML model having lowest values for the prefilling quantisation bitwidth and decoding quantisation bitwidth.

3. The method of claim 1 or claim 2, wherein selecting at least one decoding quantisation bitwidth comprises:selecting a first decoding quantisation bitwidth for a first sequence of the decoding phase; andselecting a second decoding quantisation bitwidth for a second sequence of the decoding phase;wherein the second decoding quantisation bitwidth is lower than the first decoding quantisation bitwidth.

4. The method of claim 3, further comprising:setting a switching point which defines an end of the first sequence whereby when the ML model is being used on the user device, the switching point defines when to switch from the first decoding quantisation bitwidth to the second decoding bitwidth.

5. The method of any one of the preceding claims, further comprising:setting a plurality of switching points, wherein each switching point defines an end of a sequence of the decoding phase, whereby the decoding phase comprises a plurality of time-ordered sequences andwherein selecting at least one decoding quantisation bitwidth comprises selecting a different decoding quantisation bitwidth for each one of the time-ordered sequences to obtain a plurality of decoding quantisation bitwidths which decrease in size in time-order across the decoding phase.

6. The method of claim 4 or claim 5, further comprising:setting the or each switching point and selecting the at least one decoding quantisation bitwidth prior to outputting the quantisation schedule.

7. The method of claim 6, wherein outputting the ML model with the quantisation schedule comprises outputting a quantisation schedule which comprises a static decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point.

8. The method of claim 6 or claim 7, further comprising setting the or each switching point and selecting the at least one decoding quantisation bitwidth by:selecting an initial decoding quantisation bitwidth from the set of weight quantisation bitwidths;selecting at least one switching point from a set of possible switching points wherein each switching point defines an end of a sequence of the decoding phase;selecting, from the set of weight quantisation bitwidths, an associated decoding quantisation bitwidth for each obtained at least one switching point;quantising the plurality of weights of the ML model using the selected initial decoding quantisation bitwidth, the selected at least one switching point and the selected subsequent decoding quantisation bitwidths to generate a variant ML model;repeating the selecting and quantising steps to generate a set of variant ML models each of which have different combinations of initial decoding quantisation bitwidth, at least one switching point and subsequent decoding quantisation bitwidths;evaluating, using a calibration dataset, the generated set of variant ML models; andselecting the at least one switching point and obtained initial and subsequent decoding quantisation bitwidths from the identified variant ML model having a highest performance.

9. The method of claim 4 or claim 5, further comprising:outputting a dynamic scheduler when outputting the ML model with the quantisation schedule, wherein, when using the ML model on the user device, the dynamic scheduler determines, for each input to the ML model, a decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point.

10. The method of claim 9, further comprising training the dynamic scheduler using a training dataset prior to outputting the dynamic scheduler, wherein the training dataset comprises a plurality of prefilling stage outputs each of which is labelled with a decoding quantisation schedule.

11. The method of claim 10, further comprising creating the training dataset by:receiving a plurality of inputs for the ML model,for each received input:generating a prefilling phase output by running the input through the prefill phase of the ML model;obtaining a set of decoding quantisation schedules, wherein each decoding quantisation schedule comprises at least one switching point and an associated decoding quantisation bitwidth for each at least one switching point;quantising the plurality of weights of the ML model using each decoding quantisation schedule in the set of decoding quantisation schedules to generate a set of variant ML models;evaluating, using the prefilling phase output, the generated set of variant ML models;selecting the decoding quantisation schedule from the identified variant ML model having a highest performance;labelling the prefilling phase output with the selected decoding quantisation schedule; andoutputting the prefilling phase outputs and associated labels as the training dataset.

12. The method of any one of claims 1 to 11, further comprising outputting the ML model as a plurality of separate models each having separate model weights and each having a different quantisation which is defined by the quantisation schedule.

13. The method of any one of claims 1 to 11, further comprising outputting the ML model as a unified set of model weights.

14. A method of using an autoregressive machine learning, ML, model on a user device, wherein the method comprisesreceiving, at the user device, the ML model and the quantisation schedule as set out in any one of claims 1 to 8;receiving an input to be processed through the ML model;quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule;quantising the decoding phase of the ML model using the at least one decoding quantisation bitwidth defined in the quantisation schedule; and processing, using the quantised prefill and decoding phases, the input to generate an output from the ML model.

15. A method of using an autoregressive machine learning, ML, model on a user device, wherein the method comprisesreceiving, at the user device, the ML model, the quantisation schedule and the dynamic scheduler as set out in any one of claims 9 to 11;receiving an input to be processed through the ML model;quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule;generating a prefilling phase output by running the input through the quantised prefill phase;inputting the prefilling phase output to the dynamic scheduler to generate a decoding quantisation schedule;quantising the decoding phase of the ML model using the decoding quantisation schedule;processing, using the quantised decoding phase, the prefilling phase output to generate an output from the ML model.

16. The method of any one of the preceding claims, wherein the autoregressive machine learning, ML, model is a large language model.

17. A computer-readable storage medium comprising instructions which, when executed by at least one processor, causes the at least one processor to individually or collectively carry out the method of claims 1 to 16.

18. An electronic device for optimising an autoregressive machine learning, ML, model, the electronic device comprising:at least one processor coupled to memory, wherein the at least one processor is configured to individually or collectively carry out the method of claims 1 to 13 and 16.

19. A system comprisingthe electronic device of claim 18, anda user device comprising:at least one processor coupled to memory, wherein the at least one processor is for:receiving, from the electronic device, the optimised ML model and the quantisation schedule;quantising the ML model according to the quantisation schedule and processing an input using the quantised ML model to generate an output.

20. The system of claim 19, wherein the user device is a resource-constrained device.AMENDMENTS TO THE CLAIMS HAVE BEEN FILED AS FOLLOWS:-CLAIMS1. A computer-implemented method for optimising an autoregressive machine learning, ML, model before deployment to a user device, wherein the autoregressive ML model comprises a plurality of model weights which are arranged in a prefilling phase in which an input to the model is pre-processed and a decoding phase in which the pre-processed input is processed through multiple iterations to generate an output, the method comprising:receiving performance metrics for the ML model, the performance metrics comprising at least one of an output metric for the output from the ML model and a set of weight quantisation bitwidths which are supported by the user device;selecting, based on the received performance metrics, a prefilling quantisation bitwidth for the prefilling phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the prefilling phase is converted into an integer having the selected prefilling quantisation bitwidth;selecting, based on the received performance metrics, at least one decoding quantisation bitwidth for the decoding phase of the ML model, wherein when the ML model is being used on the user device, each of the model weights in the decoding phase is converted into an integer having the selected at least one decoding quantisation bitwidth; andoutputting, to a user device, the ML model with a quantisation schedule which specifies the selected prefilling quantisation and the at least one decoding quantisation bitwidth whereby by using the quantisation schedule, the ML model has been optimised when used on the user device;wherein the selected prefilling quantisation bitwidth is greater than or equal to the at least one decoding quantisation bitwidths.

2. The method of claim 1, wherein selecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth comprisesselecting, from the set of weight quantisation bitwidths, multiple pairs of bitwidths with each pair comprising a prefilling quantisation bitwidth and a decoding quantisation bitwidth;generating multiple variant ML models, wherein each variant ML model is generated by quantising the plurality of weights of the ML model using one of the selected pairs of bitwidths;evaluating, using a calibration dataset, the generated set of variant ML models;identifying the variant ML models which meet the output metric; andselecting the prefilling quantisation bitwidth and the at least one decoding quantisation bitwidth from the identified variant ML model having lowest values for the prefilling quantisation bitwidth and decoding quantisation bitwidth.

3. The method of claim 1 or claim 2, wherein selecting at least one decoding quantisation bitwidth comprises:selecting a first decoding quantisation bitwidth for a first sequence of the decoding phase; andselecting a second decoding quantisation bitwidth for a second sequence of the decoding phase;wherein the second decoding quantisation bitwidth is lower than the first decoding quantisation bitwidth.

4. The method of claim 3, further comprising:setting a switching point which defines an end of the first sequence whereby when the ML model is being used on the user device, the switching point defines when to switch from the first decoding quantisation bitwidth to the second decoding bitwidth.

5. The method of any one of the preceding claims, further comprising:setting a plurality of switching points, wherein each switching point defines an end of a sequence of the decoding phase, whereby the decoding phase comprises a plurality of time-ordered sequences andwherein selecting at least one decoding quantisation bitwidth comprises selecting a different decoding quantisation bitwidth for each one of the time-ordered sequences to obtain a plurality of decoding quantisation bitwidths which decrease in size in time-order across the decoding phase.

6. The method of claim 4 or claim 5, further comprising:setting the or each switching point and selecting the at least one decoding quantisation bitwidth prior to outputting the quantisation schedule.

7. The method of claim 6, wherein outputting the ML model with the quantisation schedule comprises outputting a quantisation schedule which comprises a static decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point.

8. The method of claim 6 or claim 7, further comprising setting the or each switching point and selecting the at least one decoding quantisation bitwidth by:selecting an initial decoding quantisation bitwidth from the set of weight quantisation bitwidths;selecting at least one switching point from a set of possible switching points wherein each switching point defines an end of a sequence of the decoding phase;selecting, from the set of weight quantisation bitwidths, an associated decoding quantisation bitwidth for each obtained at least one switching point;quantising the plurality of weights of the ML model using the selected initial decoding quantisation bitwidth, the selected at least one switching point and the selected subsequent decoding quantisation bitwidths to generate a variant ML model;repeating the selecting and quantising steps to generate a set of variant ML models each of which have different combinations of initial decoding quantisation bitwidth, at least one switching point and subsequent decoding quantisation bitwidths;evaluating, using a calibration dataset, the generated set of variant ML models; andselecting the at least one switching point and obtained initial and subsequent decoding quantisation bitwidths from the identified variant ML model having a highest performance.

9. The method of claim 4 or claim 5, further comprising:outputting a dynamic scheduler when outputting the ML model with the quantisation schedule, wherein, when using the ML model on the user device, the dynamic scheduler determines, for each input to the ML model, a decoding quantisation schedule specifying the or each switching point together with the decoding quantisation bitwidths which are to be applied before and after the or each switching point.

10. The method of claim 9, further comprising training the dynamic scheduler using a training dataset prior to outputting the dynamic scheduler, wherein the training dataset comprises a plurality of prefilling stage outputs each of which is labelled with a decoding quantisation schedule.

11. The method of claim 10, further comprising creating the training dataset by:receiving a plurality of inputs for the ML model,for each received input:generating a prefilling phase output by running the input through the prefill phase of the ML model;obtaining a set of decoding quantisation schedules, wherein each decoding quantisation schedule comprises at least one switching point and an associated decoding quantisation bitwidth for each at least one switching point;quantising the plurality of weights of the ML model using each decoding quantisation schedule in the set of decoding quantisation schedules to generate a set of variant ML models;evaluating, using the prefilling phase output, the generated set of variant ML models;selecting the decoding quantisation schedule from the identified variant ML model having a highest performance;labelling the prefilling phase output with the selected decoding quantisation schedule; andoutputting the prefilling phase outputs and associated labels as the training dataset.

12. The method of any one of claims 1 to 11, further comprising outputting the ML model as a plurality of separate models each having separate model weights and each having a different quantisation which is defined by the quantisation schedule.

13. The method of any one of claims 1 to 11, further comprising outputting the ML model as a unified set of model weights.

14. A method of using an autoregressive machine learning, ML, model on a user device, wherein the method comprisesreceiving, at the user device, the ML model and the quantisation schedule as set out in any one of claims 1 to 8;receiving an input to be processed through the ML model;quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule;quantising the decoding phase of the ML model using the at least one decoding quantisation bitwidth defined in the quantisation schedule; and processing, using the quantised prefill and decoding phases, the input to generate an output from the ML model.

15. A method of using an autoregressive machine learning, ML, model on a user device, wherein the method comprisesreceiving, at the user device, the ML model, the quantisation schedule and the dynamic scheduler as set out in any one of claims 9 to 11;receiving an input to be processed through the ML model;quantising the prefill phase of the ML model using the prefill quantisation bitwidth defined in the quantisation schedule;generating a prefilling phase output by running the input through the quantised prefill phase;inputting the prefilling phase output to the dynamic scheduler to generate a decoding quantisation schedule;quantising the decoding phase of the ML model using the decoding quantisation schedule;processing, using the quantised decoding phase, the prefilling phase output to generate an output from the ML model.

16. The method of any one of the preceding claims, wherein the autoregressive machine learning, ML, model is a large language model.

17. A computer-readable storage medium comprising instructions which, when executed by at least one processor, causes the at least one processor to individually or collectively carry out the method of claims 1 to 16.

18. An electronic device for optimising an autoregressive machine learning, ML, model, the electronic device comprising:at least one processor coupled to memory, wherein the at least one processor is configured to individually or collectively carry out the method of claims 1 to 13 and 16.

19. A system comprisingthe electronic device of claim 18, anda user device comprising:at least one processor coupled to memory, wherein the at least one processor is for:receiving, from the electronic device, the optimised ML model and the quantisation schedule;quantising the ML model according to the quantisation schedule and processing an input using the quantised ML model to generate an output.

20. The system of claim 19, wherein the user device is a resource-constrained device.