大模型推理后端底座:高并发场景下的请求调度与显存管理实战

cover

一、GPU 显存墙与推理排队:大模型后端的性能瓶颈

大模型应用后端面临的核心挑战与传统 Web 服务截然不同——瓶颈不在 CPU,不在网络 I/O,而在 GPU 显存。一个 70B 参数的模型,仅模型权重就需要约 140GB 显存(FP16 精度),加上 KV Cache 和中间激活值,单次推理可能消耗 160-200GB 显存。即使使用 A100-80GB,也需要 2-4 张卡才能承载一个模型实例。

显存是有限且有状态的资源。当一个推理请求正在执行时,GPU 显存被其 KV Cache 占据,后续请求必须等待显存释放后才能开始推理。这导致了一个独特的问题:推理请求的排队等待时间远超实际推理时间。在某智能客服的线上监控中,平均推理时间为 1.2s,但平均排队等待时间高达 4.8s——用户感知到的 80% 延迟来自排队而非计算。

更严峻的是高并发场景下的显存碎片化问题。不同请求的输入长度差异巨大(10-4000 token),对应的 KV Cache 大小从 20MB 到 8GB 不等。频繁分配和释放不同大小的 KV Cache 块,导致显存碎片化严重——总空闲显存足够,但无法找到连续的大块显存来服务长文本请求。实测数据显示,在混合长短文本的负载下,显存碎片率可达 30%,相当于每 3 张 A100 就有 1 张的显存被浪费。

二、推理调度架构:连续批处理与显存池化的协同机制

解决大模型推理后端性能瓶颈的关键,在于将"逐请求串行推理"升级为"连续批处理 + 显存池化"的协同调度架构。

