1
# KV Cache 显存占用计算
2
def kv_cache_memory(
3
    num_layers: int,       # 模型层数，如 80
4
    num_heads: int,        # 注意力头数，如 64
5
    head_dim: int,         # 头维度，如 128
6
    seq_len: int,          # 序列长度，如 2048
7
    batch_size: int,       # 批大小
8
    dtype_bytes: float = 2, # FP16 每参数 2 字节
9
) -> int:
10
    """
11
    KV Cache 显存 = 2 (K+V) × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_bytes
12
    """
13
    per_token_per_layer = 2 * num_heads * head_dim * dtype_bytes  # K + V
14
    total = num_layers * per_token_per_layer * seq_len * batch_size
15
    return total
16

17
# LLaMA-2 70B 示例
18
# 80 层, 64 头, 128 维, FP16
19
memory = kv_cache_memory(
20
    num_layers=80, num_heads=64, head_dim=128,
21
    seq_len=2048, batch_size=1
22
)
23
# 约 2.5 GB / 序列
24

25
# 当 batch_size = 32, seq_len = 2048 时
26
# KV Cache 约 80 GB —— 超出 A100-80G 全部显存

1
# Block 结构示意
2
class KVBlock:
3
    """
4
    一个 Block 存储固定数量 token 的 KV Cache
5
    Block 大小 = block_size × num_heads × head_dim × 2 (K+V) × dtype_size
6
    """
7
    block_size: int = 16  # 每个 Block 存储 16 个 token 的 KV
8
    # 物理 Block 在 GPU 显存中可以不连续存储
9

10
class BlockTable:
11
    """
12
    类似操作系统的页表
13
    记录逻辑 Block 到物理 Block 的映射
14
    """
15
    # 请求的逻辑 Block 0 → 物理 Block 42
16
    # 请求的逻辑 Block 1 → 物理 Block 7
17
    # 请求的逻辑 Block 2 → 物理 Block 123
18
    # 物理位置不必连续！
19
    mappings: list[int]

1
def paged_attention(
2
    query: Tensor,           # [num_heads, head_dim]
3
    key_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
4
    value_cache: Tensor,     # [num_blocks, block_size, num_heads, head_dim]
5
    block_table: Tensor,     # [max_num_blocks_per_seq] 逻辑→物理映射
6
    seq_len: int,
7
    block_size: int = 16,
8
):
9
    """
10
    PagedAttention: 在分块的 KV Cache 上执行注意力计算
11

12
    关键：通过 block_table 间接访问物理 KV Block，
13
    实现"逻辑连续但物理不连续"的内存布局
14
    """
15
    num_blocks = (seq_len + block_size - 1) // block_size
16
    output = torch.zeros(num_heads, head_dim)
17

18
    # 累积 softmax 的在线计算（类似 Flash Attention）
19
    max_score = float('-inf')
20
    sum_exp = 0.0
21

22
    for block_idx in range(num_blocks):
23
        # 通过 block_table 获取物理 Block 编号
24
        physical_block = block_table[block_idx]
25

26
        # 从物理 Block 中读取 K, V
27
        k_block = key_cache[physical_block]    # [block_size, num_heads, head_dim]
28
        v_block = value_cache[physical_block]  # [block_size, num_heads, head_dim]
29

30
        # 计算注意力分数
31
        scores = torch.matmul(query, k_block.transpose(-1, -2)) / math.sqrt(head_dim)
32

33
        # 在线 softmax 更新（处理最后一个 Block 的部分填充）
34
        valid_tokens = min(block_size, seq_len - block_idx * block_size)
35
        scores = scores[:valid_tokens]
36

37
        block_max = scores.max()
38
        new_max = max(max_score, block_max)
39

40
        # 重缩放之前的累积结果
41
        sum_exp = sum_exp * math.exp(max_score - new_max)
42
        new_sum = torch.exp(scores - new_max).sum()
43
        output = output * sum_exp * math.exp(max_score - new_max)
44
        output += torch.matmul(
45
            torch.exp(scores - new_max).unsqueeze(-2), v_block[:valid_tokens]
46
        ).squeeze(-2)
47

48
        sum_exp += new_sum
49
        max_score = new_max
50

51
    output /= sum_exp
52
    return output

1
class ContinuousBatcher:
2
    """
3
    连续批处理器：在每次迭代中动态管理请求队列
4
    """
5
    def __init__(self, model, block_manager, max_batch_size):
6
        self.model = model
7
        self.block_manager = block_manager
8
        self.max_batch_size = max_batch_size
9
        self.running_requests = []  # 当前正在处理的请求
10
        self.waiting_queue = []     # 等待队列
