vLLM

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

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.

see also introduction slides

tldr: Bottleneck at AsyncLogitProcessor and Scheduling layer, given that this is row-wise operations ²

---
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

The following document describes and summarizes existing works in vLLM to improve general guided decoding performance. ³

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 ⁴
• 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 structured decoding from scheduler, given that we can compute token bitmask and broadcast towards GPU workers

• p1: Implement jump-ahead decoding

@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 xgrammar

questions

• 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 $\mu s$

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) $\to$ add approximately 4% of overhead
• use a comprehensive regex to guide the decoding phase, instead of employing multiple concatenated regex ⁵

Coalescence

intuition: Instead of expanding to $n$ 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:

import outlines.fsm as fsm
from outlines.fsm.regex import make_deterministic_fsm, create_fsm_index_tokenizer
 
new_fsm, _ = make_deterministic_fsm(fsm)
idx, _ = create_fsm_index_tokenizer(new_fsm, tokenizer)

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"]

idx_with_tokens = {
  state: {tokenizer.tokenizer.decode([key]): value for key, value in transitions.items()}
  for state, transitions in idx.items()
}

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 result name

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:

simplified_index = {
    0: {'{"': 2},
    2: {"name": 6},
    6: {'":"': 9},
    9: {'Paul': 14, 'John': 14},
    14: {'","': 17},
    17: {'age': 20},
    20: {'":': 22},
    22: {'20': 24, '30': 24},
    24: {'}': 25},
}

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 $\mathcal{F} \subset \mathcal{P}(\mathcal{V})$, where $\mathcal{P}$ is the power set operator, be subset of multi-token string that ends with tokens $\text{EOS} \in \mathcal{V}$.
• Text generation tasks is to draw samples from $\mathcal{F}$

Notable sampling methods include greedy decoding (generate tokens recursively with highest probability tokens), beam search (but using heuristic to find the mode of distribution) ⁶

A pseudocode for sampling procedure is as follow:

$"\\begin{algorithm}\n\\caption{LLM token sampling}\n\\begin{algorithmic}\n\\Function{sample}{$L$}\n \\State $s \\gets ()$\n \\For{$i \\gets 1, L$}\n \\State $\\alpha \\gets \\text{LM}(s, \\theta)$\n \\State Sample $s \\sim \\text{Categorical}(\\alpha)$\n \\If{$s = \\text{EOS}$}\n \\State \\textbf{break}\n \\EndIf\n \\State $s \\gets \\text{append}(s, s)$\n \\EndFor\n \\State \\Return $s$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"$

Algorithm 2 LLM token sampling

1:function sample($L$)

2:$s \gets ()$

3:for $i \gets 1, L$ do

4:$\alpha \gets \text{LM}(s, \theta)$

5:Sample $s \sim \text{Categorical}(\alpha)$

6:if $s = \text{EOS}$ then

7:break

8:end if

9:$s \gets \text{append}(s, s)$

10:end for

11:

12:return $s$

13:end function

Given that we are dealing with finite discrete distribution, we can then compute an un-normalized conditional distribution by applying a boolean mask $m: \mathcal{P}(\mathcal{V}) \to \{0,1\}^N$, which restricts the support of original distribution:

$$\begin{aligned} \alpha &= \text{LM}(\tilde{S_t}, \theta) \\ \tilde{\alpha} &= m(\tilde{S_t}) \odot \alpha \\ \tilde{s_{t+1}} &\approx \text{Categorial}(\tilde{\alpha}) \end{aligned}$$

augmentation upon sampling algorithm

$"\\begin{algorithm}\n\\caption{token sampling with masking}\n\\begin{algorithmic}\n\\Function{sample}{$L$}\n \\State $s \\gets ()$\n \\For{$i \\gets 1, L$}\n \\State $\\alpha \\gets \\text{LM}(s, \\theta)$\n \\State Construct the mask m($s$)\n \\State $\\tilde{\\alpha} \\gets m \\odot \\alpha$\n \\State Sample $\\tilde{s} \\sim \\text{Categorical}(\\tilde{\\alpha})$\n \\If{$\\tilde{s} = \\text{EOS}$}\n \\State \\textbf{break}\n \\EndIf\n \\State $s \\gets \\text{append}(s, \\tilde{s})$\n \\EndFor\n \\State \\Return $s$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"$

Algorithm 3 token sampling with masking

1:function sample($L$)

2:$s \gets ()$

3:for $i \gets 1, L$ do

4:$\alpha \gets \text{LM}(s, \theta)$

5:Construct the mask m($s$)

6:$\tilde{\alpha} \gets m \odot \alpha$

