# 导入必要的库
from vllm import LLM, SamplingParams

# 初始化模型
llm = LLM(model="/opt/data/models/Llama-3.2-1B-Instruct")

# 设置生成参数
sampling_params = SamplingParams(
    max_tokens=100,         # 最大生成长度
    temperature=0.7,        # 控制随机性
    top_p=0.9,              # 控制取样范围
    stop=["</s>", "\n\n"]   # 停止标记
)

# 准备输入提示
prompts = [
    "介绍一下北京的旅游景点",
    "写一首关于春天的诗",
]

# 生成文本
outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    print(f"Prompt: {output.prompt}")
    print(f"Generated: {output.outputs[0].text}")
    print("-" * 50)

def __init__(
    self,
    model_config: ModelConfig,
    scheduler_config: SchedulerConfig,
    device_config: Optional[DeviceConfig] = None,
    parallel_config: Optional[ParallelConfig] = None,
    **kwargs,
) -> None:
    # 1. 设置基本配置
    self.model_config = model_config
    self.scheduler_config = scheduler_config
    if device_config is None:
        device_config = DeviceConfig()
    self.device_config = device_config
    if parallel_config is None:
        parallel_config = ParallelConfig()
    self.parallel_config = parallel_config
    
    # 2. 初始化策略配置
    self.num_nodes = parallel_config.world_size // parallel_config.tensor_parallel_size
    self.distributed_init_method = kwargs.get("distributed_init_method", None)
    self.max_logprobs = kwargs.get("max_logprobs", None)
    
    # 3. 初始化组件
    self.tokenizer = self._init_tokenizer()
    self.schedulers = self._init_schedulers()
    self.request_tracker = RequestTracker()
    
    # 4. 初始化 Worker 池
    self.worker_pool = WorkerPool(
        model_config=model_config,
        parallel_config=parallel_config,
        scheduler_config=scheduler_config,
        device_config=device_config,
        local_rank=parallel_config.local_rank,
    )
    
    # 5. 启动引擎线程
    self._start_background_thread()

def _init_tokenizer(self) -> PreTrainedTokenizer:
    """初始化 tokenizer"""
    # 加载 tokenizer
    tokenizer = get_tokenizer(
        model_name=self.model_config.model,
        tokenizer_name=self.model_config.tokenizer,
        tokenizer_mode=self.model_config.tokenizer_mode,
        trust_remote_code=self.model_config.trust_remote_code,
        revision=self.model_config.tokenizer_revision,
    )
    
    # 配置特殊 token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    # 检查并处理特殊 token ID
    if (self.model_config.chat_template is not None and
            self.model_config.chat_template != "none"):
        tokenizer.chat_template = self.model_config.chat_template
        
    return tokenizer

def _init_schedulers(self) -> List[Scheduler]:
    """初始化调度器"""
    # 创建多个调度器实例，每个对应一个模型副本
    schedulers = []
    for rank in range(self.num_nodes):
        # 创建调度器实例
        scheduler = Scheduler(
            # 配置调度器
            scheduler_config=self.scheduler_config,
            max_num_batched_tokens=self.worker_pool.max_num_batched_tokens,
            max_num_seqs=self.worker_pool.max_num_seqs,
            device=f"cuda:{self.parallel_config.local_rank}" if torch.cuda.is_available() else "cpu",
        )
        schedulers.append(scheduler)
    return schedulers

class WorkerPool:
    def __init__(
        self,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
        device_config: DeviceConfig,
        local_rank: int,
    ) -> None:
        # 1. 设置基本配置
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
        self.device_config = device_config
        self.local_rank = local_rank
        
        # 2. 初始化其他参数
        self.max_num_seqs = scheduler_config.max_num_seqs
        self.max_num_batched_tokens = scheduler_config.max_num_batched_tokens
        
        # 3. 初始化 Worker
        self.workers: List[Worker] = []
        self.model_replicas = []
        
        # 4. 加载模型并创建 Worker
        self._initialize_workers()

def _initialize_workers(self) -> None:
    """初始化 Worker"""
    # 创建模型副本
    self.model_replicas = self._initialize_model_replicas()
    
    # 为每个模型副本创建一个 Worker
    for i, model_replica in enumerate(self.model_replicas):
        worker = Worker(
            model_config=self.model_config,
            model_replica=model_replica,
            # ... 其他参数
        )
        self.workers.append(worker)

def add_request(
    self,
    request_id: str,
    prompt: Optional[str],
    sampling_params: SamplingParams,
    prompt_token_ids: Optional[List[int]] = None,
    arrival_time: Optional[float] = None,
) -> Optional[int]:
    """添加一个新请求"""
    # 1. 记录请求到达时间
    if arrival_time is None:
        arrival_time = time.time()
    
    # 2. 处理输入
    if prompt is None and prompt_token_ids is None:
        raise ValueError("Either prompt or prompt_token_ids must be provided")
    
    # 3. 如果需要，使用 tokenizer 处理输入文本
    if prompt_token_ids is None:
        assert prompt is not None
        prompt_token_ids = self.tokenizer.encode(prompt)
    
    # 4. 创建请求并添加到队列
    request = Request(
        request_id=request_id,
        prompt=prompt,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        arrival_time=arrival_time,
    )
    
    # 5. 将请求添加到请求跟踪器
    self.request_tracker.add_request(request)
    
    # 6. 将请求添加到调度器
    self.schedulers[0].add_request(request)
    
    return len(prompt_token_ids)

