vLLM推理引擎深度解析：从PagedAttention原理到高吞吐量部署的工程实践

一、引言：大模型推理的效率瓶颈

大语言模型在服务部署时面临严峻的推理效率挑战。自回归生成需要逐token输出，每个token生成都依赖前序所有token的KV Cache。以LLaMA-2 70B为例，单个请求在2048 token上下文下需要约40GB显存存储KV Cache，远超单GPU容量。传统推理框架（如HuggingFace Transformers）存在三大瓶颈：KV Cache显存碎片化导致利用率低、批处理策略僵化导致GPU利用率低、注意力计算存在大量冗余内存读写。vLLM由加州大学伯克利分校团队开发，通过PagedAttention、连续批处理和优化的CUDA kernel，将推理吞吐量提升了2-4倍，成为大模型部署的事实标准。本文将深入解析vLLM的核心技术原理，并提供Python实现。

二、KV Cache与显存管理

自回归生成的核心瓶颈是KV Cache。每生成一个token，需要将前序所有token的Key和Value向量缓存，避免重复计算。对于一个12层、hidden_size=4096、32头的70B模型，单个token的KV Cache大小为 $2 \times 12 \times 4096 \times 32 \times 2 = 6MB$（fp16）。2048 token的上下文需要12GB KV Cache，一个batch的32个请求需要384GB——即使A100 80GB也远远不够。

传统KV Cache管理为每个请求预分配最大长度的连续内存，导致严重的内存碎片和浪费：短序列浪费大量预分配空间，内存碎片使得无法容纳更多请求。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Dict, Tuple, Optional, Set
from dataclasses import dataclass, field
import math
import os

@dataclass
class KVCacheEntry:
    """单个token的KV Cache条目"""
    layer_idx: int
    key: torch.Tensor
    value: torch.Tensor

@dataclass 
class RequestState:
    """请求状态"""
    request_id: int
    prompt_tokens: List[int]
    generated_tokens: List[int] = field(default_factory=list)
    max_length: int = 2048
    is_finished: bool = False
    
    @property
    def all_tokens(self) -> List[int]:
        return self.prompt_tokens + self.generated_tokens
    
    @property
    def current_length(self) -> int:
        return len(self.all_tokens)


class TraditionalKVCacheManager:
    """传统KV Cache管理器（预分配连续内存）"""
    
    def __init__(self, n_layers: int, n_heads: int, head_dim: int,
                 max_batch_size: int = 32, max_seq_len: int = 2048):
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        
        # 预分配固定大小的KV Cache
        # 每个请求分配max_seq_len的空间
        self.cache_size_per_request = max_seq_len
        
        # 模拟预分配（实际中是GPU显存）
        self.allocated = {}  # request_id -> allocated_size
        self.total_allocated = 0
        self.total_used = 0
        self.wasted = 0
    
    def allocate(self, request_id: int, prompt_len: int) -> bool:
        """为请求分配KV Cache"""
        if len(self.allocated) >= self.max_batch_size:
            return False
        
        # 预分配max_seq_len的空间
        allocated_size = self.max_seq_len
        self.allocated[request_id] = allocated_size
        self.total_allocated += allocated_size
        self.total_used += prompt_len
        self.wasted += (allocated_size - prompt_len)
        
        return True
    
    def free(self, request_id: int):
        """释放请求的KV Cache"""
        if request_id in self.allocated:
            self.total_allocated -= self.allocated[request_id]
            del self.allocated[request_id]
    
    def get_stats(self) -> Dict[str, float]:
        """获取统计信息"""
        utilization = self.total_used / max(self.total_allocated, 1) * 100
        waste_ratio = self.wasted / max(self.total_allocated, 1) * 100
        return {
            'total_allocated': self.total_allocated,
            'total_used': self.total_used,
            'wasted': self.wasted,
            'utilization': utilization,
            'waste_ratio': waste_ratio,
            'num_requests': len(self.allocated)
        }


