xFormers Transformer组件库

import xformers
print(f"xFormers version: {xformers.__version__}")

# 验证内存高效注意力
import torch
from xformers.ops import memory_efficient_attention

q = torch.randn(2, 512, 8, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 512, 8, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 512, 8, 64, device='cuda', dtype=torch.float16)
out = memory_efficient_attention(q, k, v)
print(f"Output shape: {out.shape}")  # torch.Size([2, 512, 8, 64])

from xformers.ops import memory_efficient_attention

out = memory_efficient_attention(
    query,            # 查询张量，形状 [B, Mq, Hq, K]
    key,              # 键张量，形状 [B, Mkv, Hkv, K]
    value,            # 值张量，形状 [B, Mkv, Hkv, Kv]
    attn_bias=None,   # 注意力偏置（支持因果掩码、位置偏置等）
    p=0.0,            # dropout 概率
    scale=None,       # 缩放因子，默认 1/sqrt(K)
)

import torch
from xformers.ops import memory_efficient_attention

B, M, H, K = 2, 512, 8, 64
q = torch.randn(B, M, H, K, device='cuda', dtype=torch.float16)
k = torch.randn(B, M, H, K, device='cuda', dtype=torch.float16)
v = torch.randn(B, M, H, K, device='cuda', dtype=torch.float16)

# 基本使用
out = memory_efficient_attention(q, k, v)
print(f"Basic output: {out.shape}")

# 因果注意力
from xformers.ops import LowerTriangularMask
out_causal = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())
print(f"Causal output: {out_causal.shape}")

# 带 dropout
out_dropout = memory_efficient_attention(q, k, v, p=0.1, attn_bias=LowerTriangularMask())
print(f"Dropout output: {out_dropout.shape}")

from xformers.ops import (
    LowerTriangularMask,        # 因果掩码（下三角）
    LowerTriangularMaskWithTensorBias,  # 因果掩码 + 自定义偏置
    AttentionBias,              # 基类
)

from xformers.ops import memory_efficient_attention, LowerTriangularMask

# 因果掩码
out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

from xformers.ops import memory_efficient_attention, LowerTriangularMaskWithTensorBias

# 因果掩码 + 位置偏置
bias = torch.randn(1, 1, M, M, device='cuda', dtype=torch.float16) * 0.01
out = memory_efficient_attention(
    q, k, v,
    attn_bias=LowerTriangularMaskWithTensorBias(bias)
)

from xformers.ops import (
    memory_efficient_attention,              # 自动选择最优后端
    scaled_dot_product_attention,            # 标准SDPA实现
)

# 查看当前使用的后端
from xformers.ops import memory_efficient_attention

# 使用 attn_bias 参数控制后端行为
# 不传 attn_bias 时自动选择最优后端
out = memory_efficient_attention(q, k, v)

# 也可以使用 PyTorch 2.0+ 的原生 SDPA
import torch.nn.functional as F
out_native = F.scaled_dot_product_attention(
    q.transpose(1, 2),  # PyTorch 使用 BHSD 布局
    k.transpose(1, 2),
    v.transpose(1, 2),
    is_causal=True,
)

from xformers.ops import SwiGLU

# 创建 SwiGLU 层
swiglu = SwiGLU(
    in_features=4096,    # 输入维度
    hidden_features=11008,  # 隐藏层维度（通常约为 in_features 的 8/3）
    out_features=4096,   # 输出维度
    bias=True,           # 是否使用偏置
)

# 使用
x = torch.randn(2, 128, 4096, device='cuda', dtype=torch.float16)
out = swiglu(x)
print(f"SwiGLU output: {out.shape}")  # [2, 128, 4096]

from xformers.ops import RotEmbOp

# RoPE 操作符
rotary_emb = RotEmbOp(
    dim=64,           # 应用RoPE的维度
    seq_len=2048,     # 最大序列长度
    base=10000,       # 频率基数θ
)

