Methods and systems for tensor memory management and for serving large language models

Abstract

There is provided a method, system and non-transitory storage medium for serving a large language model (LLM) engine by allocating memory on demand. A virtual memory space with virtual pages is initialized for a tensor. A first request with first input tokens is received, added to the tensor by being allocated to first virtual pages and mapped to first page frames in physical memory. The tensor is provided to the LLM engine for processing. A second request with second input tokens is received, and a first output token is received from the LLM engine. The first output token is added to the tensor by being allocated to a second virtual page contiguous to the first virtual pages and mapped to a second page frame. The second input tokens are allocated to second virtual pages and mapped to second page frames in physical memory. The tensor is provided for further processing.

Description

METHODS AND SYSTEMS FOR TENSOR MEMORY MANAGEMENT AND FOR SERVING LARGE LANGUAGE MODELSTECHNICAL FIELD

[0001] The present technology relates to artificial intelligence in general and more specifically to methods, systems, and non-transitory storage mediums for tensor memory management and for serving large language models (LLMs) .BACKGROUND

[0002] The advent of large language models (LLMs) has significantly advanced the field of artificial intelligence (AI) and natural language processing (NLP) . These models, characterized by their extensive training on diverse and voluminous datasets, have demonstrated remarkable capabilities in generating human-like text, understanding context, and performing complex language tasks. The deployment of such models in various applications, including but not limited to chatbots, content generation, translation services, and information retrieval systems, has necessitated the development of efficient and effective methods for serving these models to end-users.

[0003] Traditionally, serving language models involved the use of relatively smaller models that could be deployed on individual servers or clusters with moderate computational resources. However, the exponential increase in the size and complexity of modern LLMs has introduced several challenges in their deployment and serving. These challenges include the need for substantial computational power, significant memory resources, and efficient data handling mechanisms to ensure low latency and high throughput during inference.

[0004] Running these applications incurs significant costs due to the requirement for numerous hardware accelerators, such as GPUs. Estimates suggest that processing a large language model (LLM) request can be up to ten times more expensive than a traditional keyword query. Due to these high costs, enhancing throughput-and thereby reducing the cost per request-of LLM serving systems is becoming increasingly crucial.

[0005] Large language models (LLMs) typically rely on an autoregressive Transformer model. This model generates tokens sequentially, one at a time, based on the input (prompt) and the previously generated tokens. For each request, this process is repeated until a termination token is produced. This sequential generation makes the workload memory-bound, leading to underutilization of GPU computation power and limiting the serving throughput.

[0006] Throughput can be improved by batching multiple requests together. However, efficiently managing the memory space for each request in a batch is essential.

[0007] Around 30%of the memory is used to store the dynamic states of the requests. For Transformers, these states consist of the key and value tensors associated with the attention mechanism, commonly referred to as the key-value (KV) cache.

[0008] Efficiently managing the KV Cache memory can improve the throughput of LLM serving system, thus increasing the device utilization and reducing the cost.

[0009] Another challenge of serving LLMs is related to the hardware and frameworks used. For example, for Huawei’s Ascend NPU, the Deep learning framework is MindSpore, which mainly runs in static graph mode, thus requiring its input to be static-shaped. While in case of LLM, whose input (sentences) usually varies in length, the input will be padded to a predefined maximum length. The same applies to the KV Cache, which results in memory waste due to storage of the paddings.SUMMARY

[0010] Developers of the present technology have appreciated that there is a need for more efficient key-value (KV) cache and low-rank adaptation (LoRA) management methods, to allocate high bandwidth memory (HBM) only when used by processing units, which would avoid memory waste caused by padded tensors, and avoid inefficiencies caused by sparse tensor memory read in the paged attention algorithms in state-of-the-art LLM frameworks such as vLLM.

[0011] Developers of the present technology have recognized that there is a need to create a new library that manages tensor memory by application programming interfaces (APIs) other than PyTorch or MindSpore. Such a library should support efficient  scaling (i.e., scale up and scale down) , save memory for paddings, and where the tensor is stored in contiguous memory addresses from the point of view of PyTorch or MindSpore frameworks, on which operators can work easily.

[0012] Thus, one or more embodiments of the present disclosure are directed to methods and systems for tensor memory management and for serving large language models.

[0013] In accordance with a broad aspect of the present technology, there is provided a method for serving a large language model (LLM) engine by allocating memory on demand, the method being executed by at least one processing unit, the at least one processing unit being operatively connected to a physical memory comprising a plurality of page frames, the method comprising: initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages, receiving a first request comprising a first sequence of input tokens, adding the first sequence of input tokens to the tensor, said adding comprising: allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages, and mapping the first set of virtual pages to a first set of page frames in the physical memory, providing the tensor to the LLM engine for processing, receiving a second request comprising a second sequence of input tokens, receiving, from the LLM engine, a first output token for the first sequence of tokens, adding the first output token to the tensor, said adding comprising: allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory, and mapping the second page to a second page frame in the physical memory, adding the second sequence of input tokens to the tensor, said adding comprising: allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory, and mapping the second set of virtual pages to a second set of page frames in the physical memory, and providing the tensor to the LLM engine for further processing.

[0014] In one or more embodiments of the method, said providing the tensor to the LLM engine for processing comprises: accessing decoding layers of the LLM engine, generating, via the decoding layers, based on the first sequence of input tokens, a first set of query vectors, a first set of key vectors, and a first set of value vectors, adding the first set of query vectors, the first set of key vectors and the first set of value vectors to the tensor, said adding comprising: allocating the first set of query vectors, the first set of key vectors and the first set of value vectors to a third set of virtual pages, mapping the third set of virtual pages to a third set of page frames, and generating the first output token based on the first set of query vectors, the first set of key vectors and the first set of value vectors to receive the first output token.

[0015] In one or more embodiments of the method, the second set of virtual pages are contiguous to the second virtual page in the virtual memory.

[0016] In one or more embodiments of the method, at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.

[0017] In one or more embodiments of the method, said adding the first sequence of input tokens to the tensor comprises adding a set of padded tokens, said allocating the first sequence of input tokens in the tensor to the first set of virtual pages comprises allocating the set of padded tokens in the first set of virtual pages, and said mapping the first set of virtual pages to the first set of page frames in the physical memory comprises mapping the set of padded tokens to a single page frame in the first set of page frames.

[0018] In one or more embodiments of the method, said mapping is performed using the mmap API.

[0019] In one or more embodiments of the method, the method further comprises: upon reaching a stopping criterion for at least one of the first request and the second request: unmapping page frames associated with the at least one of the first request and the second request from the physical memory, and deallocating virtual pages associated with the at least one of the first request and the second request from the virtual memory.

[0020] In one or more embodiments of the method, the physical memory comprises a high-bandwidth memory (HBM) of a parallel processing unit (PPU) .

[0021] In one or more embodiments of the method, the PPU comprises one of: a graphics processing unit (GPU) , a neural processing unit (NPU) , and a tensor processing unit (TPU) .