def schedule(self) -> List[SequenceGroupMetadata]:
    """调度请求执行"""
    # 1. 获取系统状态
    now = time.time()
    
    # 2. 从等待队列和运行中队列选择可处理的序列组
    seq_groups = self._get_schedulable_groups()
    
    # 3. 根据调度策略对序列组排序
    seq_groups = self._sort_by_strategy(seq_groups)
    
    # 4. 根据资源约束选择可处理的序列组
    scheduled_groups = self._select_best_fitting_groups(seq_groups)
    
    # 5. 生成调度元数据
    metadata = self._create_schedule_metadata(scheduled_groups)
    
    return metadata

def execute_model(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    kv_caches: List[torch.Tensor],
    block_tables: torch.Tensor,
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
    """执行模型前向传播"""
    # 1. 准备输入
    batch_size = input_ids.shape[0]
    
    # 2. 如果已启用 CUDA 图捕获并且批大小匹配，使用 CUDA 图执行
    if (self.model_captured and batch_size in self.graph_runners and
            not self.disable_graphs):
        # 使用预先捕获的 CUDA 图执行模型
        return self._execute_with_graph(batch_size, input_ids, positions,
                                       kv_caches, block_tables)
    else:
        # 使用标准方式执行模型
        return self._forward_helper(input_ids, positions, kv_caches,
                                    block_tables)

def allocate(self, seq_id: int, num_tokens: int) -> List[int]:
    """为序列分配 KV 缓存块"""
    # 1. 计算需要的块数
    num_blocks = (num_tokens + self.block_size - 1) // self.block_size
    
    # 2. 分配物理块
    block_indices = []
    for _ in range(num_blocks):
        block_idx = self.block_allocator.allocate()
        block_indices.append(block_idx)
    
    # 3. 更新映射表
    self.block_tables[seq_id] = block_indices
    
    return block_indices

def _process_model_outputs(
    self,
    seq_group: SequenceGroup,
    outputs: ModelOutput,
    metadata: SequenceGroupMetadata,
) -> None:
    """处理模型输出"""
    # 1. 处理 logits 和采样结果
    logprobs = outputs.logprobs
    next_tokens = outputs.next_tokens
    
    # 2. 为每个序列更新状态
    for seq_id, seq in seq_group.seqs.items():
        # 获取对应的 token 和概率
        next_token = next_tokens[seq_id]
        logprob = logprobs[seq_id] if logprobs is not None else None
        
        # 更新序列状态
        seq.append_token_id(next_token, logprob)
        
        # 检查是否达到终止条件
        if self._is_finished(seq, next_token):
            seq.status = SequenceStatus.FINISHED

# 传统 KV 缓存的内存占用
memory_per_seq = seq_len * num_layers * 2 * hidden_dim * dtype_size

class PagedAttention:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.block_tables = {}  # 序列到物理块的映射
        self.physical_blocks = []  # 实际存储块
        
    def allocate_block(self, seq_id: int):
        """为序列分配新块"""
        if seq_id not in self.block_tables:
            self.block_tables[seq_id] = []
            
        # 分配新物理块
        block_id = self._find_free_block()
        self.block_tables[seq_id].append(block_id)
        
    def get_kv_cache(self, seq_id: int, positions: torch.Tensor):
        """获取缓存内容"""
        if seq_id not in self.block_tables:
            return None
            
        # 计算物理块索引和块内位置
        block_indices, block_offsets = self._compute_indices(positions)
        physical_blocks = [self.block_tables[seq_id][idx] for idx in block_indices]
        
        # 获取缓存内容
        return self._gather_blocks(physical_blocks, block_offsets)

class BlockAllocator:
    def __init__(self, num_blocks: int):
        self.num_blocks = num_blocks
        self.free_blocks = set(range(num_blocks))
        self.used_blocks = set()
        
    def allocate(self) -> int:
        """分配一个空闲块"""
        if not self.free_blocks:
            raise RuntimeError("No free blocks available")
        
        block_id = next(iter(self.free_blocks))
        self.free_blocks.remove(block_id)
        self.used_blocks.add(block_id)
        return block_id
        
    def free(self, block_id: int) -> None:
        """释放一个块"""
        self.used_blocks.remove(block_id)
        self.free_blocks.add(block_id)

class BlockTable:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.tables = {}  # seq_id -> List[block_id]
        
    def add_block(self, seq_id: int, block_id: int) -> None:
        """为序列添加一个块"""
        if seq_id not in self.tables:
            self.tables[seq_id] = []
        self.tables[seq_id].append(block_id)
        
    def get_physical_blocks(self, seq_id: int, positions: torch.Tensor) -> Tuple[List[int], torch.Tensor]:
        """获取物理块 ID 和偏移量"""
        block_ids = self.tables[seq_id]
        block_indices = positions // self.block_size
        block_offsets = positions % self.block_size
        physical_blocks = [block_ids[idx.item()] for idx in block_indices]
        return physical_blocks, block_offsets

# CUDA 核心伪代码
@cuda.kernel
def paged_attention_kernel(
    q: Tensor,             # [batch_size, num_heads, head_dim]
    k_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    v_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    block_tables: Tensor,  # [batch_size, max_blocks]
    output: Tensor         # [batch_size, num_heads, head_dim]
):
    # 获取线程索引
    batch_idx = cuda.blockIdx.x
    head_idx = cuda.blockIdx.y
    
    # 获取该序列的块表
    seq_blocks = block_tables[batch_idx]
    
    # 本地计算缓冲区
    local_k = shared_memory[...]
    local_v = shared_memory[...]
    
    # 加载查询向量
    query = q[batch_idx, head_idx]
    
    # 对每个块执行注意力计算
    for i in range(len(seq_blocks)):
        block_id = seq_blocks[i]
        
        # 加载 KV 缓存到共享内存
        local_k.copy_from(k_cache[block_id])
        local_v.copy_from(v_cache[block_id])
        
        # 计算注意力分数
        scores = compute_attention(query, local_k)
        
        # 应用注意力权重
        output[batch_idx, head_idx] += apply_attention(scores, local_v)

class ContinuousBatcher:
    def __init__(self):
        self.active_sequences = {}  # 活跃序列集合
        self.waiting_sequences = Queue()  # 等待队列
        
    def add_sequence(self, seq_id, prompt):
        """添加新序列到等待队列"""
        self.waiting_sequences.put((seq_id, prompt))
        
    def step(self, model):
        """执行一步批处理"""
        # 1. 调度可执行序列
        batch = self._schedule_batch()
        
        # 2. 执行模型推理
        outputs = model.forward(batch.input_ids)
        
        # 3. 处理结果
        for seq_id, output in zip(batch.seq_ids, outputs):
            self._process_output(seq_id, output)
            
        # 4. 移除已完成序列
        self._cleanup_finished_sequences()
        
    def _schedule_batch(self):
        """调度当前步骤的批次"""
        batch = Batch()
        
        # 添加活跃序列
        for seq_id, seq in self.active_sequences.items():
            if not seq.is_finished:
                batch.add(seq_id, seq.get_next_input())
                
        # 尝试添加新序列（如果有空间）
        while len(batch) < self.max_batch_size and not self.waiting_sequences.empty():
            seq_id, prompt = self.waiting_sequences.get()
            self.active_sequences[seq_id] = Sequence(prompt)
            batch.add(seq_id, self.active_sequences[seq_id].get_next_input())
            
        return batch

class Scheduler:
    def __init__(self, max_num_seqs: int, max_num_tokens: int):
        self.max_num_seqs = max_num_seqs  # 最大序列数
        self.max_num_tokens = max_num_tokens  # 最大 token 数
        self.waiting = []  # 等待队列
        self.running = {}  # 运行中的序列
        self.swapped = {}  # 被换出的序列
        
    def schedule(self) -> List[SequenceGroup]:
        """调度下一批序列"""
        scheduled = []
        
        # 优先调度运行中的序列
        running_ready = [seq for seq in self.running.values() if seq.is_ready()]
        scheduled.extend(running_ready)
        
        # 如果有剩余资源，从等待队列调度新序列
        remaining_capacity = self._get_remaining_capacity(scheduled)
        if remaining_capacity.can_schedule():
            new_seqs = self._schedule_from_waiting(remaining_capacity)
            scheduled.extend(new_seqs)
            
        return scheduled
        
    def add_request(self, request):
        """添加新请求到等待队列"""
        self.waiting.append(request)
        
    def abort_request(self, request_id):
        """中止请求"""
        # 从所有队列中移除
        # ...
        
    def _schedule_from_waiting(self, capacity):
        """从等待队列调度序列"""
        scheduled = []
        remaining_seqs = capacity.max_seqs
        remaining_tokens = capacity.max_tokens
        
        # 使用优先级排序等待队列
        self.waiting.sort(key=lambda x: x.priority, reverse=True)
        
        for seq in self.waiting:
            # 检查是否有足够资源
            if remaining_seqs > 0 and remaining_tokens >= seq.token_count:
                scheduled.append(seq)
                self.waiting.remove(seq)
                self.running[seq.id] = seq
                
                # 更新剩余容量
                remaining_seqs -= 1
                remaining_tokens -= seq.token_count
                
        return scheduled

class Engine:
    def __init__(self, model, scheduler):
        self.model = model
        self.scheduler = scheduler
        self.running_sequences = {}
        
    def add_request(self, request):
        """添加新请求"""
        self.scheduler.add_request(request)
        
    def step(self):
        """执行一步批处理"""
        # 1. 调度序列
        scheduled = self.scheduler.schedule()
        if not scheduled:
            return None
            
        # 2. 准备批处理输入
        batch = self._prepare_batch(scheduled)
        
        # 3. 执行模型推理
        outputs = self.model.forward(batch.input_ids, batch.attention_mask)
        
        # 4. 处理输出
        self._process_outputs(outputs, batch)
        
        # 5. 返回生成的 token
        return self._get_outputs()
        
    def _prepare_batch(self, scheduled):
        """准备批处理输入"""
        batch = Batch()
        
        for seq in scheduled:
            if seq.id in self.running_sequences:
                # 继续生成
                input_id = self.running_sequences[seq.id].last_token
            else:
                # 新序列，处理提示
                input_ids = self._tokenize(seq.prompt)
                self.running_sequences[seq.id] = RunningSequence(seq.id, input_ids)
                
            batch.add(seq)
            
        return batch
        
    def _process_outputs(self, outputs, batch):
        """处理模型输出"""
        for i, seq_id in enumerate(batch.seq_ids):
            # 获取下一个 token
            next_token = outputs.logits[i].argmax(-1).item()
            
            # 更新序列
            self.running_sequences[seq_id].append_token(next_token)
            
            # 检查是否完成
            if self._is_finished(seq_id, next_token):
                seq = self.running_sequences.pop(seq_id)
                self.scheduler.finish_sequence(seq_id)

def iterative_schedule(self):
    """迭代调度以最大化 GPU 利用率"""
    # 初始化批次
    batch = Batch()
    
    # 迭代添加序列，直到达到资源限制
    while True:
        # 尝试添加下一个序列
        next_seq = self._get_next_candidate()
        if next_seq is None:
            break
            
        # 计算添加后的资源使用
        new_resource_usage = self._calculate_resource_usage(batch, next_seq)
        
        # 如果超出资源限制，停止添加
        if new_resource_usage > self.max_resource_limit:
            break
            
        # 添加到批次
        batch.add(next_seq)
    
    return batch

def prefetch_next_batch(self):
    """预取下一批次数据"""
    # 预测下一步可能调度的序列
    predicted_next = self.scheduler.peek_next_batch()
    
    # 异步预加载数据
    asyncio.create_task(self._async_prefetch(predicted_next))
    
def _async_prefetch(self, sequences):
    """异步预加载数据"""
    for seq in sequences:
        if seq.id not in self.prefetched_data:
            # 预处理输入
            processed_data = self._preprocess(seq)
            
            # 存储预处理结果
            self.prefetched_data[seq.id] = processed_data

def priority_schedule(self):
    """基于优先级的调度"""
    # 计算每个请求的综合优先级分数
    for req in self.waiting:
        # 基础优先级
        priority = req.priority
        
        # 等待时间因子
        wait_time = time.time() - req.arrival_time
        wait_factor = min(1.0, wait_time / self.max_wait_time)
        
        # 综合分数 (混合优先级和等待时间)
        req.score = priority * 0.7 + wait_factor * 0.3
    
    # 按分数排序
    self.waiting.sort(key=lambda x: x.score, reverse=True)
    
    # 调度顶部序列
    return self.waiting[:self.max_batch_size]

# vLLM 引擎主循环
def _run_engine(self):
    """vLLM 引擎主循环"""
    while not self._stop_event.is_set():
        # 调度序列
        scheduled_seq_groups = self.scheduler.schedule()
        if not scheduled_seq_groups:
            time.sleep(0.001)  # 避免空转
            continue
            
        # 准备批次
        batch = self._prepare_batch(scheduled_seq_groups)
        
        # 执行模型前向传播
        outputs = self.model_runner.execute_model(
            batch.input_ids,
            batch.positions,
            batch.kv_caches,
            batch.block_tables,
        )
        
        # 后处理结果
        self._process_model_outputs(outputs, batch)
        
        # 检查完成的序列
        self._check_completion(batch.seq_groups)

# 传统 CUDA 执行方式
for i in range(iterations):
    # 每次迭代都需要 CPU-GPU 通信和调度开销
    kernel1<<<grid, block>>>(args1);  # CPU 调度 GPU 执行
    kernel2<<<grid, block>>>(args2);  # CPU 调度 GPU 执行

# CUDA 图执行方式
# 1. 捕获阶段（只执行一次）
cudaStreamBeginCapture(stream);
kernel1<<<grid, block, 0, stream>>>(args1);
kernel2<<<grid, block, 0, stream>>>(args2);
cudaGraph_t graph;
cudaStreamEndCapture(stream, &graph);

# 2. 实例化图（只执行一次）
cudaGraphExec_t graphExec;
cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0);

