02-大模型面试题

"""
Self-Attention 实现
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SelfAttention(nn.Module):
    """单头自注意力"""

    def __init__(self, d_model, d_k=None):
        super().__init__()
        self.d_k = d_k or d_model

        self.W_q = nn.Linear(d_model, self.d_k, bias=False)
        self.W_k = nn.Linear(d_model, self.d_k, bias=False)
        self.W_v = nn.Linear(d_model, self.d_k, bias=False)

    def forward(self, x, mask=None):
        """
        Args:
            x: (batch, seq_len, d_model)
            mask: (batch, seq_len, seq_len) or (seq_len, seq_len)
        """
        Q = self.W_q(x)  # (batch, seq_len, d_k)
        K = self.W_k(x)
        V = self.W_v(x)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch, seq_len, seq_len)
        scores = scores / math.sqrt(self.d_k)

        # 应用 mask (因果注意力或 padding mask)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Softmax
        attention_weights = F.softmax(scores, dim=-1)

        # 加权求和
        output = torch.matmul(attention_weights, V)

        return output, attention_weights


class MultiHeadAttention(nn.Module):
    """多头注意力"""

    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # 线性变换并分头
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k)

        # 转置: (batch, heads, seq_len, d_k)
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 注意力计算
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention = F.softmax(scores, dim=-1)
        context = torch.matmul(attention, V)

        # 合并多头
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_len, self.d_model)

        # 输出投影
        output = self.W_o(context)

        return output

"""
位置编码实现
"""

# 1. 正弦位置编码
class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]


# 2. 旋转位置编码 (RoPE)
class RotaryPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=2048, base=10000):
        super().__init__()
        self.d_model = d_model

        # 计算频率
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)

        # 预计算
        self._build_cache(max_len)

    def _build_cache(self, seq_len):
        t = torch.arange(seq_len, device=self.inv_freq.device)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())

    def _rotate_half(self, x):
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, q, k, seq_len):
        cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0)
        sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)

        q_embed = (q * cos) + (self._rotate_half(q) * sin)
        k_embed = (k * cos) + (self._rotate_half(k) * sin)

        return q_embed, k_embed


# 3. ALiBi
class ALiBiAttention(nn.Module):
    def __init__(self, num_heads):
        super().__init__()
        # 每个头有不同的斜率
        slopes = self._get_slopes(num_heads)
        self.register_buffer('slopes', slopes)

    def _get_slopes(self, n_heads):
        """计算 ALiBi 斜率"""
        def get_slopes_power_of_2(n):
            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
            ratio = start
            return [start * ratio ** i for i in range(n)]

        if math.log2(n_heads).is_integer():
            return torch.tensor(get_slopes_power_of_2(n_heads))
        else:
            closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
            return torch.tensor(
                get_slopes_power_of_2(closest_power_of_2) +
                get_slopes_power_of_2(2 * closest_power_of_2)[0::2][:n_heads - closest_power_of_2]
            )

    def forward(self, attention_scores, seq_len):
        # 构建位置偏置矩阵
        positions = torch.arange(seq_len, device=attention_scores.device)
        relative_positions = positions.unsqueeze(0) - positions.unsqueeze(1)

        # 应用斜率
        alibi = self.slopes.unsqueeze(1).unsqueeze(1) * relative_positions.abs().unsqueeze(0)

        return attention_scores - alibi

"""
Flash Attention 简化实现 (展示核心思想)
"""

def flash_attention_forward(Q, K, V, block_size=64):
    """
    Flash Attention 前向传播

    核心: 分块计算 + Online Softmax
    """
    batch, heads, seq_len, d = Q.shape

    # 输出和统计量
    O = torch.zeros_like(Q)
    L = torch.zeros(batch, heads, seq_len, 1, device=Q.device)  # log-sum-exp
    M = torch.full((batch, heads, seq_len, 1), float('-inf'), device=Q.device)  # max

    # 分块
    num_blocks = (seq_len + block_size - 1) // block_size

    for block_j in range(num_blocks):
        # K, V 块
        j_start = block_j * block_size
        j_end = min(j_start + block_size, seq_len)

        K_j = K[:, :, j_start:j_end, :]
        V_j = V[:, :, j_start:j_end, :]

        for block_i in range(num_blocks):
            # Q 块
            i_start = block_i * block_size
            i_end = min(i_start + block_size, seq_len)

            Q_i = Q[:, :, i_start:i_end, :]

            # 计算注意力分数
            S_ij = torch.matmul(Q_i, K_j.transpose(-2, -1)) / math.sqrt(d)

            # 因果 mask
            if j_start > i_end - 1:  # 这个 K 块完全在 Q 块之后
                continue

            # Online Softmax 更新
            M_i = M[:, :, i_start:i_end, :]
            L_i = L[:, :, i_start:i_end, :]
            O_i = O[:, :, i_start:i_end, :]

            # 新的 max
            M_ij = S_ij.max(dim=-1, keepdim=True).values
            M_new = torch.maximum(M_i, M_ij)

            # 更新 exp 和 sum
            exp_old = torch.exp(M_i - M_new)
            exp_new = torch.exp(S_ij - M_new)
            L_new = exp_old * L_i + exp_new.sum(dim=-1, keepdim=True)

            # 更新输出
            O_new = (exp_old * O_i + torch.matmul(exp_new, V_j)) / L_new

            # 写回
            M[:, :, i_start:i_end, :] = M_new
            L[:, :, i_start:i_end, :] = L_new
            O[:, :, i_start:i_end, :] = O_new

    return O


# 实际使用 Flash Attention
from flash_attn import flash_attn_func

def use_flash_attention(q, k, v, causal=True):
    """
    使用 Flash Attention 库

    要求:
    - q, k, v: (batch, seq_len, heads, head_dim)
    - 半精度 (fp16/bf16)
    """
    output = flash_attn_func(
        q, k, v,
        causal=causal,
        softmax_scale=1.0 / math.sqrt(q.shape[-1])
    )
    return output

"""
不同架构的实现对比
"""

# GPT 风格 (Decoder-Only)
class GPTBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = MLP(d_model, d_ff)

    def forward(self, x):
        # Pre-LN 架构
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)

    def forward(self, x):
        seq_len = x.size(1)
        # 因果 mask
        mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
        return self.mha(x, mask=mask)


# BERT 风格 (Encoder-Only)
class BERTBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, num_heads)
        self.ln1 = nn.LayerNorm(d_model)
        self.mlp = MLP(d_model, d_ff)
        self.ln2 = nn.LayerNorm(d_model)

    def forward(self, x, attention_mask=None):
        # Post-LN 架构 (原始 BERT)
        attn_out, _ = self.attn(x, x, x, key_padding_mask=attention_mask)
        x = self.ln1(x + attn_out)
        x = self.ln2(x + self.mlp(x))
        return x


# T5 风格 (Encoder-Decoder)
class T5Block(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, is_decoder=False):
        super().__init__()
        self.is_decoder = is_decoder

        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ln1 = nn.LayerNorm(d_model)

        if is_decoder:
            self.cross_attn = MultiHeadAttention(d_model, num_heads)
            self.ln2 = nn.LayerNorm(d_model)

        self.mlp = MLP(d_model, d_ff)
        self.ln_final = nn.LayerNorm(d_model)

    def forward(self, x, encoder_output=None, causal_mask=None):
        # Self-Attention
        x = x + self.self_attn(self.ln1(x), mask=causal_mask)

        # Cross-Attention (仅 Decoder)
        if self.is_decoder and encoder_output is not None:
            x = x + self.cross_attn(
                self.ln2(x),
                encoder_output,
                encoder_output
            )

        # FFN
        x = x + self.mlp(self.ln_final(x))
        return x

"""
MoE 实现
"""

class Router(nn.Module):
    """Top-K 路由器"""

    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        """
        Args:
            x: (batch, seq_len, d_model)

        Returns:
            gates: (batch, seq_len, top_k) - 权重
            indices: (batch, seq_len, top_k) - 专家索引
        """
        # 计算路由概率
        logits = self.gate(x)  # (batch, seq_len, num_experts)

        # Top-K 选择
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)

        # Softmax 归一化 (只在选中的专家上)
        top_k_gates = F.softmax(top_k_logits, dim=-1)

        return top_k_gates, top_k_indices


class MoELayer(nn.Module):
    """MoE 层"""

    def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # 路由器
        self.router = Router(d_model, num_experts, top_k)

        # 专家 (每个专家是一个 FFN)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        batch, seq_len, d_model = x.shape

        # 路由
        gates, indices = self.router(x)  # (B, S, K), (B, S, K)

        # 展平便于处理
        x_flat = x.view(-1, d_model)  # (B*S, D)
        gates_flat = gates.view(-1, self.top_k)  # (B*S, K)
        indices_flat = indices.view(-1, self.top_k)  # (B*S, K)

        # 计算输出
        output = torch.zeros_like(x_flat)

        for k in range(self.top_k):
            expert_indices = indices_flat[:, k]  # (B*S,)
            expert_gates = gates_flat[:, k:k+1]  # (B*S, 1)

            for e in range(self.num_experts):
                # 找到路由到专家 e 的 token
                mask = (expert_indices == e)
                if mask.sum() == 0:
                    continue

                expert_input = x_flat[mask]
                expert_output = self.experts[e](expert_input)
                output[mask] += expert_gates[mask] * expert_output

        return output.view(batch, seq_len, d_model)

    def load_balance_loss(self, gates, indices):
        """负载均衡损失"""
        # 计算每个专家的负载
        num_tokens = gates.numel() // self.top_k

        # f_i: 路由到专家 i 的 token 比例
        expert_counts = torch.zeros(self.num_experts, device=gates.device)
        for e in range(self.num_experts):
            expert_counts[e] = (indices == e).sum()
        f = expert_counts / num_tokens

        # P_i: 路由到专家 i 的平均概率
        P = torch.zeros(self.num_experts, device=gates.device)
        for e in range(self.num_experts):
            mask = (indices == e)
            if mask.sum() > 0:
                P[e] = gates[mask].mean()

        # 负载均衡损失: 鼓励均匀分配
        loss = self.num_experts * (f * P).sum()

        return loss

"""
预训练配置示例
"""

# DeepSpeed 配置
deepspeed_config = {
    "train_batch_size": 2048,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 128,

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-4,
            "betas": [0.9, 0.95],
            "eps": 1e-8,
            "weight_decay": 0.1
        }
    },

    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 1e-4,
            "warmup_num_steps": 2000,
            "total_num_steps": 100000
        }
    },

    "fp16": {
        "enabled": True,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu"},
        "offload_param": {"device": "none"},
        "overlap_comm": True,
        "contiguous_gradients": True
    },

    "gradient_clipping": 1.0,

    "activation_checkpointing": {
        "partition_activations": True,
        "contiguous_memory_optimization": True
    }
}


# Megatron-LM 张量并行
class TensorParallelLinear(nn.Module):
    """张量并行线性层 (列并行)"""

    def __init__(self, in_features, out_features, tp_size, tp_rank):
        super().__init__()
        self.tp_size = tp_size
        self.tp_rank = tp_rank

        # 每个 rank 只持有 1/tp_size 的权重
        self.local_out = out_features // tp_size
        self.weight = nn.Parameter(
            torch.empty(self.local_out, in_features)
        )
        nn.init.kaiming_uniform_(self.weight)

    def forward(self, x):
        # 本地计算
        local_out = F.linear(x, self.weight)

        # AllGather 收集所有结果
        output_list = [torch.empty_like(local_out) for _ in range(self.tp_size)]
        dist.all_gather(output_list, local_out)

        return torch.cat(output_list, dim=-1)

"""
LoRA 实现
"""

class LoRALinear(nn.Module):
    """LoRA 线性层"""

    def __init__(self, original_layer, r=8, alpha=16, dropout=0.1):
        """
        Args:
            original_layer: 原始线性层
            r: LoRA 秩
            alpha: 缩放因子
            dropout: Dropout 概率
        """
        super().__init__()

        self.original = original_layer
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r

        in_features = original_layer.in_features
        out_features = original_layer.out_features

        # 冻结原始权重
        for param in self.original.parameters():
            param.requires_grad = False

        # LoRA 矩阵
        self.lora_A = nn.Parameter(torch.zeros(r, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))

        # 初始化
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)  # B 初始化为 0，确保初始输出不变

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 原始输出
        original_output = self.original(x)

        # LoRA 输出
        lora_output = self.dropout(x)
        lora_output = F.linear(lora_output, self.lora_A)
        lora_output = F.linear(lora_output, self.lora_B)

        return original_output + lora_output * self.scaling

    def merge_weights(self):
        """合并权重 (推理时使用)"""
        self.original.weight.data += self.scaling * (self.lora_B @ self.lora_A)
        return self.original


def apply_lora_to_model(model, r=8, alpha=16, target_modules=['q_proj', 'v_proj']):
    """给模型应用 LoRA"""

    for name, module in model.named_modules():
        if any(target in name for target in target_modules):
            if isinstance(module, nn.Linear):
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]

                parent = model.get_submodule(parent_name) if parent_name else model
                lora_layer = LoRALinear(module, r=r, alpha=alpha)
                setattr(parent, child_name, lora_layer)

    return model


# 使用 PEFT 库
from peft import get_peft_model, LoraConfig, TaskType

def create_lora_model(base_model):
    """使用 PEFT 创建 LoRA 模型"""

    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        bias="none"
    )

    model = get_peft_model(base_model, lora_config)
    model.print_trainable_parameters()

    return model

"""
RLHF 实现
"""

# Stage 2: 奖励模型
class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.backbone = base_model
        self.reward_head = nn.Linear(
            base_model.config.hidden_size, 1
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True
        )

        # 取最后一个 token 的隐藏状态
        last_hidden = outputs.hidden_states[-1]
        # 找到序列的最后一个非 padding token
        sequence_lengths = attention_mask.sum(dim=1) - 1
        last_token_hidden = last_hidden[
            torch.arange(last_hidden.size(0)),
            sequence_lengths
        ]

        reward = self.reward_head(last_token_hidden)
        return reward


def train_reward_model(model, dataloader, optimizer, epochs=3):
    """训练奖励模型"""
    model.train()

    for epoch in range(epochs):
        for batch in dataloader:
            chosen_ids = batch['chosen_input_ids']
            chosen_mask = batch['chosen_attention_mask']
            rejected_ids = batch['rejected_input_ids']
            rejected_mask = batch['rejected_attention_mask']

            # 计算奖励
            r_chosen = model(chosen_ids, chosen_mask)
            r_rejected = model(rejected_ids, rejected_mask)

            # Bradley-Terry 损失
            loss = -F.logsigmoid(r_chosen - r_rejected).mean()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


# Stage 3: PPO
class PPOTrainer:
    def __init__(self, policy_model, ref_model, reward_model,
                 tokenizer, kl_coef=0.1):
        self.policy = policy_model
        self.ref = ref_model
        self.reward = reward_model
        self.tokenizer = tokenizer
        self.kl_coef = kl_coef

    def compute_rewards(self, query_tensors, response_tensors):
        """计算奖励 (包含 KL 惩罚)"""
        with torch.no_grad():
            # 奖励模型分数
            full_ids = torch.cat([query_tensors, response_tensors], dim=1)
            rm_scores = self.reward(full_ids)

            # KL 散度
            policy_logits = self.policy(
                response_tensors,
                attention_mask=torch.ones_like(response_tensors)
            ).logits

            ref_logits = self.ref(
                response_tensors,
                attention_mask=torch.ones_like(response_tensors)
            ).logits

            kl_div = F.kl_div(
                F.log_softmax(policy_logits, dim=-1),
                F.softmax(ref_logits, dim=-1),
                reduction='batchmean'
            )

        rewards = rm_scores - self.kl_coef * kl_div
        return rewards

    def ppo_step(self, queries, responses, old_logprobs, advantages):
        """PPO 更新步骤"""
        # 计算新的 log 概率
        new_logits = self.policy(responses).logits
        new_logprobs = F.log_softmax(new_logits, dim=-1)

        # 概率比
        ratio = torch.exp(new_logprobs - old_logprobs)

        # PPO Clipped 目标
        clip_eps = 0.2
        pg_loss1 = -advantages * ratio
        pg_loss2 = -advantages * torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
        pg_loss = torch.max(pg_loss1, pg_loss2).mean()

        return pg_loss


# DPO (Direct Preference Optimization) - RLHF 的简化替代
def dpo_loss(policy_model, ref_model, chosen_ids, rejected_ids, beta=0.1):
    """
    DPO 直接优化偏好，无需单独的奖励模型

    Loss = -log σ(β × (log π(y_w|x)/π_ref(y_w|x) - log π(y_l|x)/π_ref(y_l|x)))
    """
    # 计算 log 概率
    policy_chosen_logps = get_logprobs(policy_model, chosen_ids)
    policy_rejected_logps = get_logprobs(policy_model, rejected_ids)

    with torch.no_grad():
        ref_chosen_logps = get_logprobs(ref_model, chosen_ids)
        ref_rejected_logps = get_logprobs(ref_model, rejected_ids)

    # DPO 损失
    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)

    loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()

    return loss

"""
KV Cache 实现
"""

class KVCache:
    """基础 KV Cache"""

    def __init__(self, num_layers, max_batch, max_seq_len,
                 num_heads, head_dim, device, dtype=torch.float16):
        self.num_layers = num_layers
        self.max_seq_len = max_seq_len

        # 预分配缓存
        cache_shape = (num_layers, 2, max_batch, num_heads, max_seq_len, head_dim)
        self.cache = torch.zeros(cache_shape, device=device, dtype=dtype)

        # 当前序列长度
        self.seq_len = 0

    def update(self, layer_idx, k, v):
        """更新缓存"""
        # k, v: (batch, heads, seq_len, head_dim)
        new_seq_len = k.shape[2]

        self.cache[layer_idx, 0, :, :, self.seq_len:self.seq_len + new_seq_len, :] = k
        self.cache[layer_idx, 1, :, :, self.seq_len:self.seq_len + new_seq_len, :] = v

        self.seq_len += new_seq_len

    def get(self, layer_idx):
        """获取缓存的 K, V"""
        k = self.cache[layer_idx, 0, :, :, :self.seq_len, :]
        v = self.cache[layer_idx, 1, :, :, :self.seq_len, :]
        return k, v


class PagedKVCache:
    """分页 KV Cache (vLLM 风格)"""

    def __init__(self, num_layers, num_heads, head_dim,
                 block_size=16, num_blocks=1000, device='cuda'):
        self.block_size = block_size
        self.num_blocks = num_blocks

        # 物理块 (预分配的内存池)
        block_shape = (num_layers, 2, num_heads, block_size, head_dim)
        self.blocks = torch.zeros(num_blocks, *block_shape, device=device)

        # 块分配表
        self.free_blocks = list(range(num_blocks))
        self.block_tables = {}  # request_id -> [block_ids]

    def allocate(self, request_id, num_tokens):
        """为请求分配块"""
        num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size

        if len(self.free_blocks) < num_blocks_needed:
            raise MemoryError("No free blocks")

        allocated = [self.free_blocks.pop() for _ in range(num_blocks_needed)]
        self.block_tables[request_id] = allocated

        return allocated

    def free(self, request_id):
        """释放请求的块"""
        if request_id in self.block_tables:
            self.free_blocks.extend(self.block_tables.pop(request_id))

    def update(self, request_id, layer_idx, k, v, position):
        """更新指定位置的 KV"""
        block_idx = position // self.block_size
        block_offset = position % self.block_size

        physical_block = self.block_tables[request_id][block_idx]

        self.blocks[physical_block, layer_idx, 0, :, block_offset, :] = k
        self.blocks[physical_block, layer_idx, 1, :, block_offset, :] = v


# GQA (Grouped-Query Attention)
class GroupedQueryAttention(nn.Module):
    """分组查询注意力"""

    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.num_groups = num_heads // num_kv_heads

        self.head_dim = d_model // num_heads

        self.q_proj = nn.Linear(d_model, num_heads * self.head_dim)
        self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim)
        self.o_proj = nn.Linear(num_heads * self.head_dim, d_model)

    def forward(self, x, kv_cache=None, position=0):
        batch, seq_len, _ = x.shape

        q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)

        # 更新 KV Cache
        if kv_cache is not None:
            k, v = kv_cache.update_and_get(k, v, position)

        # 扩展 K, V 以匹配 Q 的头数
        k = k.repeat_interleave(self.num_groups, dim=2)
        v = v.repeat_interleave(self.num_groups, dim=2)

        # 注意力计算
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)

        output = output.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.o_proj(output)

"""
投机解码实现
"""

class SpeculativeDecoder:
    """投机解码器"""

    def __init__(self, target_model, draft_model, tokenizer,
                 num_speculative_tokens=4):
        self.target = target_model
        self.draft = draft_model
        self.tokenizer = tokenizer
        self.K = num_speculative_tokens

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens):
        """投机生成"""
        generated = input_ids.clone()

        while generated.shape[1] - input_ids.shape[1] < max_new_tokens:
            # 1. Draft 阶段: 小模型生成 K 个候选
            draft_tokens = self._draft_tokens(generated)

            # 2. Verify 阶段: 大模型验证
            accepted_tokens, next_token = self._verify_tokens(
                generated, draft_tokens
            )

            # 3. 更新生成序列
            generated = torch.cat([generated, accepted_tokens, next_token], dim=1)

        return generated

    def _draft_tokens(self, prefix):
        """用小模型生成候选 token"""
        draft_tokens = []
        draft_input = prefix.clone()

        for _ in range(self.K):
            logits = self.draft(draft_input).logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)

            # 采样
            token = torch.multinomial(probs, num_samples=1)
            draft_tokens.append(token)
            draft_input = torch.cat([draft_input, token], dim=1)

        return torch.cat(draft_tokens, dim=1)  # (batch, K)

    def _verify_tokens(self, prefix, draft_tokens):
        """用大模型验证"""
        batch_size = prefix.shape[0]

        # 拼接 prefix 和 draft tokens
        full_input = torch.cat([prefix, draft_tokens], dim=1)

        # 大模型前向 (一次性计算所有位置的概率)
        target_logits = self.target(full_input).logits

        # 小模型概率 (重新计算以获取完整分布)
        draft_logits = self.draft(full_input).logits

        # 验证每个位置
        accepted = []
        prefix_len = prefix.shape[1]

        for i in range(self.K):
            pos = prefix_len + i
            draft_token = draft_tokens[:, i:i+1]

            # 目标模型在 pos-1 位置对 pos token 的概率
            target_probs = F.softmax(target_logits[:, pos-1, :], dim=-1)
            draft_probs = F.softmax(draft_logits[:, pos-1, :], dim=-1)

            p = target_probs.gather(1, draft_token).squeeze(-1)
            q = draft_probs.gather(1, draft_token).squeeze(-1)

            # 接受概率
            accept_prob = torch.minimum(torch.ones_like(p), p / (q + 1e-10))
            accept_mask = torch.rand_like(accept_prob) < accept_prob

            if accept_mask.all():
                accepted.append(draft_token)
            else:
                # 某个位置被拒绝，从调整分布中采样
                adjusted_probs = torch.clamp(target_probs - draft_probs, min=0)
                adjusted_probs = adjusted_probs / adjusted_probs.sum(dim=-1, keepdim=True)
                next_token = torch.multinomial(adjusted_probs, num_samples=1)
                break
        else:
            # 所有候选都被接受，从大模型采样下一个
            next_probs = F.softmax(target_logits[:, -1, :], dim=-1)
            next_token = torch.multinomial(next_probs, num_samples=1)

        accepted_tokens = torch.cat(accepted, dim=1) if accepted else torch.empty(batch_size, 0, dtype=torch.long, device=prefix.device)

        return accepted_tokens, next_token

"""
RAG 系统实现
"""

class AdvancedRAG:
    """高级 RAG 系统"""

    def __init__(self, embedding_model, llm, vector_store):
        self.embedding_model = embedding_model
        self.llm = llm
        self.vector_store = vector_store
        self.bm25 = None

    def index_documents(self, documents, chunk_size=512, overlap=50):
        """索引文档"""
        chunks = []

        for doc in documents:
            # 语义分块
            doc_chunks = self._semantic_chunking(doc, chunk_size, overlap)
            chunks.extend(doc_chunks)

        # 向量化
        embeddings = self.embedding_model.encode([c['text'] for c in chunks])

        # 存储
        self.vector_store.add(embeddings, chunks)

        # 构建 BM25 索引
        from rank_bm25 import BM25Okapi
        tokenized = [c['text'].split() for c in chunks]
        self.bm25 = BM25Okapi(tokenized)
        self.bm25_chunks = chunks

    def _semantic_chunking(self, document, chunk_size, overlap):
        """语义分块"""
        # 按段落分割
        paragraphs = document['text'].split('\n\n')

        chunks = []
        current_chunk = ""

        for para in paragraphs:
            if len(current_chunk) + len(para) < chunk_size:
                current_chunk += para + "\n\n"
            else:
                if current_chunk:
                    chunks.append({
                        'text': current_chunk.strip(),
                        'metadata': document.get('metadata', {})
                    })
                current_chunk = para + "\n\n"

        if current_chunk:
            chunks.append({
                'text': current_chunk.strip(),
                'metadata': document.get('metadata', {})
            })

        return chunks

    def retrieve(self, query, top_k=5, use_hybrid=True):
        """混合检索"""
        results = []

        # 向量检索
        query_embedding = self.embedding_model.encode([query])[0]
        vector_results = self.vector_store.search(query_embedding, top_k=top_k)

        if use_hybrid and self.bm25:
            # BM25 检索
            tokenized_query = query.split()
            bm25_scores = self.bm25.get_scores(tokenized_query)
            bm25_top_k = sorted(
                enumerate(bm25_scores),
                key=lambda x: x[1],
                reverse=True
            )[:top_k]

            # RRF 融合
            results = self._reciprocal_rank_fusion(
                vector_results, bm25_top_k
            )
        else:
            results = vector_results

        return results

    def _reciprocal_rank_fusion(self, vector_results, bm25_results, k=60):
        """倒数排名融合"""
        scores = {}

        for rank, (doc_id, score) in enumerate(vector_results):
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)

        for rank, (doc_idx, bm25_score) in enumerate(bm25_results):
            doc_id = f"bm25_{doc_idx}"
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)

        # 排序
        sorted_results = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return sorted_results

    def query_expansion(self, query):
        """查询扩展"""
        prompt = f"""Given the following query, generate 3 alternative versions that might help find relevant information:

Query: {query}

Alternative queries:
1."""

        expanded = self.llm.generate(prompt, max_tokens=100)
        queries = [query] + [q.strip() for q in expanded.split('\n') if q.strip()]
        return queries

    def generate_answer(self, query, contexts, with_citations=True):
        """生成答案"""
        # 构建上下文
        context_str = "\n\n".join([
            f"[{i+1}] {ctx['text']}"
            for i, ctx in enumerate(contexts)
        ])

        if with_citations:
            prompt = f"""Based on the following context, answer the question. Include citations [1], [2], etc. when using information from the context.

Context:
{context_str}

Question: {query}

Answer with citations:"""
        else:
            prompt = f"""Based on the following context, answer the question.

Context:
{context_str}

Question: {query}

Answer:"""

        answer = self.llm.generate(prompt, max_tokens=500)

        # 验证答案
        if with_citations:
            answer = self._validate_citations(answer, len(contexts))

        return answer

    def _validate_citations(self, answer, num_contexts):
        """验证引用"""
        import re
        citations = re.findall(r'\[(\d+)\]', answer)

        # 检查引用是否有效
        invalid_citations = [c for c in citations if int(c) > num_contexts]
        if invalid_citations:
            # 移除无效引用
            for c in invalid_citations:
                answer = answer.replace(f'[{c}]', '')

        return answer

"""
Agent 系统实现
"""

from dataclasses import dataclass
from typing import List, Dict, Any, Callable, Optional
from enum import Enum

class ActionStatus(Enum):
    SUCCESS = "success"
    FAILED = "failed"
    PENDING = "pending"

@dataclass
class Tool:
    name: str
    description: str
    function: Callable
    parameters: Dict[str, Any]

@dataclass
class Action:
    tool_name: str
    parameters: Dict[str, Any]
    thought: str

@dataclass
class Observation:
    result: Any
    status: ActionStatus
    error: Optional[str] = None


class ReActAgent:
    """ReAct 风格的 Agent"""

    def __init__(self, llm, tools: List[Tool], max_iterations=10):
        self.llm = llm
        self.tools = {t.name: t for t in tools}
        self.max_iterations = max_iterations
        self.memory = []

    def run(self, task: str) -> str:
        """执行任务"""
        self.memory = [{"role": "user", "content": task}]

        for i in range(self.max_iterations):
            # 思考下一步
            action = self._think()

            if action is None:
                # 任务完成
                return self._extract_final_answer()

            # 执行动作
            observation = self._execute(action)

            # 更新记忆
            self.memory.append({
                "role": "assistant",
                "content": f"Thought: {action.thought}\nAction: {action.tool_name}\nAction Input: {action.parameters}"
            })
            self.memory.append({
                "role": "observation",
                "content": f"Observation: {observation.result}"
            })

            # 检查是否失败
            if observation.status == ActionStatus.FAILED:
                # 尝试恢复
                recovery_action = self._recover_from_failure(observation.error)
                if recovery_action:
                    continue
                else:
                    return f"Task failed: {observation.error}"

        return "Max iterations reached"

    def _think(self) -> Optional[Action]:
        """生成下一步动作"""
        tools_desc = "\n".join([
            f"- {t.name}: {t.description}"
            for t in self.tools.values()
        ])

        prompt = f"""You are a helpful assistant that can use tools to complete tasks.