def test_traditional_cache():
    """测试传统KV Cache管理"""
    manager = TraditionalKVCacheManager(
        n_layers=32, n_heads=32, head_dim=128,
        max_batch_size=32, max_seq_len=2048
    )
    
    # 模拟请求（长度不一）
    requests = [
        (1, 50),   # 短请求
        (2, 100),
        (3, 30),    # 很短
        (4, 500),
        (5, 80),
    ]
    
    for req_id, prompt_len in requests:
        manager.allocate(req_id, prompt_len)
    
    stats = manager.get_stats()
    print("传统KV Cache管理:")
    print(f"  总分配: {stats['total_allocated']:,} tokens")
    print(f"  实际使用: {stats['total_used']:,} tokens")
    print(f"  浪费: {stats['wasted']:,} tokens")
    print(f"  利用率: {stats['utilization']:.1f}%")
    print(f"  浪费率: {stats['waste_ratio']:.1f}%")
    
    # 计算实际显存
    per_token_bytes = 2 * 32 * 32 * 128 * 2  # K+V, layers, heads, dim, fp16
    total_bytes = stats['total_allocated'] * per_token_bytes
    used_bytes = stats['total_used'] * per_token_bytes
    wasted_bytes = stats['wasted'] * per_token_bytes
    
    print(f"\n显存占用 (fp16):")
    print(f"  总分配: {total_bytes / 1024**3:.2f} GB")
    print(f"  实际使用: {used_bytes / 1024**3:.2f} GB")
    print(f"  浪费: {wasted_bytes / 1024**3:.2f} GB")

if __name__ == "__main__":
    test_traditional_cache()

传统KV Cache管理的利用率通常只有10-30%，大量显存被浪费。vLLM的PagedAttention通过将KV Cache分页管理解决了这一问题。

三、PagedAttention：虚拟内存思想应用于KV Cache

PagedAttention的核心思想借鉴操作系统的虚拟内存管理。将KV Cache分为固定大小的block（类似内存页），每个block存储固定数量token的KV。逻辑上连续的KV Cache在物理上可以分散在不连续的block中，通过block table映射逻辑地址到物理地址。

这种设计带来三大优势：消除内存碎片（block大小固定，可以填充任何空隙）、支持变长序列（按需分配block，不多分配）、实现内存共享（不同请求可以共享相同前缀的block，如系统prompt）。

@dataclass
class BlockEntry:
    """KV Cache块"""
    block_id: int
    tokens: List[int] = field(default_factory=list)
    ref_count: int = 0  # 引用计数（用于共享）
    
    @property
    def num_tokens(self) -> int:
        return len(self.tokens)
    
    @property
    def is_full(self) -> int:
        return len(self.tokens) >= self.block_size


