35. Continuous Batching

# 实时负载监测器核心组件
class DynamicLoadMonitor:
    def __init__(self):
        self.gpu_utilization = 0.0
        self.memory_utilization = 0.0
        self.cpu_utilization = 0.0
        self.request_queue_length = 0
        self.current_batch_size = 0
        self.average_request_length = 0
        self.request_arrival_rate = 0.0
        
        # 历史数据
        self.history_data = []
        self.max_history_size = 1000
    
    def update_load(self, gpu_util, memory_util, cpu_util, queue_length, batch_size, avg_req_len, arrival_rate):
        """更新负载数据"""
        self.gpu_utilization = gpu_util
        self.memory_utilization = memory_util
        self.cpu_utilization = cpu_util
        self.request_queue_length = queue_length
        self.current_batch_size = batch_size
        self.average_request_length = avg_req_len
        self.request_arrival_rate = arrival_rate
        
        # 记录历史数据
        current_time = time.time()
        self.history_data.append({
            "timestamp": current_time,
            "gpu_util": gpu_util,
            "memory_util": memory_util,
            "cpu_util": cpu_util,
            "queue_length": queue_length,
            "batch_size": batch_size,
            "avg_req_len": avg_req_len,
            "arrival_rate": arrival_rate
        })
        
        # 限制历史数据大小
        if len(self.history_data) > self.max_history_size:
            self.history_data = self.history_data[-self.max_history_size:]
    
    def get_load_features(self):
        """获取负载特征"""
        return {
            "gpu_util": self.gpu_utilization,
            "memory_util": self.memory_utilization,
            "cpu_util": self.cpu_utilization,
            "queue_length": self.request_queue_length,
            "current_batch_size": self.current_batch_size,
            "avg_req_len": self.average_request_length,
            "arrival_rate": self.request_arrival_rate,
            "time_of_day": time.localtime().tm_hour
        }
    
    def get_recent_load_trend(self, window=100):
        """获取最近的负载趋势"""
        if len(self.history_data) < window:
            window = len(self.history_data)
        
        recent_data = self.history_data[-window:]
        
        # 计算趋势
        gpu_trend = (recent_data[-1]["gpu_util"] - recent_data[0]["gpu_util"]) / window
        memory_trend = (recent_data[-1]["memory_util"] - recent_data[0]["memory_util"]) / window
        queue_trend = (recent_data[-1]["queue_length"] - recent_data[0]["queue_length"]) / window
        arrival_trend = (recent_data[-1]["arrival_rate"] - recent_data[0]["arrival_rate"]) / window
        
        return {
            "gpu_trend": gpu_trend,
            "memory_trend": memory_trend,
            "queue_trend": queue_trend,
            "arrival_trend": arrival_trend
        }

# 自适应批次大小调整器核心组件
class AdaptiveBatchSizeAdjuster:
    def __init__(self):
        self.load_monitor = DynamicLoadMonitor()
        self.batch_size_predictor = BatchSizePredictor()
        self.min_batch_size = 1
        self.max_batch_size = 1024
        self.current_batch_size = 64
        self.adjustment_interval = 0.1  # 调整间隔（秒）
        self.last_adjustment_time = 0
    
    def update_load(self, gpu_util, memory_util, cpu_util, queue_length, batch_size, avg_req_len, arrival_rate):
        """更新负载数据"""
        self.load_monitor.update_load(gpu_util, memory_util, cpu_util, queue_length, batch_size, avg_req_len, arrival_rate)
    
    def adjust_batch_size(self):
        """调整批次大小"""
        current_time = time.time()
        
        # 检查是否需要调整
        if current_time - self.last_adjustment_time < self.adjustment_interval:
            return self.current_batch_size
        
        # 获取负载特征
        load_features = self.load_monitor.get_load_features()
        load_trend = self.load_monitor.get_recent_load_trend()
        
        # 合并特征
        features = {**load_features, **load_trend}
        
        # 预测最优批次大小
        predicted_batch_size = self.batch_size_predictor.predict(features)
        
        # 确保在合理范围内
        predicted_batch_size = max(self.min_batch_size, min(self.max_batch_size, predicted_batch_size))
        
        # 平滑调整
        self.current_batch_size = int(
            self.current_batch_size * 0.7 + predicted_batch_size * 0.3
        )
        
        # 更新调整时间
        self.last_adjustment_time = current_time
        
        return self.current_batch_size
    
    def get_current_batch_size(self):
        """获取当前批次大小"""
        return self.current_batch_size