[0022] In accordance with a broad aspect of the present technology, there is provided a system for serving a large language model (LLM) engine by allocating physical memory on demand, the system comprising: at least one processor, a non-transitory storage medium operatively connected to the at least one processor, the non-transitory storage medium storing computer-readable instructions thereon, and at least one parallel processing unit (PPU) connected to the non-transitory storage medium, the PPU comprising: a plurality of processing cores and a PPU memory, the PPU memory comprising a plurality of page frames. The at least one processor, upon executing the computer-readable instructions, is configured for: initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages, receiving a first request comprising a first sequence of input tokens, adding the first sequence of input tokens to the tensor, said adding comprising: allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages, and mapping the first set of virtual pages to a first set of page frames in the plurality of page frames, providing the tensor to the LLM engine for processing using the at least one PPU, receiving a second request comprising a second sequence of input tokens, receiving, from the LLM engine, a first output token for the first sequence of tokens, adding the first output token to the tensor, said adding comprising: allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory, and mapping the second page to a second page frame in the physical memory, adding the second sequence of input tokens to the tensor, said adding comprising: allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory, and mapping the second set of virtual pages to a second set of page frames in the physical memory, and providing the tensor to the LLM engine for further processing.

[0023] In one or more embodiments of the system, the at least one PPU comprises one of: a graphics processing unit (GPU) , a neural processing unit (NPU) , and a tensor processing unit (TPU) .

[0024] In one or more embodiments of the system, at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.

[0025] In one or more embodiments of the system, said providing the tensor to the LLM engine for processing comprises: accessing decoding layers of the LLM engine, generating, via the decoding layers, based on the first sequence of input tokens, a first set of query vectors, a first set of key vectors, and a first set of value vectors, adding the first set of query vectors, the first set of key vectors and the first set of value vectors to the tensor, said adding comprising: allocating the first set of query vectors, the first set of key vectors and the first set of value vectors to a third set of virtual pages, mapping the third set of virtual pages to a third set of page frames, and generating the first output token based on the first set of query vectors, the first set of key vectors and the first set of value vectors to receive the first output token.

[0026] In one or more embodiments of the system, the second set of virtual pages are contiguous to the second virtual page in the virtual memory.

[0027] In one or more embodiments of the system, at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.

[0028] In one or more embodiments of the system, said adding the first sequence of input tokens to the tensor comprises adding a set of padded tokens, said allocating the first sequence of input tokens in the tensor to the first set of virtual pages comprises allocating the set of padded tokens in the first set of virtual pages, and said mapping the first set of virtual pages to the first set of page frames in the physical memory comprises mapping the set of padded tokens to a single page frame in the first set of page frames.

[0029] In one or more embodiments of the system, said mapping is performed using the mmap API.

[0030] In one or more embodiments of the system, the at least one processor is further configured for: upon reaching a stopping criterion for at least one of the first request and the second request: unmapping page frames associated with the at least one of the  first request and the second request from the physical memory, and deallocating virtual pages associated with the at least one of the first request and the second request from the virtual memory.

[0031] In one or more embodiments of the system, the at least one PPU comprises a plurality of PPUs, and the system further comprises a scheduler configured for receiving available memory from each respective PPU memory from the plurality of PPUs and scheduling processing of the first request and the second request by respective ones of the plurality of PPUs.

[0032] In accordance with a broad aspect of the present technology, there is provided a non-transitory storage medium comprising computer-readable instructions, when executed by at least one processing operatively connected to a physical memory comprising a plurality of page frames, cause execution of a method comprising: initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages, receiving a first request comprising a first sequence of input tokens, adding the first sequence of input tokens to the tensor, said adding comprising: allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages, and mapping the first set of virtual pages to a first set of page frames in the physical memory, providing the tensor to the LLM engine for processing, receiving a second request comprising a second sequence of input tokens, receiving, from the LLM engine, a first output token for the first sequence of tokens, adding the first output token to the tensor, said adding comprising: allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory, and mapping the second page to a second page frame in the physical memory, adding the second sequence of input tokens to the tensor, said adding comprising: allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory, and mapping the second set of virtual pages to a second set of page frames in the physical memory, and providing the tensor to the LLM engine for further processing.BRIEF DESCRIPTION OF THE DRAWINGS

[0033] FIG. 1A illustrates a schematic diagram of a network environment in accordance with non-limiting embodiments of the present technology.

[0034] FIG. 1B illustrates a schematic diagram of an exemplary computing device used in the network environment of FIG. 1A in accordance with non-limiting embodiments of the present technology.

[0035] FIG. 2 illustrates a schematic diagram of exemplary components of an inference server and a set of parallel processing units (PPUs) used in the network environment of FIG. 1A in accordance with non-limiting embodiments of the present technology.

[0036] FIG. 3 illustrates a schematic diagram in accordance with non-limiting embodiments of the present technology.

[0037] FIG. 4 illustrates a schematic diagram of a scheduler and a plurality of workers in an LLM inference engine in accordance with non-limiting embodiments of the present technology.

[0038] FIG. 5 illustrates a non-limiting example of KV cache management of two requests at two different iterations from the point of view of the tensor virtual memory and from the point of view of the physical memory in accordance with non-limiting embodiments of the present technology.

[0039] FIG. 6 illustrates of padding of the tensor from the point of view of the virtual memory and from the point of view of the physical memory in accordance with non-limiting embodiments of the present technology.

[0040] FIG. 7 and FIG. 8 illustrate a flowchart of a method for serving a LLM engine by allocating memory on demand, the method being executed in accordance with non-limiting embodiments of the present technology.DETAILED DESCRIPTION

[0041] One or more embodiments of the present technology provide for memory-efficient tensor memory management methods and systems which can be implemented in a large language model (LLM) serving framework.

[0042] One or more embodiments of the present technology also provide for a novel LLM serving system architecture according to the tensor memory management methods and systems.

[0043] One or more embodiments of the present technology may be implemented for management of key-value (KV) cache and low-rank adaptation (LoRA) by allocating high-bandwidth memory (HBM) on demand. More specifically, one or more embodiments of the present technology use the concept of virtual memory in operating systems to map tensors comprising tokens and intermediate states computed by LLMs to physical memory such that from the point of view of the machine learning (ML) frameworks, related tokens and states are located in contiguous memory address spaces, which does not require redefining operators and using block tables.

[0044] One or more embodiments of the present technology support scale up and down, save memory for paddings tensors and use contiguous memory address from the point of view of the machine leaning frameworks, which can be easily managed by operators (e.g., attention operators) without requiring redesign of operator architecture and use of block tables.

[0045] FIG. 1A illustrates an exemplary network environment 100, in accordance with one or more examples in the present disclosure. Machine learning techniques implementing the framework disclosed herein may take place in the exemplary network environment 100. Network environments suitable for use in implementing embodiments of the disclosure may include a plurality of client devices 120, plurality of servers 130, and / or other device types.