class PagedAttentionManager:
    """PagedAttention KV Cache管理器"""
    
    def __init__(self, block_size: int = 16, 
                 total_blocks: int = 1000,
                 n_layers: int = 32, n_heads: int = 32, head_dim: int = 128):
        self.block_size = block_size  # 每个block存储的token数
        self.total_blocks = total_blocks
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = head_dim
        
        # 物理block池
        self.blocks: Dict[int, BlockEntry] = {
            i: BlockEntry(block_id=i) for i in range(total_blocks)
        }
        
        # 空闲block列表
        self.free_blocks: List[int] = list(range(total_blocks))
        
        # 每个请求的block table
        # request_id -> List[block_id]
        self.block_tables: Dict[int, List[int]] = {}
        
        # 请求状态
        self.request_states: Dict[int, RequestState] = {}
        
        # 统计
        self.allocated_blocks = 0
        self.shared_blocks = 0
    
    def allocate_request(self, request_id: int, prompt_tokens: List[int],
                        max_length: int = 2048) -> bool:
        """为新请求分配KV Cache"""
        if not self.free_blocks:
            return False
        
        # 计算需要的block数
        num_tokens = len(prompt_tokens)
        num_blocks_needed = math.ceil(num_tokens / self.block_size)
        
        if num_blocks_needed > len(self.free_blocks):
            return False
        
        # 分配block
        allocated = []
        for _ in range(num_blocks_needed):
            block_id = self.free_blocks.pop(0)
            self.blocks[block_id].ref_count = 1
            allocated.append(block_id)
        
        # 填充token到block
        token_idx = 0
        for block_id in allocated:
            block = self.blocks[block_id]
            end_idx = min(token_idx + self.block_size, num_tokens)
            block.tokens = prompt_tokens[token_idx:end_idx]
            token_idx = end_idx
        
        # 记录block table
        self.block_tables[request_id] = allocated
        self.request_states[request_id] = RequestState(
            request_id=request_id,
            prompt_tokens=prompt_tokens,
            max_length=max_length
        )
        
        self.allocated_blocks += num_blocks_needed
        
        return True
    
    def append_token(self, request_id: int, token: int):
        """为请求追加一个token"""
        if request_id not in self.block_tables:
            return
        
        request = self.request_states[request_id]
        request.generated_tokens.append(token)
        
        # 找到最后一个block
        block_table = self.block_tables[request_id]
        last_block_id = block_table[-1]
        last_block = self.blocks[last_block_id]
        
        if last_block.num_tokens < self.block_size:
            # 当前block还有空间
            last_block.tokens.append(token)
        else:
            # 需要新block
            if not self.free_blocks:
                request.is_finished = True
                return
            
            new_block_id = self.free_blocks.pop(0)
            self.blocks[new_block_id].ref_count = 1
            self.blocks[new_block_id].tokens = [token]
            block_table.append(new_block_id)
            self.allocated_blocks += 1
    
    def free_request(self, request_id: int):
        """释放请求的KV Cache"""
        if request_id not in self.block_tables:
            return
        
        for block_id in self.block_tables[request_id]:
            block = self.blocks[block_id]
            block.ref_count -= 1
            if block.ref_count <= 0:
                block.tokens = []
                self.free_blocks.append(block_id)
                self.allocated_blocks -= 1
        
        del self.block_tables[request_id]
        del self.request_states[request_id]
    
    def share_prefix(self, request_id: int, prefix_tokens: List[int],
                    existing_request_id: int = None) -> bool:
        """共享前缀KV Cache（如系统prompt）"""
        if existing_request_id and existing_request_id in self.block_tables:
            # 从已有请求共享前缀block
            existing_table = self.block_tables[existing_request_id]
            
            # 找到共享的前缀block数
            num_shared_blocks = 0
            shared_tokens = 0
            for block_id in existing_table:
                block = self.blocks[block_id]
                if shared_tokens + len(block.tokens) <= len(prefix_tokens):
                    num_shared_blocks += 1
                    shared_tokens += len(block.tokens)
                else:
                    break
            
            if num_shared_blocks > 0:
                # 共享这些block
                shared_block_ids = existing_table[:num_shared_blocks]
                for bid in shared_block_ids:
                    self.blocks[bid].ref_count += 1
                
                self.block_tables[request_id] = list(shared_block_ids)
                self.allocated_blocks += num_shared_blocks
                self.shared_blocks += num_shared_blocks
                
                # 剩余的token需要新分配
                remaining_tokens = prefix_tokens[shared_tokens:]
                if remaining_tokens:
                    num_new_blocks = math.ceil(len(remaining_tokens) / self.block_size)
                    token_idx = 0
                    for _ in range(num_new_blocks):
                        if not self.free_blocks:
                            return False
                        block_id = self.free_blocks.pop(0)
                        self.blocks[block_id].ref_count = 1
                        end_idx = min(token_idx + self.block_size, len(remaining_tokens))
                        self.blocks[block_id].tokens = remaining_tokens[token_idx:end_idx]
                        token_idx = end_idx
                        self.block_tables[request_id].append(block_id)
                        self.allocated_blocks += 1
                
                self.request_states[request_id] = RequestState(
                    request_id=request_id,
                    prompt_tokens=prefix_tokens
                )
                return True
        
        return False
    
    def get_stats(self) -> Dict[str, float]:
        """获取统计信息"""
        total_capacity = self.total_blocks * self.block_size
        used_capacity = sum(
            len(self.blocks[bid].tokens) 
            for bid in range(self.total_blocks) 
            if self.blocks[bid].ref_count > 0
        )
        
        # 去重计算（共享block只算一次）
        unique_used = sum(
            len(self.blocks[bid].tokens)
            for bid in range(self.total_blocks)
            if self.blocks[bid].ref_count > 0
        )
        
        logical_used = sum(
            len(self.blocks[bid].tokens) * self.blocks[bid].ref_count
            for bid in range(self.total_blocks)
            if self.blocks[bid].ref_count > 0
        )
        
        return {
            'total_blocks': self.total_blocks,
            'free_blocks': len(self.free_blocks),
            'allocated_blocks': self.allocated_blocks,
            'shared_blocks': self.shared_blocks,
            'block_size': self.block_size,
            'total_capacity': total_capacity,
            'used_capacity': unique_used,
            'logical_used': logical_used,
            'utilization': unique_used / total_capacity * 100,
            'num_requests': len(self.block_tables)
        }


