推理优化可视化实验室

def generate(model, prompt_ids, max_new_tokens):
    # Prefill: 整个 prompt 一次前向传播 (计算密集)
    logits = model.forward(prompt_ids)
    next_token = sample(logits[-1])
    kv_cache = model.kv_cache  # prefill 阶段保存

    # Decode: 逐 token 生成 (内存密集, KV-Cache 读取主导)
    output = [next_token]
    for _ in range(max_new_tokens - 1):
        logits = model.forward([next_token], kv_cache=kv_cache)
        next_token = sample(logits[-1])
        output.append(next_token)
    return output

class KVCache:
    def __init__(self, n_layers, n_kv_heads, head_dim, max_seq):
        self.cache_k = torch.zeros(n_layers, max_seq, n_kv_heads, head_dim)
        self.cache_v = torch.zeros(n_layers, max_seq, n_kv_heads, head_dim)
        self.pos = 0

    def update(self, layer, new_k, new_v):
        # 写入新 token 的 K, V 到当前位置
        self.cache_k[layer, self.pos] = new_k
        self.cache_v[layer, self.pos] = new_v

    def get(self, layer):
        # 读取到当前位置的全部缓存 (decode 瓶颈所在)
        return self.cache_k[layer, :self.pos], self.cache_v[layer, :self.pos]

class BlockManager:
    def __init__(self, num_blocks, block_size=16):
        self.block_size = block_size
        self.free_blocks = list(range(num_blocks))
        self.block_table = {}  # req_id -> [block_ids]

    def allocate(self, req_id, num_tokens):
        blocks = self.block_table.setdefault(req_id, [])
        needed = (num_tokens + self.block_size - 1) // self.block_size
        while len(blocks) < needed and self.free_blocks:
            blocks.append(self.free_blocks.pop())

    def free(self, req_id):
        for b in self.block_table.get(req_id, []):
            self.free_blocks.append(b)
        self.block_table[req_id] = []

class ContinuousBatchScheduler:
    def __init__(self, engine, max_batch_size):
        self.engine = engine
        self.max_batch = max_batch_size
        self.running, self.waiting = [], []

    def step(self):
        # 移除已完成的请求
        self.running = [r for r in self.running if not r.finished]
        # 动态补充新请求到批次中
        while self.waiting and len(self.running) < self.max_batch:
            self.running.append(self.waiting.pop(0))
        # 对所有运行中的请求执行一次 decode 迭代
        for req in self.running:
            self.engine.decode_step(req)

def speculative_decode(target, draft, prompt, max_tokens, K=4):
    tokens = list(prompt)
    for _ in range(max_tokens):
        # 1. 草稿模型快速生成 K 个候选 token
        draft_tokens = draft.generate(tokens, n=K)
        # 2. 目标模型一次前向验证全部 K 个 token
        target_logits = target.forward(tokens + draft_tokens)
        # 3. 接受匹配的 token, 拒绝第一个不匹配的
        accepted = 0
        for i, dt in enumerate(draft_tokens):
            if argmax(target_logits[len(tokens)+i]) == dt:
                accepted += 1
            else:
                break
        # 4. 添加接受的 token + 1 个修正 token
        tokens += draft_tokens[:accepted]
        if accepted < K:
            tokens += [argmax(target_logits[len(tokens)+accepted])]
    return tokens

def quantize_weights(w, n_bits=8):
    """对称量化: FP16 -> INT8/INT4"""
    scale = w.abs().max() / (2 ** (n_bits - 1) - 1)
    q = torch.clamp((w / scale).round(),
                    -(2 ** (n_bits - 1)), 2 ** (n_bits - 1) - 1)
    return q.to(torch.int8), scale  # 存储量化值和缩放因子

# GPTQ: 二阶 Hessian 信息, 逐列量化补偿
# AWQ: 激活感知, 保护显著权重通道
# GGUF: k-quant 格式, 支持 CPU/GPU 混合推理
# SmoothQuant: 将激活难度迁移至权重, 适合 W8A8

def flash_attention(Q, K, V, block_size=64):
    """Flash Attention: 分块 + 在线 softmax, IO 复杂度 O(N^2/sqrt(M))"""
    N = Q.shape[0]
    O = torch.zeros_like(Q)
    for i in range(0, N, block_size):        # 遍历 Q 块
        Qi = Q[i:i+block_size]
        Oi, m, l = torch.zeros_like(Qi), torch.full((block_size,), -1e9), torch.zeros(block_size)
        for j in range(0, N, block_size):    # 遍历 K/V 块 (在 SRAM 中)
            S = Qi @ K[j:j+block_size].T / (Q.shape[-1] ** 0.5)
            m_new = torch.max(m, S.max(dim=-1).values)
            P = torch.exp(S - m_new.unsqueeze(-1))
            l = l * torch.exp(m - m_new) + P.sum(dim=-1)
            Oi = Oi * torch.exp(m - m_new).unsqueeze(-1) + P @ V[j:j+block_size]
            m = m_new
        O[i:i+block_size] = Oi / l.unsqueeze(-1)
    return O

import torch.distributed as dist

class TensorParallelLinear(nn.Module):
    """张量并行: 按输出维度切分权重矩阵到各 GPU"""
    def __init__(self, in_features, out_features, rank, world_size):
        super().__init__()
        self.local_out = out_features // world_size
        self.weight = nn.Parameter(torch.randn(self.local_out, in_features))
        self.rank, self.world_size = rank, world_size

    def forward(self, x):
        local_out = F.linear(x, self.weight)       # 各 GPU 计算自己的分片
        gathered = [torch.zeros_like(local_out)
                    for _ in range(self.world_size)]
        dist.all_gather(gathered, local_out)        # All-Gather 合并
        return torch.cat(gathered, dim=-1)

# Pipeline Parallel: 按层切分, P2P 传递激活值
# Data Parallel: 完整模型复制, 梯度 All-Reduce