# 批次大小预测器
class BatchSizePredictor:
    def __init__(self):
        # 初始化预测模型
        self.model = self._load_model()
    
    def _load_model(self):
        """加载预测模型"""
        # 这里使用简化的模型，实际可能使用更复杂的机器学习模型
        return SimpleBatchSizePredictor()
    
    def predict(self, features):
        """预测最优批次大小"""
        # 使用模型预测
        return self.model.predict(features)

# 简单批次大小预测器
class SimpleBatchSizePredictor:
    def predict(self, features):
        """简单的批次大小预测"""
        # 基于GPU利用率和队列长度的简单预测
        gpu_util = features["gpu_util"]
        queue_length = features["queue_length"]
        avg_req_len = features["avg_req_len"]
        
        # 基础批次大小
        base_batch_size = 64
        
        # 根据GPU利用率调整
        if gpu_util < 0.5:
            base_batch_size *= 1.5
        elif gpu_util > 0.8:
            base_batch_size *= 0.8
        
        # 根据队列长度调整
        if queue_length > 100:
            base_batch_size *= 1.2
        elif queue_length < 10:
            base_batch_size *= 0.9
        
        # 根据请求长度调整
        if avg_req_len > 512:
            base_batch_size *= 0.8
        elif avg_req_len < 64:
            base_batch_size *= 1.2
        
        return int(base_batch_size)

# 面向MoE模型的Continuous Batching优化核心组件
class MoEContinuousBatcher:
    def __init__(self, num_experts, expert_capacity):
        self.num_experts = num_experts
        self.expert_capacity = expert_capacity
        self.expert_queues = [[] for _ in range(num_experts)]  # 每个专家的请求队列
        self.expert_batches = [[] for _ in range(num_experts)]  # 每个专家的当前批次
        self.request_expert_mapping = {}  # request_id -> list of expert_ids
        self.expert_load = [0 for _ in range(num_experts)]  # 每个专家的当前负载
    
    def add_request(self, request, expert_ids):
        """添加请求到专家队列"""
        request_id = request.request_id
        self.request_expert_mapping[request_id] = expert_ids
        
        # 将请求添加到对应的专家队列
        for expert_id in expert_ids:
            self.expert_queues[expert_id].append(request)
    
    def remove_request(self, request_id):
        """从专家队列中移除请求"""
        if request_id not in self.request_expert_mapping:
            return
        
        expert_ids = self.request_expert_mapping[request_id]
        
        # 从对应的专家队列和批次中移除请求
        for expert_id in expert_ids:
            # 从队列中移除
            self.expert_queues[expert_id] = [req for req in self.expert_queues[expert_id] if req.request_id != request_id]
            
            # 从批次中移除
            self.expert_batches[expert_id] = [req for req in self.expert_batches[expert_id] if req.request_id != request_id]
        
        # 清理映射
        del self.request_expert_mapping[request_id]
    
    def build_expert_batches(self):
        """为每个专家构建批次"""
        for expert_id in range(self.num_experts):
            # 如果当前批次未满，从队列中添加请求
            while len(self.expert_batches[expert_id]) < self.expert_capacity and self.expert_queues[expert_id]:
                # 从队列中取出一个请求
                request = self.expert_queues[expert_id].pop(0)
                
                # 添加到批次
                self.expert_batches[expert_id].append(request)
        
        return self.expert_batches
    
    def update_expert_load(self, expert_id, load):
        """更新专家负载"""
        if 0 <= expert_id < self.num_experts:
            self.expert_load[expert_id] = load
    
    def get_expert_load(self):
        """获取所有专家的负载"""
        return self.expert_load
    
    def get_balanced_batch(self):
        """获取负载平衡的批次"""
        # 找出负载最低的专家
        min_load_expert = min(range(self.num_experts), key=lambda x: self.expert_load[x])
        
        return self.expert_batches[min_load_expert]
    
    def clear_batches(self):
        """清空所有批次"""
        self.expert_batches = [[] for _ in range(self.num_experts)]
    
    def get_num_pending_requests(self):
        """获取待处理的请求数量"""
        return sum(len(queue) for queue in self.expert_queues)