Available tools:
{tools_desc}

Based on the conversation history, decide the next action.
If the task is complete, respond with "Final Answer: [your answer]"
Otherwise, respond with:
Thought: [your reasoning]
Action: [tool name]
Action Input: [parameters as JSON]

History:
{self._format_memory()}

Next step:"""

        response = self.llm.generate(prompt, max_tokens=500)

        return self._parse_action(response)

    def _parse_action(self, response: str) -> Optional[Action]:
        """解析 LLM 响应"""
        if "Final Answer:" in response:
            return None

        # 提取 Thought, Action, Action Input
        import re
        import json

        thought_match = re.search(r'Thought:\s*(.+?)(?=Action:|$)', response, re.DOTALL)
        action_match = re.search(r'Action:\s*(\w+)', response)
        input_match = re.search(r'Action Input:\s*({.+})', response, re.DOTALL)

        if not action_match:
            return None

        thought = thought_match.group(1).strip() if thought_match else ""
        tool_name = action_match.group(1)

        try:
            parameters = json.loads(input_match.group(1)) if input_match else {}
        except json.JSONDecodeError:
            parameters = {}

        return Action(tool_name=tool_name, parameters=parameters, thought=thought)

    def _execute(self, action: Action) -> Observation:
        """执行动作"""
        tool = self.tools.get(action.tool_name)

        if not tool:
            return Observation(
                result=None,
                status=ActionStatus.FAILED,
                error=f"Unknown tool: {action.tool_name}"
            )

        try:
            result = tool.function(**action.parameters)
            return Observation(result=result, status=ActionStatus.SUCCESS)
        except Exception as e:
            return Observation(
                result=None,
                status=ActionStatus.FAILED,
                error=str(e)
            )

    def _recover_from_failure(self, error: str) -> Optional[Action]:
        """从失败中恢复"""
        prompt = f"""The previous action failed with error: {error}