# 3. 重复执行（高效）
for i in range(iterations):
    # 无需 CPU 干预，极低开销
    cudaGraphLaunch(graphExec, stream);

@torch.inference_mode()
def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
    """捕获模型执行为 CUDA 图以加速推理"""
    # 1. 准备捕获环境
    self._check_if_cuda_graphs_supported()
    # 保存当前内存状态，在捕获后释放临时内存
    mempool = torch.cuda.graph_pool_handle()
    self.graph_memory_pool = mempool
    self.dummy_inputs = {}
    self.graph_runners = {}
    
    # 2. 获取批处理大小
    batch_sizes = self._get_graph_batch_sizes()
    
    # 3. 为每个批处理大小捕获图
    for batch_size in batch_sizes:
        # 准备虚拟模型输入（重要！）
        dummy_inputs = self._prepare_dummy_inputs(batch_size)
        self.dummy_inputs[batch_size] = dummy_inputs
        
        # 捕获模型前向传播
        with self.attn_state.graph_capture(batch_size), \
             cuda_graph_context(self.device) as context:
            # 执行捕获
            outputs = self._forward_helper(**dummy_inputs)
            
        # 创建图运行器
        self.graph_runners[batch_size] = CUDAGraphRunner(
            context.graph, dummy_inputs, outputs
        )
    
    # 4. 设置已捕获标志
    self.model_captured = True