def test_paged_attention():
    """测试PagedAttention"""
    manager = PagedAttentionManager(
        block_size=16,
        total_blocks=100,
        n_layers=32, n_heads=32, head_dim=128
    )
    
    # 分配请求
    print("=== PagedAttention管理 ===")
    
    # 请求1：长prompt
    manager.allocate_request(1, list(range(50)))
    print(f"请求1 (50 tokens) 分配完成")
    
    # 请求2：短prompt
    manager.allocate_request(2, list(range(30)))
    print(f"请求2 (30 tokens) 分配完成")
    
    # 请求3：共享请求1的前缀
    manager.share_prefix(3, list(range(50)), existing_request_id=1)
    print(f"请求3 (共享请求1前缀) 分配完成")
    
    # 追加token
    manager.append_token(1, 100)
    manager.append_token(2, 200)
    manager.append_token(3, 300)
    print(f"各请求追加1个token")
    
    stats = manager.get_stats()
    print(f"\nPagedAttention统计:")
    print(f"  总block数: {stats['total_blocks']}")
    print(f"  空闲block: {stats['free_blocks']}")
    print(f"  已分配block: {stats['allocated_blocks']}")
    print(f"  共享block: {stats['shared_blocks']}")
    print(f"  利用率: {stats['utilization']:.1f}%")
    print(f"  逻辑使用: {stats['logical_used']} tokens")
    print(f"  物理使用: {stats['used_capacity']} tokens (去重)")
    print(f"  节省(共享): {stats['logical_used'] - stats['used_capacity']} tokens")
    
    # 对比传统方案
    print("\n=== 对比 ===")
    traditional = TraditionalKVCacheManager(
        n_layers=32, n_heads=32, head_dim=128,
        max_batch_size=32, max_seq_len=2048
    )
    traditional.allocate(1, 50)
    traditional.allocate(2, 30)
    traditional.allocate(3, 50)
    
    t_stats = traditional.get_stats()
    print(f"传统方案利用率: {t_stats['utilization']:.1f}%")
    print(f"PagedAttention利用率: {stats['utilization']:.1f}%")
    print(f"PagedAttention额外支持共享: {stats['shared_blocks']} blocks")

if __name__ == "__main__":
    test_paged_attention()

PagedAttention将KV Cache利用率从传统方案的10-30%提升到接近100%，同时支持前缀共享，进一步节省显存。系统prompt可以在多个请求间共享，一个1000 token的系统prompt在100个并发请求中可以节省约200GB的KV Cache显存。

四、连续批处理

传统批处理采用静态批处理：等待一批请求全部到达后开始处理，等待一批全部完成后才接收新请求。这种方式导致GPU空闲——短请求完成后要等待长请求完成才能释放资源。连续批处理（Continuous Batching，又称Iteration-Level Batching）在每个token生成步骤后检查是否有请求完成，完成则立即释放资源并接收新请求，实现GPU持续高利用率。

