while True:
    # CPU 阶段：准备批次
    batch = scheduler.prepare_batch(request_queue)
    
    # GPU 阶段：执行推理
    outputs = model.forward(batch)
    
    # CPU 阶段：处理输出，更新缓存
    scheduler.update_cache_and_queue(outputs)

import torch

# 创建两个非默认 CUDA Stream
compute_stream = torch.cuda.Stream()
prepare_stream = torch.cuda.Stream()

# 在不同 Stream 中并行执行操作
with torch.cuda.stream(compute_stream):
    result = model(input_tensor)

with torch.cuda.stream(prepare_stream):
    next_input = next_batch.cuda()

# 等待所有 Stream 完成
torch.cuda.synchronize()

# 创建高优先级计算 Stream 和低优先级准备 Stream
compute_stream = torch.cuda.Stream(priority=-1)  # 高优先级
prepare_stream = torch.cuda.Stream(priority=0)   # 默认优先级

# 高优先级 Stream 会在资源竞争时获得优先处理

# 分配 pinned memory 用于 CPU-GPU 数据传输
cpu_data = torch.empty(batch_size, seq_len, dtype=torch.long).pin_memory()

# 异步传输到 GPU
with torch.cuda.stream(prepare_stream):
    gpu_data = cpu_data.cuda(non_blocking=True)

def stream_synchronization_example():
    compute_stream = torch.cuda.Stream()
    copy_stream = torch.cuda.Stream()
    
    # 异步数据传输
    with torch.cuda.stream(copy_stream):
        gpu_data = cpu_data.cuda(non_blocking=True)
    
    # 计算依赖于传输的数据
    with torch.cuda.stream(compute_stream):
        # 确保数据传输完成后再进行计算
        compute_stream.wait_stream(copy_stream)
        result = model(gpu_data)
    
    return result

import threading
import queue
import torch
from typing import List, Dict, Any

class AsyncContinuousBatcher:
    def __init__(self, model, num_prepare_threads=4):
        self.model = model
        self.num_prepare_threads = num_prepare_threads
        
        # CUDA Streams
        self.compute_stream = torch.cuda.Stream()
        self.prepare_stream = torch.cuda.Stream()
        
        # 双缓冲区
        self.compute_buffer = BatchBuffer()
        self.prepare_buffer = BatchBuffer()
        
        # 请求队列
        self.request_queue = queue.Queue()
        self.output_queue = queue.Queue()
        
        # 控制变量
        self.running = False
        self.lock = threading.Lock()
        
    def prepare_next_batch(self, buffer: 'BatchBuffer'):
        """在 prepare_stream 中准备下一批次"""
        with torch.cuda.stream(self.prepare_stream):
            # 从请求队列获取新请求
            new_requests = []
            while not self.request_queue.empty():
                try:
                    new_requests.append(self.request_queue.get_nowait())
                except queue.Empty:
                    break
            
            # 构建批次
            buffer.fill(new_requests)
            
            # 更新 KV Cache
            buffer.update_kv_cache()
    
    def compute_current_batch(self, buffer: 'BatchBuffer'):
        """在 compute_stream 中执行当前批次"""
        with torch.cuda.stream(self.compute_stream):
            # 执行模型推理
            outputs = self.model.forward(buffer.get_inputs())
            
            # 返回结果
            return outputs
    
    def run_async(self):
        """主异步执行循环"""
        self.running = True
        
        # 预热：准备第一个批次
        self.prepare_buffer.fill_initial_requests()
        
        while self.running:
            # 交换缓冲区
            with self.lock:
                compute_buf = self.compute_buffer
                prepare_buf = self.prepare_buffer
                self.compute_buffer, self.prepare_buffer = \
                    self.prepare_buffer, self.compute_buffer
            
            # 并行执行：GPU 计算当前批次，CPU 准备下一批次
            compute_future = threading.Thread(
                target=self.compute_current_batch, 
                args=(compute_buf,)
            )
            prepare_future = threading.Thread(
                target=self.prepare_next_batch, 
                args=(prepare_buf,)
            )
            
            compute_future.start()
            prepare_future.start()
            
            # 等待计算完成
            compute_future.join()
            
            # 处理输出
            self.process_outputs(compute_future.result)
    
    def process_outputs(self, outputs):
        """处理模型输出，更新请求状态"""
        for output in outputs:
            request_id = output['request_id']
            token = output['token']
            
            # 更新请求状态
            self.update_request_state(request_id, token)
            
            # 如果请求完成，发送结果
            if output['finished']:
                self.send_response(request_id, output['full_text'])