[0046] Components of a network environment may communicate with each other via one or more network (s) 110, which may be wired, wireless, or both. By way of example, network 110 may include one or more Wide Area Networks (WANs) , one or more Local Area Networks (LANs) , one or more public networks such as the Internet, and / or one or more private networks. Where the network includes a wireless telecommunications network, components such as a base station, a communications tower, access points, or other components may provide wireless connectivity.

[0047] The plurality of client devices 120 comprises a first client device 122, a second client device 124 and a third client device 126. It will be appreciated that the plurality of client devices 120 may include a different number of client devices, such as dozens, hundreds or thousands of client devices.

[0048] The first client device 122, the second client device 124 and the third client device 126 may include at least some of the components, features, and functionality of the example computer system 150 of FIG. 1B. By way of example and not limitation, a given client device 122, 124, 126 of the plurality of client devices 120 may be embodied as a Personal Computer (PC) , a laptop computer, a mobile device, a smartphone, a tablet computer, a virtual reality headset, a video player, a video camera, a vehicle, a virtual machine, a drone, a robot, a handheld communications device, a vehicle computer system, an embedded system controller, a workstation, an edge device, any combination of these delineated devices, or any other suitable device.

[0049] FIG. 1B illustrates a block diagram of an exemplary computer system 150 configured to implement various functions according to one or more embodiments in the present disclosure. In some examples, computer system 150 may be implemented in the plurality of client devices 120 or the plurality of servers 130 in network environment 100 as shown in FIG. 1A. One or more computing systems 150, one or more client devices 120, one or more servers 130, or the combination thereof may form a processing system to perform the processes in the present disclosure.

[0050] As shown in FIG. 1B, computer system 150 may include one or more processor (s) 160, a communication interface 164, a memory 162, one or more storage unit (s) 166, and a display 190. Processor (s) 160 may be configured to perform the operations in accordance with the instructions stored in memory 162. Processor (s) 160 may include any appropriate type of general-purpose or special-purpose microprocessor (e.g., a CPU or GPU, respectively) , digital signal processor, microcontroller, or the like. Memory 162 may be configured to store computer-readable instructions that, when executed by processor (s) 160, can cause processor (s) 160 to perform various operations disclosed herein. Memory 162 may be any non-transitory type of mass storage, such as volatile or non-volatile, magnetic, semiconductor-based, tape-based, optical, removable, non-removable, or other type of storage device or tangible computer-readable medium including, but not limited to, a read-only memory (ROM) , a flash memory, a dynamic random-access memory (RAM) , and / or a static RAM. Storage unit (s) 166 may include hard disk drives (HDDs) , solid-sate drives (SSDs) configured for high-capacity storage.

[0051] Various processes / flowcharts described in the present disclosure may be realized in instructions stored in memory 162, when executed by processor (s) 160.

[0052] Communication interface 164 may be configured to communicate information between computer system 150 and other devices or systems, such as client devices 120 and / or the plurality of servers 130 as shown in FIG. 1A. For example, communication  interface 164 may include an integrated services digital network (ISDN) card, a cable modem, a satellite modem, or a modem to provide a data communication connection. As another example, communication interface 164 may include a local area network (LAN) card to provide a data communication connection to a compatible LAN. As a further example, communication interface 164 may include a high-speed network adapter such as a fiber optic network adaptor, 10G Ethernet adaptor, or the like. Wireless links can also be implemented by communication interface 164. In such an implementation, communication interface 164 can send and receive electrical, electromagnetic or optical signals that carry digital data streams representing various types of information via a network. The network can typically include a cellular communication network (e.g., 4G, 5G) , a Wireless Local Area Network (WLAN) , a Wide Area Network (WAN) , or the like.

[0053] Communication interface 164 may also include various I / O devices such as a keyboard, a mouse, a touchpad, a touch screen, a microphone, a camera, a biosensor, etc. A user may input data to computer system 150 (e.g., a terminal device) through communication interface 164.

[0054] Display 190 may be integrated as part of computer system 150 or may be provided as a separate device communicatively coupled to computer system 150. Display 190 may include a display device such as a liquid crystal display (LCD) , a light emitting diode display (LED) , an organic LED (OLED) , a plasma display, or any other type of display, and provide a graphical user interface (GUI) presented on the display for user input and data depiction. In some embodiments, display 190 may be integrated as part of communication interface 164.

[0055] Turning back to FIG. 1A, Each of the plurality of client devices 120 is configured to transmit respective requests to the plurality of servers 130 over the network 110 via gateway 115. Each respective request may include a request to perform a machine learning (ML) related task by the plurality of servers 130. A given request, also referred to as prompt, may include input data (e.g., text, image, video, audio or combination thereof) .

[0056] The gateway 115 is configured to act as an entry point for client requests from the plurality of client devices 120 by routing client requests to the appropriate backend services in the plurality of servers 330. The gateway 115 is configured to handle authentication, load balancing, and to ensure secure and efficient communication between client devices 120 and the ML models being served in the plurality of servers 130.

[0057] The plurality of servers 130 are configured to implement a ML serving engine that receives, distributes, and processes the requests for execution by one or more ML models (e.g., LLMs) to obtain respective outputs for transmission to the plurality of client devices 120.

[0058] In one or more embodiments, the plurality of servers 130 are configured to process, schedule, and encode each request into a sequence of input tokens, apply one or more ML models (e.g., LLMs) on the sequence of input tokens to generate a sequence of output tokens, and decode the sequence of output tokens to obtain an output which is transmitted back to the respective client devices of the plurality of client devices 130 via gateway 115.

[0059] In one or more embodiments, the ML models executed by the plurality of servers 130 are large language models (LLMs) which may have a transformer-based architecture (e.g. Pangu, GPT series, BERT, Llama, T5, RoBERTa, and XL net) or non-transformer architecture (e.g., RNN, CNN, Hybrid model and attention mechanism models) .

[0060] The plurality of servers 130 are configured to execute a deep learning framework (not shown) . Non-limiting examples of deep learning frameworks include PyTorch (META) , TensorFlow (GOOGLE) , and MindSpore (HUAWEI) .

[0061] The plurality of servers 130 comprises one or more inference servers 132 configured to execute ML models and generate predictions based on requests, one or more data servers 134 configured to manage, store, and provide access to the data required for model execution and inference, and one or more application servers 136 configured to handle application logic, manage user interactions, and coordinate communication between the inference servers 132 and data servers 134.

[0062] It will be appreciated that the plurality of servers 130 may implement a plurality of abstraction layers for deployment and operationalization of ML models such as LLMs. As a non-limiting example, the abstraction layers may include one or more of: hardware abstraction layer (s) , resource management layer (s) , system software layer (s) , framework layer (s) , middleware layer (s) , inference engine layer (s) , application layer (s) , monitoring and management layer (s) and security layer (s) .

[0063] In addition to FIG. 1A and FIG. 1B, reference is also made to FIG. 2, which illustrates hardware and software components of the one or more inference servers 132 and a parallel processing unit (PPU) cluster 138. The one or more inference servers 132 are operatively connected to the PPU cluster 138.