# 结合Speculative Decoding的Continuous Batching融合机制核心组件
class SpeculativeContinuousBatcher:
    def __init__(self):
        self.continuous_batcher = ContinuousBatcher()
        self.speculative_decoder = SpeculativeDecoder()
        self.request_speculative_info = {}  # request_id -> speculative_info
    
    def add_request(self, request):
        """添加请求"""
        self.continuous_batcher.add_request(request)
        
        # 初始化Speculative Decoding信息
        self.request_speculative_info[request.request_id] = {
            "predicted_tokens": [],
            "correct_predicted_tokens": 0,
            "total_predicted_tokens": 0,
            "prediction_length": 4  # 默认预测长度
        }
    
    def remove_request(self, request_id):
        """移除请求"""
        self.continuous_batcher.remove_request(request_id)
        
        # 清理Speculative Decoding信息
        if request_id in self.request_speculative_info:
            del self.request_speculative_info[request_id]
    
    async def execute_batch(self):
        """执行批次"""
        # 获取当前批次
        batch = self.continuous_batcher.get_current_batch()
        
        if not batch:
            return []
        
        # 对批次中的请求进行Speculative Decoding
        speculative_results = await self._perform_speculative_decoding(batch)
        
        # 处理Speculative Decoding结果
        final_results = await self._process_speculative_results(batch, speculative_results)
        
        # 更新批次
        self.continuous_batcher.update_batch(batch, final_results)
        
        return final_results
    
    async def _perform_speculative_decoding(self, batch):
        """对批次中的请求进行Speculative Decoding"""
        speculative_results = []
        
        for request in batch:
            request_id = request.request_id
            spec_info = self.request_speculative_info[request_id]
            
            # 动态调整预测长度
            self._adjust_prediction_length(request_id)
            
            # 执行Speculative Decoding
            predicted_tokens, prediction_length = await self.speculative_decoder.predict(
                request, spec_info["prediction_length"]
            )
            
            # 更新Speculative Decoding信息
            spec_info["predicted_tokens"] = predicted_tokens
            spec_info["total_predicted_tokens"] += prediction_length
            
            speculative_results.append({
                "request_id": request_id,
                "predicted_tokens": predicted_tokens,
                "prediction_length": prediction_length
            })
        
        return speculative_results
    
    async def _process_speculative_results(self, batch, speculative_results):
        """处理Speculative Decoding结果"""
        final_results = []
        
        for request, spec_result in zip(batch, speculative_results):
            request_id = request.request_id
            spec_info = self.request_speculative_info[request_id]
            
            # 验证预测的Token
            correct_tokens, actual_tokens = await self._verify_predicted_tokens(request, spec_result["predicted_tokens"])
            
            # 更新正确预测的Token数量
            spec_info["correct_predicted_tokens"] += correct_tokens
            
            # 如果有正确的预测Token，直接使用
            if correct_tokens > 0:
                final_results.append({
                    "request_id": request_id,
                    "generated_tokens": actual_tokens[:correct_tokens],
                    "status": "partial" if len(actual_tokens) < request.max_tokens else "completed"
                })
            else:
                # 没有正确的预测Token，使用正常的Continuous Batching
                actual_tokens = await self.continuous_batcher.generate_tokens([request])
                final_results.append({
                    "request_id": request_id,
                    "generated_tokens": actual_tokens,
                    "status": "partial" if len(actual_tokens) < request.max_tokens else "completed"
                })
        
        return final_results
    
    async def _verify_predicted_tokens(self, request, predicted_tokens):
        """验证预测的Token"""
        # 这里简化实现，实际需要调用模型进行验证
        # 假设前N个Token是正确的
        correct_tokens = min(2, len(predicted_tokens))  # 简化假设
        actual_tokens = predicted_tokens[:correct_tokens] + ["actual_token"]  # 简化假设
        
        return correct_tokens, actual_tokens
    
    def _adjust_prediction_length(self, request_id):
        """动态调整预测长度"""
        spec_info = self.request_speculative_info[request_id]
        
        # 计算预测准确率
        if spec_info["total_predicted_tokens"] > 0:
            accuracy = spec_info["correct_predicted_tokens"] / spec_info["total_predicted_tokens"]
        else:
            accuracy = 0.5
        
        # 根据准确率调整预测长度
        if accuracy > 0.8:
            spec_info["prediction_length"] = min(8, spec_info["prediction_length"] + 1)
        elif accuracy < 0.4:
            spec_info["prediction_length"] = max(2, spec_info["prediction_length"] - 1)