class BatchBuffer:
    """批次缓冲区，管理输入数据和 KV Cache"""
    
    def __init__(self):
        self.inputs = None
        self.kv_cache = None
        self.metadata = {}
        self.batch_size = 0
    
    def fill(self, requests: List[Dict[str, Any]]):
        """填充批次数据"""
        if not requests:
            return
        
        # 组织输入数据
        input_ids = []
        attention_masks = []
        request_ids = []
        
        for req in requests:
            input_ids.append(req['input_ids'])
            attention_masks.append(req['attention_mask'])
            request_ids.append(req['request_id'])
        
        # 批量处理
        self.inputs = {
            'input_ids': torch.stack(input_ids),
            'attention_mask': torch.stack(attention_masks)
        }
        self.metadata['request_ids'] = request_ids
        self.batch_size = len(requests)
    
    def update_kv_cache(self):
        """更新 KV Cache"""
        # 实现 KV Cache 管理逻辑
        pass
    
    def get_inputs(self):
        """获取批次输入"""
        return self.inputs

class DoubleBufferManager:
    def __init__(self):
        self.buffers = [BatchBuffer(), BatchBuffer()]
        self.current_compute_idx = 0
        self.current_prepare_idx = 1
        self.switch_lock = threading.Lock()
    
    def get_compute_buffer(self):
        """获取用于计算的缓冲区"""
        return self.buffers[self.current_compute_idx]
    
    def get_prepare_buffer(self):
        """获取用于准备的缓冲区"""
        return self.buffers[self.current_prepare_idx]
    
    def switch_buffers(self):
        """切换缓冲区角色"""
        with self.switch_lock:
            self.current_compute_idx, self.current_prepare_idx = \
                self.current_prepare_idx, self.current_compute_idx
    
    def process_compute_output(self, output, prepare_buffer):
        """将计算输出传递给准备缓冲区"""
        # 将预测结果传递给准备缓冲区
        prepare_buffer.process_prev_output(output)

def handle_edge_cases(self):
    """处理异步批处理的边界情况"""
    
    # 1. 初始启动：第一个批次需要预热
    if self.is_first_batch():
        self.prepare_first_batch_sync()
    
    # 2. 请求队列为空：保持缓冲区不为空
    if self.request_queue.empty():
        # 可以选择等待新请求或执行空批次
        if self.should_wait_for_requests():
            self.wait_for_requests()
        else:
            # 执行空批次以保持流水线运行
            self.execute_empty_batch()
    
    # 3. 批次大小变化：动态调整缓冲区
    def adjust_batch_size(self, new_size):
        if new_size != self.current_batch_size:
            # 重新分配缓冲区内存
            self.reallocate_buffers(new_size)
    
    # 4. 错误恢复：异常情况下重置缓冲区
    def reset_on_error(self):
        self.compute_buffer.reset()
        self.prepare_buffer.reset()
        self.scheduler.reset()

class MemoryOptimizedAsyncBatcher(AsyncContinuousBatcher):
    def __init__(self, model, memory_budget_gb=80):
        super().__init__(model)
        self.memory_budget_bytes = memory_budget_gb * 1024**3
        self.current_memory_usage = 0
    
    def optimize_memory_allocation(self, batch_size):
        """根据批次大小优化内存分配"""
        # 计算所需内存
        estimated_memory = self.estimate_batch_memory(batch_size)
        
        # 如果超出预算，减少批次大小
        if self.current_memory_usage + estimated_memory > self.memory_budget_bytes:
            max_batch_size = self.calculate_max_batch_size()
            if batch_size > max_batch_size:
                batch_size = max_batch_size
                print(f"Reducing batch size to {batch_size} due to memory constraints")
        
        return batch_size
    
    def reuse_buffers(self):
        """重用缓冲区内存而不是重新分配"""
        # 清空现有缓冲区而不是释放内存
        self.compute_buffer.clear()
        self.prepare_buffer.clear()
        
        # 重置元数据但保留底层张量
        self.compute_buffer.reset_metadata()
        self.prepare_buffer.reset_metadata()

import time
import torch
from contextlib import contextmanager