class ContinuousBatchingScheduler:
    """连续批处理调度器"""
    
    def __init__(self, max_batch_size: int = 32, 
                 max_total_tokens: int = 8192):
        self.max_batch_size = max_batch_size
        self.max_total_tokens = max_total_tokens
        
        # 运行队列
        self.running_queue: List[RequestState] = []
        # 等待队列
        self.waiting_queue: List[RequestState] = []
        # 已完成
        self.completed: List[RequestState] = []
    
    def add_request(self, request: RequestState):
        """添加新请求到等待队列"""
        self.waiting_queue.append(request)
    
    def schedule(self) -> List[RequestState]:
        """调度一批请求"""
        # 从运行队列中移除已完成的
        still_running = []
        for req in self.running_queue:
            if req.is_finished:
                self.completed.append(req)
            else:
                still_running.append(req)
        self.running_queue = still_running
        
        # 计算当前batch的总token数
        current_tokens = sum(req.current_length for req in self.running_queue)
        available_slots = self.max_batch_size - len(self.running_queue)
        available_tokens = self.max_total_tokens - current_tokens
        
        # 从等待队列中添加新请求
        while (self.waiting_queue and 
               available_slots > 0 and available_tokens > 0):
            req = self.waiting_queue.pop(0)
            
            # 检查是否可以加入
            req_tokens = req.current_length
            if req_tokens <= available_tokens:
                self.running_queue.append(req)
                available_slots -= 1
                available_tokens -= req_tokens
            else:
                # 放回等待队列
                self.waiting_queue.insert(0, req)
                break
        
        return self.running_queue
    
    def step(self, generate_fn) -> Dict[str, int]:
        """执行一步生成"""
        batch = self.schedule()
        
        if not batch:
            return {'batch_size': 0, 'total_tokens': 0}
        
        # 为batch中的每个请求生成一个token
        total_tokens_before = sum(req.current_length for req in batch)
        
        for req in batch:
            # 模拟token生成
            new_token = generate_fn(req)
            req.generated_tokens.append(new_token)
            
            # 检查完成条件
            if (req.current_length >= req.max_length or 
                new_token == 2):  # EOS token
                req.is_finished = True
        
        total_tokens_after = sum(req.current_length for req in batch)
        
        return {
            'batch_size': len(batch),
            'total_tokens': total_tokens_after,
            'tokens_generated': total_tokens_after - total_tokens_before + len(batch),
            'completed': sum(1 for r in batch if r.is_finished)
        }
    
    def get_stats(self) -> Dict[str, int]:
        """获取统计信息"""
        return {
            'running': len(self.running_queue),
            'waiting': len(self.waiting_queue),
            'completed': len(self.completed),
            'total_tokens': sum(r.current_length for r in self.running_queue)
        }


def test_continuous_batching():
    """测试连续批处理"""
    scheduler = ContinuousBatchingScheduler(
        max_batch_size=4, max_total_tokens=2048
    )
    
    # 添加请求
    for i in range(10):
        req = RequestState(
            request_id=i,
            prompt_tokens=list(range(20 + i * 5)),
            max_length=50 + i * 10
        )
        scheduler.add_request(req)
    
    # 模拟生成
    import random
    random.seed(42)
    
    def mock_generate(req):
        # 模拟token生成
        if random.random() < 0.1:  # 10%概率生成EOS
            return 2
        return random.randint(3, 1000)
    
    step = 0
    while scheduler.running_queue or scheduler.waiting_queue:
        result = scheduler.step(mock_generate)
        step += 1
        
        if step % 5 == 0 or result['completed'] > 0:
            stats = scheduler.get_stats()
            print(f"Step {step}: batch={result['batch_size']}, "
                  f"running={stats['running']}, waiting={stats['waiting']}, "
                  f"completed={stats['completed']}")
        
        if step > 200:
            print("达到最大步数，停止")
            break
    
    final_stats = scheduler.get_stats()
    print(f"\n最终统计: 完成 {final_stats['completed']} 个请求")

if __name__ == "__main__":
    test_continuous_batching()

五、优化的注意力计算

vLLM使用自定义CUDA kernel实现PagedAttention。在注意力计算时，直接从分散的物理block中读取KV，避免了传统方式中将KV Cache拼接到连续内存的开销。这需要特殊的内存访问模式和kernel设计。