class CUDAGraphRunner:
    """管理 CUDA 图的捕获和执行"""
    
    def __init__(self, graph, static_inputs, static_outputs):
        # 存储图和静态内存引用
        self.graph = graph
        self.static_inputs = static_inputs
        self.static_outputs = static_outputs
        
    def capture(self, **inputs):
        """执行并捕获计算图"""
        # 使用指定输入执行一次前向传播，让 CUDA 记录操作序列
        with torch.cuda.graph(self.graph):
            # 更新静态输入
            for name, tensor in inputs.items():
                self.static_inputs[name].copy_(tensor)
                
            # 捕获前向传播
            self.static_outputs = self.forward_fn(**self.static_inputs)
            
        return self
        
    def forward(self, **inputs):
        """使用CUDA图执行前向传播"""
        # 更新输入数据
        for name, tensor in inputs.items():
            self.static_inputs[name].copy_(tensor)
            
        # 执行图 - 高效！
        self.graph.replay()
        
        return self.static_outputs

def _prepare_dummy_inputs(self, batch_size: int) -> Dict[str, torch.Tensor]:
    """准备用于图捕获的虚拟输入"""
    # 1. 创建固定形状的输入张量
    input_ids = torch.zeros(
        batch_size, dtype=torch.long, device=self.device
    )
    
    # 2. 分配位置 ID
    positions = torch.zeros(
        batch_size, dtype=torch.long, device=self.device
    )
    
    # 3. 创建注意力掩码
    attention_mask = torch.ones(
        batch_size, self.max_context_len, dtype=torch.bool, device=self.device
    )
    
    # 4. 准备 KV 缓存和块表
    # 此处省略复杂的 KV 缓存准备逻辑...
    
    return {
        "input_ids": input_ids,
        "positions": positions,
        "attention_mask": attention_mask,
        "kv_caches": kv_caches,
        "block_tables": block_tables,
    }

