#import "templates/template.typ": *
#show: doc => slides(doc)

#show link: l => underline(l)

#show figure: it => {
  show figure.caption: set text(7pt)
  it
}

#slide(
  title: "About this paper", pad(
    x: 2em,
  )[
  #figure(
    image("images/header.png")
  )
  - This paper is presented at SOSP 2023.
  - PagedAttention, which is proposed in this paper, is the core of #link("https://github.com/vllm-project/vllm")[vLLM], the one of the state-of-the-art LLM serving systems.
])

== Preliminaries

#slide(
  title: "KV Cache", pad(
    x: 2em,
  )[
  #figure(image("images/KVCache.png"), 
  caption: "The mechanism of KV Cache. Cited from https://resources.nvidia.com/ja-jp-llm-developer-day/mastering-llm-techniques-inference-optimization")
])

#slide(
  title: "batching techniques for serving LLMs", pad(
    x: 2em,
  )[
  - Because requests for a LLM share the same model weights, #underline("batching can be applied to amortize the overhead of loading the model weights.")
  - The batching techniques for serving LLMs are categorized into two types.
    - request-level batching
    #figure(
      image("images/request-level.png", height: 30%)
    )
    - token-level batching
    #figure(
      image("images/token-level.png", height: 30%)
    )
])

== Background

#slide(
  title: "memory-bound throughput", pad(
    x: 2em,
  )[
  - The throughput of these systems is memory-bound.
    - These systems process multiple requests concurrently.
    - The memory consumption make a barrier to the scalability.
    - Less memory consumption means more requests can be processed concurrently.
  #figure(
    image("images/batch-memory.png", height: 50%), caption: "The memory consumption of LLMs served by existing systems and vLLM with different batch sizes."
  )
])

#slide(
  title: "memory distribution", pad(
    x: 2em,
  )[
  - The memory distribution of the LLM is like below. (The size is the example of LLMs with 13B parameters.)
    - Parameters: 65% (26GB)
    - KV Cache: 30% (12GB)
    - Others (e.g., activation): 5% (2GB)
  #figure(
    image("images/memory-dist.png", height: 60%), caption: "The memory distribution of LLMs with 13B parameters.",
  )
])

#slide(
  title: "", pad(
    x: 2em,
  )[
  - The model parameters are constant during inference phase and the activations only occupy a small fraction of the GPU memory.
    - KV Cache is the main bottleneck of the memory consumption.
  #figure(
    image("images/memory-dist-inst.png", height: 60%), caption: "The memory distribution of LLMs with 13B parameters.",
  )
])

#slide(
  title: "existing ways to manage the KV Cache", pad(
    x: 2em
  )[
  The existing way to manage KV Cache is as follows:
  + *pre-allocate* the whole size of KV Cache in *coutiguous* virtual memory. (It is corresponding to the size of the model's maximum output length.)
  + pre-fill the KV Cache with the users' prompts.
  + iterate decoding and append the corresponding results to the KV Cache one by one.

  #figure(
    image("images/Inference-steps.png", height: 60%), caption: "Cited from https://aijishu.com/a/1060000000458185",
  )
])

#slide(
  title: "existing ways to manage the KV Cache", pad(
    x: 2em
  )[
  This way has three types of memory wastes:
  + *internal-fragmentation*: the pre-allocated KV Cache is not fully used since the final output length is not always the maximum.
  + *external-fragmentation*: memory allocators often cause external fragmentation. (e.g., Buddy memory allocator)
  + *reserved* slots: the KV Cache is pre-allocated with the maximum output length, though the output length increases one by one.
  #figure(
    image("images/existing-kvmanage.png", width: 100%), caption: "KV cache memory management in existing systems.",
  )
])

#slide(
  title: "other challenging points", pad(
    x: 2em
  )[
  LLM serving systems offer a range of decoding algorithms, such as:
  - *greedy decoding*
  - *beam search*
  - *top-k sampling*
  #text(blue)[The KV cache corresponding to the user's prompt can be shared] to minimize the memory consumption.
  
  On the other hand, the KV cache during the autoregressive generation phase should remain unshared due to the different sample results and their dependencies on context and positions.
])

== Method