[0064] The one or more inference servers 132 are configured to execute an LLM inference engine 200 and a set of APIs 210.

[0065] In one or more embodiments, the one or more inference servers 132 are configured to process requests by batches to achieve higher processor utilization by exploiting the parallel computing units in the PPU cluster 138. The one or more inference servers 132 are configured to receive multiple requests each associated with an input token sequence. The one or more inference servers 132 iteratively apply the LLM (not shown) of the LLM inference engine 200 to the batch of requests to generate the output tokens for the requests together. In some embodiments, batching for the LLM model is made possible by grouping requests that have the same length of input token sequences together or at each iteration, padding shorter requests such that the shorter requests have the same length as the longest request in the batch.

[0066] For each request, the LLM inference engine 200 is configured to generate a sequence of output tokens in response to the sequence of input prompt tokens. The LLM inference engine 200 executes a prefill phase and a decode phase.

[0067] During the prefill phase, also referred to as prompt phase, the LLM inference engine 200 receives a request comprising a list of input prompt tokens and computes the probability of a first new token. The LLM inference engine 200 generates the key vectors and values vector (intermediate states) which are used to generate the first new token. A self-attention layer first applies linear transformations on each position of the input tokens to obtain the query, key and value vectors. The key vectors and value vectors are stored in a tensor in a cache memory of the PPU cluster 138, also referred to as KV cache.

[0068] Each new subsequent token depends on all the previous tokens. Since all input prompt tokens are known and stored in the KV cache, the computation is parallelized in the PPU cluster 138 by using matrix-matrix multiplication operations.

[0069] During the decode phase, also referred to as autoregressive generation phase, the LLM inference engine 200 generates the remaining tokens autoregressively in a sequential manner until a stopping criterion is met. At iteration t, the LLM takes a new token as an input and computes the probability of the subsequent token with the key vectors and the value vectors from the previous iterations output states (keys and values) .

[0070] In one or more embodiments, the set of APIs 210 may be provided by the application servers 136. In one or more embodiments, the one or more inference servers 132 may implement one or more of data parallelism, tensor parallelism, and pipeline parallelism to process input requests.

[0071] The set of APIs 210 may include one or more APIs to manage software and hardware components required to execute the LLM inference engine 200. As a non-limiting example, the set of APIs 210 may include resource management APIs, performance monitoring APIs, hardware configuration APIs, load balancing APIs, firmware and driver management APIs, job scheduling APIs and scalability APIs.

[0072] In the context of the present disclosure, the set of APIs 210 includes at least tensor memory management APIs (not shown in FIG. 2) and high-bandwidth memory (HBM) management APIs (not shown in FIG. 2) , which will be described in more detail hereinafter. The LLM inference engine 200 is configured to use the tensor memory management and HBM management APIs to allocate HBM memory on-demand by using virtual memory for the tensors in the KV cache. In some embodiments, for frameworks requiring static-shaped inputs using padded tensors, the LLM inference engine 200 is configured to use the tensor memory management and HBM management APIs to map padded tensor tokens to the same physical memory page, thereby saving resources.

[0073] The one or more inference servers 132 are connected to the parallel processing unit (PPU) cluster 138 (shown in FIG. 2) comprising a plurality of PPUs (i.e., hardware accelerators) for handling large scale computations required by the LLM inference engine 200. It will be appreciated that execution of ML models, such as LLMs including transformer models, involves a significant number of operations, such as tensor multiplication between input data and high-dimensional weight tensors that can be computationally intensive. The PPUs of PPU cluster 138 are optimized to perform these operations efficiently by parallel processing, which provides significant improvements in latency or throughput when the number of parameters in the ML models are large.

[0074] The one or more inference servers 132 comprise one or more general purpose processing units 160 (e.g., CPUs) , one or more memories 162, one or more communication interfaces 164, and one or more storage units 166, which may be similar to the components of computer system 150. As a non-limiting example, the one or more inference servers 132 may be implemented as one of the HUAWEI Atlas 800 inference server and the NVIDIA DGX series server.

[0075] In one or more embodiments, the PPU cluster 138, which may also be referred to as accelerator cluster 138 or execution engine 138, may be integrated with the one or more inference servers 132.

[0076] In some instances, the one or more inference server 132 and the PPU cluster 138 are configured to process workloads that require processing capabilities of more than one PPU. Thus, the one or more inference servers 132 are configured to distribute the workload across multiple hardware accelerators in the PPU cluster 138.

[0077] In the example shown in FIG. 2, the PPU cluster 138 comprises a first PPU 182, a second PPU 184, and a third PPU 186. It should be noted that the PPU cluster 138 may include a different number of PPUs.

[0078] Each given PPU 182, 184, 186 may be implemented as one or more parallel processing units (PPU) , also referred to as accelerators, such as graphics processing units (GPUs) , neural processing units (NPUs) , and tensor processing units (TPUs) for executing specialized ML computation tasks. As a non-limiting example, each given PPU 182, 184, 186 may be implemented as one of the HUAWEI Ascend Series NPU, the NVIDIA A100 Tensor Core GPU, and the GOOGLE TPU.

[0079] In one or more embodiments, each given PPU 182, 184, 186 in the PPU cluster 138 is configured to execute a respective instance of a tensor memory manager (not shown in FIG. 2) .

[0080] FIG. 2 shows components of the first PPU 182, which comprise one or more specialized processing cores (PCs) 192 (e.g., GPU cores) operatively connected to a PPU memory 194 (e.g., GPU memory) . It should be understood that the second PPU 184 and the third PPU 186 may comprise similar components. The specialized PCs 192 may for example include streaming multiprocessors (SMs) , and the PPU memory 194 may include shared memory (e.g., SRAM) and high-bandwidth memory (HBM) (e.g., SDRAM) . The one or more inference servers 132 may send instructions and data to a given PPU 182, 184, 186 and receive data therefrom using one or more APIs from the set of APIs 210.

[0081] The one or more PCs 192 are configured to: (i) execute a respective portion of model parameters of the ML model (i.e., LLM) ; (ii) perform matrix multiplication, apply activation functions, attention calculations and / or other types of calculations according to the configuration of the respective portion of the model parameters of the ML model stored in the PPU memory 194; and (iii) implement communication mechanisms to transmit intermediate results to other PPUs in the PPU cluster 138.

[0082] The PPU memory 194 is configured to: (i) store a respective portion of model parameters of the ML model including weights matrices; (ii) store dynamic states of requests processed by the ML model, including activations, attention scores and hidden states; and (iii) store query, key and value tensors in a KV cache; and (iv) buffer data for transmission to other PPUs in the PPU cluster 138.

[0083] In the context of the present disclosure, the PPU memory 194 is also referred to as physical memory, the physical memory being configured for storing tensor elements in page frames having respective addresses. The page frames are generally divided into fixed size pages. The page frames of the PPU memory 194 may be mapped to one or more virtual pages in virtual memory using a page table.

