import flash_attn_interface
flash_attn_interface.flash_attn_func(...)

from flash_attn import (
    flash_attn_func,              # 标准 Q, K, V 分离输入
    flash_attn_qkvpacked_func,    # QKV 打包输入 (更快)
    flash_attn_varlen_func,       # 变长序列支持
    flash_attn_with_kvcache,      # KV Cache 推理加速
)

import torch
from flash_attn import flash_attn_func

# 参数设定
batch_size = 2
seqlen = 1024
nheads = 32         # Query 的注意力头数
nheads_k = 8        # Key/Value 的注意力头数 (GQA: nheads 必须能被 nheads_k 整除)
headdim = 128       # 每个注意力头的维度 (支持最大 256)

# 创建输入张量
# 注意: 张量 layout 为 (batch_size, seqlen, nheads, headdim)
q = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
k = torch.randn(batch_size, seqlen, nheads_k, headdim, device='cuda', dtype=torch.float16)
v = torch.randn(batch_size, seqlen, nheads_k, headdim, device='cuda', dtype=torch.float16)

# 基础调用 (非因果注意力)
output = flash_attn_func(q, k, v)
# output shape: (batch_size, seqlen, nheads, headdim)

# 因果注意力 (自回归模型使用)
output_causal = flash_attn_func(q, k, v, causal=True)

# 带滑动窗口的局部注意力 (如 Mistral 7B)
output_local = flash_attn_func(q, k, v, causal=True, window_size=(512, 0))

print(f"输出形状: {output.shape}")  # torch.Size([2, 1024, 32, 128])

import torch
from flash_attn import flash_attn_qkvpacked_func

batch_size = 2
seqlen = 1024
nheads = 32
headdim = 128

# QKV 打包为一个张量, shape: (batch_size, seqlen, 3, nheads, headdim)
qkv = torch.randn(batch_size, seqlen, 3, nheads, headdim, device='cuda', dtype=torch.float16)

# 因果注意力
output = flash_attn_qkvpacked_func(qkv, causal=True)
print(f"输出形状: {output.shape}")  # torch.Size([2, 1024, 32, 128])

import torch
from flash_attn import flash_attn_varlen_func

# 假设 batch 中有 3 个序列，长度分别为 100, 200, 150
seqlens = [100, 200, 150]
total_tokens = sum(seqlens)  # 450
nheads = 32
nheads_k = 8
headdim = 128

# 所有序列拼接成一个张量, shape: (total_tokens, nheads, headdim)
q = torch.randn(total_tokens, nheads, headdim, device='cuda', dtype=torch.float16)
k = torch.randn(total_tokens, nheads_k, headdim, device='cuda', dtype=torch.float16)
v = torch.randn(total_tokens, nheads_k, headdim, device='cuda', dtype=torch.float16)

# 累积序列长度 (batch_size + 1,), dtype 必须是 int32
cu_seqlens_q = torch.tensor([0, 100, 300, 450], device='cuda', dtype=torch.int32)
cu_seqlens_k = torch.tensor([0, 100, 300, 450], device='cuda', dtype=torch.int32)

max_seqlen_q = max(seqlens)
max_seqlen_k = max(seqlens)

output = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=max_seqlen_q,
    max_seqlen_k=max_seqlen_k,
    causal=True,
)
print(f"输出形状: {output.shape}")  # torch.Size([450, 32, 128])

import torch
from flash_attn import flash_attn_with_kvcache

batch_size = 4
max_seqlen = 2048     # 预分配的最大序列长度
nheads = 32
nheads_k = 8
headdim = 128

# 预分配 KV Cache
k_cache = torch.zeros(batch_size, max_seqlen, nheads_k, headdim, device='cuda', dtype=torch.float16)
v_cache = torch.zeros(batch_size, max_seqlen, nheads_k, headdim, device='cuda', dtype=torch.float16)

# 模拟已填充了 100 个 token 后, 新增 1 个 token
cache_seqlens = torch.full((batch_size,), 100, device='cuda', dtype=torch.int32)

# 新增的 query, key, value (seqlen=1 对应一个新 token)
q_new = torch.randn(batch_size, 1, nheads, headdim, device='cuda', dtype=torch.float16)
k_new = torch.randn(batch_size, 1, nheads_k, headdim, device='cuda', dtype=torch.float16)
v_new = torch.randn(batch_size, 1, nheads_k, headdim, device='cuda', dtype=torch.float16)

# 执行注意力计算，同时更新 KV Cache
output = flash_attn_with_kvcache(
    q_new,
    k_cache,
    v_cache,
    k=k_new,
    v=v_new,
    cache_seqlens=cache_seqlens,
    causal=True,
)
print(f"输出形状: {output.shape}")  # torch.Size([4, 1, 32, 128])
# k_cache 和 v_cache 已在位置 100 处被原地更新

:: Windows 上安装
set DISTUTILS_USE_SDK=1
pip install flash-attn --no-build-isolation

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

batch_size, seqlen, nheads, headdim = 2, 1024, 32, 128

# ========== PyTorch SDPA ==========
# 注意: SDPA 的张量 layout 为 (batch, nheads, seqlen, headdim)
q_sdpa = torch.randn(batch_size, nheads, seqlen, headdim, device='cuda', dtype=torch.float16)
k_sdpa = torch.randn(batch_size, nheads, seqlen, headdim, device='cuda', dtype=torch.float16)
v_sdpa = torch.randn(batch_size, nheads, seqlen, headdim, device='cuda', dtype=torch.float16)

output_sdpa = F.scaled_dot_product_attention(
    q_sdpa, k_sdpa, v_sdpa,
    is_causal=True,
)
# output shape: (batch_size, nheads, seqlen, headdim)

# ========== Flash Attention ==========
# 注意: Flash Attention 的张量 layout 为 (batch, seqlen, nheads, headdim)
q_fa = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
k_fa = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
v_fa = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)

output_fa = flash_attn_func(
    q_fa, k_fa, v_fa,
    causal=True,
)
# output shape: (batch_size, seqlen, nheads, headdim)