11

12
    def step(self):
13
        """执行一步调度"""
14
        # 1. 移除已完成的请求，释放 KV Cache Block
15
        for req in self.running_requests:
16
            if req.is_finished():
17
                self.block_manager.free(req)
18
                self.running_requests.remove(req)
19

20
        # 2. 从等待队列中填入新请求
21
        while (len(self.running_requests) < self.max_batch_size
22
               and self.waiting_queue):
23
            new_req = self.waiting_queue.pop(0)
24
            # 为新请求分配 KV Cache Block
25
            if self.block_manager.can_allocate(new_req):
26
                self.block_manager.allocate(new_req)
27
                self.running_requests.append(new_req)
28

29
        # 3. 对当前所有 running 请求执行一次前向传播
30
        # 注意：不同请求处于不同阶段（prefill/decode）
31
        output = self.model.forward(self.running_requests)
32

33
        # 4. 每个请求独立处理自己的输出
34
        for req, out in zip(self.running_requests, output):
35
            req.append_token(out.next_token)

1
class PrefixCacheManager:
2
    """
3
    管理共享前缀的 KV Cache Block
4
    使用 Copy-on-Write 避免数据竞争
5
    """
6
    def __init__(self, block_size=16):
7
        self.block_size = block_size
8
        self.prefix_blocks = {}  # prefix_hash → [physical_block_ids]
9
        self.ref_counts = {}     # physical_block_id → 引用计数
10

11
    def get_or_create_prefix(self, prefix_tokens):
12
        """获取或创建共享前缀 Block"""
13
        prefix_hash = hash(tuple(prefix_tokens))
14

15
        if prefix_hash in self.prefix_blocks:
16
            # 前缀已存在，增加引用计数
17
            blocks = self.prefix_blocks[prefix_hash]
18
            for block_id in blocks:
19
                self.ref_counts[block_id] += 1
20
            return blocks
21
        else:
22
            # 首次创建，计算 KV Cache 并存储
23
            blocks = self._compute_and_store(prefix_tokens)
24
            self.prefix_blocks[prefix_hash] = blocks
25
            for block_id in blocks:
26
                self.ref_counts[block_id] = 1
27
            return blocks
28

29
    def release_prefix(self, prefix_hash):
30
        """释放前缀引用，引用计数为 0 时回收"""
31
        blocks = self.prefix_blocks[prefix_hash]
32
        for block_id in blocks:
33
            self.ref_counts[block_id] -= 1
34
            if self.ref_counts[block_id] == 0:
35
                self._free_block(block_id)

1
| 场景               | System Prompt 长度 | 并发请求数 | 无缓存 KV 显存 | 有缓存 KV 显存 | 节省比例 |
2
|--------------------|--------------------|------------|----------------|----------------|----------|
3
| ChatGPT 风格       | 500 tokens         | 20         | 10.0 GB        | 2.0 GB         | 80%      |
4
| RAG 应用           | 2000 tokens        | 10         | 20.0 GB        | 6.0 GB         | 70%      |
5
| Agent 多工具       | 3000 tokens        | 8          | 24.0 GB        | 5.0 GB         | 79%      |
6
| 代码助手           | 1000 tokens        | 30         | 30.0 GB        | 4.0 GB         | 87%      |

1
# vLLM 的 PagedAttention CUDA Kernel 设计要点
2
"""
3
1. 合并内存访问（Coalesced Memory Access）
4
   - Block 内的 KV 数据在 GPU 显存中连续存储
5
   - 线程束（Warp）可以高效地合并读取
6

7
2. 减少 Block Table 查找
8
   - Block Table 常驻 GPU 显存
9
   - 查找开销远小于连续内存带来的收益
10

11
3. 支持 Tensor Parallelism
12
   - KV Cache 按 Head 维度切分到多 GPU
13
   - 每个 GPU 只需存储部分 Head 的 KV Cache
14

15
4. 变长序列支持
16
   - 同一 Batch 中不同序列长度不同
17
   - 通过 seq_len 数组标识每个序列的有效长度
18
"""

1
| 系统                    | 模型        | GPU      | 吞吐量 (req/s) | vs HF 加速比 |
2
|------------------------|-------------|----------|----------------|--------------|
3
| HuggingFace Transformers| LLaMA-13B  | A100     | 2.5            | 1.0x         |
4
| TGI (HuggingFace)      | LLaMA-13B  | A100     | 17.0           | 6.8x         |
5
| vLLM                   | LLaMA-13B  | A100     | 60.0           | 24x          |
6
| vLLM                   | LLaMA-70B  | 4xA100   | 8.5            | —            |
7
| TensorRT-LLM           | LLaMA-70B  | 4xA100   | 7.2            | —            |

