34. Token 级调度原理

# Token依赖分析核心组件
class TokenDependencyAnalyzer:
    def __init__(self):
        self.dependency_graphs = {}  # request_id -> dependency_graph
    
    def analyze_dependencies(self, request_id, tokens):
        """分析Token之间的依赖关系"""
        if request_id not in self.dependency_graphs:
            self.dependency_graphs[request_id] = {}
        
        graph = self.dependency_graphs[request_id]
        
        # 构建Token依赖图
        for i, token in enumerate(tokens):
            if i == 0:
                # 第一个Token没有前驱依赖
                graph[token] = set()
            else:
                # 当前Token依赖于前一个Token
                graph[token] = {tokens[i-1]}
        
        return graph
    
    def get_dependency_chain(self, request_id, token):
        """获取Token的依赖链"""
        if request_id not in self.dependency_graphs:
            return []
        
        graph = self.dependency_graphs[request_id]
        if token not in graph:
            return []
        
        chain = [token]
        current = token
        
        # 回溯依赖链
        while graph[current]:
            prev_token = next(iter(graph[current]))
            chain.insert(0, prev_token)
            current = prev_token
        
        return chain
    
    def calculate_dependency_depth(self, request_id, token):
        """计算Token的依赖深度"""
        chain = self.get_dependency_chain(request_id, token)
        return len(chain)
    
    def cleanup_dependencies(self, request_id):
        """清理请求的依赖关系"""
        if request_id in self.dependency_graphs:
            del self.dependency_graphs[request_id]

# 动态优先级管理器核心组件
class DynamicPriorityManager:
    def __init__(self):
        self.priorities = {}  # (request_id, token) -> priority
        self.request_priorities = {}  # request_id -> base_priority
        self.dependency_analyzer = TokenDependencyAnalyzer()
    
    def set_base_priority(self, request_id, priority):
        """设置请求的基础优先级"""
        self.request_priorities[request_id] = priority
    
    def calculate_token_priority(self, request_id, token, tokens_generated, total_tokens):
        """计算Token的动态优先级"""
        # 获取基础优先级
        base_priority = self.request_priorities.get(request_id, 0)
        
        # 计算依赖深度
        dependency_depth = self.dependency_analyzer.calculate_dependency_depth(request_id, token)
        
        # 计算生成进度比例
        progress_ratio = tokens_generated / total_tokens if total_tokens > 0 else 0
        
        # 计算优先级分数
        # 优先级 = 基础优先级 + 依赖深度权重 + 进度权重
        priority = (
            base_priority * 0.5 +
            dependency_depth * 0.3 +
            (1 - progress_ratio) * 0.2
        )
        
        self.priorities[(request_id, token)] = priority
        return priority
    
    def get_token_priority(self, request_id, token):
        """获取Token的优先级"""
        return self.priorities.get((request_id, token), 0)
    
    def get_highest_priority_tokens(self, candidate_tokens, limit=100):
        """获取最高优先级的Token"""
        # 计算每个Token的优先级
        token_priorities = []
        for request_id, token in candidate_tokens:
            priority = self.get_token_priority(request_id, token)
            token_priorities.append((priority, request_id, token))
        
        # 按优先级排序
        token_priorities.sort(reverse=True, key=lambda x: x[0])
        
        # 返回前N个最高优先级的Token
        return [(request_id, token) for _, request_id, token in token_priorities[:limit]]
    
    def cleanup_priorities(self, request_id):
        """清理请求的优先级信息"""
        # 清理Token优先级
        tokens_to_remove = [key for key in self.priorities if key[0] == request_id]
        for key in tokens_to_remove:
            del self.priorities[key]
        
        # 清理请求基础优先级
        if request_id in self.request_priorities:
            del self.request_priorities[request_id]
        
        # 清理依赖关系
        self.dependency_analyzer.cleanup_dependencies(request_id)