def execute_model(self, input_ids, positions, kv_caches, block_tables):
    """执行模型前向传播"""
    batch_size = input_ids.shape[0]
    
    # 如果已捕获且批处理大小匹配，使用 CUDA 图
    if self.model_captured and batch_size in self.graph_runners:
        # 准备输入
        inputs = {
            "input_ids": input_ids,
            "positions": positions,
            "kv_caches": kv_caches,
            "block_tables": block_tables,
        }
        
        # 使用图运行器 - 高性能路径
        return self.graph_runners[batch_size].forward(**inputs)
        
    else:
        # 标准执行路径
        return self._forward_helper(
            input_ids, positions, kv_caches, block_tables
        )

# 在 model_runner.py 的 execute_model 方法中
def execute_model(self, *args, **kwargs):
    """执行模型推理，自动选择最佳路径"""
    batch_size = kwargs["input_ids"].shape[0]
    
    # 选择执行路径
    if (self.model_captured and 
        batch_size in self.graph_runners and 
        not self.config.disable_cuda_graphs):
        # CUDA 图路径 - 高性能
        start_time = time.time()
        outputs = self._execute_with_graph(batch_size, *args, **kwargs)
        if self.profiler:
            self.profiler.record_cuda_graph_execution(time.time() - start_time)
        return outputs
    else:
        # 常规执行路径
        return self._execute_standard(*args, **kwargs)

# 未使用 KV 缓存的自回归生成
for i in range(max_seq_len):
    # 对所有已生成的 token 重新计算 K 和 V
    keys = self.compute_keys(input_ids[:, :i+1])    # 重复计算
    values = self.compute_values(input_ids[:, :i+1])  # 重复计算
    
    # 计算下一个 token 的输出
    next_token_logits = self.compute_attention(
        query=self.compute_query(input_ids[:, i:i+1]),
        keys=keys,
        values=values
    )
    
    # 采样下一个 token
    next_token = sample(next_token_logits)
    input_ids = torch.cat([input_ids, next_token], dim=1)

# 使用 KV 缓存的自回归生成
# 初始化缓存
key_cache = []
value_cache = []

for i in range(max_seq_len):
    if i == 0:
        # 处理提示词
        keys = self.compute_keys(input_ids)
        values = self.compute_values(input_ids)
        # 存入缓存
        key_cache.append(keys)
        value_cache.append(values)
    else:
        # 只计算最新 token 的 K 和 V
        new_key = self.compute_keys(input_ids[:, -1:])
        new_value = self.compute_values(input_ids[:, -1:])
        # 添加到缓存
        key_cache.append(new_key)
        value_cache.append(new_value)
    
    # 使用完整缓存计算注意力
    next_token_logits = self.compute_attention(
        query=self.compute_query(input_ids[:, -1:]),
        keys=torch.cat(key_cache, dim=1),
        values=torch.cat(value_cache, dim=1)
    )
    
    # 采样下一个 token
    next_token = sample(next_token_logits)
    input_ids = torch.cat([input_ids, next_token], dim=1)

class KVCacheManager:
    def __init__(self, num_layers, num_heads, head_size, block_size):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_size = head_size
        self.block_size = block_size
        
        # 物理缓存存储
        self.key_blocks = torch.empty(
            (0, block_size, num_heads, head_size),
            dtype=torch.float16, device="cuda"
        )
        self.value_blocks = torch.empty(
            (0, block_size, num_heads, head_size),
            dtype=torch.float16, device="cuda"
        )
        
        # 块分配器
        self.block_allocator = BlockAllocator()
        
        # 序列到块的映射
        self.block_tables = {}

def allocate_kv_cache(self, seq_id: int, prompt_len: int) -> None:
    """为新序列分配 KV 缓存"""
    # 计算需要的块数
    num_blocks = ceil_div(prompt_len, self.block_size)
    
    # 分配物理块
    block_ids = []
    for _ in range(num_blocks):
        block_id = self.block_allocator.allocate()
        if block_id >= len(self.key_blocks):
            # 缓存扩容
            self._extend_cache(block_id + 1 - len(self.key_blocks))
        block_ids.append(block_id)
    
    # 更新块表
    self.block_tables[seq_id] = block_ids

def free_kv_cache(self, seq_id: int) -> None:
    """释放序列的 KV 缓存"""
    if seq_id not in self.block_tables:
        return
    
    # 获取分配的块
    block_ids = self.block_tables[seq_id]
    
    # 释放块
    for block_id in block_ids:
        self.block_allocator.free(block_id)
    
    # 移除映射
    del self.block_tables[seq_id]

def _extend_cache(self, num_blocks: int) -> None:
    """扩展 KV 缓存"""
    current_size = len(self.key_blocks)
    new_size = current_size + num_blocks
    
    # 创建新存储空间
    new_key_blocks = torch.empty(
        (new_size, self.block_size, self.num_heads, self.head_size),
        dtype=torch.float16, device="cuda"
    )
    new_value_blocks = torch.empty(
        (new_size, self.block_size, self.num_heads, self.head_size),
        dtype=torch.float16, device="cuda"
    )
    
    # 复制现有数据
    if current_size > 0:
        new_key_blocks[:current_size] = self.key_blocks
        new_value_blocks[:current_size] = self.value_blocks
        
    # 替换缓存
    self.key_blocks = new_key_blocks
    self.value_blocks = new_value_blocks