[0084] As stated previously, the LLM inference engine 200 is configured to sample and generate new tokens one by one, where the generation process of each token depends on the previous tokens in the sequence. Each output token needs to know all the previous iterations’ output states (i.e., key and value vectors stored in the KV tensor in PPU memory 194) . Thus, during the decode phase, each new output token depends on the KV tensors of the input tokens computed at the prefill phase, and the KV tensors computed until the current time step.

[0085] The LLM inference engine 200 caches the key and value vectors of each existing tokens in the KV cache of PPU memory 194 for generating future tokens in the sequence. At each new iteration, when new elements are computed, the LLM adds the computed key values to the KV cache to be used in the subsequent iteration. The decode phase stops when the sequence of output tokens reaches a maximum length or when an end of sequence token is reached.

[0086] In one or more embodiments, the layers of the LLM model in LLM inference engine 200 are sharded on the PPU cluster 138, such that each of the first PPU 182, the second PPU 184, and the third PPU 186 executes a respective subset of layers (i.e., model parameters) of the LLM.

[0087] In one or more embodiments, the LLM inference engine 200 is configured to manage one or more caches on the local memory in the inference server 132 and the PPU cluster 138 necessary for executing calculations in the layers of the LLM executed by the LLM inference engine 200.

[0088] As a non-limiting example, the LLM inference engine 200 is configured to store a key cache tensor and a value cache tensor for a request until the request has been completed. In response to receiving one or more new requests (requests for which the encoding phase has not been processed) , the LLM inference engine 200 is configured to allocate memory to each request for maintaining the key cache tensor and the value cache tensor for the request. For each iteration of the decoding phase, the LLM inference engine 200 may add the key tensor and the value tensor for the iteration to the respective key cache and value cache allocated for the request and retrieve the key cache and value cache for the request to execute an operation. Upon completing one or more requests, the LLM inference engine 200 may provide the outputs for the requests to an appropriate component executed by the plurality of servers 130 and free the allocated cache memory for the completed requests, such that the freed memory can be used for other requests.

[0089] FIG. 3 shows a non-limiting example of a LLM serving system 300 implementing tensor memory management according to one or more embodiments of the present technology.

[0090] The LLM serving system 300 may be executed within the plurality of servers 130.

[0091] The LLM serving system 300 is configured to receive a plurality of requests from the client devices 120 (shown in FIG. 1A) . Each request includes a list of input prompt tokens. In some embodiments, the LLM serving system 300 is configured to pad the tokens of the requests such that each request has the same length.

[0092] The LLM serving system 300 comprises a LLM inference engine 302 and a tensor memory manager 320. The LLM inference engine 302 is an embodiment of the LLM inference engine 200.

[0093] The LLM inference engine 302 is configured to batch multiple requests by grouping the requests and handling the requests as a single batch to improve usage of computing resources. It will be appreciated that batching is possible because the requests share the same LLM model weights.

[0094] The LLM serving system 300 has access to one or more HBM management APIs 330 to manage HBM of PPU memory 194 of each given PPU 182, 184, 186 in the PPU cluster 138 of FIG. 2. The HBM management API 330s are used to monitor and manage HBM memory in the PPU cluster 138.

[0095] In one or more embodiments, the one or more HBM managements APIs 330 are configured to: (i) receive current availability and usage statistics of the PPU memory 194, including total available memory, used memory, and free memory in the PPU cluster 138; (ii) allocate, reallocate and deallocate memory in the PPU memory 194 in the PPU cluster 138; (iii) perform additional memory management and optimization strategies.

[0096] Non-limiting examples of a HBM management APIs 138 include NVIDIA (nvidia-smi) , AMD ROCm-SMI, NVIDIA Nsight Systems, NVIDIA CUDA Runtime &Driver APIs, HUAWEI Ascend (npu-smi) , and HUAWEI Ascend CANN Runtime APIs.

[0097] The tensor memory manager 320 is implemented on top of the HBM management APIs 330.

[0098] The tensor memory manager 320 provides a virtual tensor API 322 used by the LLM inference engine 302. The LLM inference engine 302 uses the virtual tensor API 322 to create tensors for its KV caches, and calls the APIs to notify the tensor memory manager 320 when it needs to allocate, deallocate, and / or increase the memory. The tensor memory manager 320 determines the allocation, deallocation, and / or scaling of memory for tensors of the LLM, and executes the HBM management APIs 330.

[0099] The tensor memory manager 320 is configured to: (i) reserve a virtual memory address space comprising a plurality of virtual pages and use it to initialize a tensor; (ii) create physical memory page frames and map the page frames to the virtual memory pages addresses when needed to store the key vectors and value vectors using a page table; (iii) unmap the physical  memory page frames from the virtual memory addresses and release the memory page frame upon reaching a stopping criterion; and (iv) map padded tokens in the tensor to a single page frame in the PPU memory 194.

[0100] In one or more embodiments, each PPU 182, 184, 186 in the PPU cluster 138 executes a respective instance of the tensor memory manager 320.

[0101] It will be appreciated that the tensor memory manager 320 using the HBM management APIs 330 enables storing sequences of tokens into contiguous (virtual) memory spaces from the point of view of the LLM inference engine 302. Thus, the LLM inference engine 302 does not need specialized adaptation to read from discontiguous memory spaces, thereby rendering the inference process more efficient than with frameworks storing sequence of tokens in dis-contiguous memory spaces, such as the vLLM framework using the custom PagedAttention operator. The tensor memory manager 320 uses a page table to map respective virtual pages addresses of virtual memory to corresponding page frames addresses in physical memory (i.e., PPU memory 194) .

[0102] The tensor memory manager 320 is configured to support multi-program multi-data (MPMD) parallelization patterns. In such embodiments implementing a distributed computing configuration with multiple PPUs, each PPU 182, 184, 186 has a respective tensor memory manager 320, which may be programmed differently to fulfill heterogenous computing tasks.

[0103] In some embodiments, the tensor memory manager 320 is configured to manage low-rank adaptation (LoRA) , which freezes pre-trained model weights and injects trainable rank decomposition matrices in layers of the LLM architecture. LoRA adaptors can be injected into the LLM for different types of custom applications. With a request for a certain application, the corresponding LoRA weights may be loaded into the HBM of the PPU memory 192 to be executed. The allocation and deallocation of LoRA tensor memory can also be managed by the tensor memory manager 320.

[0104] The tensor memory manager 320 is configured to access the mmap API. The tensor memory manager 320 is configured to use the mmap API to map virtual pages in virtual memory to page frames in physical memory (e.g., PPU memory 194) .

[0105] It will be appreciated that the tensor memory manager 320 does not require pre-allocation of physical memory. The tensor memory manager 320 is configured to allocate or reserve virtual memory pages for the tensor and map the virtual memory pages to page frames in the PPU memory 194 upon use.

