# 标准 Attention (PyTorch 实现)
import torch
import torch.nn.functional as F

def standard_attention(Q, K, V, scale=None):
    """
    Q, K, V: (batch, seq_len, d_k)
    """
    d_k = Q.size(-1)
    if scale is None:
        scale = d_k ** -0.5
    
    # Step 1: 计算注意力分数
    scores = torch.matmul(Q, K.transpose(-2, -1)) * scale
    # O(N²) 内存！
    
    # Step 2: Softmax
    attn_weights = F.softmax(scores, dim=-1)
    # O(N²) 内存！
    
    # Step 3: 加权求和
    outputs = torch.matmul(attn_weights, V)
    
    return outputs

# Flash Attention 的核心思想
def flash_attention_tiled(Q, K, V, block_size=64):
    """
    Flash Attention 使用 Tiling 技术:
    1. 将 Q, K, V 分成 blocks
    2. 逐 block 计算，避免 Materialization
    3. 在 SRAM 上计算，结果流式写回 HBM
    """
    batch_size, seq_len, d_k = Q.shape
    
    # 初始化输出和归一化因子
    outputs = torch.zeros_like(Q)
    l = torch.zeros((batch_size, seq_len, 1))  # 归一化因子
    m = torch.full((batch_size, seq_len, 1), -float('inf'))  # 最大值
    
    # 分块处理
    for i in range(0, seq_len, block_size):
        # 加载一个 Q block 到 SRAM
        Q_block = Q[:, i:i+block_size, :]
        
        for j in range(0, seq_len, block_size):
            # 加载 K, V block 到 SRAM
            K_block = K[:, j:j+block_size, :]
            V_block = V[:, j:j+block_size, :]
            
            # 在 SRAM 上计算（远快于 HBM）
            block_scores = torch.matmul(Q_block, K_block.transpose(-2, -1))
            block_scores = block_scores / (d_k ** 0.5)
            
            # Safe softmax
            block_max = block_scores.amax(dim=-1, keepdim=True)
            block_scores_minus_max = block_scores - block_max
            
            # 更新最大值和归一化因子
            new_m = torch.maximum(m[:, i:i+block_size], block_max)
            
            # ... 完整的 safe softmax 计算
            # ...
            
    return outputs

# 方式一：pip 安装（推荐）
pip install flash-attn

# 方式二：从源码安装
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention

# 安装 PyTorch 扩展
pip install .

# RTX 3090 / A100 (sm_80/86)
pip install flash-attn --no-build-isolation

# H100 (sm_90)
pip install flash-attn --no-build-isolation --index-url https://wheels.flash-attention.com/3.4/

# 使用 Docker
docker run --gpus all -it ghcr.io/dao-aeilab/flash-attention:latest

import torch
from flash_attn import flash_attn_func

# 检查版本
print(flash_attn_func.__module__)  # 应该输出 flash_attn

# 检查 CUDA 可用性
print(torch.cuda.is_available())  # True
print(torch.cuda.get_device_name(0))  # NVIDIA A100-SXM4-80GB

import torch
from flash_attn import flash_attn_func, flash_attn_qkvpacked_func

# 标准调用
Q = torch.randn(2, 4, 64, dtype=torch.float16, device='cuda')
K = torch.randn(2, 4, 64, dtype=torch.float16, device='cuda')
V = torch.randn(2, 4, 64, dtype=torch.float16, device='cuda')

# Flash Attention 前向计算
output = flash_attn_func(Q, K, V, dropout_p=0.0, softmax_mode='expander')
print(output.shape)  # torch.Size([2, 4, 64])

# 误差对比（应该 < 1e-3）
def standard_attention(Q, K, V):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    return torch.matmul(torch.softmax(scores, dim=-1), V)

standard_out = standard_attention(Q, K, V)
diff = (output.float() - standard_out.float()).abs().max()
print(f"Max difference: {diff.item():.6f}")  # 应该 < 1e-3

from flash_attn import flash_attn_qkvpacked_func

# QKV 打包格式 (更高效)
qkv = torch.randn(2, 4, 96, dtype=torch.float16, device='cuda')
# 96 = 3 * 32 (Q, K, V 各 32 维)

# 一次调用完成 QKV 分割 + Attention
output = flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_mode='expander')

from flash_attn import flash_attn_func