def get_kv_cache(self, seq_id: int, positions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """获取特定位置的 KV 缓存"""
    if seq_id not in self.block_tables:
        return None, None
    
    # 获取块表
    block_table = self.block_tables[seq_id]
    
    # 计算块索引和偏移
    block_indices = positions // self.block_size
    block_offsets = positions % self.block_size
    
    # 构建批索引张量
    indices_tensor = torch.empty(
        (len(positions), 3), dtype=torch.long, device="cuda"
    )
    
    # 填充索引 [block_idx, block_offset, head_idx]
    for i, (block_idx, offset) in enumerate(zip(block_indices, block_offsets)):
        indices_tensor[i, 0] = block_table[block_idx.item()]
        indices_tensor[i, 1] = offset
    
    # 提取缓存
    keys = torch.index_select(
        self.key_blocks.view(-1, self.num_heads, self.head_size),
        0, indices_tensor[:, 0] * self.block_size + indices_tensor[:, 1]
    )
    values = torch.index_select(
        self.value_blocks.view(-1, self.num_heads, self.head_size),
        0, indices_tensor[:, 0] * self.block_size + indices_tensor[:, 1]
    )
    
    return keys, values

# CUDA 内核（伪代码）
@cuda.jit
def paged_attention_kernel(
    query,           # [batch_size, num_heads, head_size]
    key_cache,       # [num_blocks, block_size, num_heads, head_size]
    value_cache,     # [num_blocks, block_size, num_heads, head_size]
    block_tables,    # [batch_size, max_blocks_per_seq]
    seq_lens,        # [batch_size]
    output           # [batch_size, num_heads, head_size]
):
    # 获取线程索引
    batch_idx = cuda.blockIdx.x
    head_idx = cuda.blockIdx.y
    
    # 获取序列的块表和长度
    seq_block_table = block_tables[batch_idx]
    seq_len = seq_lens[batch_idx]
    
    # 获取当前查询向量
    q = query[batch_idx, head_idx]
    
    # 计算注意力分数和加权和
    acc = cuda.local.array(shape=(HEAD_SIZE,), dtype=float32)
    for i in range(HEAD_SIZE):
        acc[i] = 0.0
    
    # 遍历序列的所有块
    for block_idx in range(ceil_div(seq_len, BLOCK_SIZE)):
        physical_block = seq_block_table[block_idx]
        block_len = min(BLOCK_SIZE, seq_len - block_idx * BLOCK_SIZE)
        
        # 计算注意力分数
        scores = cuda.local.array(shape=(BLOCK_SIZE,), dtype=float32)
        for i in range(block_len):
            k = key_cache[physical_block, i, head_idx]
            scores[i] = dot_product(q, k)
        
        # 应用 softmax
        softmax_normalization(scores, block_len)
        
        # 计算加权和
        for i in range(block_len):
            v = value_cache[physical_block, i, head_idx]
            for j in range(HEAD_SIZE):
                acc[j] += scores[i] * v[j]
    
    # 写回输出
    for i in range(HEAD_SIZE):
        output[batch_idx, head_idx, i] = acc[i]

def fork_sequence(
    self, src_seq_id: int, dst_seq_id: int, fork_pos: int
) -> None:
    """在指定位置分叉序列，共享 KV 缓存"""
    if src_seq_id not in self.block_tables:
        return
    
    # 获取源序列的块表
    src_blocks = self.block_tables[src_seq_id]
    
    # 计算分叉位置对应的块
    fork_block_idx = fork_pos // self.block_size
    
    # 共享前 fork_block_idx 个块
    shared_blocks = src_blocks[:fork_block_idx + 1]
    
    # 复制源块引用
    if dst_seq_id in self.block_tables:
        # 清理可能的现有块
        self.free_kv_cache(dst_seq_id)
    
    # 创建新的块表
    self.block_tables[dst_seq_id] = shared_blocks.copy()
    
    # 增加共享块的引用计数
    for block_id in shared_blocks:
        self.block_allocator.increment_ref_count(block_id)

def swap_out(self, seq_id: int) -> Dict:
    """将序列缓存交换出 GPU"""
    if seq_id not in self.block_tables:
        return None
    
    # 获取块表
    block_ids = self.block_tables[seq_id]
    
    # 准备 CPU 存储
    cpu_key_cache = []
    cpu_value_cache = []
    
    # 复制到 CPU
    for block_id in block_ids:
        cpu_key_cache.append(
            self.key_blocks[block_id].detach().cpu()
        )
        cpu_value_cache.append(
            self.value_blocks[block_id].detach().cpu()
        )
    
    # 释放 GPU 块
    self.free_kv_cache(seq_id)
    
    # 返回 CPU 缓存
    return {
        "block_table": block_ids,
        "key_cache": cpu_key_cache,
        "value_cache": cpu_value_cache
    }

def swap_in(self, seq_id: int, cpu_cache: Dict) -> None:
    """将序列缓存交换回 GPU"""
    # 提取 CPU 缓存
    block_table = cpu_cache["block_table"]
    cpu_key_cache = cpu_cache["key_cache"]
    cpu_value_cache = cpu_cache["value_cache"]
    
    # 分配新块
    new_block_ids = []
    for i in range(len(block_table)):
        # 分配块
        block_id = self.block_allocator.allocate()
        if block_id >= len(self.key_blocks):
            self._extend_cache(block_id + 1 - len(self.key_blocks))
        
        # 将数据复制回 GPU
        self.key_blocks[block_id].copy_(
            cpu_key_cache[i].to("cuda")
        )
        self.value_blocks[block_id].copy_(
            cpu_value_cache[i].to("cuda")
        )
        
        new_block_ids.append(block_id)
    
    # 更新块表
    self.block_tables[seq_id] = new_block_ids

def step(self):
    """执行一步生成"""
    # 1. 调度序列
    scheduled = self.scheduler.schedule()
    
    # 2. 准备批处理输入
    batch = self._prepare_batch(scheduled)
    
    # 3. 为新序列分配 KV 缓存
    for seq_id in batch.new_seq_ids:
        self.cache_manager.allocate_kv_cache(
            seq_id, batch.prompt_lens[seq_id]
        )
    
    # 4. 执行模型前向传播
    outputs = self.model_runner.forward(
        batch.input_ids,
        batch.positions,
        batch.kv_caches,
        batch.block_tables
    )
    
    # 5. 处理输出并更新序列状态
    for i, seq_id in enumerate(batch.seq_ids):
        # 采样下一个 token
        next_token = self._sample_next_token(outputs, i)
        
        # 添加到序列
        self.sequences[seq_id].append_token(next_token)
        
        # 检查是否完成
        if self._is_finished(seq_id):
            # 释放 KV 缓存
            self.cache_manager.free_kv_cache(seq_id)

class ModelSharding:
    def __init__(self, model, num_gpus):
        self.num_gpus = num_gpus
        self.model = model
        self.tp_size = num_gpus  # 张量并行度
        
    def shard_model(self):
        """将模型分片到多个 GPU"""
        # 1. 识别可并行层
        tp_modules = []
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Embedding)):
                tp_modules.append((name, module))
        
        # 2. 为每个分片创建子模型
        self.devices = [f"cuda:{i}" for i in range(self.tp_size)]
        self.shards = [copy.deepcopy(self.model) for _ in range(self.tp_size)]
        
        # 3. 分割权重
        for name, module in tp_modules:
            if isinstance(module, nn.Linear):
                # 沿输出维度分割线性层
                self._shard_linear_layer(name)
            elif isinstance(module, nn.Embedding):
                # 沿词汇表维度分割嵌入层
                self._shard_embedding_layer(name)