[0106] The LLM inference engine 302 comprises a scheduler 304, a KV cache 308, one or more worker (s) 306, an attention operator 310, metadata definition 312, and tensor definition 314.

[0107] In one or more embodiments, each of the worker (s) 306 is a software component associated with a respective PPU 182, 184, 186 in the PPU cluster 138. The workers 306 are software components implemented by the distributed computing framework (e.g. Ray) . Each of the worker (s) 306 is configured to manage the respective PPU 182, 184, 186 and to execute programs on the respective PPU 182, 184, 186 including a respective instance of the tensor memory manager 320 and the HBM management API 330.

[0108] The scheduler 304 is configured to schedule tasks by grouping multiple requests into batches to take advantage of the parallel processing capabilities of the PPUs 182, 184, 186 in the PPU cluster 138 (e.g., GPUs) .

[0109] In one or more embodiments, the scheduler 304 is configured to: (i) receive a plurality of requests; (ii) queue the plurality of requests; (ii) receive, from each of the worker (s) 306, a respective available HBM; (iii) group the plurality of requests into batches; (iv) allocate, for each request, based on the respective available HBM, a respective KV cache 308; (v) prepare the input tokens for each batch; (vi) cause and coordinate execution of the model computations by the workers 306; (vii) aggregate the output of workers 306.

[0110] In some embodiments, each of the worker (s) 306 is configured to: (i) receive input tokens and control messages from the scheduler 304; (ii) compute, based on the input tokens and the control messages, intermediate states comprising key-value pairs; (iii) update the KV cache 308; (iv) transmit intermediate data to subsequent worker in the workers 306.

[0111] The metadata definition 312 comprises the KV indices of the key-value pairs stored in the tensor in the KV cache 308 during the inference process. It should be noted that the size of the KV indices is constant.

[0112] The tensor definition 314 comprises definitions and structures of tensors used in the LLM, including the shapes, data types, and storage formats of tensors.

[0113] The attention operator 310, also known as attention kernel 310, is configured to execute attention mechanisms within the LLM. The attention mechanisms include calculating attention scores, and calculating a relevance of each portion of the input token  sequence with respect to each position in the output token sequence. The attention mechanism is configured to transform the input sequence of tokens into matrices of queries (Q) , matrices of keys (K) and matrices of values (V) using learned weight matrices for queries (WQ) , learned weight matrices for keys (WK) and learned weight matrices for values (WV) . The attention mechanism may then use a self-attention layer to compute an attention score by multiplying query vectors at one position with all preceding key vectors and compute the output as a weighted average over the value vectors. In the context of the present disclosure, key matrices and value matrices are stored in contiguous memory spaces from the point of view of the attention operator 310.

[0114] The KV cache 308 is configured to store tensors comprising key vectors and value vectors of tokens computed and used by each worker 306.

[0115] FIG. 4 shows an embodiment of a scheduler 404 with workers 410, 420. Each worker 410, 420 executes a respective tensor memory manager instance 412, 414 and a respective model shard 414, 424.

[0116] FIG. 5 illustrates a non-limiting example of KV cache management of two requests at two different iterations from the point of view of the tensor using virtual memory 500 and from the point of view of the physical memory 520. In one or more embodiments, the physical memory 520 corresponds to the PPU memory 194 of FIG. 2.

[0117] The physical memory 520 includes a plurality of page frames 522, 524, 526, 528 and 530. It should be noted that the plurality of page frames are not pre-allocated.

[0118] A tensor is initialized by allocating a plurality of virtual pages, each virtual page being associated with a respective virtual memory page address (not shown) in the virtual memory 500.

[0119] At the current iteration, the tensor stores a first sequence of tokens ( ‘7’ ) in a first virtual page 506. The first virtual page 506 is mapped by the tensor memory manager 320 to the physical memory 520. It will be appreciated that a page table is used to map the virtual page addresses to corresponding physical memory page frame addresses.

[0120] From the point of view of the physical memory 520 comprising the plurality of page frames 522, 524, 526, 528 and 530, the first sequence of tokens in the first virtual page 506 (token ‘7’ ) is mapped to the third page frame 526, with page frames 522, 524, 528 and 530 being free. In one or more embodiments, the first virtual page 506 is mapped to the third page frame 526 using the mmap API.

[0121] An attention operator 572 is applied on the first sequence of tokens in the first virtual page 506 to obtain the new token stored in second virtual page 508 at the next iteration.

[0122] At the next iteration, the tensor stores the first sequence of tokens ( ‘7’ ) in the first virtual page 506 and the new token (‘8’ ) generated using the attention operator 572 in a second virtual page 508. The new token is stored contiguous to the first sequence of tokens 506, i.e., the second virtual page 508 is adjacent to the first virtual page 506.

[0123] At the same time, a second request comprising a second sequence of tokens ( ‘1’ , ’ 2’ , ‘3’ ) is received and stored in the tensor. The tensor memory manager 320 allocates a second set of virtual pages comprising a third virtual page 532, a fourth virtual page 534 and a fifth virtual page 536 to each respective token ‘1’ , ’ 2’ , ’ 3’ in the second sequence of tokens. Tokens in the second sequence of tokens are contiguous to each other in the virtual memory 500, and the second sequence of tokens is contiguous to the first sequence of tokens from the point of view of virtual memory 500.

[0124] From the point of view of the physical memory 520, the second sequence of tokens ( ‘1’ , ’ 2’ , ’ 3’ ) is mapped to the first page frame 522, the second page frame 524 and the fifth page frame 530, respectively while the first sequence of tokens is mapped to the third page frame 526 and the new token is cached into fourth page frame 528. Thus, it can be seen that while the first sequence of tokens ( ‘7’ ) , the new token ( ‘8’ ) and the second sequence of tokens ( ‘1’ , ’ 2’ , ’ 3’ ) are discontinuous from the point of view of the physical memory 520, the first sequence of tokens stored in first virtual page 506 and the new token stored in the second virtual page 508 are contiguous from the point of view of the virtual memory 500, which facilitates computations.

[0125] FIG. 6 illustrates a non-limiting example of padding of the tensor in virtual memory 600 to physical memory 650 in accordance with implementations of the present technology. The padding of the tensor may be required for frameworks running in static graph mode, which require static-shaped inputs.

[0126] The tensor in virtual memory 600 stores tokens ‘1’ , ‘2’ and ‘3’ respectively in the first virtual page or block 602, the second virtual page 604, and the third virtual page 608. The remaining virtual pages 650, 612, 614, 616, 618 and 620 in the tensor in the virtual memory 600 are prefilled or padded for future computations.

[0127] From the point of view of the physical memory 660, the first virtual page 602, the second virtual page 604, and the third virtual page 608 are mapped respectively to the first page frame 622, the second page frame 624 and the third page frame 626 in the physical memory 660. The remaining virtual pages 610, 612, 614, 616, 618 and 620 from the tensor in virtual memory 600, corresponding to a padded token ( ‘p’ ) , are mapped to a single page frame 640 in the physical memory 600, while the remaining page frames 634 (not separately numbered) are free.