#slide(
  title: "system overview of vLLM", pad(
    x: 2em
  )[
  - PagedAttention divides the KV Cache into multiple #text(font: "Times New Roman", style: "italic")[pages].
  - Each page is allocated only when it is needed.
  - The pages are managed by a centralized scheduler.
    - Because the scheduler works in user-space, it causes several problems. (about that later)
  #figure(
    image("images/vllm-overview.png", height: 60%), caption: "The system overview of PagedAttention."
  )
])

#slide(
  title: "PagedAttention", pad(
    x: 2em
  )[
  vLLM works with PagedAttention as follows:
  + As in OS's virtual memory management, PagedAttention does not require reserving the memory for the maximum possible generated sequence length initially.
  + In the first autoregressive decoding step, vLLM generates the new token with the PagedAttention algorithm on physical blocks 7 and 1.
  + At the second decoding step, as the last logical block is full, vLLM stores the newly generated KV cache in a new logical block (physical block 3).
  #figure(
    image("images/pagedAttention.png", height: 50%), caption: "Block table translation in vLLM."
  )
])

#slide(
  title: "Application to Other Decoding Senarios - Parallel sampling - ", pad(
    x: 2em
  )[
  In LLM-based program assistants, an LLM generates multiple completions for a single input prompt in parallel. Users can choose a favorite output.
  #figure(
    image("images/parallel-sampling.png", height: 50%), caption: "Parallel sampling example."
  )
  In parallel sampling, the KV cache corresponding to the user's prompt can be shared to minimize the memory consumption.
])

#slide(
  title: "Application to Other Decoding Senarios - Parallel sampling - ", pad(
    x: 2em
  )[
  When generating multiple texts, when a different token is output for the first time, the corresponding page is duplicated for each token, creating a separate page for each (#text(blue)[copy-on-write]). 
  
  Since one page can be shared across multiple contexts, #text(blue)[reference counting] is implemented to manage memory.
  #figure(
    image("images/parallel-sampling.png", height: 50%), caption: "Parallel sampling example."
  )
])

#slide(
  title: "Application to Other Decoding Senarios - Beam search - ", pad(
    x: 2em
  )[
  In LLM tasks like machine translation, the users expect the top-k most appropriate translations output by the LLM.
  #figure(
    image("images/beam-search.png", height: 50%), caption: "Beam search example."
  )
  Unlike parallel decoding, beam search facilities sharing not only the blocks corresponding to the user's prompt but also other blocks across different candidates.
])

#slide(
  title: "Application to Other Decoding Senarios - Beam search - ", pad(
    x: 2em
  )[
  Previous LLM serving systems require frequent memory copies of the KV cache across the beam candidates. 

  For example, in the case shown in the figure, after the dotted line, candidate 3 (Block 12) would need to copy a large portion of candidate 2's KV cache to continue generation.

  #figure(
    image("images/beam-search.png", height: 50%), caption: "Beam search example."
  )
  
  This frequent memory copy overhead is significantly reduced by vLLM's physical block sharing.
])

#slide(
  title: "Scheduling and Preemption", pad(
    x: 2em
  )[
  As the number of requests and their outputs grow, systems can run out of the GPU's physical memory to store the newly generated KV cache.
  When the request traffic surpasses the system's capacity, it must prioritize a subset of requests. 

  LLM serving systems face a unique challenge in the scheduling and preemption in LLM serving systems:
  #text(blue)[the resulting output lengths are not known a priori, contignent on both the input prompt and the model.]

  There are two classic questions that these LLM serving systems must answer:
  + What blocks should it evict?
  + How to recover evicted blocks if needed again?
])

#slide(
  title: "What blocks should it evict?", pad(
    x: 2em
  )[
  Since we know that all blocks of a sequence are accessed together, #text(blue)[all-or-nothing eviction policy] is applied to the blocks.

  Furthermore, multiple sequences within one request (e.g., beam candidates in one beam search request) are gang-scheduled as a sequence group.
])

#slide(
  title: "How to recover evicted blocks if needed again?", pad(
    x: 2em
  )[
  There are two ways to recover evicted blocks:
  + *Swapping*: This is the classic technique used by most virtual memory implementations which copy the evicted pages to a swap space on the disk. In vLLM, evicted blocks are copied to the CPU memory. Besides the GPU block allocator, vLLM includes a CPU block allocator to manage the physical blocks swapped to CPU RAM.
  + *Recomputation*: This is a technique that recomputes the evicted blocks when needed. Note that #text(blue)[recomputation latency is significantly lower than the original latency, as the tokens generated at decoding can be concatenated with the original user prompt as a new prompt.]
])