def _shard_linear_layer(self, name):
    """分割线性层权重"""
    original_weight = self.model.get_submodule(name).weight
    original_shape = original_weight.shape
    
    # 沿输出维度分割
    split_size = original_shape[0] // self.tp_size
    for i in range(self.tp_size):
        # 计算分片范围
        start_idx = i * split_size
        end_idx = (i + 1) * split_size if i < self.tp_size - 1 else original_shape[0]
        
        # 分配分片
        shard = self.shards[i].get_submodule(name)
        with torch.no_grad():
            shard.weight = nn.Parameter(
                original_weight[start_idx:end_idx].clone().to(self.devices[i])
            )
        
        # 更新形状
        shard.in_features = original_shape[1]
        shard.out_features = end_idx - start_idx

def tp_forward(self, input_tensor):
    """使用张量并行执行前向传播"""
    # 1. 分发输入到各个设备
    input_shards = []
    for i in range(self.tp_size):
        input_shards.append(input_tensor.to(self.devices[i]))
    
    # 2. 并行执行每个分片
    output_shards = []
    for i in range(self.tp_size):
        with torch.cuda.device(self.devices[i]):
            output_shards.append(self.shards[i](input_shards[i]))
    
    # 3. 合并结果
    if isinstance(output_shards[0], torch.Tensor):
        # 对于分割的线性层，沿输出维度合并
        return torch.cat(output_shards, dim=-1)
    else:
        # 对于其他层，可能需要特殊处理
        return self._merge_custom_outputs(output_shards)

def load_quantized_model(model_path, quantization="int8"):
    """加载量化模型"""
    if quantization == "int8":
        # 使用 INT8 量化
        return load_int8_model(model_path)
    elif quantization == "int4":
        # 使用 INT4 量化
        return load_int4_model(model_path)
    elif quantization == "awq":
        # 使用 AWQ 量化
        return load_awq_model(model_path)
    else:
        # 默认加载全精度模型
        return load_full_precision_model(model_path)

def load_int8_model(model_path):
    """加载 INT8 量化模型"""
    # 1. 加载原始模型配置和权重
    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_config(config)
    state_dict = torch.load(f"{model_path}/pytorch_model.bin")
    
    # 2. 对每个权重进行量化
    quantized_state_dict = {}
    for name, param in state_dict.items():
        # 跳过某些不应量化的层
        if "ln" in name or "embedding" in name:
            quantized_state_dict[name] = param
        else:
            # 执行 INT8 量化
            quantized_param = quantize_to_int8(param)
            quantized_state_dict[name] = quantized_param
    
    # 3. 加载量化后的权重
    model.load_state_dict(quantized_state_dict)
    
    # 4. 注册自定义内核用于量化计算
    replace_with_quantized_kernels(model)
    
    return model

class MixedPrecisionManager:
    def __init__(self, model):
        self.model = model
        
    def apply_mixed_precision(self):
        """应用混合精度策略"""
        # 分析层的重要性
        layer_importance = self._analyze_layer_importance()
        
        # 应用不同精度
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear):
                layer_name = name.split(".")[-2]  # 获取层名称
                if layer_name in layer_importance:
                    importance = layer_importance[layer_name]
                    
                    # 根据重要性选择精度
                    if importance > 0.8:
                        # 重要层使用更高精度
                        module.to(torch.float16)
                    elif importance > 0.5:
                        # 中等重要性使用 INT8
                        quantize_module_to_int8(module)
                    else:
                        # 低重要性使用 INT4
                        quantize_module_to_int4(module)

class ActivationCheckpointing:
    def __init__(self, model):
        self.model = model
        
    def apply_checkpointing(self):
        """对模型应用激活检查点"""
        for name, module in self.model.named_children():
            if isinstance(module, TransformerLayer):
                # 包装 Transformer 层使用检查点
                self.model._modules[name] = torch.utils.checkpoint.checkpoint_wrapper(module)