[0128] The tensor in virtual memory 600 is provided to the LLM 650 for computing the next token.

[0129] At the next iteration, the tensor 600 stores tokens ‘1’ , ‘2’ , ‘3’ respectively in the first virtual page 602, the second virtual page 604, and the third virtual page 608, and new token ‘4’ computed from the previous tokens is stored in the fourth virtual page 610, the remaining virtual pages 612, 614, 616, 618 and 620 in the virtual memory 600 are prefilled or padded for future computations.

[0130] In the physical memory 660, the first virtual page 602, the second virtual page 604, and the third virtual page 608 are mapped respectively to the first page frame 622, the second page frame 624, and the third page frame 626, with the new token ‘4’ in the fourth virtual page 610 of the tensor in virtual memory 600 mapped to the fourth page frame 628.

[0131] The remaining blocks 612, 614, 616, 618 and 620 from the tensor 600, corresponding to a padded value, are mapped to a single tenth page frame 640 in the physical memory 610, while the remaining physical page frames 636 (not separately numbered) are free.

[0132] FIG. 7 and FIG. 8 illustrate a flowchart of a method 700 for serving a LLM engine by allocating memory on demand, the method being executed in accordance with non-limiting embodiments of the present technology.

[0133] The method 700 may be executed by a computing device such as computing device 150, the one or more inference servers 132 and / or the PPU cluster 138.

[0134] In one or more embodiments, the method 700 is stored in the form of computer-readable instructions in a non-transitory storage medium. At least one processing unit (e.g., CPU, GPU, NPU, or TPU) , upon loading the computer-readable instructions from the non-transitory storage medium, is configured or operable to execute method 700.

[0135] In some embodiments, the method 700 is executed by an API such as the tensor memory manager 320. The tensor memory manager 320 has access to the mmap API.

[0136] In other embodiments, the method 700 may be executed within the LLM serving system 300.

[0137] The method 700 begins at processing step 702.

[0138] According to processing step 702, the at least one processing unit initializes a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages.

[0139] The at least one processing unit reserves a virtual memory address space by allocating virtual addresses of the plurality of virtual pages to the tensor. The plurality of virtual pages include a contiguous range of virtual memory addresses to be used for storing the KV cache.

[0140] According to processing step 704, the at least one processing unit receives a first request comprising a first sequence of input tokens.

[0141] In some embodiments, the first request may originate from a client device (e.g. first client device 122) , may be received via gateway 115 and may be processed to obtain the first sequence of input tokens.

[0142] According to processing step 706, the at least one processing unit adds the first sequence of input tokens to the tensor, said adding comprising: allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages, and mapping the first set of virtual pages to a first set of page frames in the physical memory.

[0143] The at least one processing unit allocates the first sequence of input tokens to the virtual addresses of the first set of virtual pages in the plurality of virtual pages of the memory.

[0144] The at least one processing unit allocates the first set of page frames (i.e., physical pages) which are mapped to the first set of virtual pages. The at least one processing unit uses a page table to translate virtual memory addresses to physical memory addresses. In one or more embodiments, the at least one processing unit uses the mmap API to perform the mapping.

[0145] It will be appreciated that page frames in the first set of page frames may not be contiguous in physical memory.

[0146] In one or more embodiments, the physical memory may be the PPU memory 192 (e.g., GPU memory) . In one or more other embodiments, the physical memory may be a CPU memory.

[0147] According to processing step 708, the at least one processing unit provides the tensor to the LLM engine for processing.

[0148] In one or more embodiments, the processing of the tensor by the LLM engine comprises: accessing decoding layers of the LLM engine, generating, via the decoding layers, based on the first sequence of input tokens, a first set of query vectors, a first set of key vectors, and a first set of value vectors, adding the first set of query vectors, the first set of key vectors and the first set of value vectors to the tensor. The adding comprises allocating the first set of query vectors, the first set of key vectors and the first set of value vectors to a third set of virtual pages in virtual memory, and mapping the third set of virtual pages to a third set of page frames. The LLM engine generates the first output token based on the first set of query vectors, the first set of key vectors and the first set of value vectors to obtain the first output token. The LLM engine applies operators (e.g., attention kernels) on the first set of query vectors, the first set of key vectors and the first set of value vectors to generate the first output token. The page table is updated to add the mapping between the third set of virtual pages and the third set of page frames.

[0149] It will be appreciated that since the first sequence of input tokens is stored in contiguous memory from the point of view of the LLM engine, the operator (s) used for processing the first sequence of input tokens do not need to be specifically adapted, in contrast with techniques such as vLLM which require modifying operators to perform computations over non-contiguous memory.

[0150] According to processing step 710, the at least one processing unit receives a second request comprising a second sequence of input tokens.

[0151] In some embodiments, the second request may originate from a client device (e.g. second client device 124 or first client device 122) and may be processed to obtain the second sequence of input tokens. As a non-limiting example, the second request may be received via gateway 115 and processed by the plurality of servers 130 to obtain the second sequence of input tokens.

[0152] According to processing step 712, the at least one processing unit receives, from the LLM engine, a first output token for the first sequence of tokens.

[0153] According to processing step 714, the at least one processing unit adds the first output token to the tensor, said adding comprising: allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory, and mapping the second page to a second page frame in the physical memory.

[0154] The first output token is added to a second virtual page located adjacent to the first set of virtual pages. The page table is updated to add the mapping between the second virtual page and the second page frame in the physical memory.

[0155] In one or more embodiments, the at least one processing unit uses the mmap API to perform the mapping.

[0156] According to processing step 716, the at least one processing unit adds the second sequence of input tokens to the tensor, with the adding comprising: allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory, and mapping the second set of virtual pages to a second set of page frames in the physical memory.

[0157] The at least one processing unit allocates the second sequence of input tokens to the virtual addresses of the second set of virtual pages in the plurality of virtual pages of the memory. The at least one processing unit maps the second set of virtual pages to a second set of page frames in the physical memory.

[0158] The second set of virtual pages are contiguous in virtual memory. In the physical memory, at least some page frames in the second set of page frames are discontiguous.

[0159] According to processing step 718, the at least one processing unit provides the tensor to the LLM engine for further processing.

[0160] In some embodiments, processing steps 704 to 716 may be repeated for incoming requests until a stopping criterion is met.

[0161] Upon reaching the stopping criterion (e.g., end token) for at least one of the first request and the second request, the at least one processing unit causes deallocation of the page frames associated with the at least one of the first request and the second request, and unmaps the virtual pages associated with the at least one of the first request and the second request from the virtual memory.

[0162] The method 700 then ends.

Claims