class PagedAttention(nn.Module):
    """PagedAttention的Python实现（实际中用CUDA kernel）"""
    
    def __init__(self, n_heads: int, head_dim: int, block_size: int = 16):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.block_size = block_size
        
        # QKV投影
        self.q_proj = nn.Linear(head_dim * n_heads, head_dim * n_heads, bias=False)
        self.k_proj = nn.Linear(head_dim * n_heads, head_dim * n_heads, bias=False)
        self.v_proj = nn.Linear(head_dim * n_heads, head_dim * n_heads, bias=False)
        self.o_proj = nn.Linear(head_dim * n_heads, head_dim * n_heads, bias=False)
    
    def forward(self, hidden_states: torch.Tensor,
               kv_cache_blocks: Dict[int, torch.Tensor],
               block_table: List[int],
               seq_lens: List[int]) -> torch.Tensor:
        """
        PagedAttention前向传播
        
        参数:
            hidden_states: (batch, 1, hidden_size) 当前token
            kv_cache_blocks: block_id -> (2, block_size, n_heads, head_dim) KV block
            block_table: List[block_id] 逻辑block到物理block的映射
            seq_lens: 每个请求的序列长度
        """
        batch_size = hidden_states.size(0)
        
        # 计算Q, K, V
        q = self.q_proj(hidden_states)  # (batch, 1, hidden)
        k_new = self.k_proj(hidden_states)
        v_new = self.v_proj(hidden_states)
        
        # Reshape为多头
        q = q.view(batch_size, self.n_heads, self.head_dim)
        k_new = k_new.view(batch_size, self.n_heads, self.head_dim)
        v_new = v_new.view(batch_size, self.n_heads, self.head_dim)
        
        # 将新的K, V写入KV Cache block
        # 实际中由PagedAttention manager处理
        
        outputs = []
        for b in range(batch_size):
            seq_len = seq_lens[b]
            
            # 收集KV Cache（通过block table）
            keys = []
            values = []
            
            num_blocks = math.ceil(seq_len / self.block_size)
            for block_idx in range(num_blocks):
                if block_idx < len(block_table):
                    block_id = block_table[block_idx]
                    if block_id in kv_cache_blocks:
                        block_kv = kv_cache_blocks[block_id]
                        # block_kv: (2, block_size, n_heads, head_dim)
                        keys.append(block_kv[0])  # K
                        values.append(block_kv[1])  # V
            
            if keys:
                all_keys = torch.cat(keys, dim=0)[:seq_len]  # (seq_len, n_heads, head_dim)
                all_values = torch.cat(values, dim=0)[:seq_len]
            else:
                all_keys = k_new[b:b+1]
                all_values = v_new[b:b+1]
            
            # 添加新的K, V
            all_keys = torch.cat([all_keys, k_new[b:b+1]], dim=0)
            all_values = torch.cat([all_values, v_new[b:b+1]], dim=0)
            
            # 多头注意力
            q_b = q[b]  # (n_heads, head_dim)
            
            # (n_heads, head_dim) x (seq_len, n_heads, head_dim) -> (n_heads, seq_len)
            attn_scores = torch.einsum('hd,lhd->hl', q_b, all_keys)
            attn_scores = attn_scores / math.sqrt(self.head_dim)
            
            # 因果掩码（只关注前面的token）
            actual_len = all_keys.size(0)
            causal_mask = torch.ones(actual_len, dtype=torch.bool, device=q.device)
            
            attn_weights = F.softmax(attn_scores, dim=-1)
            
            # (n_heads, seq_len) x (seq_len, n_heads, head_dim) -> (n_heads, head_dim)
            out = torch.einsum('hl,lhd->hd', attn_weights, all_values)
            outputs.append(out)
        
        output = torch.stack(outputs)  # (batch, n_heads, head_dim)
        output = output.reshape(batch_size, -1)
        output = self.o_proj(output)
        
        return output


def test_paged_attention_layer():
    """测试PagedAttention层"""
    n_heads = 8
    head_dim = 64
    hidden_size = n_heads * head_dim
    
    attn = PagedAttention(n_heads, head_dim, block_size=16)
    
    # 模拟输入
    batch_size = 2
    hidden_states = torch.randn(batch_size, 1, hidden_size)
    
    # 模拟KV Cache blocks
    kv_blocks = {}
    for i in range(5):
        kv_blocks[i] = torch.randn(2, 16, n_heads, head_dim)
    
    block_table = [0, 1, 2, 3, 4]
    seq_lens = [50, 30]
    
    output = attn(hidden_states, kv_blocks, block_table, seq_lens)
    
    print(f"输入: {hidden_states.shape}")
    print(f"输出: {output.shape}")
    print(f"KV Cache blocks: {len(kv_blocks)} blocks x 16 tokens = {len(kv_blocks)*16} tokens容量")

if __name__ == "__main__":
    test_paged_attention_layer()

六、vLLM服务架构

