def create_causal_mask(seq_len, device):
    """创建因果掩码"""
    # 创建上三角矩阵，对角线以上为 True
    mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
    # 将 True 转换为 -inf
    mask = mask.masked_fill(mask == 1, float('-inf'))
    return mask

# 使用示例
seq_len = 4
mask = create_causal_mask(seq_len, 'cuda')
# tensor([[  0., -inf, -inf, -inf],
#         [  0.,   0., -inf, -inf],
#         [  0.,   0.,   0., -inf],
#         [  0.,   0.,   0.,   0.]])

# 注意力稀疏性验证代码
def analyze_attention_sparsity(attn_weights, thresholds=[0.01, 0.001, 0.0001]):
    """分析注意力矩阵的稀疏性"""
    results = {}
    total_elements = attn_weights.numel()

    for thresh in thresholds:
        sparse_count = (attn_weights < thresh).sum().item()
        sparsity = sparse_count / total_elements
        results[thresh] = sparsity

    return results

import torch
import torch.nn.functional as F
import math

class ScaledDotProductAttention(torch.nn.Module):
    """Scaled Dot-Product Attention 完整实现"""

    def __init__(self, dropout=0.0):
        super().__init__()
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None, return_weights=False):
        """
        Args:
            query: [batch, num_heads, seq_len_q, head_dim]
            key:   [batch, num_heads, seq_len_k, head_dim]
            value: [batch, num_heads, seq_len_k, head_dim]
            mask:  [batch, 1, seq_len_q, seq_len_k] or broadcastable

        Returns:
            output: [batch, num_heads, seq_len_q, head_dim]
            attn_weights (optional): [batch, num_heads, seq_len_q, seq_len_k]
        """
        d_k = query.shape[-1]

        # Step 1: 计算注意力分数
        # [batch, heads, seq_q, head_dim] @ [batch, heads, head_dim, seq_k]
        # -> [batch, heads, seq_q, seq_k]
        scores = torch.matmul(query, key.transpose(-2, -1))

        # Step 2: 缩放
        scores = scores / math.sqrt(d_k)

        # Step 3: 应用掩码（如果有）
        if mask is not None:
            # mask 中 True 的位置会被设为 -inf
            scores = scores.masked_fill(mask, float('-inf'))

        # Step 4: Softmax 归一化
        attn_weights = F.softmax(scores, dim=-1)

        # Step 5: Dropout（训练时）
        attn_weights = self.dropout(attn_weights)

        # Step 6: 加权求和
        # [batch, heads, seq_q, seq_k] @ [batch, heads, seq_k, head_dim]
        # -> [batch, heads, seq_q, head_dim]
        output = torch.matmul(attn_weights, value)

        if return_weights:
            return output, attn_weights
        return output


class MultiHeadSelfAttention(torch.nn.Module):
    """多头自注意力完整实现"""

    def __init__(self, hidden_size, num_heads, dropout=0.0):
        super().__init__()
        assert hidden_size % num_heads == 0

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        # 投影层
        self.q_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)

        self.attention = ScaledDotProductAttention(dropout)

    def forward(self, hidden_states, attention_mask=None, use_cache=False, past_kv=None):
        """
        Args:
            hidden_states: [batch, seq_len, hidden_size]
            attention_mask: [batch, 1, seq_len, seq_len]
            use_cache: 是否返回 KV Cache
            past_kv: (past_key, past_value) 之前缓存的 KV

        Returns:
            output: [batch, seq_len, hidden_size]
            present_kv (optional): (key, value) 当前的 KV
        """
        batch_size, seq_len, _ = hidden_states.shape

        # 投影
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        # 重塑为多头形式
        query = self._reshape_for_multihead(query, batch_size)
        key = self._reshape_for_multihead(key, batch_size)
        value = self._reshape_for_multihead(value, batch_size)

        # 如果有过去的 KV，拼接
        if past_kv is not None:
            past_key, past_value = past_kv
            key = torch.cat([past_key, key], dim=2)
            value = torch.cat([past_value, value], dim=2)

        # 计算注意力
        attn_output = self.attention(query, key, value, mask=attention_mask)

        # 重塑回原始形状
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.hidden_size)

        # 输出投影
        output = self.o_proj(attn_output)

        if use_cache:
            return output, (key, value)
        return output

    def _reshape_for_multihead(self, x, batch_size):
        """[batch, seq, hidden] -> [batch, heads, seq, head_dim]"""
        x = x.view(batch_size, -1, self.num_heads, self.head_dim)
        return x.transpose(1, 2)


# 使用示例
def demo_attention():
    batch_size = 2
    seq_len = 10
    hidden_size = 256
    num_heads = 8

    # 创建模型
    attn = MultiHeadSelfAttention(hidden_size, num_heads)

    # 创建输入
    x = torch.randn(batch_size, seq_len, hidden_size)

    # 创建因果掩码
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
    mask = mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq, seq]

    # 前向传播
    output, (key, value) = attn(x, attention_mask=mask, use_cache=True)

    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"KV Cache shapes: K={key.shape}, V={value.shape}")

if __name__ == "__main__":
    demo_attention()

def generate_with_kv_cache(model, prompt_tokens, max_new_tokens=100):
    """使用 KV Cache 进行自回归生成"""

    # Prefill 阶段：处理 prompt
    hidden = model.embed(prompt_tokens)  # [1, prompt_len, hidden]
    output, kv_cache = model.attention(hidden, use_cache=True)

    # Decode 阶段：逐 token 生成
    generated = []
    for _ in range(max_new_tokens):
        # 只处理最后一个 token
        last_hidden = hidden[:, -1:, :]  # [1, 1, hidden]

        # 使用缓存的 KV
        output, kv_cache = model.attention(
            last_hidden,
            use_cache=True,
            past_kv=kv_cache
        )

        # 预测下一个 token
        next_token = model.lm_head(output[:, -1, :]).argmax(dim=-1)
        generated.append(next_token)

        # 更新 hidden（仅用于下一次迭代的 embedding）
        hidden = model.embed(next_token.unsqueeze(1))

        if next_token == EOS_TOKEN:
            break

    return generated