# 简化的Speculative Decoder
class SpeculativeDecoder:
    async def predict(self, request, prediction_length):
        """执行Speculative Decoding"""
        # 简化实现，生成随机预测Token
        predicted_tokens = [f"pred_token_{i}" for i in range(prediction_length)]
        return predicted_tokens, prediction_length

# Block级内存管理器核心组件
class BlockMemoryManager:
    def __init__(self, num_blocks, block_size, device):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.device = device
        
        # 初始化Block内存池
        self.memory_pool = self._create_memory_pool()
        
        # 空闲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_info
    
    def _create_memory_pool(self):
        """创建Block内存池"""
        # 为每个Block分配内存
        # 这里简化实现，实际使用CUDA内存分配
        return [torch.empty(self.block_size, dtype=torch.float16, device=self.device) for _ in range(self.num_blocks)]
    
    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,
                "active": True
            }
        
        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_memory_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,
            "block_size": self.block_size
        }

# Continuous Batching配置示例

continuous_batching:
  # 基础配置
  enabled: true
  min_batch_size: 1
  max_batch_size: 1024
  batch_build_interval: 0.001  # 批次构建间隔（秒）
  
  # 动态负载监测配置
  dynamic_load_monitoring:
    enabled: true
    monitoring_interval: 0.1  # 监测间隔（秒）
    gpu_util_threshold: 0.8  # GPU利用率阈值
    memory_util_threshold: 0.8  # 内存利用率阈值
  
  # 自适应批次大小配置
  adaptive_batch_size:
    enabled: true
    adjustment_interval: 0.1  # 调整间隔（秒）
    use_ml_model: false  # 是否使用机器学习模型
    smoothing_factor: 0.7  # 平滑因子
  
  # MoE模型配置
  moe:
    enabled: false
    num_experts: 8
    expert_capacity: 32
    load_balancing_strategy: "round_robin"  # 可选：round_robin, least_load
  
  # Speculative Decoding融合配置
  speculative_decoding:
    enabled: false
    prediction_length: 4  # 默认预测长度
    min_prediction_length: 2
    max_prediction_length: 8
    accuracy_threshold: 0.5  # 预测准确率阈值
  
  # 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_batch_building: true
    enable_prefetch: true
    enable_batch_fusion: true
  
  # 监控和日志配置
  monitoring:
    enabled: true
    metrics_enabled: true
    tracing_enabled: false
    logging_enabled: true
    log_level: "info"
    log_format: "json"

# Continuous Batching性能测试命令示例

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

# 测试不同批次大小的性能
for batch_size in 32 64 128 256 512 1024; do
    python -m vllm.entrypoints.benchmark_continuous_batching \
        --model meta-llama/Llama-2-7b-hf \
        --num-requests 500 \
        --concurrency 100 \
        --input-len 128 \
        --output-len 128 \
        --batch-size $batch_size \
        --enable-adaptive-batch-size false;
done

# 测试自适应批次大小的影响
for adaptive_enabled in true false; do
    python -m vllm.entrypoints.benchmark_continuous_batching \
        --model meta-llama/Llama-2-7b-hf \
        --num-requests 500 \
        --concurrency 100 \
        --input-len 128 \
        --output-len 128 \
        --enable-adaptive-batch-size $adaptive_enabled;
done

# 测试MoE模型的性能
python -m vllm.entrypoints.benchmark_continuous_batching \
    --model deepseek-ai/DeepSeek-MoE-16B-base \
    --num-requests 1000 \
    --concurrency 100 \
    --input-len 128 \
    --output-len 128 \
    --enable-continuous-batching true \
    --enable-moe-optimization true