class vLLMEngine:
    """vLLM推理引擎（简化版）"""
    
    def __init__(self, model_config: Dict, 
                 block_size: int = 16,
                 max_batch_size: int = 32,
                 gpu_memory_utilization: float = 0.9):
        self.model_config = model_config
        self.block_size = block_size
        self.max_batch_size = max_batch_size
        self.gpu_memory_utilization = gpu_memory_utilization
        
        # 组件
        self.kv_cache_manager = PagedAttentionManager(
            block_size=block_size,
            total_blocks=10000,  # 根据GPU显存计算
            n_layers=model_config.get('n_layers', 32),
            n_heads=model_config.get('n_heads', 32),
            head_dim=model_config.get('head_dim', 128)
        )
        
        self.scheduler = ContinuousBatchingScheduler(
            max_batch_size=max_batch_size,
            max_total_tokens=8192
        )
        
        # 模型（实际中加载真实模型）
        self.model = None
        
        # 统计
        self.total_tokens_generated = 0
        self.total_requests = 0
        self.start_time = None
    
    def add_request(self, request_id: int, prompt: List[int],
                   max_length: int = 2048, 
                   temperature: float = 0.7,
                   top_p: float = 0.9):
        """添加生成请求"""
        request = RequestState(
            request_id=request_id,
            prompt_tokens=prompt,
            max_length=max_length
        )
        request.temperature = temperature
        request.top_p = top_p
        
        self.scheduler.add_request(request)
        self.total_requests += 1
    
    def step(self) -> Dict[str, int]:
        """执行一步推理"""
        batch = self.scheduler.schedule()
        
        if not batch:
            return {'batch_size': 0}
        
        results = {}
        
        for req in batch:
            # 模拟token生成
            token = self._generate_token(req)
            req.generated_tokens.append(token)
            
            # 更新KV Cache
            self.kv_cache_manager.append_token(req.request_id, token)
            
            # 检查完成
            if req.current_length >= req.max_length or token == 2:
                req.is_finished = True
                results[req.request_id] = req.all_tokens
            
            self.total_tokens_generated += 1
        
        return {
            'batch_size': len(batch),
            'completed': sum(1 for r in batch if r.is_finished),
            'total_tokens': self.total_tokens_generated
        }
    
    def _generate_token(self, request: RequestState) -> int:
        """生成单个token（模拟）"""
        import random
        # 模拟：10%概率EOS，否则随机token
        if random.random() < 0.05:
            return 2  # EOS
        return random.randint(3, 1000)
    
    def run(self, max_steps: int = 1000):
        """运行引擎"""
        self.start_time = time.time()
        
        for step in range(max_steps):
            result = self.step()
            
            if step % 50 == 0:
                stats = self.get_stats()
                print(f"Step {step}: batch={result['batch_size']}, "
                      f"completed={result.get('completed', 0)}, "
                      f"tokens={result['total_tokens']}, "
                      f"running={stats['running']}, waiting={stats['waiting']}")
            
            # 所有请求完成
            if (not self.scheduler.running_queue and 
                not self.scheduler.waiting_queue):
                print(f"所有请求完成于 step {step}")
                break
        
        elapsed = time.time() - self.start_time
        throughput = self.total_tokens_generated / elapsed if elapsed > 0 else 0
        
        print(f"\n=== 最终统计 ===")
        print(f"总请求数: {self.total_requests}")
        print(f"总生成token: {self.total_tokens_generated}")
        print(f"总耗时: {elapsed:.2f}s")
        print(f"吞吐量: {throughput:.0f} tokens/s")
    
    def get_stats(self) -> Dict[str, int]:
        """获取统计"""
        return self.scheduler.get_stats()


import time

def test_vllm_engine():
    """测试vLLM引擎"""
    model_config = {
        'n_layers': 32,
        'n_heads': 32,
        'head_dim': 128,
        'vocab_size': 32000
    }
    
    engine = vLLMEngine(
        model_config=model_config,
        block_size=16,
        max_batch_size=8
    )
    
    # 添加请求
    for i in range(20):
        prompt = list(range(20 + i * 3))
        engine.add_request(
            request_id=i,
            prompt=prompt,
            max_length=50 + i * 5
        )
    
    # 运行
    engine.run(max_steps=500)