== Evaluation

#slide(
  title: "Model and server configurations", pad(
    x: 2em
  )[
  - Model:
    - OPT (with 13B, 66B, 175B parameters)
    - LLaMA (with 13B parameters)
  - Server: A2 instance with NVIDIA A100 GPUs on GCP
  
  The detailed configurations are like below.
  #figure(
    image("images/eval-conf.png", height:50%), caption: "The detailed configurations of the model and server."
  )
])

#slide(
  title: "Workloads", pad(
    x: 2em
  )[
  Datasets:
  - ShareGPT: collection of user-shared cconversations with ChatGPT.
  - Alpaca: instruction dataset generated by GPT-3.5 with self-instruct.

  ShareGPT dataset has 8.4 times longer input prompts and 5.8 times longer outputs on average than the Alpaca dataset.
  #figure(
    image("images/dataset-dist.png", height: 40%), caption: "The distribution of the dataset."
  )

  These datasets do not include timestamps, so the authors added timestamps using Poisson distribution to simulate the real-world scenario.
])

#slide(
  title: "Baselines and Key metrics", pad(
    x: 2em
  )[
  *Baselines*:
  - #link("https://github.com/NVIDIA/FasterTransformer")[FasterTransformer]: distributed inference engine highly optimized for latency. Its development has transitioned to #link("https://github.com/NVIDIA/TensorRT-LLM/tree/release/0.5.0")[TensorRT-LLM].
  - #link("https://arxiv.org/pdf/2306.02707")[Orca]: a state-of-the-art at that time LLM serving system, which proposed the *iteration-level batching* technique.

  *Key metrics*:
  #text(blue)[The mean of every request's end-to-end latency divided by its output length.] A high-throughput serving system should have a low normalized latency against high request rates.
  #figure(
    image("images/e2e-latency.png", height: 40%), caption: "The end-to-end latency of the baselines."
  )
])

#slide(
  title: "Baselines and Key metrics", pad(
    x: 2em
  )[
  They implemented three version of Orca as baselines:
  + *Orca (Oracle)*: It has the knowledge of the length of the outputs that will be actually generated for the requests. This shows upper-bound performance of Orca, #text(red)[which is infeasible in real-world scenarios].
  + *Orca (Pow2)*: It over-reserves the space for outputs by at most 2x.
  + *Orca (Max)*: It always reserves the maximum possible output length.

  #table(
    columns: (auto, auto, auto),
    table.header(
      [], [*Actual output length* (examples)], [*Reserved memory size*]
    ),
    [*Orca (Oracle)*], [125 tokens], [125 tokens],
    [*Orca (Pow2)*], [125 tokens], [250 tokens],
    [*Orca (Max)*], [125 tokens], [2048 tokens]
  )
])

#slide(
  title: "Basic sampling", pad(
    x: 2em
  )[
  #figure(
    image("images/single-sequence-generation.png", height:80%), caption: "The end-to-end latency with OPT models on the ShareGPT and Alpaca dataset"
  )
])

#slide(
  title: "Basic sampling", pad(
    x: 2em
  )[
  #figure(
    image("images/single-sequence-generation-explosion.png", height: 60%), caption: "The end-to-end latency with OPT models on the ShareGPT and Alpaca dataset"
  )
  As the request rate increases, the latency initially increases at a gradual pace but then suddenly explodes.
  This is because the request rate surpasses the capacity of the serving system, the queue length continues to grow infinitely and so does the latency of the requests.
])

#slide(
  title: "Basic sampling", pad(
    x: 2em
  )[
  #figure(
    image("images/single-sequence-generation-less-pronounced.png", height: 60%), caption: "Average number of batched requests when serving OPT-13B for the ShareGPT (2 reqs/s) and Alpaca (30 reqs/s) datasets."
  )
  vLLM's advantage over Orca (Oracle) is less pronounced in (f).
  This is because the model and server configuration for OPT-175B allows for #text(blue)[large GPU memory space] available to store KV cache, while #text(blue)[the Alpaca dataset has short sequences.]

  These models can afford to reserve the maximum possible output length without running out of memory.
])