import torch
from xformers.helpers.rope import (
    apply_rotary_emb,
    get_rotary_embedding,
)

# 获取旋转位置编码
d_head = 64
max_seq_len = 2048
freqs = get_rotary_embedding(max_seq_len, d_head, device='cuda')

# 应用旋转位置编码
q = torch.randn(2, 512, 8, 64, device='cuda', dtype=torch.float16)
positions = torch.arange(512, device='cuda')
q_rotated = apply_rotary_emb(q, freqs, positions)

from xformers.components.attention import LocalAttention

local_attn = LocalAttention(
    window_size=256,      # 注意力窗口大小
    attention_dropout=0.1,
    causal=True,
)

# 使用
q = torch.randn(2, 1024, 8, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device='cuda', dtype=torch.float16)
out = local_attn(q, k, v)

from xformers.components.attention import RandomAttention

random_attn = RandomAttention(
    r=0.1,                # 采样比例
    attention_dropout=0.1,
)

out = random_attn(q, k, v)

from xformers.components.attention import BlockSparseAttention

# 定义稀疏模式
# 1 表示计算该块，0 表示跳过
layout = torch.randint(0, 2, (16, 16), device='cuda')
# 确保因果性：上三角为0
layout = layout.tril()

sparse_attn = BlockSparseAttention(
    layout=layout,        # 稀疏布局矩阵
    block_size=64,        # 块大小
    attention_dropout=0.0,
)

out = sparse_attn(q, k, v)

import torch
import torch.nn as nn
from xformers.ops import memory_efficient_attention, LowerTriangularMask

class XFormersAttention(nn.Module):
    """使用xFormers的注意力层"""
    def __init__(self, d_model, num_heads, causal=True):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.causal = causal

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, S, D = x.shape
        H, K = self.num_heads, self.head_dim

        q = self.q_proj(x).view(B, S, H, K)
        k = self.k_proj(x).view(B, S, H, K)
        v = self.v_proj(x).view(B, S, H, K)

        attn_bias = LowerTriangularMask() if self.causal else None
        out = memory_efficient_attention(q, k, v, attn_bias=attn_bias)

        out = out.reshape(B, S, D)
        return self.out_proj(out)


class XFormersTransformerBlock(nn.Module):
    """使用xFormers组件的Transformer块"""
    def __init__(self, d_model, num_heads, d_ff=None, causal=True):
        super().__init__()
        d_ff = d_ff or 4 * d_model
        self.attn = XFormersAttention(d_model, num_heads, causal=causal)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

import torch
import torch.nn as nn
from xformers.ops import SwiGLU

class LLaMAFFN(nn.Module):
    """LLaMA风格的前馈网络，使用SwiGLU激活"""
    def __init__(self, d_model, d_ff=None, multiple_of=256):
        super().__init__()
        # LLaMA中d_ff通常是d_model的约8/3倍，并对齐到multiple_of
        if d_ff is None:
            d_ff = int(2 * 4 * d_model / 3)
            d_ff = multiple_of * ((d_ff + multiple_of - 1) // multiple_of)

        self.swiglu = SwiGLU(
            in_features=d_model,
            hidden_features=d_ff,
            out_features=d_model,
            bias=False,  # LLaMA不使用偏置
        )

    def forward(self, x):
        return self.swiglu(x)

# 使用
ffn = LLaMAFFN(d_model=4096).cuda().half()
x = torch.randn(2, 512, 4096, device='cuda', dtype=torch.float16)
out = ffn(x)
print(f"LLaMA FFN output: {out.shape}")

from transformers import AutoModelForCausalLM

# 方式1: 通过attn_implementation参数启用
model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-350m",
    torch_dtype=torch.float16,
    device_map="auto",
    attn_implementation="sdpa",  # 使用缩放点积注意力（xFormers后端）
)