# 跨请求Token复用管理器核心组件
class CrossRequestTokenReuseManager:
    def __init__(self, max_cache_size=100000):
        self.token_cache = {}  # token_hash -> (token_data, usage_count, last_access_time)
        self.max_cache_size = max_cache_size
        self.lock = threading.RLock()
    
    def get_token_hash(self, token):
        """计算Token的哈希值"""
        # 使用SHA-256计算Token的哈希值
        return hashlib.sha256(token.encode()).hexdigest()
    
    def cache_token(self, token, token_data):
        """缓存Token数据"""
        with self.lock:
            token_hash = self.get_token_hash(token)
            
            # 更新缓存
            self.token_cache[token_hash] = (
                token_data,
                1,  # 初始使用计数
                time.time()  # 最后访问时间
            )
            
            # 如果缓存超过最大大小，清理最久未使用的Token
            if len(self.token_cache) > self.max_cache_size:
                self._cleanup_old_tokens()
    
    def get_cached_token(self, token):
        """获取缓存的Token数据"""
        with self.lock:
            token_hash = self.get_token_hash(token)
            
            if token_hash not in self.token_cache:
                return None
            
            # 更新使用计数和最后访问时间
            token_data, usage_count, _ = self.token_cache[token_hash]
            self.token_cache[token_hash] = (
                token_data,
                usage_count + 1,
                time.time()
            )
            
            return token_data
    
    def _cleanup_old_tokens(self):
        """清理最久未使用的Token"""
        # 按最后访问时间排序
        sorted_tokens = sorted(
            self.token_cache.items(),
            key=lambda x: x[1][2]  # 按最后访问时间排序
        )
        
        # 清理多余的Token，保留90%的容量
        tokens_to_remove = len(self.token_cache) - int(self.max_cache_size * 0.9)
        for i in range(tokens_to_remove):
            token_hash, _ = sorted_tokens[i]
            del self.token_cache[token_hash]
    
    def invalidate_token(self, token):
        """使Token缓存失效"""
        with self.lock:
            token_hash = self.get_token_hash(token)
            if token_hash in self.token_cache:
                del self.token_cache[token_hash]
    
    def cleanup(self):
        """清理所有缓存"""
        with self.lock:
            self.token_cache.clear()
    
    def get_cache_stats(self):
        """获取缓存统计信息"""
        with self.lock:
            total_size = len(self.token_cache)
            total_usage = sum(usage_count for _, usage_count, _ in self.token_cache.values())
            avg_usage = total_usage / total_size if total_size > 0 else 0
            
            return {
                "total_size": total_size,
                "total_usage": total_usage,
                "avg_usage": avg_usage,
                "max_size": self.max_cache_size
            }

# Token调度器中的跨请求Token共享逻辑
class TokenScheduler:
    def __init__(self):
        self.cross_request_token_manager = CrossRequestTokenReuseManager()
        # 其他初始化...
    
    async def schedule_tokens(self):
        """调度Token执行"""
        while True:
            # 获取候选Token
            candidate_tokens = self._get_candidate_tokens()
            
            # 检查跨请求Token复用
            tokens_to_execute = []
            reused_tokens = []
            
            for request_id, token in candidate_tokens:
                # 检查Token是否可复用
                cached_token = self.cross_request_token_manager.get_cached_token(token)
                if cached_token:
                    # Token可复用，直接使用缓存结果
                    reused_tokens.append((request_id, token, cached_token))
                else:
                    # Token不可复用，需要执行
                    tokens_to_execute.append((request_id, token))
            
            # 处理复用的Token
            for request_id, token, cached_data in reused_tokens:
                await self._process_reused_token(request_id, token, cached_data)
            
            # 调度需要执行的Token
            if tokens_to_execute:
                # 根据优先级排序Token
                prioritized_tokens = self.dynamic_priority_manager.get_highest_priority_tokens(
                    tokens_to_execute, limit=self._max_batch_size
                )
                
                # 执行Token生成
                results = await self._execute_tokens(prioritized_tokens)
                
                # 缓存生成的Token结果
                for (request_id, token), result in zip(prioritized_tokens, results):
                    self.cross_request_token_manager.cache_token(token, result)
                    await self._process_generated_token(request_id, token, result)
            
            # 等待下一次调度
            await asyncio.sleep(0.001)
    
    async def _process_reused_token(self, request_id, token, cached_data):
        """处理复用的Token"""
        # 更新请求状态
        request = self.request_manager.get_request(request_id)
        if request:
            # 更新生成的Token
            request.generated_tokens.append(token)
            request.tokens_generated += 1
            
            # 检查请求是否完成
            if request.tokens_generated >= request.max_tokens:
                await self._complete_request(request_id)
            else:
                # 生成下一个Token候选
                await self._generate_next_token_candidate(request_id)
    
    async def _execute_tokens(self, tokens):
        """执行Token生成"""
        # 构建执行批次
        batch = self._build_execution_batch(tokens)
        
        # 执行模型前向计算
        results = await self.model_runner.run_forward(batch)
        
        return results