# 使用 cuSeDensify 模式（处理稀疏/不规则注意力）
# 适用于不确定长度的序列
Q = torch.randn(2, None, 64, dtype=torch.float16, device='cuda')
K = torch.randn(2, None, 64, dtype=torch.float16, device='cuda')
V = torch.randn(2, None, 64, dtype=torch.float16, device='cuda')
cu_seqlens_q = torch.tensor([0, 2, 4], dtype=torch.int32, device='cuda')  # 每个序列的长度
cu_seqlens_k = torch.tensor([0, 2, 4], dtype=torch.int32, device='cuda')

output = flash_attn_func(
    Q, K, V,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=4,
    max_seqlen_k=4
)

# 使用 Hugging Face Transformers (自动集成 Flash Attention)
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "meta-llama/Llama-2-7b-hf"

# Flash Attention 自动启用（当可用时）
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    attn_implementation="flash_attention_2"  # 明确指定
)

# 推理
inputs = tokenizer("Hello, world!", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=100)

# xFormers 也支持 Flash Attention
from xformers.ops import memory_efficient_attention

output = memory_efficient_attention(
    Q, K, V,
    attn_bias=None,
    p=0.0,
    scale=None
)

# NVIDIA Megatron-LM
from megatron.core.models.gpt.gpt_model import GPTModel

model = GPTModel(
    args,
    num_tokentypes=0,
    parallel_output=False
)

# 在配置中启用 Flash Attention
# megatron_config = {
#     "attention": "flash",
#     "num_attention_heads": 32,
#     "hidden_size": 4096
# }

import torch
import time
from flash_attn import flash_attn_func

def benchmark_attention(seq_len, batch_size=4, heads=16, head_dim=64):
    Q = torch.randn(batch_size, seq_len, heads * head_dim, 
                    dtype=torch.float16, device='cuda')
    K = torch.randn(batch_size, seq_len, heads * head_dim,
                    dtype=torch.float16, device='cuda')
    V = torch.randn(batch_size, seq_len, heads * head_dim,
                    dtype=torch.float16, device='cuda')
    
    # Warmup
    for _ in range(10):
        _ = flash_attn_func(Q, K, V)
    
    # Benchmark
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(100):
        _ = flash_attn_func(Q, K, V)
    torch.cuda.synchronize()
    
    elapsed = (time.time() - start) / 100 * 1000  # ms
    return elapsed

# 测试不同序列长度
for seq_len in [512, 1024, 2048, 4096]:
    ms = benchmark_attention(seq_len)
    print(f"Seq len {seq_len}: {ms:.2f} ms")

# 安装 Flash Attention-2
pip install flash-attn --no-build-isolation

# 验证版本
python -c "import flash_attn; print(flash_attn.__version__)"

from flash_attn import flash_attn_func

# FA2 的 API 与 FA1 相同
output = flash_attn_func(Q, K, V, dropout_p=0.0, softmax_mode='expander')

# LLaMA 使用 Flash Attention
# 模型配置中指定
config = {
    "hidden_size": 4096,
    "num_attention_heads": 32,
    "num_key_value_heads": 32,
    "attn_implementation": "flash_attention_2"
}

# Mistral 7B
from transformers import MistralForCausalLM

model = MistralForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.float16
)

# CodeLlama (LLaMA + 代码能力)
from transformers import CodeLlamaForCausalLM

model = CodeLlamaForCausalLM.from_pretrained(
    "codellama/CodeLlama-7b-hf",
    attn_implementation="flash_attention_2"
)

import torch
import torch.nn as nn
from flash_attn import flash_attn_func

class FlashAttentionLayer(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = 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.head_dim).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Flash Attention
        attn_output = flash_attn_func(Q, K, V, dropout_p=0.0, softmax_mode='expander')
        
        # 恢复形状
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        return self.W_o(attn_output)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from flash_attn import flash_attn_func

# 分布式训练示例
model = FlashAttentionModel()
model = DDP(model)

for batch in dataloader:
    optimizer.zero_grad()
    
    Q, K, V = model(batch)
    output = flash_attn_func(Q, K, V)
    
    loss.backward()
    optimizer.step()

@article{dao2022flashattention,
  title={FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness},
  author={Dao, Tri},
  journal={Advances in Neural Information Processing Systems},
  year={2022}
}

@article{dao2023flashattention2,
  title={FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author={Dao, Tri},
  journal={arXiv preprint arXiv:2307.08691},
  year={2023}
}