class PerformanceProfiler:
    def __init__(self):
        self.events = []
    
    @contextmanager
    def record_event(self, name):
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        
        start_event.record()
        yield
        end_event.record()
        
        self.events.append((name, start_event, end_event))
    
    def get_metrics(self):
        """获取性能指标"""
        torch.cuda.synchronize()
        
        metrics = {}
        for name, start, end in self.events:
            metrics[f"{name}_time_ms"] = start.elapsed_time(end)
        
        return metrics

# 性能分析示例
profiler = PerformanceProfiler()

with profiler.record_event("async_batching"):
    async_batcher.run_async()

with profiler.record_event("sync_batching"):
    sync_batcher.run_sync()

metrics = profiler.get_metrics()
print(f"Async time: {metrics['async_batching_time_ms']:.2f} ms")
print(f"Sync time: {metrics['sync_batching_time_ms']:.2f} ms")
print(f"Speedup: {metrics['sync_batching_time_ms']/metrics['async_batching_time_ms']:.2f}x")

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from datetime import datetime, timedelta

def visualize_timeline(cpu_spans, gpu_spans, title="CPU-GPU Timeline"):
    """可视化 CPU 和 GPU 的时间线"""
    fig, ax = plt.subplots(figsize=(15, 4))
    
    # 绘制 CPU 活动
    for start, end, label in cpu_spans:
        duration = (end - start).total_seconds() * 1000  # 转换为毫秒
        start_ms = (start - cpu_spans[0][0]).total_seconds() * 1000
        rect = patches.Rectangle((start_ms, 1.2), duration, 0.6, 
                                linewidth=1, edgecolor='red', facecolor='lightcoral', 
                                label='CPU Work' if label == 'work' else '')
        ax.add_patch(rect)
    
    # 绘制 GPU 活动
    for start, end, label in gpu_spans:
        duration = (end - start).total_seconds() * 1000
        start_ms = (start - gpu_spans[0][0]).total_seconds() * 1000
        rect = patches.Rectangle((start_ms, 0.2), duration, 0.6, 
                                linewidth=1, edgecolor='green', facecolor='lightgreen',
                                label='GPU Work' if label == 'compute' else '')
        ax.add_patch(rect)
    
    # 计算重叠度
    total_time = max(max((end - start).total_seconds() * 1000 
                        for start, end, _ in cpu_spans + gpu_spans))
    overlap_time = calculate_overlap_time(cpu_spans, gpu_spans)
    overlap_percentage = (overlap_time / total_time) * 100
    
    ax.set_xlim(0, total_time)
    ax.set_ylim(0, 2)
    ax.set_xlabel('Time (ms)')
    ax.set_yticks([0.5, 1.5])
    ax.set_yticklabels(['GPU', 'CPU'])
    ax.set_title(f'{title}\nOverlap: {overlap_percentage:.1f}% ({overlap_time:.1f}ms/{total_time:.1f}ms)')
    ax.legend()
    
    plt.tight_layout()
    return fig

def calculate_overlap_time(cpu_spans, gpu_spans):
    """计算 CPU 和 GPU 的重叠时间"""
    overlap = 0
    for cpu_start, cpu_end, _ in cpu_spans:
        for gpu_start, gpu_end, _ in gpu_spans:
            # 计算时间段重叠
            latest_start = max(cpu_start, gpu_start)
            earliest_end = min(cpu_end, gpu_end)
            if latest_start < earliest_end:
                overlap += (earliest_end - latest_start).total_seconds() * 1000
    return overlap

from vllm import LLM, SamplingParams

# 启用异步调度器（vLLM 0.13.0+）
llm = LLM(
    model="meta-llama/Llama-3.1-70B",
    enable_async_scheduler=True,  # 启用异步调度
    async_scheduler_config={
        "num_prepare_threads": 4,           # CPU 准备线程数
        "enable_double_buffering": True,    # 启用双缓冲
        "cpu_gpu_overlap_target": 0.9,     # 目标重叠度
        "prepare_timeout_ms": 1000,        # 准备超时时间
        "max_prepare_batch_size": 64,      # 最大准备批次大小
    },
    gpu_memory_utilization=0.9,
    tensor_parallel_size=4,  # 4 GPU 并行
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=2048,
)

# 性能对比：异步模式 vs 同步模式
prompts = [
    "Explain asynchronous continuous batching in LLM inference.",
    "What are the benefits of CPU-GPU overlap in inference?",
    # ... 更多提示
]

# 异步模式下性能提升 20-30%
outputs = llm.generate(prompts, sampling_params)

import sglang as sgl