# 方式2: 设置环境变量
# XFORMERS_MORE_ADVISORY=1 会启用xFormers的更多优化
import os
os.environ["XFORMERS_MORE_ADVISORY"] = "1"

import torch
from xformers.ops import memory_efficient_attention, LowerTriangularMask

def grouped_query_attention_xformers(q, k, v, num_q_heads, num_kv_heads, causal=True):
    """
    使用xFormers实现GQA
    q: [B, S, num_q_heads, K]
    k: [B, S, num_kv_heads, K]
    v: [B, S, num_kv_heads, Kv]
    """
    attn_bias = LowerTriangularMask() if causal else None
    # xFormers的memory_efficient_attention原生支持不同头数的Q和K/V
    out = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
    return out

# 示例: 32个查询头，8个KV头
B, S, num_q_heads, num_kv_heads, K = 2, 512, 32, 8, 64
q = torch.randn(B, S, num_q_heads, K, device='cuda', dtype=torch.float16)
k = torch.randn(B, S, num_kv_heads, K, device='cuda', dtype=torch.float16)
v = torch.randn(B, S, num_kv_heads, K, device='cuda', dtype=torch.float16)

out = grouped_query_attention_xformers(q, k, v, num_q_heads, num_kv_heads)
print(f"GQA output: {out.shape}")  # [2, 512, 32, 64]

import torch
from xformers.ops import memory_efficient_attention, LowerTriangularMask, AttentionBias

def inference_with_kv_cache(q_new, k_cache, v_cache, cache_len):
    """
    使用KV Cache进行推理
    q_new: [B, 1, H, K] - 当前步的查询
    k_cache: [B, max_len, H, K] - K缓存（已填充cache_len个位置）
    v_cache: [B, max_len, H, K] - V缓存
    """
    # 只使用缓存中已填充的部分
    k_used = k_cache[:, :cache_len]
    v_used = v_cache[:, :cache_len]

    # 单步生成时无需因果掩码（query只有1个位置）
    out = memory_efficient_attention(q_new, k_used, v_used)
    return out

# 示例
B, H, K, max_len = 1, 32, 64, 2048
k_cache = torch.zeros(B, max_len, H, K, device='cuda', dtype=torch.float16)
v_cache = torch.zeros(B, max_len, H, K, device='cuda', dtype=torch.float16)

# 模拟: 已有128个token在缓存中
cache_len = 128
k_cache[:, :cache_len] = torch.randn(B, cache_len, H, K, device='cuda', dtype=torch.float16)
v_cache[:, :cache_len] = torch.randn(B, cache_len, H, K, device='cuda', dtype=torch.float16)

# 生成新token
q_new = torch.randn(B, 1, H, K, device='cuda', dtype=torch.float16)
out = inference_with_kv_cache(q_new, k_cache, v_cache, cache_len)
print(f"Inference output: {out.shape}")

# xFormers内部的后端选择逻辑（简化版）
def memory_efficient_attention(query, key, value, attn_bias=None, p=0.0, scale=None):
    # 1. 检查输入是否满足FlashAttention要求
    if _can_use_flash(query, key, value, attn_bias):
        return _flash_attention(query, key, value, attn_bias, p, scale)

    # 2. 检查cutlass后端
    if _can_use_cutlass(query, key, value, attn_bias):
        return _cutlass_attention(query, key, value, attn_bias, p, scale)

    # 3. 检查triton后端
    if _can_use_triton(query, key, value, attn_bias):
        return _triton_attention(query, key, value, attn_bias, p, scale)

    # 4. 退回到PyTorch原生实现
    return _torch_attention(query, key, value, attn_bias, p, scale)

# 标准实现（未融合）- 3次内核调用
def swiglu_unfused(x, w1, b1, w2, b2, w3, b3):
    gate = F.linear(x, w1, b1)       # 内核1: 线性变换1
    up = F.linear(x, w2, b2)          # 内核2: 线性变换2
    mid = F.silu(gate) * up           # 内核3: silu + 逐元素乘
    out = F.linear(mid, w3, b3)       # 内核4: 输出投影
    return out