# 智能上下文切换管理器核心组件
class IntelligentContextSwitchManager:
    def __init__(self):
        self.switch_history = []  # 上下文切换历史记录
        self.switch_prediction_model = self._load_prediction_model()
        self.optimal_switch_threshold = 0.5  # 初始切换阈值
    
    def _load_prediction_model(self):
        """加载上下文切换预测模型"""
        # 这里使用简化的模型，实际可能使用更复杂的机器学习模型
        return SimpleSwitchPredictor()
    
    def predict_switch_overhead(self, current_request, next_request):
        """预测上下文切换的开销"""
        # 提取特征
        features = self._extract_switch_features(current_request, next_request)
        
        # 使用模型预测开销
        overhead = self.switch_prediction_model.predict(features)
        
        return overhead
    
    def _extract_switch_features(self, current_request, next_request):
        """提取上下文切换特征"""
        return {
            "current_request_length": len(current_request.generated_tokens),
            "next_request_length": len(next_request.generated_tokens),
            "current_request_remaining": current_request.max_tokens - current_request.tokens_generated,
            "next_request_remaining": next_request.max_tokens - next_request.tokens_generated,
            "same_model": current_request.model_name == next_request.model_name,
            "current_request_priority": current_request.priority,
            "next_request_priority": next_request.priority,
            "current_load": self._get_current_system_load()
        }
    
    def should_switch_context(self, current_request, next_request, current_batch):
        """决定是否应该切换上下文"""
        # 预测切换开销
        switch_overhead = self.predict_switch_overhead(current_request, next_request)
        
        # 计算继续执行当前请求的收益
        continue_benefit = self._calculate_continue_benefit(current_request, current_batch)
        
        # 计算切换到下一个请求的收益
        switch_benefit = self._calculate_switch_benefit(next_request, current_batch)
        
        # 计算净收益
        net_benefit = switch_benefit - continue_benefit - switch_overhead
        
        # 根据净收益和阈值决定是否切换
        return net_benefit > self.optimal_switch_threshold
    
    def _calculate_continue_benefit(self, request, current_batch):
        """计算继续执行当前请求的收益"""
        # 简单的收益计算，实际可能更复杂
        return len(current_batch) * 0.1  # 继续执行当前批次的收益
    
    def _calculate_switch_benefit(self, request, current_batch):
        """计算切换到下一个请求的收益"""
        # 简单的收益计算，实际可能更复杂
        return request.priority * 0.5  # 优先级越高，收益越大
    
    def update_switch_threshold(self, actual_overhead, predicted_overhead):
        """根据实际切换开销更新切换阈值"""
        # 计算预测误差
        error = abs(actual_overhead - predicted_overhead)
        
        # 调整切换阈值
        if actual_overhead > predicted_overhead:
            # 实际开销比预测大，提高阈值
            self.optimal_switch_threshold += error * 0.1
        else:
            # 实际开销比预测小，降低阈值
            self.optimal_switch_threshold -= error * 0.05
        
        # 确保阈值在合理范围内
        self.optimal_switch_threshold = max(0.1, min(1.0, self.optimal_switch_threshold))
    
    def record_switch(self, current_request, next_request, actual_overhead):
        """记录上下文切换"""
        # 预测的开销
        predicted_overhead = self.predict_switch_overhead(current_request, next_request)
        
        # 记录切换历史
        self.switch_history.append({
            "timestamp": time.time(),
            "current_request_id": current_request.request_id,
            "next_request_id": next_request.request_id,
            "predicted_overhead": predicted_overhead,
            "actual_overhead": actual_overhead
        })
        
        # 限制历史记录大小
        if len(self.switch_history) > 1000:
            self.switch_history = self.switch_history[-1000:]
        
        # 更新切换阈值
        self.update_switch_threshold(actual_overhead, predicted_overhead)
    
    def _get_current_system_load(self):
        """获取当前系统负载"""
        # 简单的负载计算，实际可能更复杂
        return len(self.request_manager.get_active_requests()) / 1000.0