def custom_checkpointing(self):
    """实现自定义粒度的检查点"""
    checkpoint_layers = []
    
    # 为每个 Transformer 层选择检查点策略
    for i, layer in enumerate(self.model.layers):
        if i % 2 == 0:  # 偶数层：细粒度检查点
            # 单独对 attention 和 mlp 应用检查点
            layer.attention = torch.utils.checkpoint.checkpoint_wrapper(
                layer.attention, 
                preserve_rng_state=False
            )
            layer.mlp = torch.utils.checkpoint.checkpoint_wrapper(
                layer.mlp,
                preserve_rng_state=False
            )
        else:  # 奇数层：整层检查点
            checkpoint_layers.append(layer)
    
    # 对选择的层应用检查点
    for layer in checkpoint_layers:
        layer_index = self.model.layers.index(layer)
        self.model.layers[layer_index] = torch.utils.checkpoint.checkpoint_wrapper(
            layer, 
            preserve_rng_state=False
        )

class GPUMemoryManager:
    def __init__(self, device, max_memory_fraction=0.9):
        self.device = device
        self.max_memory = int(torch.cuda.get_device_properties(device).total_memory 
                              * max_memory_fraction)
        self.allocated = 0
        self.buffers = {}
        self.free_buffers = {}
        
    def allocate(self, size, dtype):
        """分配内存缓冲区"""
        # 检查现有的空闲缓冲区
        buffer_key = (size, dtype)
        if buffer_key in self.free_buffers and self.free_buffers[buffer_key]:
            # 重用现有缓冲区
            buffer = self.free_buffers[buffer_key].pop()
            return buffer
        
        # 检查内存限制
        bytes_size = size * dtype_size(dtype)
        if self.allocated + bytes_size > self.max_memory:
            # 尝试回收内存
            self._garbage_collect()
            if self.allocated + bytes_size > self.max_memory:
                raise RuntimeError(f"Out of memory: trying to allocate {bytes_size} bytes")
            
            # 分配新缓冲区
            buffer = torch.empty(size, dtype=dtype, device=self.device)
            self.allocated += bytes_size
            buffer_id = id(buffer)
            self.buffers[buffer_id] = (buffer, size, dtype, bytes_size)
            
            return buffer
            
        def free(self, buffer):
            """释放缓冲区（移至空闲池）"""
            buffer_id = id(buffer)
            if buffer_id in self.buffers:
                _, size, dtype, _ = self.buffers[buffer_id]
                buffer_key = (size, dtype)
                
                # 添加到空闲池
                if buffer_key not in self.free_buffers:
                    self.free_buffers[buffer_key] = []
                self.free_buffers[buffer_key].append(buffer)
        
        def _garbage_collect(self):
            """强制垃圾回收释放内存"""
            # 清除超过阈值的空闲缓冲区
            for buffer_key in list(self.free_buffers.keys()):
                if len(self.free_buffers[buffer_key]) > 5:  # 保留一定数量的缓冲区
                    # 释放多余缓冲区
                    excess_buffers = self.free_buffers[buffer_key][5:]
                    self.free_buffers[buffer_key] = self.free_buffers[buffer_key][:5]
                    
                    # 更新分配跟踪
                    for buffer in excess_buffers:
                        buffer_id = id(buffer)
                        if buffer_id in self.buffers:
                            _, _, _, bytes_size = self.buffers[buffer_id]
                            self.allocated -= bytes_size
                            del self.buffers[buffer_id]
                    
                    # 明确删除引用
                    del excess_buffers
            
            # 强制 PyTorch 缓存清理
            torch.cuda.empty_cache()

class MemoryMonitor:
    def __init__(self, device, warning_threshold=0.9, critical_threshold=0.95):
        self.device = device
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self.last_check_time = 0
        self.check_interval = 1.0  # 秒
        
    def check_memory(self):
        """检查内存使用情况并作出响应"""
        current_time = time.time()
        if current_time - self.last_check_time < self.check_interval:
            return MemoryStatus.OK
            
        # 更新检查时间
        self.last_check_time = current_time
        
        # 获取当前内存使用
        used_memory = torch.cuda.memory_allocated(self.device)
        total_memory = torch.cuda.get_device_properties(self.device).total_memory
        memory_fraction = used_memory / total_memory
        
        # 根据使用率确定状态
        if memory_fraction > self.critical_threshold:
            return MemoryStatus.CRITICAL
        elif memory_fraction > self.warning_threshold:
            return MemoryStatus.WARNING
        else:
            return MemoryStatus.OK
    
    def get_recommended_actions(self, status):
        """根据内存状态提供建议操作"""
        if status == MemoryStatus.CRITICAL:
            return [
                "swap_out_inactive_sequences",
                "reduce_batch_size",
                "force_garbage_collect"
            ]
        elif status == MemoryStatus.WARNING:
            return [
                "delay_new_requests",
                "compress_kv_cache"
            ]
        else:
            return []

def adaptive_sequence_scheduling(self):
    """根据内存使用自适应调整调度"""
    # 检查内存状态
    memory_status = self.memory_monitor.check_memory()
    actions = self.memory_monitor.get_recommended_actions(memory_status)
    
    # 应用推荐操作
    if "reduce_batch_size" in actions:
        # 临时降低批大小
        self.current_max_batch_size = max(1, self.max_batch_size // 2)
    else:
        # 恢复正常批大小
        self.current_max_batch_size = self.max_batch_size
        
    if "swap_out_inactive_sequences" in actions:
        # 交换出不活跃序列
        inactive_seqs = self._identify_inactive_sequences()
        for seq_id in inactive_seqs:
            self.swap_out_sequence(seq_id)
            
    if "force_garbage_collect" in actions:
        # 强制垃圾回收
        torch.cuda.empty_cache()
        gc.collect()