See also Paged Attention (Kwon et al., 2023)
KV-Compress
variable compression rates per attention head
source: github
idea.
Look at past attention weights for each pair of key and value vectors (a measure of the degree with which that KV’s representation has been queried during past attention operations)
Then select the KV with the least attention to evict
Think of LFU (least frequency used) cache management policy
the KV cache for each sequence in a particular layer is allocated on the GPU as a # attention heads sequence length tensor.
Lien vers l'originalImportant
total memory allocation scales with the maximum sequence length for all attention heads of the KV cache
Notes
A variation of Ada-SnapKV
Motivation:
- group-query-compression: compress KV-cache of GQA without repeating it into the dimension of query heads.
- Modified
PagedAttention
that compute against KV-cache (contains variable numbers of KVs per head)
For vLLM, each cache block stores KV for every attention head of every layer
For KV-Compress, each block only holds KVs for a single head. Block tables are expanded so that unique block for each specific KV head and layer can be retrieved
Query-Group Compression (QGC)
KV compression algorithm doesn’t have GQA design in mind.
- Pyramid-KV cache and compress KV after repetition for alignment with query tensors
- Redundancy in cache before compression
modification of eviction-based methods per groups
Block layout and allocation
idea: adapt PagedAttention to page out cache on a per-head, per-layer–as well as per sequence–basis
explanation
A simplified example with two KV heads and a block size of two:
- KV metrics are visualized for a given cache state, highlighting blocks of a particular sequence in the decoding batch that is scheduled to evict two blocks.
- Logical indices are displayed under the corresponding metrics slot.
Evict from Paged KV cache
need to evict KV blocks instead of evict single KV attention
automatic prefix caching
excerpt from github
block manager and evictor
see also: v2 and v1, benchmark
Reasoning for v2:
- support sliding windows attention
- lookahead slot for speculative decoding
speculative decoding
See slides
Speculative execution for LLMs is an excellent inference-time optimization.
— Andrej Karpathy (@karpathy) 31 août 2023
It hinges on the following unintuitive observation: forwarding an LLM on a single input token takes about as much time as forwarding an LLM on K input tokens in a batch (for larger K than you might… https://t.co/FiwTwqsfho
- not all parameters are required for generations tokens
- constraints tokens with low information-density
Ideas
Uses a small cheap “draft model” to generate candidate K tokens ⇒ feed back to the large models in a batch
- have a sort of sampling logics to get the probability of the next token, then forward passing for all later tokens.
continuous batching
(Yu et al., 2022) solves the static batching to reduce cost and improve throughput by appending requests continuously into existing KV cache 1
Lien vers l'original
guided decoding
- not supported from
SamplingParams
- requires support batch/async logits processing
- engine will die if failed
Benchmark script: vllm-project/vllm#10046
overhead
Currently logit_processor are happening in frontend, so we should move this to model_executor layers
waterfall
see also introduction slides
tldr: Bottleneck at AsyncLogitProcessor
and Scheduling
layer, given that this is row-wise operations 2
---
title: Initialization flow
---
graph TB
subgraph Engine
AsyncLLMEngine[AsyncLLMEngine]
end
subgraph Executors
GPU[GPUExecutorAsync]
TPU[TPUExecutorAsync]
XPU[XPUExecutorAsync]
end
subgraph Workers
GPUWorker[GPUWorker]
end
subgraph Model Runners
EmbeddingModelRunner[EmbeddingModelRunner]
GPUModelRunner[GPUModelRunner]
end
subgraph Control Plane
Scheduling[Scheduling]
SequenceGroup[SequenceGroup]
KVCache[KVCache]
Executors
end
AsyncLLMEngine --> |init| C[init]
C --> |device_type=gpu| GPU
C --> |device_type=tpu| TPU
C --> |device_type=xpu| XPU
GPU --> Workers
Workers --> |model_type=decoder| GPUModelRunner
Workers --> |model_type=embeddings| EmbeddingModelRunner
GPUModelRunner --> ModelClassImpl[LlamaModelForCausalLM]
---
title: Request flow
---
graph TB
subgraph Engine
AsyncLLMEngine[AsyncLLMEngine]
end
subgraph Executors
GPU[GPUExecutorAsync]
TPU[TPUExecutorAsync]
XPU[XPUExecutorAsync]
end
subgraph Workers
GPUWorker[GPUWorker]
end
subgraph Model Runners
EmbeddingModelRunner[EmbeddingModelRunner]
GPUModelRunner[GPUModelRunner]
end
subgraph control plane
Scheduling[Scheduling]
end
Request[prompt, sampling_params] --> AsyncLLMEngine
AsyncLLMEngine --> |add_request_async| AsyncLogitProcessor[AsyncLogitProcessorList]
AsyncLogitProcessor --> Scheduling --> Executors
GPU --> GPUWorker --> GPUModelRunner --> |.execute_model| ModelClassImpl[LlamaModelForCausalLM]
some related items
Worker base:
vllm/worker/worker_base.py
Initialize GPU cache and sequence group in ModelRunner step
Executor will handle all KVCache, block manager, and evictor layer here during model execution
broadcast SPMD with sequence groups
proposal
The following document describes and summarizes existing works in vLLM to improve general guided decoding performance. 3
This design will largely affect how
logit_processor
are currently being handle within the vLLM architecture.Main mega thread: vllm-project/vllm#5423
Goal:
- Improve general TPS when using guided decoding.
- Standardize logit processor interface 4
- separate compute_logits and preparing logits into two separate steps
Orthogonal, but still goals:
- vllm-project/vllm#5006
- Logit processor plugins, similar to how vLLM plugins are handled. vllm-project/vllm#4769
- https://github.com/mlc-ai/xgrammar
Scope:
logit_processor
,sampling controller interface
background
reference: vllm-project/vllm#5329
Currently, generations with FSM is super slow, even with warmup steps to initialize given FSM. This behaviour is further exemplified when running with context longer than 4096 tokens.
Additionally, all outlines logit processors are considered stateful, which slows down the model executor, given in V0 logit processors are applied row-by-row blocking
Thus comparing to sglang, vLLM v0 is currently not up to par.
plan
- Implement jump-ahead decoding through a JSONWorker, we can then extend this to CFGWorker
- similar to how spec decode is currently implemented in V0
@cadedaniel: “tree scoring in [spec decode] could use the same API as multi-path jump decoding.”
How should we handle FSM per requests?
- Currently, users can specify different schemas per request, which means the FSM will be compiled per request. This is suboptimal because it slows down general TTFT.
- For most use cases, we should assume JSON schema similar to how the system prompt is currently being handled (pass during server init)
Why should we follow the plugins system?
- If going with the best options, then what is the reasoning behind supporting different backends?
- Agree for extensibility, but seems to add additional overhead.
appendix.
The following includes background information about guided generations.
batched constrained decoding using pushdown automaton
Implemented in mlc-ai/xgrammar
Quote
calculate adaptive token bit-mask per batch
IMPORTANT
operating on string level, not
token_id
GrammarMatcher
⇒ FSM in xgrammarquestions
- byte-level automaton
overhead of token_id ⇒ string
Token for context-independent tokens vs dependent tokens within the generation masks
async pre-compile
synchronize apply mask for CPU → GPU?
How do we apply said masks to GPU block? Zero-overhead generations?
worst-case scenario for grammar compilation?
mask gen overhead: 36
time linearly increase for batch size?
parallelize for compilation.
do we need to parallelize on vLLM?
no, xgrammar parallelize it, with
pthread
shape of masks?
bitmask, tensors of vocab size ⇒ concat with recast ⇒ GPU
supported tokenizers?
GLM yet to be supported (Nov 22nd)
Given that detokenizer is in a separate process with vLLM, then can we stops duplicating this process?
Currently with
xgrammar
: detokenizer included in mask generations.token_id ⇒ tokens
future plans
- Function calling support
- Support more grammar (CFG, Python grammar)
compressed FSM for jump-ahead tokens.
Implemented in (Zheng et al., 2024)
Method 1: FSM-based decoding
intuition: Using FSM (Willard & Louf, 2023) to guide generations by increasing logit bias for tokens that conform to given JSON schema. This allows us to track the current state during decoding and filter out invalid tokens by applying logit bias to the output.
limitation: we can see that given construction of FSM requires token-level access, it can only transition the state by only one token at a time, resulting in slow decoding.
Method 2: Interleaved-based
intuition: breaks down JSON schemas, each containing either a chunk prefill part or constrained decoding part. They are then executed interleaved by inference system. Faster than per-token decoding given that chunked prefill components can process multiple tokens per forward pass
See also https://github.com/guidance-ai/guidance#guidance-acceleration using llama.cpp as backend.
limitation:
- interleaved-based require custom syntax, making it less expressive compared to regex.
- struggles to deal with tokenization boundaries due to conflicts between decode and chunked prefill segments.
- frequent communications between interpreter and back-end adds additional overhead.
Method 3: Jump-Forward Decoding with compressed FSM
tokenization boundary handling
During decoding, it is preferred to combine multiple characters into a single tokens.
For example, when decoding
"Hello"
in context of JSON decoding, LLM might output the following token"
,He
,llo
,",
This may cause some strange behaviour if we combine the last
"
with,
(this regex"[\w\d\s]*"
with the last,
will lead to endless decoding because this token",
is not valid even if the LM wants to stop.)Fix:
- implement re-tokenization mechanism during jump-forward phase (append string instead of the tokens, followed with re-tokenization of the entire text) add approximately 4% of overhead
- use a comprehensive regex to guide the decoding phase, instead of employing multiple concatenated regex 5
Coalescence
intuition: Instead of expanding to state, we can compress certain chunks into one state to reduce the size of said FSM.
figure 1: initial FSM state
figure 2: compressed FSM state
A way to adapt character regex to work with tokens in
outlines
:stateDiagram-v2 [*] --> InputPrompt: Start state "input prompt" as InputPrompt state "next-token probability distribution" as GetProb state "valid tokens" as ListTokens { [*] --> CheckTransitions CheckTransitions --> FilterTokens: Get index[0].keys() FilterTokens --> [*] } state "Sample Token" as SampleToken state "Update FSM State" as UpdateState InputPrompt --> GetProb: "model.generate" GetProb --> ListTokens: Get next-token distribution ListTokens --> SampleToken: Use filtered token list SampleToken --> UpdateState: Selected token X UpdateState --> [*]: new_state = index[0]["X"]
example
stateDiagram-v2 direction LR 0 --> 2: n 0 --> 1: t 1 --> 2: a 2 --> 4: na 2 --> 3: a 3 --> 5: am 4 --> 6: me 5 --> 6: me 2 --> 6: name 6 --> 7: e 6 --> 8: c 7 --> 9: p 8 --> 9: p 9 --> 11: Paul 9 --> 12: Pa 9 --> 10: Jo 11 --> 13: aul 12 --> 14: ul 10 --> 26: o 26 --> 27: h 27 --> 14: n 13 --> 14: l 14 --> 16: s 14 --> 15: s 15 --> 17: s 16 --> 17: s 17 --> 18: a 17 --> 19: ag 18 --> 20: ge 19 --> 20: e 20 --> 21: 30 20 --> 22: 20 21 --> 24: 2 22 --> 24: 2 22 --> 23: 3 24 --> 25: 0 25 --> [*]
note: each state of FSM represents a forward pass to the LM. In vanilla generation, this is essentially necessary. Thus there is no added overhead of FSM for controlling the generated outputs.
From state 2-6, we observer that there are eight different paths to get the same generations of
name
. We probably don’t need to do this, given that it will all give us resultname
But suffice to say, we can hijack this behaviour to accelerate generations by append either of the following tokens word to currently generated sequence:
- [”name”]
- [”n”, “a”, “m”, “e”]
- [”na”, “m”, “e”]
- [”nam”, “e”]
- [”n”, “am”, “e”]
- [”n”, “ame”]
- [”na”, “me”]
- [”n”, “a”, “me”]
A simplified index can be shown as:
That’s at least a 5x speedup over structured generations, given that out of the 9 tokens, two states are single-state transitions. Therefore we only need to call the model twice!!
difference in sampling distribution
All these paths lead to the same string and the same speedup, however they lead to potentially very different states for the LLM when it reaches state 6. That is, the strings are the same, but each path leads to a different conditional probability distribution in stage 6.
Guided generations with FSM.
(Willard & Louf, 2023), implemented at https://github.com/dottxt-ai/outlines
assumption: we are building against autoregressive transformers models
- Let , where is the power set operator, be subset of multi-token string that ends with tokens .
- Text generation tasks is to draw samples from
Notable sampling methods include greedy decoding (generate tokens recursively with highest probability tokens), beam search (but using heuristic to find the mode of distribution) 6
A pseudocode for sampling procedure is as follow:
Given that we are dealing with finite discrete distribution, we can then compute an un-normalized conditional distribution by applying a boolean mask , which restricts the support of original distribution:
augmentation upon sampling algorithm
finite automaton
We define a finite-state machine, given by 7 where character comprising the strings in are drawn from , i.e:
> FSM making for regular expression
([0-9]*)?\.?[0-9]*
example illustration
For simplicity, let the vocabulary consists of strings
- generations start: FSM in state 0, so it masks “A”, since it wouldn’t accepted by the FSM. Then we only sample ”.”, “42”, “.2”, “1” in this case
- if we sample “.2” then we advance the FSM to state 3. In this case. only “42” and “1” are valid completions, so we mask other values before sampling. If we sample “1” instead, then we advance FSM to state 1, in which case ”.”, “.42”, “.2”, and “1” are valid completions
determinism
Looping through the vocabulary is still the biggest issue. For that, we preprocess the vocabulary using Regex’s FSM and build a index. Thus a proceeding for producing matches starting at any point in the FSM is required.
We define finding sub-sequences of FSM that accept string as follow:
We can then define construction of
Lien vers l'original
Bibliographie
- Yu, G.-I., Jeong, J. S., Kim, G.-W., Kim, S., & Chun, B.-G. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 521–538. https://www.usenix.org/conference/osdi22/presentation/yu
- Lew, A. K., Zhi-Xuan, T., Grand, G., & Mansinghka, V. K. (2023). Sequential Monte Carlo Steering of Large Language Models using Probabilistic Programs. arXiv preprint arXiv:2306.03081 arxiv
- Willard, B. T., & Louf, R. (2023). Efficient Guided Generation for Large Language Models. arXiv preprint arXiv:2307.09702 arxiv
- Zheng, L., Yin, L., Xie, Z., Sun, C., Huang, J., Yu, C. H., Cao, S., Kozyrakis, C., Stoica, I., Gonzalez, J. E., Barrett, C., & Sheng, Y. (2024). SGLang: Efficient Execution of Structured Language Model Programs. arXiv preprint arXiv:2312.07104 arxiv
- Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles.
Remarque
-
The paper and presentation for the paper. Most notable open source implementation is vLLM.
p/s: Actually, I think first implemented in huggingface/tgi ↩
-
Current implementation of logits processor mandates that we gather all logits from hidden state, scale if needed then apply the processors.
reference: vllm-project/vllm#5329
Note that there is also vllm-project/vllm#5006 that improves vLLM’s own Outlines implementations of the FSM where it halves memory transition from Python list to Tensors ↩
-
Benchmark script can be found at vllm-project/vllm#10046.
Current RFC vllm-project/vllm#5423 ↩
-
vllm-project/vllm#6273 proposed a sampling controller interface, but @cadedaniel shares some concerns wrt fast-forward tokens ↩
-
this phenomena is also known as coalescence in structured generations, where it exploit deterministic structures in desired outputs to skip expensive forward pass ↩
-
(Lew et al., 2023) recently proposes a sequential Monte Carlo steering. The idea is to classify causal generations as a posteriori inference problem in a class of discrete probabilistic sequence models.
See also Feynman-Kac transformers models ↩
-
- is a finite set of states
- is a finite alphabet
- is the transition function
- is the start state
- is the set of all accepted states.