# 序列 A 的 Block Table
seq_a_blocks = [3, 7, 1, 9]  # 逻辑顺序：Block 3 → 7 → 1 → 9

# 序列 B 的 Block Table
seq_b_blocks = [5, 2, 8]  # 逻辑顺序：Block 5 → 2 → 8

# KV 缓存的物理存储
k_cache: [num_blocks, num_kv_heads, head_size/x, block_size, x]
v_cache: [num_blocks, num_kv_heads, head_size, block_size]

# 典型配置（A100 80GB）
num_blocks = 32768  # 可动态调整
block_size = 16     # 每个 block 存储 16 个 token

def paged_attention(
    q: torch.Tensor,           # [batch, num_heads, head_size]
    k_cache: torch.Tensor,     # [num_blocks, num_kv_heads, head_size/x, block_size, x]
    v_cache: torch.Tensor,     # [num_blocks, num_kv_heads, head_size, block_size]
    block_tables: torch.Tensor, # [batch, max_num_blocks_per_seq]
    seq_lens: torch.Tensor,    # [batch]
    scale: float,
) -> torch.Tensor:
    """
    PagedAttention 前向传播
    """
    batch_size, num_heads, head_size = q.shape
    output = torch.zeros_like(q)
    
    for i in range(batch_size):
        # 获取当前序列的 block table
        seq_block_table = block_tables[i]
        seq_len = seq_lens[i]
        
        # 计算需要访问的 block 数量
        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
        
        # 遍历所有 block
        for j in range(num_blocks):
            block_id = seq_block_table[j]
            block_start = j * BLOCK_SIZE
            block_end = min(block_start + BLOCK_SIZE, seq_len)
            actual_block_len = block_end - block_start
            
            # 从物理 block 中加载 K, V
            k_block = k_cache[block_id, :, :, :actual_block_len, :]
            v_block = v_cache[block_id, :, :, :actual_block_len]
            
            # 计算注意力
            qk = torch.matmul(q[i], k_block.transpose(-1, -2)) * scale
            attn_weights = torch.softmax(qk, dim=-1)
            output[i] += torch.matmul(attn_weights, v_block)
    
    return output

class BlockAllocator:
    def __init__(self, num_blocks, block_size, dtype):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.free_blocks = list(range(num_blocks))
        self.allocated_blocks = set()
        
        # 预分配 GPU 内存
        self.k_cache = torch.empty(
            (num_blocks, num_kv_heads, head_size, block_size),
            dtype=dtype,
            device='cuda'
        )
        self.v_cache = torch.empty(
            (num_blocks, num_kv_heads, head_size, block_size),
            dtype=dtype,
            device='cuda'
        )
    
    def allocate(self, num_blocks_needed):
        """分配指定数量的 block"""
        if len(self.free_blocks) < num_blocks_needed:
            raise RuntimeError("Out of memory")
        
        allocated = self.free_blocks[:num_blocks_needed]
        self.free_blocks = self.free_blocks[num_blocks_needed:]
        self.allocated_blocks.update(allocated)
        return allocated
    
    def free(self, block_ids):
        """释放 block"""
        for block_id in block_ids:
            self.allocated_blocks.remove(block_id)
            self.free_blocks.append(block_id)

# 缓存池
prompt_cache = {
    "prompt_hash_1": [block_1, block_2, block_3],
    "prompt_hash_2": [block_5, block_6],
}

# 新请求到达时
if prompt_hash in prompt_cache:
    # 直接复用已计算的 KV 缓存
    block_table = prompt_cache[prompt_hash].copy()
else:
    # 预填充并缓存
    block_table = prefill_and_cache(prompt)

from vllm import LLM, SamplingParams

# 创建 LLM 实例
llm = LLM(model="meta-llama/Llama-2-70b-chat-hf")

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=512,
    top_p=0.9
)

# 生成
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
]
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)