# xFormers融合实现 - 更少的内核调用和内存访问
# SwiGLU类内部将前两步融合，减少中间张量的HBM读写

# RoPE的向量化实现（xFormers内部优化版）
def apply_rotary_emb(x, freqs, positions):
    """
    x: [B, S, H, D] 输入张量
    freqs: [max_seq_len, D/2] 预计算的频率
    positions: [S] 位置索引
    """
    D = x.shape[-1]
    # 将x分成奇偶两组（对应2D子空间的两个分量）
    x1 = x[..., :D//2]  # 偶数维度
    x2 = x[..., D//2:]  # 奇数维度

    # 获取当前位置的旋转参数
    cos = torch.cos(freqs[positions])  # [S, D/2]
    sin = torch.sin(freqs[positions])  # [S, D/2]

    # 应用旋转
    out1 = x1 * cos - x2 * sin
    out2 = x1 * sin + x2 * cos

    return torch.cat([out1, out2], dim=-1)

# 检查版本兼容性
import xformers
import torch

print(f"PyTorch version: {torch.__version__}")
print(f"xFormers version: {xformers.__version__}")

# xFormers版本应与PyTorch版本匹配
# 不匹配可能导致运行时错误或性能下降

import os

# 强制使用特定后端
os.environ["XFORMERS_MEM_EFF_ATTENTION"] = "cutlass"  # 强制cutlass

# 启用调试信息
os.environ["XFORMERS_MORE_ADVISORY"] = "1"

# 禁用xFormers，退回PyTorch原生实现
os.environ["DISABLE_XFORMERS"] = "1"

from xformers.ops import LowerTriangularMask

# 正确：复用同一个mask对象
causal_mask = LowerTriangularMask()
out1 = memory_efficient_attention(q1, k1, v1, attn_bias=causal_mask)
out2 = memory_efficient_attention(q2, k2, v2, attn_bias=causal_mask)

# 错误：偏置张量形状不匹配
# bias = torch.randn(B, H, M, M)  # 形状可能不匹配

# xFormers / FlashAttention 布局
q_bshd = torch.randn(B, S, H, D, device='cuda', dtype=torch.float16)

# PyTorch 标准布局
q_bhsd = q_bshd.transpose(1, 2)

# 从 PyTorch 布局转换为 xFormers 布局
q_for_xformers = q_bhsd.transpose(1, 2)
out = memory_efficient_attention(q_for_xformers, k_for_xformers, v_for_xformers)

# PyTorch 2.0+ SDPA（推荐用于新项目）
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# xFormers（更灵活，支持更多特性）
from xformers.ops import memory_efficient_attention, LowerTriangularMask
out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

import torch.distributed as dist
from xformers.ops import memory_efficient_attention

# xFormers 的 memory_efficient_attention 本身不处理分布式
# 但可以在每个 GPU 上独立运行，配合 DDP/FSDP 使用

# 序列并行（需要配合其他库）
# xFormers 本身不提供 Ring Attention 实现
# 对于超长序列的分布式训练，建议使用 FlashAttention 的 Ring Attention

import torch
from xformers.ops import memory_efficient_attention

# 推理时使用 torch.no_grad() 和 torch.inference_mode()
with torch.inference_mode():
    out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

# 使用 KV Cache 时，注意增量更新
# xFormers 没有内置的 KV Cache 管理，需要手动管理
# 推荐在推理场景使用 FlashAttention 的 flash_attn_with_kvcache

from torch.cuda.amp import autocast

# xFormers 在 fp16/bf16 下性能最佳
with autocast(dtype=torch.float16):
    out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

# bf16 在训练中更稳定（Ampere+架构）
with autocast(dtype=torch.bfloat16):
    out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())