1
| 系统          | 有效 KV Cache 比例 | GPU 显存利用率 | 最大并发数 |
2
|---------------|-------------------|---------------|-----------|
3
| HuggingFace   | ~20%              | ~25%          | 10        |
4
| TGI           | ~55%              | ~60%          | 28        |
5
| vLLM          | >95%              | >90%          | 54        |

1
| 场景                    | 输入长度 | 输出长度 | vLLM 吞吐优势 |
2
|------------------------|---------|---------|---------------|
3
| 短对话 (Chatbot)        | 128     | 64      | 24x vs HF     |
4
| 长文档摘要              | 2048    | 256     | 20x vs HF     |
5
| 代码生成                | 512     | 512     | 22x vs HF     |
6
| RAG 检索增强            | 1024    | 128     | 18x vs HF     |
7
| 共享 System Prompt      | 500     | 200     | 30x vs HF     |

1
# vLLM LoRA Serve 示例
2
from vllm import LLM, SamplingParams
3

4
llm = LLM(
5
    model="meta-llama/Llama-2-7b-hf",
6
    enable_lora=True,
7
    max_loras=4,            # 最多同时加载 4 个 LoRA
8
    max_lora_rank=16,
9
)
10

11
# 不同的请求可以使用不同的 LoRA
12
outputs = llm.generate(
13
    [
14
        {
15
            "prompt": "翻译为英文",
16
            "lora_request": LoRARequest("translator", 1, "/path/to/translator"),
17
        },
18
        {
19
            "prompt": "写代码",
20
            "lora_request": LoRARequest("coder", 2, "/path/to/coder"),
21
        },
22
        {
23
            "prompt": "通用对话",
24
            "lora_request": None,  # 不使用 LoRA
25
        },
26
    ]
27
)

1
| 特性              | vLLM          | TensorRT-LLM | HuggingFace TGI | SGLang        |
2
|------------------|---------------|--------------|-----------------|---------------|
3
| PagedAttention   | 是 原生        | 否 自有方案   | 否               | 是 RadixAttn   |
4
| Continuous Batch | 是             | 是            | 是               | 是             |
5
| Prefix Caching   | 是             | 是            | 否               | 是 自动        |
6
| Speculative Dec  | 是             | 是            | 否               | 是             |
7
| LoRA 动态加载     | 是             | 否            | 是               | 是             |
8
| 多模态           | 是             | 是            | 是               | 是             |
9
| 开源协议         | Apache 2.0    | Apache 2.0   | Apache 2.0      | Apache 2.0    |
10
| 易用性           | 高 Python     | 中 C++       | 高 Docker       | 高 Python     |
11
| 性能             | 高            | 极高          | 中              | 高            |

1
# 安装
2
# pip install vllm
3

4
from vllm import LLM, SamplingParams
5

6
# 基础用法
7
llm = LLM(model="meta-llama/Llama-2-7b-hf")
8
sampling_params = SamplingParams(
9
    temperature=0.7,
10
    top_p=0.9,
11
    max_tokens=256,
12
)
13
outputs = llm.generate(["你好，请介绍一下自己"], sampling_params)
14
for output in outputs:
15
    print(output.outputs[0].text)

1
# 启动 OpenAI 兼容的 API 服务器
2
python -m vllm.entrypoints.openai.api_server \
3
    --model meta-llama/Llama-2-7b-hf \
4
    --host 0.0.0.0 \
5
    --port 8000 \
6
    --tensor-parallel-size 2 \
7
    --gpu-memory-utilization 0.9 \
8
    --max-model-len 4096 \
9
    --enable-prefix-caching
10

11
# 客户端调用（与 OpenAI API 完全兼容）
12
import openai
13
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
14
response = client.chat.completions.create(
15
    model="meta-llama/Llama-2-7b-hf",
16
    messages=[{"role": "user", "content": "你好"}],
17
)

1
| 参数                      | 说明                         | 推荐值          |
2
|--------------------------|------------------------------|----------------|
3
| gpu-memory-utilization   | GPU 显存利用率上限            | 0.85 - 0.95    |
4
| max-model-len            | 最大序列长度                  | 按需设置        |
5
| tensor-parallel-size     | 张量并行度                    | GPU 数量        |
6
| block-size               | KV Cache Block 大小           | 16（默认）      |
7
| swap-space               | CPU Swap 空间大小（GB）       | 4 - 8          |
8
| enable-prefix-caching    | 是否启用前缀缓存              | True            |
9
| max-num-seqs             | 最大并发序列数                | 64 - 256       |