1.A method for serving a large language model (LLM) engine by allocating memory on demand, the method being executed by at least one processing unit, the at least one processing unit being operatively connected to a physical memory comprising a plurality of page frames, the method comprising:initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages;receiving a first request comprising a first sequence of input tokens;adding the first sequence of input tokens to the tensor, said adding comprising:allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages; andmapping the first set of virtual pages to a first set of page frames in the physical memory;providing the tensor to the LLM engine for processing;receiving a second request comprising a second sequence of input tokens;receiving, from the LLM engine, a first output token for the first sequence of tokens;adding the first output token to the tensor, said adding comprising:allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory; andmapping the second page to a second page frame in the physical memory;adding the second sequence of input tokens to the tensor, said adding comprising:allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory; andmapping the second set of virtual pages to a second set of page frames in the physical memory; andproviding the tensor to the LLM engine for further processing.2.The method of claim 1, wherein said providing the tensor to the LLM engine for processing comprises:accessing decoding layers of the LLM engine;generating, via the decoding layers, based on the first sequence of input tokens, a first set of query vectors, a first set of key vectors, and a first set of value vectors;adding the first set of query vectors, the first set of key vectors and the first set of value vectors to the tensor, said adding comprising:allocating the first set of query vectors, the first set of key vectors and the first set of value vectors to a third set of virtual pages;mapping the third set of virtual pages to a third set of page frames; andgenerating the first output token based on the first set of query vectors, the first set of key vectors and the first set of value vectors to receive the first output token.3.The method of claim 1 or 2, wherein the second set of virtual pages are contiguous to the second virtual page in the virtual memory.4.The method of any one of claims 1 to 3, wherein at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.5.The method of any one of claims 1 to 4, whereinsaid adding the first sequence of input tokens to the tensor comprises adding a set of padded tokens; wherein said allocating the first sequence of input tokens in the tensor to the first set of virtual pages comprises allocating the set of padded tokens in the first set of virtual pages; and whereinsaid mapping the first set of virtual pages to the first set of page frames in the physical memory comprises mapping the set of padded tokens to a single page frame in the first set of page frames.6.The method of any one of claims 1 to 5, wherein said mapping is performed using the mmap API.7.The method of any one of claims 1 to 6, further comprising:upon reaching a stopping criterion for at least one of the first request and the second request:unmapping page frames associated with the at least one of the first request and the second request from the physical memory; anddeallocating virtual pages associated with the at least one of the first request and the second request from the virtual memory.8.The method of any one of claims 1 to 7, wherein the physical memory comprises a high-bandwidth memory (HBM) of a parallel processing unit (PPU) .9.The method of claim 8, wherein the PPU comprises one of: a graphics processing unit (GPU) , a neural processing unit (NPU) , and a tensor processing unit (TPU) .10.A system for serving a large language model (LLM) engine by allocating physical memory on demand, the system comprising:at least one processor;a non-transitory storage medium operatively connected to the at least one processor, the non-transitory storage medium storing computer-readable instructions thereon; andat least one parallel processing unit (PPU) connected to the non-transitory storage medium, the PPU comprising:a plurality of processing cores and a PPU memory, the PPU memory comprising a plurality of page frames;the at least one processor, upon executing the computer-readable instructions, being configured for:initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages;receiving a first request comprising a first sequence of input tokens;adding the first sequence of input tokens to the tensor, said adding comprising:allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages; andmapping the first set of virtual pages to a first set of page frames in the plurality of page frames;providing the tensor to the LLM engine for processing using the at least one PPU;receiving a second request comprising a second sequence of input tokens;receiving, from the LLM engine, a first output token for the first sequence of tokens;adding the first output token to the tensor, said adding comprising:allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory; andmapping the second page to a second page frame in the physical memory;adding the second sequence of input tokens to the tensor, said adding comprising:allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory; andmapping the second set of virtual pages to a second set of page frames in the physical memory; andproviding the tensor to the LLM engine for further processing.11.The system of claim 10, wherein the at least one PPU comprises one of: a graphics processing unit (GPU) , a neural processing unit (NPU) , and a tensor processing unit (TPU) .12.The system of claim 10 or 11, wherein at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.13.The system of any one of claims 10 to 12, wherein said providing the tensor to the LLM engine for processing comprises:accessing decoding layers of the LLM engine;generating, via the decoding layers, based on the first sequence of input tokens, a first set of query vectors, a first set of key vectors, and a first set of value vectors;adding the first set of query vectors, the first set of key vectors and the first set of value vectors to the tensor, said adding comprising:allocating the first set of query vectors, the first set of key vectors and the first set of value vectors to a third set of virtual pages;mapping the third set of virtual pages to a third set of page frames; andgenerating the first output token based on the first set of query vectors, the first set of key vectors and the first set of value vectors to receive the first output token.14.The system of any one of claims 10 to 13, wherein the second set of virtual pages are contiguous to the second virtual page in the virtual memory.15.The system of any one of claims 10 to 14, wherein at least two page frames in the first set of page frames, the second page frame and the second set of page frames are discontiguous.16.The system of any one of claims 10 to 15, whereinsaid adding the first sequence of input tokens to the tensor comprises adding a set of padded tokens; wherein said allocating the first sequence of input tokens in the tensor to the first set of virtual pages comprises allocating the set of padded tokens in the first set of virtual pages; and whereinsaid mapping the first set of virtual pages to the first set of page frames in the physical memory comprises mapping the set of padded tokens to a single page frame in the first set of page frames.17.The system of any one of claims 11 to 16, wherein said mapping is performed using the mmap API.18.The system of any one of claims 10 to 17, wherein the at least one processor is further configured for:upon reaching a stopping criterion for at least one of the first request and the second request:unmapping page frames associated with the at least one of the first request and the second request from the physical memory; anddeallocating virtual pages associated with the at least one of the first request and the second request from the virtual memory.19.The system of any one of claims 10 to 18, whereinthe at least one PPU comprises a plurality of PPUs; and whereinthe system further comprises a scheduler configured for receiving available memory from each respective PPU memory from the plurality of PPUs and scheduling processing of the first request and the second request by respective ones of the plurality of PPUs.20.A non-transitory storage medium comprising computer-readable instructions, when executed by at least one processing operatively connected to a physical memory comprising a plurality of page frames, cause execution of a method comprising:initializing a virtual memory space for a tensor, the virtual memory space comprising a plurality of virtual pages;receiving a first request comprising a first sequence of input tokens;adding the first sequence of input tokens to the tensor, said adding comprising:allocating the first sequence of input tokens in the tensor to a first set of virtual pages in the plurality of virtual pages; andmapping the first set of virtual pages to a first set of page frames in the physical memory;providing the tensor to the LLM engine for processing;receiving a second request comprising a second sequence of input tokens;receiving, from the LLM engine, a first output token for the first sequence of tokens;adding the first output token to the tensor, said adding comprising:allocating the first output token to a second virtual page contiguous to the first set of virtual pages in the virtual memory; andmapping the second page to a second page frame in the physical memory;adding the second sequence of input tokens to the tensor, said adding comprising:allocating the second sequence of input tokens to a second set of virtual pages in the virtual memory; andmapping the second set of virtual pages to a second set of page frames in the physical memory; andproviding the tensor to the LLM engine for further processing.

Images