# Block分配器核心组件
class BlockAllocator:
    def __init__(self, num_blocks, block_size):
        self.num_blocks = num_blocks
        self.block_size = block_size
        
        # 初始化空闲Block列表
        self.free_blocks = list(range(num_blocks))
        
        # 记录Block的分配情况
        self.block_allocation = {}
        # request_id -> list of block indices
        
        # 记录Block的使用情况
        self.block_usage = {}
        # block_index -> usage information
    
    def allocate_blocks(self, request_id, num_blocks_needed):
        """分配Block"""
        if num_blocks_needed <= 0:
            return []
        
        # 检查是否有足够的空闲Block
        if len(self.free_blocks) < num_blocks_needed:
            # 尝试回收一些Block
            self._reclaim_blocks()
            
            # 再次检查
            if len(self.free_blocks) < num_blocks_needed:
                return []
        
        # 分配Block
        allocated_blocks = self.free_blocks[:num_blocks_needed]
        self.free_blocks = self.free_blocks[num_blocks_needed:]
        
        # 更新分配记录
        if request_id not in self.block_allocation:
            self.block_allocation[request_id] = []
        self.block_allocation[request_id].extend(allocated_blocks)
        
        # 更新使用记录
        for block_idx in allocated_blocks:
            self.block_usage[block_idx] = {
                "request_id": request_id,
                "allocated_time": time.time(),
                "last_used_time": time.time(),
                "usage_count": 0
            }
        
        return allocated_blocks
    
    def free_blocks(self, request_id):
        """释放请求的Block"""
        if request_id not in self.block_allocation:
            return
        
        # 释放所有分配给该请求的Block
        for block_idx in self.block_allocation[request_id]:
            # 标记为空闲
            self.free_blocks.append(block_idx)
            
            # 清理使用记录
            if block_idx in self.block_usage:
                del self.block_usage[block_idx]
        
        # 清理分配记录
        del self.block_allocation[request_id]
    
    def update_block_usage(self, block_idx):
        """更新Block的使用情况"""
        if block_idx in self.block_usage:
            self.block_usage[block_idx]["last_used_time"] = time.time()
            self.block_usage[block_idx]["usage_count"] += 1
    
    def _reclaim_blocks(self):
        """回收Block"""
        # 简单的回收策略，实际可能更复杂
        # 这里采用最久未使用的策略
        current_time = time.time()
        
        # 找出最久未使用的Block
        lru_blocks = sorted(
            self.block_usage.items(),
            key=lambda x: x[1]["last_used_time"]
        )
        
        # 回收最多10%的Block
        num_to_reclaim = max(1, int(self.num_blocks * 0.1))
        
        for block_idx, usage_info in lru_blocks[:num_to_reclaim]:
            request_id = usage_info["request_id"]
            
            # 从请求的Block列表中移除
            if request_id in self.block_allocation:
                if block_idx in self.block_allocation[request_id]:
                    self.block_allocation[request_id].remove(block_idx)
            
            # 标记为空闲
            self.free_blocks.append(block_idx)
            
            # 清理使用记录
            del self.block_usage[block_idx]
    
    def get_allocation_stats(self):
        """获取分配统计信息"""
        allocated_blocks = sum(len(blocks) for blocks in self.block_allocation.values())
        free_blocks = len(self.free_blocks)
        
        return {
            "total_blocks": self.num_blocks,
            "allocated_blocks": allocated_blocks,
            "free_blocks": free_blocks,
            "utilization": allocated_blocks / self.num_blocks if self.num_blocks > 0 else 0
        }

# Token级调度配置示例

token_scheduling:
  # 基础配置
  enabled: true
  max_batch_size: 1024
  scheduling_interval: 0.001  # 调度间隔（秒）
  
  # 动态优先级配置
  dynamic_priority:
    enabled: true
    base_priority_weight: 0.5
    dependency_depth_weight: 0.3
    progress_weight: 0.2
  
  # 跨请求Token复用配置
  cross_request_token_reuse:
    enabled: true
    max_cache_size: 100000
    cache_ttl: 3600  # 缓存过期时间（秒）
    token_hash_algorithm: "sha256"
  
  # 智能上下文切换配置
  intelligent_context_switching:
    enabled: true
    optimal_switch_threshold: 0.5
    switch_prediction_model: "linear"
    record_switch_history: true
    max_switch_history: 1000
  
  # Block分配配置
  block_allocation:
    num_blocks: 1024
    block_size: 16  # Token数量
    reclaim_strategy: "lru"  # 可选：lru, fifo, random
    reclaim_ratio: 0.1  # 每次回收的比例
  
  # 性能优化配置
  performance_optimization:
    enabled: true
    enable_async_scheduling: true
    enable_batch_processing: true
    enable_memory_pinning: true
  
  # 监控和日志配置
  monitoring:
    enabled: true
    metrics_enabled: true
    tracing_enabled: false
    logging_enabled: true
    log_level: "info"
    log_format: "json"

# Token级调度性能测试命令示例

# 测试基础性能
python -m vllm.entrypoints.benchmark_token_scheduling \
    --model meta-llama/Llama-2-7b-hf \
    --num-requests 1000 \
    --concurrency 100 \
    --input-len 128 \
    --output-len 128 \
    --enable-token-scheduling true

# 测试不同调度粒度的性能
for granularity in "token" "request" "hybrid"; do
    python -m vllm.entrypoints.benchmark_token_scheduling \
        --model meta-llama/Llama-2-7b-hf \
        --num-requests 500 \
        --concurrency 100 \
        --input-len 128 \
        --output-len 128 \
        --scheduling-granularity $granularity;
done

# 测试跨请求Token复用的影响
for reuse_enabled in true false; do
    python -m vllm.entrypoints.benchmark_token_scheduling \
        --model meta-llama/Llama-2-7b-hf \
        --num-requests 500 \
        --concurrency 100 \
        --input-len 128 \
        --output-len 128 \
        --enable-cross-request-token-reuse $reuse_enabled;
done