flowchart TB
    subgraph Gateway[请求接入层]
        Q[请求队列<br/>优先级排序]
        TC[Token 计数器<br/>预估 KV Cache 大小]
    end

    subgraph Scheduler[调度引擎]
        CB[连续批处理<br/>Continuous Batching]
        MP[显存池化管理<br/>PagedAttention]
        PA[前缀缓存<br/>Prefix Caching]
    end

    subgraph GPU[GPU 推理层]
        R1[推理实例-1<br/>A100 #0-#1]
        R2[推理实例-2<br/>A100 #2-#3]
    end

    Q -->|请求入队| TC
    TC -->|预估显存需求| CB
    CB -->|批次组装| MP
    MP -->|显存块分配| PA
    PA -->|KV Cache 复用| R1
    PA -->|KV Cache 复用| R2

    R1 -->|推理完成| CB
    R2 -->|推理完成| CB
    CB -->|释放显存块| MP

    MP -.->|显存不足| Q
    Q -.->|低优先级请求等待| CB

    style CB fill:#e74c3c,color:#fff
    style MP fill:#3498db,color:#fff
    style PA fill:#27ae60,color:#fff

连续批处理(Continuous Batching) 是推理调度的核心优化。传统静态批处理需要等待批次内所有请求完成后才能开始下一批,导致短文本请求被长文本请求拖慢。连续批处理则允许已完成的请求立即退出批次,新请求立即加入,实现"流水线式"推理。实测中,连续批处理将吞吐量提升了 2.3 倍,同时将短文本请求的 P99 延迟降低了 60%。

PagedAttention 显存池化 借鉴了操作系统的虚拟内存分页机制。将 KV Cache 划分为固定大小的"页"(通常 16KB),以页表形式管理逻辑地址到物理显存地址的映射。推理请求按需分配页,无需预分配最大可能的连续显存块,从根本上消除了显存碎片化问题。vLLM 的实测数据显示,PagedAttention 将显存利用率从 60% 提升至 95% 以上。

前缀缓存(Prefix Caching) 针对多轮对话场景优化。同一用户的连续对话共享系统提示词和历史上下文,这些公共前缀的 KV Cache 可以跨请求复用,避免重复计算。在智能客服场景中,前缀缓存将平均推理时间从 1.2s 降至 0.7s,降幅达 42%。

三、生产级推理调度后端代码实现

3.1 请求调度与连续批处理引擎

"""
大模型推理调度引擎 —— 连续批处理 + 优先级调度
核心逻辑:
1. 请求按优先级入队(付费用户 > 免费用户)
2. 调度器根据当前显存余量和队列深度动态组装批次
3. 已完成请求立即释放显存,新请求立即填入
"""
import asyncio
import time
from dataclasses import dataclass, field
from typing import List, Optional
from enum import Enum

class Priority(Enum):
    HIGH = 1     # 付费用户
    NORMAL = 2   # 普通用户
    LOW = 3      # 离线批处理

@dataclass
class InferenceRequest:
    request_id: str
    prompt_tokens: List[int]          # 输入 token 序列
    max_new_tokens: int               # 最大生成 token 数
    priority: Priority                # 优先级
    arrival_time: float               # 到达时间戳
    estimated_kv_pages: int = 0       # 预估 KV Cache 页数
    kv_pages_allocated: int = 0       # 已分配页数
    generated_tokens: int = 0         # 已生成 token 数
    is_completed: bool = False        # 是否完成

@dataclass
class BatchState:
    """当前执行中的批次状态"""
    requests: List[InferenceRequest] = field(default_factory=list)
    total_kv_pages_used: int = 0
    max_kv_pages: int = 0             # GPU 可用总页数

class ContinuousBatchScheduler:
    def __init__(self, max_kv_pages: int, max_batch_size: int = 32):
        self.max_kv_pages = max_kv_pages
        self.max_batch_size = max_batch_size
        self.waiting_queue: List[InferenceRequest] = []
        self.batch = BatchState(max_kv_pages=max_kv_pages)
        self._lock = asyncio.Lock()

    async def submit_request(self, req: InferenceRequest) -> asyncio.Future:
        """提交推理请求到等待队列"""
        # 预估 KV Cache 页数:输入 token + 最大输出 token
        total_tokens = len(req.prompt_tokens) + req.max_new_tokens
        req.estimated_kv_pages = self._estimate_kv_pages(total_tokens)
        req.arrival_time = time.time()

        future = asyncio.get_event_loop().create_future()
        req._future = future

        async with self._lock:
            # 按优先级插入队列(优先级相同时按到达时间排序)
            self.waiting_queue.append(req)
            self.waiting_queue.sort(
                key=lambda r: (r.priority.value, r.arrival_time)
            )
        return future

    def schedule_iteration(self) -> List[InferenceRequest]:
        """
        单次调度迭代:从等待队列中选择可加入批次的请求
        调度策略:优先填满高优先级请求,显存不足时跳过
        """
        scheduled = []
        remaining = []

        for req in self.waiting_queue:
            available_pages = self.max_kv_pages - self.batch.total_kv_pages_used
            # 检查显存和批次大小是否还有余量
            if (req.estimated_kv_pages <= available_pages and
                len(self.batch.requests) < self.max_batch_size):
                req.kv_pages_allocated = req.estimated_kv_pages
                self.batch.requests.append(req)
                self.batch.total_kv_pages_used += req.estimated_kv_pages
                scheduled.append(req)
            else:
                remaining.append(req)

        self.waiting_queue = remaining
        return scheduled

    def on_request_completed(self, req: InferenceRequest):
        """请求完成后立即释放显存页"""
        self.batch.total_kv_pages_used -= req.kv_pages_allocated
        self.batch.requests.remove(req)
        req.is_completed = True
        req._future.set_result(req)

    @staticmethod
    def _estimate_kv_pages(num_tokens: int, page_size: int = 16) -> int:
        """
        预估 KV Cache 所需页数
        每页 16KB,约可存储 16 个 token 的 KV Cache
        """
        return (num_tokens + page_size - 1) // page_size

3.2 显存池化管理与 PagedAttention 集成

"""
显存池化管理器 —— 基于 PagedAttention 的 KV Cache 分配
核心设计:固定大小页分配,消除显存碎片化
"""
from typing import Dict, List, Optional
from dataclasses import dataclass

@dataclass
class KVPage:
    page_id: int
    gpu_ptr: int          # GPU 显存物理地址
    ref_count: int = 0    # 引用计数(支持前缀缓存共享)

class PagedKVCacheManager:
    def __init__(self, total_pages: int, page_size_bytes: int = 16384):
        self.total_pages = total_pages
        self.page_size = page_size_bytes
        # 空闲页链表
        self.free_pages: List[KVPage] = [
            KVPage(page_id=i, gpu_ptr=self._allocate_gpu_memory(page_size_bytes))
            for i in range(total_pages)
        ]
        # 请求 → 已分配页的映射
        self.request_pages: Dict[str, List[KVPage]] = {}
        # 前缀缓存:prefix_hash → 页列表
        self.prefix_cache: Dict[str, List[KVPage]] = {}

    def allocate(self, request_id: str, num_pages: int) -> Optional[List[KVPage]]:
        """
        为请求分配指定数量的页
        如果空闲页不足,返回 None(调度器应将请求放回等待队列)
        """
        if len(self.free_pages) < num_pages:
            return None

        allocated = []
        for _ in range(num_pages):
            page = self.free_pages.pop()
            page.ref_count += 1
            allocated.append(page)

        self.request_pages[request_id] = allocated
        return allocated

    def free(self, request_id: str):
        """释放请求占用的所有页(引用计数归零时回收)"""
        if request_id not in self.request_pages:
            return

        for page in self.request_pages[request_id]:
            page.ref_count -= 1
            if page.ref_count == 0:
                # 无其他请求引用此页,回收至空闲池
                self.free_pages.append(page)

        del self.request_pages[request_id]

    def try_prefix_cache(self, request_id: str,
                         prefix_hash: str) -> Optional[List[KVPage]]:
        """
        尝试复用前缀缓存
        如果命中,共享页的引用计数 +1,仅分配新增 token 的页
        """
        if prefix_hash not in self.prefix_cache:
            return None

        cached_pages = self.prefix_cache[prefix_hash]
        # 共享已有页
        for page in cached_pages:
            page.ref_count += 1

        self.request_pages[request_id] = list(cached_pages)
        return cached_pages

    def _allocate_gpu_memory(self, size: int) -> int:
        """模拟 GPU 显存分配,返回物理地址"""
        # 实际实现中调用 CUDA malloc
        return id(self) + size

四、推理后端的隐性瓶颈:调度开销、显存拷贝与冷启动延迟

大模型推理后端的优化并非止步于连续批处理和显存池化,还有三类容易被忽视的隐性瓶颈。

调度器自身的计算开销。 每次调度迭代需要遍历等待队列、计算显存余量、排序优先级。当等待队列深度超过 1000 时,单次调度耗时可达 5-10ms,在 100ms 的推理步长中占比 5-10%。优化方案是将等待队列按优先级分层,高优先级队列长度控制在 100 以内,确保核心请求的调度延迟低于 1ms。

显存拷贝的带宽瓶颈。 PagedAttention 的页表管理涉及 CPU-GPU 之间的元数据同步。每个推理步骤都需要将页表从 CPU 内存拷贝到 GPU 显存,在 32 个请求、每个请求 200 页的场景下,页表大小约 200KB,拷贝耗时约 0.5ms。虽然绝对值不大,但在高频推理步骤中会累积。优化方案是将页表常驻 GPU 显存,仅同步增量变更。

模型冷启动延迟。 从模型加载到首次推理就绪,需要 30-60 秒(加载权重 + 预热)。在弹性扩缩容场景下,冷启动延迟意味着新扩容的实例在近 1 分钟内无法承接流量。解决方案是维护预热实例池——始终保持 1-2 个已加载模型但未承接流量的实例,扩容时直接将流量切至预热实例,冷启动时间降至 1 秒以内。代价是额外的 GPU 资源闲置成本。

禁用场景:当并发请求量低于 GPU 批处理阈值(通常 4-8 个请求)时,连续批处理的收益有限,反而增加了调度复杂度。此时应使用简单的串行推理模式,减少调度开销。

五、总结

大模型推理后端底座的设计核心,是将 GPU 显存作为一级调度资源来管理。连续批处理解决了串行推理的吞吐瓶颈,PagedAttention 消除了显存碎片化,前缀缓存优化了多轮对话的重复计算。三个机制协同工作,将 GPU 显存利用率从 60% 提升至 95%,吞吐量提升 2-3 倍。工程落地的关键要点:第一,请求调度必须感知显存余量,显存不足时宁可排队也不强行分配;第二,前缀缓存需要设计合理的淘汰策略(LRU 或基于访问频率),避免缓存空间被冷数据占据;第三,预热实例池是弹性扩缩容的必要投入,冷启动延迟是生产环境不可接受的。落地路线建议:先在单模型场景验证连续批处理和 PagedAttention 的基线性能,再引入前缀缓存和优先级调度,最后建设预热实例池支撑弹性扩缩容。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