How should we proceed? Options:
1. Retry with different parameters
2. Try a different approach
3. Give up

Respond with your decision and next action if applicable."""

        response = self.llm.generate(prompt, max_tokens=200)

        if "give up" in response.lower():
            return None

        return self._parse_action(response)

    def _format_memory(self) -> str:
        """格式化记忆"""
        return "\n".join([
            f"{m['role']}: {m['content']}"
            for m in self.memory
        ])

    def _extract_final_answer(self) -> str:
        """提取最终答案"""
        last_response = self.memory[-1]['content'] if self.memory else ""
        if "Final Answer:" in last_response:
            return last_response.split("Final Answer:")[-1].strip()
        return last_response


# 工具定义示例
def search_tool(query: str) -> str:
    """搜索工具"""
    # 实际实现会调用搜索 API
    return f"Search results for '{query}': ..."

def calculator_tool(expression: str) -> float:
    """计算器工具"""
    return eval(expression)

# 创建 Agent
tools = [
    Tool(
        name="search",
        description="Search the web for information",
        function=search_tool,
        parameters={"query": "string"}
    ),
    Tool(
        name="calculator",
        description="Evaluate mathematical expressions",
        function=calculator_tool,
        parameters={"expression": "string"}
    )
]

# agent = ReActAgent(llm, tools)
# result = agent.run("What is the population of China times 2?")