@sgl.function
def async_generation(s, prompt):
    s += sgl.user(prompt)
    s += sgl.assistant(sgl.gen("response", max_tokens=512))

# 启用重叠调度
state = async_generation.run(
    prompt="Explain asynchronous continuous batching",
    backend=sgl.Runtime(
        model_path="meta-llama/Llama-3.1-70B",
        enable_overlap=True,  # 启用重叠调度
        tp_size=4,
        mem_fraction_static=0.85
    )
)

def benchmark_throughput(batcher, load_levels=[10, 50, 100, 200]):
    """在不同负载下测试吞吐量"""
    results = {}
    
    for load in load_levels:
        # 生成测试请求
        test_requests = generate_test_requests(load)
        
        # 测试同步模式
        start_time = time.time()
        sync_results = sync_batcher.process_requests(test_requests)
        sync_time = time.time() - start_time
        
        # 测试异步模式
        start_time = time.time()
        async_results = async_batcher.process_requests(test_requests)
        async_time = time.time() - start_time
        
        # 计算吞吐量
        sync_tps = sum(len(r['tokens']) for r in sync_results) / sync_time
        async_tps = sum(len(r['tokens']) for r in async_results) / async_time
        
        results[load] = {
            'sync_tps': sync_tps,
            'async_tps': async_tps,
            'improvement': (async_tps - sync_tps) / sync_tps * 100
        }
    
    return results

# 基准测试结果示例
benchmark_results = benchmark_throughput(async_batcher)
for load, metrics in benchmark_results.items():
    print(f"Load {load}: Sync {metrics['sync_tps']:.1f}, "
          f"Async {metrics['async_tps']:.1f}, "
          f"Improvement {metrics['improvement']:.1f}%")

def test_double_buffer_switching():
    """测试双缓冲区切换逻辑"""
    manager = DoubleBufferManager()
    
    # 初始状态
    assert manager.current_compute_idx == 0
    assert manager.current_prepare_idx == 1
    
    # 执行一次切换
    manager.switch_buffers()
    assert manager.current_compute_idx == 1
    assert manager.current_prepare_idx == 0
    
    # 再次切换应该回到初始状态
    manager.switch_buffers()
    assert manager.current_compute_idx == 0
    assert manager.current_prepare_idx == 1

def test_stream_synchronization():
    """测试 Stream 同步机制"""
    compute_stream = torch.cuda.Stream()
    prepare_stream = torch.cuda.Stream()
    
    # 创建测试数据
    data = torch.randn(100, 100).cuda()
    
    # 在 prepare_stream 中修改数据
    with torch.cuda.stream(prepare_stream):
        modified_data = data * 2
    
    # 在 compute_stream 中等待并使用数据
    with torch.cuda.stream(compute_stream):
        compute_stream.wait_stream(prepare_stream)
        result = modified_data.sum()
    
    # 确保结果正确
    expected = (data * 2).sum()
    assert torch.allclose(result, expected)

class AsyncBatcherLogger:
    def __init__(self):
        self.logger = logging.getLogger('async_batcher')
        
    def log_buffer_switch(self, compute_buf_id, prepare_buf_id, timestamp):
        """记录缓冲区切换事件"""
        self.logger.info(f"Buffer switch at {timestamp}: "
                        f"compute={compute_buf_id}, prepare={prepare_buf_id}")
    
    def log_stream_event(self, stream_name, event_type, batch_id, timestamp):
        """记录 Stream 事件"""
        self.logger.info(f"Stream {stream_name} {event_type} for batch {batch_id} "
                        f"at {timestamp}")

class AsyncBatchingMonitor:
    def __init__(self):
        self.metrics = {
            'cpu_gpu_overlap': 0.0,
            'scheduler_latency': 0.0,
            'queue_depth': 0,
            'gpu_utilization': 0.0,
            'throughput': 0.0
        }
    
    def collect_metrics(self, async_batcher):
        """收集异步批处理监控指标"""
        # CPU-GPU 重叠度
        self.metrics['cpu_gpu_overlap'] = self.calculate_overlap()
        
        # 调度延迟
        self.metrics['scheduler_latency'] = async_batcher.scheduler.avg_latency
        
        # 队列深度
        self.metrics['queue_depth'] = async_batcher.request_queue.qsize()
        
        # GPU 利用率
        self.metrics['gpu_utilization'] = torch.cuda.utilization()
        
        # 吞吐量
        self.metrics['throughput'] = self.calculate_throughput()
        
        return self.metrics