""" Chunked Prefill + KV Cache Offload Simulation v2 改进： 1. 简化日志输出 2. 添加reduce时间 3. 计算必须等待KV load完成 """ import threading import time from dataclasses import dataclass from typing import Optional from concurrent.futures import ThreadPoolExecutor, Future # ============== 配置参数 ============== NUM_CHUNKS = 8 GPU_SLOTS = 4 # 模拟时间 (秒) TIME_COMPUTE_BLOCK = 0.10 # 计算一个attention block TIME_REDUCE = 0.03 # 两个partial result做一次reduce TIME_TRANSFER = 0.08 # 传输一个KV cache TIME_PROJ = 0.02 # projection生成KV # ============== 全局时间基准 ============== START_TIME = None def now() -> float: """返回相对开始的时间(ms)""" return (time.time() - START_TIME) * 1000 def log_compute(msg: str): """计算队列日志（无缩进）""" print(f"[{now():7.1f}ms] [COMPUTE] {msg}") def log_transfer(msg: str): """传输队列日志（缩进）""" print(f"[{now():7.1f}ms] [TRANSFER] {msg}") def log_info(msg: str): """一般信息""" print(f"[{now():7.1f}ms] {msg}") # ============== GPU Slot管理 ============== class GPUSlots: def __init__(self, num_slots: int): self.slots = [None] * num_slots # slot_id -> kv_idx self.kv_to_slot = {} # kv_idx -> slot_id self.lock = threading.Lock() # KV ready events: kv_idx -> Event self.kv_ready = {} def alloc(self, kv_idx: int) -> int: with self.lock: for sid, val in enumerate(self.slots): if val is None: self.slots[sid] = kv_idx self.kv_to_slot[kv_idx] = sid # 创建ready event if kv_idx not in self.kv_ready: self.kv_ready[kv_idx] = threading.Event() return sid raise RuntimeError(f"No free slot for KV{kv_idx}") def free(self, slot_id: int): with self.lock: kv_idx = self.slots[slot_id] if kv_idx is not None: del self.kv_to_slot[kv_idx] # 清除event if kv_idx in self.kv_ready: del self.kv_ready[kv_idx] self.slots[slot_id] = None def free_kv(self, kv_idx: int): with self.lock: if kv_idx in self.kv_to_slot: sid = self.kv_to_slot[kv_idx] self.slots[sid] = None del self.kv_to_slot[kv_idx] if kv_idx in self.kv_ready: del self.kv_ready[kv_idx] def mark_ready(self, kv_idx: int): """标记KV已就绪（load完成或proj完成）""" with self.lock: if kv_idx in self.kv_ready: self.kv_ready[kv_idx].set() def wait_ready(self, kv_idx: int): """等待KV就绪""" with self.lock: event = self.kv_ready.get(kv_idx) if event: event.wait() def has_kv(self, kv_idx: int) -> bool: with self.lock: return kv_idx in self.kv_to_slot def state(self) -> str: with self.lock: return "[" + "][".join( f"KV{v}" if v is not None else "----" for v in self.slots ) + "]" # ============== 操作执行 ============== class Executor: def __init__(self, gpu: GPUSlots): self.gpu = gpu self.compute_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="Compute") self.transfer_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="Transfer") def proj_kv(self, q_idx: int) -> Future: """Projection生成KV，返回Future""" def task(): log_compute(f"PROJ Q{q_idx}->KV{q_idx} START") time.sleep(TIME_PROJ) slot_id = self.gpu.alloc(q_idx) self.gpu.mark_ready(q_idx) # proj完成，KV立即可用 log_compute(f"PROJ Q{q_idx}->KV{q_idx} END slot={slot_id} | {self.gpu.state()}") return slot_id return self.compute_pool.submit(task) def compute_attn(self, q_idx: int, kv_indices: list) -> Future: """计算attention block，会等待所有KV就绪""" def task(): # 等待所有需要的KV就绪 for kv_idx in kv_indices: self.gpu.wait_ready(kv_idx) kv_str = ",".join(map(str, kv_indices)) log_compute(f"ATTN Q{q_idx}*KV[{kv_str}] START") time.sleep(TIME_COMPUTE_BLOCK * len(kv_indices)) log_compute(f"ATTN Q{q_idx}*KV[{kv_str}] END") return (q_idx, kv_indices) return self.compute_pool.submit(task) def reduce(self, q_idx: int, num_partials: int) -> Future: """Online softmax reduce多个partial结果""" def task(): if num_partials <= 1: return # n个partial需要n-1次两两reduce num_reduces = num_partials - 1 log_compute(f"REDUCE Q{q_idx} ({num_partials} partials) START") time.sleep(TIME_REDUCE * num_reduces) log_compute(f"REDUCE Q{q_idx} END") return self.compute_pool.submit(task) def load_kv(self, kv_idx: int) -> Future: """从CPU load KV到GPU""" def task(): if self.gpu.has_kv(kv_idx): log_transfer(f"LOAD KV{kv_idx} SKIP (already on GPU)") return kv_idx slot_id = self.gpu.alloc(kv_idx) log_transfer(f"LOAD KV{kv_idx} START -> slot{slot_id}") time.sleep(TIME_TRANSFER) self.gpu.mark_ready(kv_idx) # load完成，标记就绪 log_transfer(f"LOAD KV{kv_idx} END | {self.gpu.state()}") return kv_idx return self.transfer_pool.submit(task) def offload_kv(self, kv_idx: int) -> Future: """从GPU offload KV到CPU""" def task(): log_transfer(f"OFFLOAD KV{kv_idx} START") time.sleep(TIME_TRANSFER) self.gpu.free_kv(kv_idx) log_transfer(f"OFFLOAD KV{kv_idx} END | {self.gpu.state()}") return kv_idx return self.transfer_pool.submit(task) def shutdown(self): self.compute_pool.shutdown(wait=True) self.transfer_pool.shutdown(wait=True) # ============== 调度器 ============== def schedule_query(exe: Executor, q_idx: int): """调度单个Query的处理""" print(f"\n{'='*50}") log_info(f"===== Query {q_idx} START =====") hist_kv = list(range(q_idx)) # 历史KV: 0 ~ q_idx-1 num_partials = 0 # Phase 1: Projection生成当前KV proj_fut = exe.proj_kv(q_idx) proj_fut.result() # 等待完成 # Phase 2: 对角块计算 + 同时prefetch历史KV # 启动对角块计算 diag_fut = exe.compute_attn(q_idx, [q_idx]) num_partials += 1 # 同时prefetch历史KV (最多3个slot可用) prefetch_slots = min(len(hist_kv), GPU_SLOTS - 1) prefetch_kv = hist_kv[:prefetch_slots] prefetch_futs = [exe.load_kv(kv) for kv in prefetch_kv] # 等待对角块完成 diag_fut.result() # Phase 3: Offload当前KV释放slot offload_fut = exe.offload_kv(q_idx) # 等待prefetch完成，然后计算这批历史KV for f in prefetch_futs: f.result() if prefetch_kv: hist_fut = exe.compute_attn(q_idx, prefetch_kv) num_partials += 1 else: hist_fut = None # 等待offload完成 offload_fut.result() # Phase 4: 处理剩余历史KV remaining_kv = hist_kv[prefetch_slots:] computed_kv = prefetch_kv.copy() while remaining_kv: # 等待上一批计算完成 if hist_fut: hist_fut.result() # 释放已计算的KV for kv in computed_kv: exe.gpu.free_kv(kv) # Load下一批 batch_size = min(len(remaining_kv), GPU_SLOTS) batch_kv = remaining_kv[:batch_size] remaining_kv = remaining_kv[batch_size:] load_futs = [exe.load_kv(kv) for kv in batch_kv] for f in load_futs: f.result() # 计算这批 hist_fut = exe.compute_attn(q_idx, batch_kv) num_partials += 1 computed_kv = batch_kv # 等待最后一批计算完成 if hist_fut: hist_fut.result() # 清理GPU for kv in computed_kv: exe.gpu.free_kv(kv) # Phase 5: Reduce所有partial results reduce_fut = exe.reduce(q_idx, num_partials) reduce_fut.result() log_info(f"===== Query {q_idx} END =====") def main(): global START_TIME START_TIME = time.time() print("Chunked Prefill + KV Cache Offload Simulation v2") print(f"Config: {NUM_CHUNKS} chunks, {GPU_SLOTS} GPU slots") print(f"Time: compute={TIME_COMPUTE_BLOCK}s, transfer={TIME_TRANSFER}s, reduce={TIME_REDUCE}s") gpu = GPUSlots(GPU_SLOTS) exe = Executor(gpu) try: for q_idx in range(NUM_CHUNKS): schedule_query(exe, q_idx) print(f"\n{'='*50}") log_info(f"ALL DONE! Total: {now():.1f}ms") finally: exe.shutdown() if __name__ == "__main__": main()