大模型服务中的并发控制：从信号量到令牌桶的请求限流策略

大模型服务中的并发控制从信号量到令牌桶的请求限流策略一、LLM推理的独特并发挑战大模型推理服务与传统的Web API在后端资源模式上有本质差异(1) GPU显存是固定容量资源——每增加一个并发请求KV Cache占用成比例增长超出显存上限即OOM不存在排队等待的缓冲机制(2) 单个请求的计算粒度极大——一个生成请求可能持续数秒到数分钟取决于生成长度传统的毫秒级请求限流粒度过细(3) 请求的资源消耗严重异构——一个生成100个token的请求消耗的计算和显存远大于生成10个token的请求不能将请求数作为统一的资源度量。因此大模型服务的并发控制需要从传统的请求数限流升级为资源感知的流量整形——基于当前可用的GPU显存和计算能力动态决定是否接纳新请求、是否对现有请求降级、以及如何在请求间公平分配资源。flowchart TB A[请求到达] -- B{并发控制层} B -- C[令牌桶] C -- C1[检查可用令牌br/每个请求消耗令牌数br/∝ max_tokens] B -- D[信号量] D -- D1[检查可用槽位br/当前活跃请求数br/ max_concurrent] B -- E[显存感知] E -- E1[预估KV Cache需求br/当前空闲显存br/ 请求所需] C1 -- F{通过?} D1 -- F E1 -- F F --|是| G[请求入队执行] F --|否| H{降级策略} H --|排队等待| I[等待队列br/FIFO或Priority] H --|快速拒绝| J[HTTP 429br/ Retry-After头] H --|降级执行| K[减少max_tokensbr/使用更小的模型br/或量化版本]二、令牌桶经典算法的LLM适配令牌桶Token Bucket算法是API限流中最经典的模式。在LLM场景下需要对令牌的含义进行适配——一个令牌可以对应一个输出token、一个推理请求、或一分钟的GPU计算时间。关键是令牌消耗量要反映实际的资源成本。import time import threading from collections import deque from typing import Optional, Tuple import asyncio class TokenBucketLimiter: LLM适配的令牌桶限流器。 与标准令牌桶的关键区别 1. 令牌消耗量 ∝ max_tokens而非固定1个令牌/请求 2. 支持部分令牌——短生成可消耗少于1个令牌 3. 预热机制防止冷启动时突发流量 def __init__( self, rate: float, # 令牌生成速率令牌/秒 capacity: int, # 桶容量最大突发令牌数 min_tokens_per_request: float 1.0, # 每请求最小令牌消耗 tokens_per_output_token: float 0.01 # 每个输出token的令牌成本 ): self.rate rate self.capacity capacity self.min_tokens_per_request min_tokens_per_request self.tokens_per_output_token tokens_per_output_token self._tokens float(capacity) # 当前可用令牌 self._last_refill time.monotonic() self._lock threading.Lock() def _refill(self): 根据经过的时间补充令牌。 令牌生成速率 rate (令牌/秒) 补充量 rate * 经过的秒数 now time.monotonic() elapsed now - self._last_refill self._tokens min(self.capacity, self._tokens elapsed * self.rate) self._last_refill now def estimate_tokens_required( self, max_output_tokens: int, prompt_length: int ) - float: 预估请求所需的令牌数。 成本模型 - 基础成本: min_tokens_per_request - 输出成本: tokens_per_output_token * max_output_tokens - prompt成本: 0.001 * prompt_length可选较小权重 Args: max_output_tokens: 请求的最大生成token数 prompt_length: prompt的token长度 Returns: 预估令牌消耗 base_cost self.min_tokens_per_request output_cost self.tokens_per_output_token * max_output_tokens prompt_cost 0.001 * prompt_length # 较小权重 return base_cost output_cost prompt_cost def acquire( self, max_output_tokens: int 100, prompt_length: int 0, timeout: Optional[float] None ) - Tuple[bool, float]: 尝试获取执行请求所需的令牌。 非阻塞模式timeoutNone立即返回结果 阻塞模式timeout0等待直到获得令牌或超时 Args: max_output_tokens: 最大生成长度 prompt_length: prompt长度 timeout: 最大等待时间秒 Returns: (是否获取成功, 等待时间) tokens_required self.estimate_tokens_required( max_output_tokens, prompt_length ) start_time time.monotonic() with self._lock: while True: self._refill() if self._tokens tokens_required: self._tokens - tokens_required elapsed time.monotonic() - start_time return True, elapsed if timeout is None: return False, 0.0 elapsed time.monotonic() - start_time if elapsed timeout: return False, elapsed # 计算需要等待的时间 wait_needed (tokens_required - self._tokens) / self.rate # 在锁外等待time.sleep持有GIL但这里我们想要真正的等待 time.sleep(min(wait_needed, 0.1)) property def available_tokens(self) - float: 当前可用令牌数用于监控。 with self._lock: self._refill() return self._tokens # 异步版本用于FastAPI/Starlette class AsyncTokenBucketLimiter(TokenBucketLimiter): 异步令牌桶限流器。 用于FastAPI等异步框架的中间件。 async def async_acquire( self, max_output_tokens: int 100, prompt_length: int 0, timeout: Optional[float] 5.0 ) - Tuple[bool, float]: 异步获取令牌不阻塞事件循环。 在等待期间允许其他协程执行。 import asyncio tokens_required self.estimate_tokens_required( max_output_tokens, prompt_length ) start_time time.monotonic() while True: with self._lock: self._refill() if self._tokens tokens_required: self._tokens - tokens_required return True, time.monotonic() - start_time if timeout is not None and ( time.monotonic() - start_time timeout ): return False, time.monotonic() - start_time # 异步等待——释放事件循环 await asyncio.sleep(0.05)三、显存感知的动态并发控制单纯的请求数限流无法防止OOM——10个生成10个token的请求占用的显存远小于10个生成1000个token的请求。显存感知的并发控制通过实时追踪KV Cache使用量来决定是否接纳新请求。import torch from dataclasses import dataclass from typing import Dict, List dataclass class GPUResourceState: GPU资源的实时状态快照。 total_memory_mb: float allocated_memory_mb: float reserved_memory_mb: float free_memory_mb: float active_requests: int kv_cache_usage_mb: float class MemoryAwareScheduler: 显存感知的请求调度器。 核心逻辑 1. 追踪每个活跃请求的KV Cache占用 2. 新请求到达时预估其KV Cache需求 3. 仅当 free_memory estimated_kv_cache 时才接纳 4. 维护预占机制——批准请求时预留显存 def __init__( self, max_total_requests: int 32, gpu_memory_reserved_ratio: float 0.9, # GPU显存90%用于KV Cache safety_margin_mb: float 512 # 512MB安全边际 ): self.max_total_requests max_total_requests self.gpu_memory_reserved_ratio gpu_memory_reserved_ratio self.safety_margin_mb safety_margin_mb # 活跃请求追踪 self._active_requests: Dict[str, int] {} # req_id → seq_len self._lock threading.Lock() def estimate_kv_cache_per_token(self, num_layers: int, num_heads: int, head_dim: int, dtype_size: int 2 # FP16 2 bytes ) - float: 估算每个token的KV Cache大小。 KV Cache 2(KV) * num_layers * num_heads * head_dim * dtype_size Args: num_layers: Transformer层数 num_heads: 注意力头数 head_dim: 每头维度 dtype_size: 数据类型字节数 Returns: 每token的KV Cache大小MB kv_size_bytes ( 2 * num_layers * num_heads * head_dim * dtype_size ) return kv_size_bytes / (1024 * 1024) def get_current_kv_cache_usage_mb(self) - float: 计算当前所有活跃请求的KV Cache总量MB。 with self._lock: total_tokens sum(self._active_requests.values()) kv_per_token_mb self.estimate_kv_cache_per_token( num_layers32, num_heads32, head_dim128 ) return total_tokens * kv_per_token_mb def can_accept( self, estimated_max_tokens: int, gpu_state: GPUResourceState ) - Tuple[bool, str]: 判断是否可以接纳新请求。 检查条件按优先级 1. 活跃请求数是否已达上限 2. 预估KV Cache是否有足够显存 Args: estimated_max_tokens: 该请求的预估最大token数 gpu_state: GPU资源状态 Returns: (是否可以接纳, 拒绝原因) with self._lock: if len(self._active_requests) self.max_total_requests: return False, max_concurrent_requests_reached # 计算新请求的KV Cache需求 kv_per_token_mb self.estimate_kv_cache_per_token( num_layers32, num_heads32, head_dim128 ) estimated_kv_mb estimated_max_tokens * kv_per_token_mb # 可用显存扣除安全边际 available_mb ( gpu_state.total_memory_mb * self.gpu_memory_reserved_ratio - gpu_state.allocated_memory_mb - self.safety_margin_mb ) if estimated_kv_mb available_mb: return False, ( finsufficient_memory: need {estimated_kv_mb:.0f}MB, fhave {available_mb:.0f}MB ) return True, ok def register_request(self, request_id: str): 注册新请求预占资源。 with self._lock: self._active_requests[request_id] 0 def update_request_length(self, request_id: str, seq_len: int): 更新请求的当前序列长度。 with self._lock: if request_id in self._active_requests: self._active_requests[request_id] seq_len def release_request(self, request_id: str): 释放请求占用的资源。 with self._lock: self._active_requests.pop(request_id, None)四、多级限流的组合策略实际部署中单级限流无法满足复杂场景需求。推荐的多级组合策略级一网关层用户级别的QPM每分钟请求数限制——使用令牌桶保护下游服务不被单用户冲垮。级二服务层GPU显存感知的动态并发控制——确保不OOM是所有限流保护中优先级最高的。级三请求层基于预估token消耗的令牌分配——短生成优先于长生成长提高整体吞吐。五、总结大模型服务的并发控制是需要多级策略协同的系统工程。三个关键原则(1) 显存保护优先于吞吐优化——OOM是比延迟增加严重得多的事件恢复时间分钟级应在所有限流决策中最优先保护(2) 资源成本异构性是核心挑战——一个请求不是统一的计量单位令牌消耗应正比于max_tokens或预估的GPU时间(3) 降级优于拒绝——当无法按请求的原始参数服务时提供减少max_tokens或使用更小模型的降级选项可在不牺牲可用性的前提下保护系统。