7:Sample $\tilde{s} \sim \text{Categorical}(\tilde{\alpha})$

8:if $\tilde{s} = \text{EOS}$ then

9:break

10:end if

11:$s \gets \text{append}(s, \tilde{s})$

12:end for

13:

14:return $s$

15:end function

finite automaton

We define a finite-state machine, given by $(Q, \Sigma , \delta, q_0, F)$ ⁷ where character comprising the strings in $\mathcal{V}$ are drawn from $\Sigma$, i.e: $\mathcal{V} \in \mathcal{P}(\Sigma)$

> FSM making for regular expression ([0-9]*)?\.?[0-9]*

example illustration

For simplicity, let the vocabulary $\mathcal{V}$ consists of strings $\{A, ., 42, .2, 1\}$

• 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 $M$ that accept string $v$ as follow:

$"\\begin{algorithm}\n\\caption{Find sub-sequences of the FSM $M$ that accept the string $v$}\n\\begin{algorithmic}\n\\Function{FindSubSequences}{$M, v$}\n \\State $M = (Q, \\Sigma, \\delta, q_0, F)$\n \\State $\\texttt{res} \\gets ()$\n \\For{$r \\in \\delta^{-1}(\\cdot, v_0)$} \\Comment{$\\text{ Loop through states that read } v_0$}\n \\State $p \\gets (r)$\n \\For{$i \\gets 1, |v| - 1$} \\Comment{$\\text{ Walk the FSM}$}\n \\If{$\\delta(r, v_i) = \\emptyset$} \\Comment{$\\text{ The FSM does not read } v_i$}\n \\State $p \\gets ()$\n \\State \\textbf{break} \\Comment{$\\text{ Stop walking and try the next start state}$}\n \\EndIf\n \\State $r \\gets \\delta(r, v_i)$\n \\State $p \\gets \\text{append}(p, r)$\n \\EndFor\n \\State $\\texttt{res} \\gets \\text{append}(\\texttt{res}, p)$\n \\EndFor\n \\State \\Return $\\texttt{res}$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"$

Algorithm 4 Find sub-sequences of the FSM $M$ that accept the string $v$

1:function FindSubSequences($M, v$)

2:$M = (Q, \Sigma, \delta, q_0, F)$

3:$\texttt{res} \gets ()$

4:for $r \in \delta^{-1}(\cdot, v_0)$ do ▷$\text{ Loop through states that read } v_0$

5:$p \gets (r)$

6:for $i \gets 1, |v| - 1$ do ▷$\text{ Walk the FSM}$

7:if $\delta(r, v_i) = \emptyset$ then ▷$\text{ The FSM does not read } v_i$

8:$p \gets ()$

9:break ▷$\text{ Stop walking and try the next start state}$

10:end if

11:$r \gets \delta(r, v_i)$

12:$p \gets \text{append}(p, r)$

13:end for

14:$\texttt{res} \gets \text{append}(\texttt{res}, p)$

15:end for

16:

17:return $\texttt{res}$

18:end function

We can then define construction of $\sigma$

$"\\begin{algorithm}\n\\caption{Construct a map from FSM states to subsets of $\\mathcal{V}$}\n\\begin{algorithmic}\n\\Function{MapStatesToVocab}{$M, \\mathcal{V}$}\n \\State $M = (Q, \\Sigma, \\delta, q_0, F)$\n \\State Initialize the map $\\sigma$ with empty sets for each element in $Q$\n \\For{$v \\in \\mathcal{V}$} \\Comment{$\\text{Loop through the vocabulary}$}\n \\State $Z \\gets \\text{find\\_sub\\_sequences}(M, v)$\n \\For{$z \\in Z$} \\Comment{$\\text{Loop through state sequences accepting } v$}\n \\State $\\sigma(z_0) \\gets \\sigma(z_0) \\cup v$\n \\EndFor\n \\EndFor\n \\State \\Return $\\sigma$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"$

Algorithm 5 Construct a map from FSM states to subsets of $\mathcal{V}$

1:function MapStatesToVocab($M, \mathcal{V}$)

2:$M = (Q, \Sigma, \delta, q_0, F)$

3:Initialize the map $\sigma$ with empty sets for each element in $Q$

4:for $v \in \mathcal{V}$ do ▷$\text{Loop through the vocabulary}$

5:$Z \gets \text{find\_sub\_sequences}(M, v)$

6:for $z \in Z$ do ▷$\text{Loop through state sequences accepting } v$

7:$\sigma(z_0) \gets \sigma(z_0) \cup v$

8:end for

9:end for

10:

11:return $\sigma$

12:end function

Lien vers l'original