#slide(
  title: "Basic sampling", pad(
    x: 2em
  )[
  #figure(
    image("images/batched-requests-shareGPT.png", height: 60%), caption: "Average number of batched requests when serving OPT-13B for the ShareGPT (2 reqs/s) and Alpaca (30 reqs/s) datasets."
  )
  vLLM has the capability of batching and processing more requests concurrently than the baselines.
])

#slide(
  title: "Parallel Sampling and Beam Search", pad(
    x: 2em
  )[
  #figure(
    image("images/parallel-beam.png", height:60%), caption: "The end-to-end latency of parallel generation and beam search with OPT-13B on the Alpaca dataset."
  )
  
  With a larger number of sequences to sample or wider beam width, the performance gap between vLLM and the baselines widens. It's because the efficient memory sharing ,and management of vLLM can be more beneficial in these scenarios.
])

#slide(
  title: "Parallel Sampling and Beam Search", pad(
    x: 2em
  )[
  #figure(
    image("images/avr-amount-mem-save.png", height:60%), caption: "Average amount of memory saving from sharing KV blocks, when serving OPT-13B for the Alpaca dataset."
  )
  
  - 6.1%-9.8% memory saving on parallel sampling
  - 37.6%-55.2% memory saving on beam search
])

== Limitations

#slide(
  title: "Redundancy of PagedAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/pagedAttention-corrected.png", height: 70%), caption: "Block table translation in vLLM."
  )
  
  #text(red)[The image of PagedAttention in the paper is misleading.]
  The block table manager is in the user space, so it maps the user space block IDs to the virtual memory addresses. This may deteriorate the coalescing of memory accesses.
])

#slide(
  title: "Redundancy of PagedAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/redundancy-pagedattention.png", height:60%), caption: "PagedAttention involves two layers of memory management: one in user space and one in OS kernel space."
  )
  
  PagedAttention needs to track #text(blue)[virtual] memory addresses of KV cache blocks and pass them to the attention kernel at runtime.

  This approach requires deplicating what the operationg system already does for enabling #text(blue)[virtual-to-physical] address tranlsation.
])

#slide(
  title: "Redundancy of PagedAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/FA-Paged.png", height:60%), caption: [Overhead of PagedAttention in prefill kernels (model: Llama-3-8B, one A100 GPU). Numbers on top show overhead over the corresponding non-paged implementation of FlashAttention (FA) and FlashInfer (FI). #text(red)[Note that the FlashAttention and FlashInfer functionalities shown in the figure are features of vLLM.]]
  )

  This redundancy impacts the system's programmability, making it more challenging to correctly and efficiently implement other techniques that enhance the throughput and user experience of LLMs.

  This figure shows that PagedAttention with other techniques has a significant overhead compared to the non-paged implementation of FlashAttention and FlashInfer.
])

== Related Work

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-API.png"), caption: "vAttention API."
  )
  #link("https://arxiv.org/pdf/2405.04437")[vAttention] provides the same functionality as PagedAttention but without the redundant memory management #text(blue)[leveraging the low-level CUDA virtual memory management.]
  For this purpose, they implemented custom CUDA memory allocator APIs because the existing CUDA APIs do not support the page size granularity of 64KB, which is the optimal size for the KV cache.
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-workload-a.png"), caption: "The workload of vAttention."
  )
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-workload-b.png"), caption: "The workload of vAttention."
  )
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-workload-c.png"), caption: "The workload of vAttention."
  )
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-workload-d.png"), caption: "The workload of vAttention."
  )
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  #figure(
    image("images/vAttention-workload-e.png"), caption: "The workload of vAttention."
  )
])

#slide(
  title: "Solution to Redundancy of PagedAttention: vAttention", pad(
    x: 2em
  )[
  vAttention has the following advantages:
  - The functionality of vAttention is #text(blue)[user-friendly for programmers], as it provides the same capabilities as PagedAttention without requiring manual memory management in user space.
    - It follows that other techniques that improve the throughput and user experience of LLMs can be implemented more easily.
  - The memory management overhead is reduced by leveraging the OS's virtual memory management.

  However, it does not seem that vLLM has integrated vAttention. (#link("https://github.com/vllm-project/vllm/pull/6102")[PR \#6102])
])