if __name__ == "__main__":
    test_vllm_engine()

七、性能对比与优化策略

def benchmark_comparison():
    """性能对比"""
    # 模拟基准测试结果（基于vLLM论文数据）
    benchmarks = [
        {
            'model': 'LLaMA-7B',
            'framework': 'HuggingFace',
            'throughput': 24.0,
            'latency_ms': 141,
            'gpu_memory': '14.5 GB'
        },
        {
            'model': 'LLaMA-7B',
            'framework': 'TGI',
            'throughput': 42.0,
            'latency_ms': 85,
            'gpu_memory': '13.2 GB'
        },
        {
            'model': 'LLaMA-7B',
            'framework': 'vLLM',
            'throughput': 72.0,
            'latency_ms': 52,
            'gpu_memory': '12.8 GB'
        },
        {
            'model': 'LLaMA-13B',
            'framework': 'HuggingFace',
            'throughput': 12.0,
            'latency_ms': 280,
            'gpu_memory': '27.1 GB'
        },
        {
            'model': 'LLaMA-13B',
            'framework': 'vLLM',
            'throughput': 38.0,
            'latency_ms': 98,
            'gpu_memory': '24.5 GB'
        },
    ]
    
    print("推理框架性能对比:")
    print(f"{'模型':<15} {'框架':<15} {'吞吐量':<12} {'延迟':<12} {'显存'}")
    print("-" * 75)
    for b in benchmarks:
        print(f"{b['model']:<15} {b['framework']:<15} "
              f"{b['throughput']:.0f} tok/s   "
              f"{b['latency_ms']}ms       {b['gpu_memory']}")
    
    # 计算提升
    print("\n性能提升:")
    llama7_hf = [b for b in benchmarks if b['model'] == 'LLaMA-7B' and b['framework'] == 'HuggingFace'][0]
    llama7_vllm = [b for b in benchmarks if b['model'] == 'LLaMA-7B' and b['framework'] == 'vLLM'][0]
    
    speedup = llama7_vllm['throughput'] / llama7_hf['throughput']
    latency_reduction = (1 - llama7_vllm['latency_ms'] / llama7_hf['latency_ms']) * 100
    memory_reduction = (1 - float(llama7_vllm['gpu_memory'].split()[0]) / 
                       float(llama7_hf['gpu_memory'].split()[0])) * 100
    
    print(f"  吞吐量提升: {speedup:.1f}x")
    print(f"  延迟降低: {latency_reduction:.1f}%")
    print(f"  显存节省: {memory_reduction:.1f}%")
    
    print("\n优化策略:")
    strategies = [
        ("PagedAttention", "消除KV Cache碎片，利用率从30%提升至95%+"),
        ("连续批处理", "GPU利用率从40%提升至80%+，消除空闲等待"),
        ("前缀共享", "系统prompt在多请求间共享，节省30-50%显存"),
        ("优化的CUDA Kernel", "减少内存读写，注意力计算提速2-3x"),
        ("Tensor并行", "多GPU分摊计算，支持超大模型推理"),
        ("Speculative Decoding", "小模型草拟+大模型验证，延迟降低40%+"),
        ("量化推理", "INT8/INT4量化减少显存占用和加速计算"),
    ]
    for name, desc in strategies:
        print(f"  - {name}: {desc}")

if __name__ == "__main__":
    benchmark_comparison()

八、总结

vLLM通过PagedAttention、连续批处理和优化的CUDA kernel，将大模型推理效率提升了2-4倍，成为大模型部署的事实标准。PagedAttention借鉴操作系统虚拟内存管理思想，将KV Cache分页管理，消除了内存碎片，将显存利用率从30%提升到95%以上。连续批处理通过iteration-level调度，实现了GPU的持续高利用率。前缀共享机制使系统prompt在多个请求间共享，大幅减少显存开销。这些创新的组合效应使vLLM在相同硬件上服务更多并发请求、生成更快响应、支持更大上下文窗口。随着大模型应用规模的扩大，高效推理引擎的优化——从PagedAttention到Speculative Decoding，从量化到分布式推理——将持续成为AI基础设施的关键竞争领域。掌握vLLM的核心技术原理，对于构建高性能大模型服务系统至关重要。