#slide(
  title: "Tradeoff between Time-between-tokens Latency and Throughput", pad(
    x: 2em
  )[
  #grid(
    columns: 2,
    [
      #figure(
        image("images/tbt-throughput-tradeoff.png", height: 60%), caption: "The tradeoff between time-between-tokens latency and throughput."
      )
    ],
    [
      #figure(
        image("images/Inference-steps.png", height: 60%), caption: "Cited from https://aijishu.com/a/1060000000458185"
      )
    ]
  )

  Current LLM serving systems involve a tradeoff between throughput and latency depending on their scheduling policy.
  Prioritizing prefills optimizes throughput but sacrifices TBT (time-between-tokens) latency, while prioritizing decoding has the opposite effect.
])

#slide(
  title: "Scheduling policies", pad(
    x: 2em
  )[
  #figure(
    image("images/scheduling.png", height: 60%), caption: "Scheduling policies of different LLM serving systems. A, B, C and D represent different requests. Subscript d represents a decode iteration, p represents a full prefill and p0, p1 represent two chunked prefills of a given prompt."
  )

  - A generation stall occurs when one or more prefills are scheduled in between consecutive decode iterations of a request.
  - #link("https://www.usenix.org/system/files/osdi24-agrawal.pdf")[Sarathi-Serve] generates a schedule that eliminates generation stalls yet delivers high throughput splitting the prefills into smaller chunks. (#text(blue)[This functionality has already been merged into vLLM.])
])

#slide(
  title: "Scheduling policies", pad(
    x: 2em
  )[
  #figure(
    image("images/scheduling-algorithm.png", height: 60%)
  )

  Aformentioned LLM serving systems all use FCFS (First-Come-First-Serve) scheduling policy.
  There are other scheduling policies that can be applied to LLM serving systems.
  
  These other scheduling policies are based on the principle: #text(blue)[shorter tasks should be scheduled before longer tasks].
  (SRPT is infeasible due to the indeterminable output lengths of LLMs.)
])

#slide(
  title: "Scheduling policies", pad(
    x: 2em
  )[
  #figure(
    image("images/scheduling-algorithm.png", height: 60%)
  )

  MLFQ (Multi-Level Feedback Queue) is a CPU scheduling algorithm that assigns processes to different priority levels based on their execution behavior. It dynamically adjusts process priority, giving preference to I/O-bound or short-running tasks and demoting CPU-bound processes to lower-priority queues to improve overall system responsiveness.
])

#slide(
  title: "Scheduling policies", pad(
    x: 2em
  )[
  #figure(
    image("images/scheduling-algorithm.png", height: 60%)
  )

  #link("https://arxiv.org/pdf/2305.05920")[FastServe] proposes Skip-Join MLFQ, a scheduling policy leveraging LLMs' unique characteristics to improve throughput and latency.

  Skip-Join MLFQ #text(blue)[determine the priority of a request based on the number of tokens of the user input prompt.] It skips the requests with a large number of tokens to avoid the generation stall.
])

== Conclusion

#slide(
  title: "", pad(
    x: 2em
  )[
  - PagedAttention proposed the method to manage the KV cache efficiently inspired by the virtual memory management of OS.
  - vAttention adrressed the redundancy of PagedAttention by leveraging the low-level CUDA virtual memory management.
  - TBT latency caused by long stalls during prefill generation can be minimized by splitting the prefills into smaller chunks.
  - The scheduling policy based on the number of tokens of the user input prompt can improve the throughput and latency of LLM serving systems.
])

#slide(
  title: "My thoughts", pad(
    x: 2em
  )[
  -	The scheduling policy in LLM serving systems can be optimized by predicting the output length of requests.
	  - Speculative execution in LLMs, which runs multiple computations in parallel based on predictions from smaller draft models, helps reduce latency and improve efficiency.
	  -	Similarly, predictive methods can enhance the scheduling policy.
	-	Preemption occurs when the KV cache exceeds GPU memory. Current strategies simply evict or recompute blocks, but preemption could be improved by using KV cache compaction techniques, preserving output quality without unnecessary recomputation.
])