FlashAttention遇上旋转位置编码：RoPE是怎么跟注意力计算配合的？

某团队在昇腾NPU上跑ChatGLM-6B，用FlashAttention加速推理。跑了一段时间后发现模型的旋转位置编码（RoPE）效果不对------输入"北京是中国的首都"和"中国的首都是北京"，模型对这两个句子的注意力分布完全一样，没有体现出位置差异。

问题出在FlashAttention和RoPE的集成方式上。RoPE需要在Attention计算之前把位置信息注入到Q和K里，如果注入的位置信息和FlashAttention的在线Softmax状态不同步，会导致位置编码失效。

FlashAttention在前向上做了一次大改动（在线Softmax），RoPE如果集成方式不对，会跟FlashAttention产生冲突。今天把这个集成机制讲清楚，给出昇腾NPU上的正确实现。

先打个比方：录音里的时间戳

想象一段录音，每句话都标了时间戳："第1秒说Hello，第2秒说World"。但如果录音设备出了bug，所有时间戳都丢了，播放的时候只知道"有Hello和World"，不知道哪个先说、哪个后说------语义就完全变了。

RoPE就是给每个token加"时间戳"的机制------告诉模型"这个token在第几个位置"。FlashAttention在前向传播时改变了Softmax的计算方式，如果RoPE注入的位置信息没有跟着变，模型就会"失聪"，不知道token的顺序。

RoPE的数学原理

标准位置编码 vs 旋转位置编码

绝对位置编码（标准PE）：给每个位置i分配一个固定向量P_i，加到词嵌入上。

复制代码

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

x'_i = x_i + PE(i)  # 词嵌入 + 位置编码

旋转位置编码（RoPE）：不给每个位置分配固定向量，而是把Q和K旋转一个角度。

复制代码

RoPE的核心思想：
  两个token的注意力分数，应该跟它们的相对位置有关
  即：attention(q_i, k_j) = attention(q_i.rotate(θ_i), k_j.rotate(θ_j))
  
  数学上，这等价于：
  q_i^T k_j = (R_θ_i q_i)^T (R_θ_j k_j) = q_i^T R_θ_i^T R_θ_j k_j
            = q_i^T R_θ_{i-j} k_j

R_θ是一个旋转矩阵，把向量在复平面上旋转一个角度。

RoPE的具体实现

python 复制代码

import torch
import math

def precompute_freqs_cis(dim, end_idx, theta=10000.0):
    """预计算旋转角度"""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(end_idx)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)  # 复数

def apply_rotary_pos_emb(q, k, freqs_cis):
    """
    给Q和K应用RoPE
    
    参数：
      q: [B, H, S, D]（D是head_dim，通常是128）
      k: [B, H, S, D]
      freqs_cis: [S, D/2]（复数形式的旋转角度）
    """
    B, H, S, D = q.shape
    D_half = D // 2
    
    # 把q分成两半：[B, H, S, D] → [B, H, S, D/2, 2]
    q_complex = torch.view_as_complex(q.float().reshape(B, H, S, D_half, 2))
    k_complex = torch.view_as_complex(k.float().reshape(B, H, S, D_half, 2))
    
    # 逐位置旋转
    # freqs_cis: [S, D/2] → unsqueeze → [1, 1, S, D/2]
    q_rotated = q_complex * freqs_cis.unsqueeze(0).unsqueeze(0)
    k_rotated = k_complex * freqs_cis.unsqueeze(0).unsqueeze(0)
    
    # 转回实数：[B, H, S, D/2] → [B, H, S, D]
    q_out = torch.view_as_real(q_rotated).flatten(-2).to(q.dtype)
    k_out = torch.view_as_real(k_rotated).flatten(-2).to(k.dtype)
    
    return q_out, k_out

# 示例
seq_len = 4096
head_dim = 128
batch_size = 1
num_heads = 32

q = torch.randn(batch_size, num_heads, seq_len, head_dim, device='npu', dtype=torch.float16)
k = torch.randn(batch_size, num_heads, seq_len, head_dim, device='npu', dtype=torch.float16)

# 预计算旋转角度
freqs_cis = precompute_freqs_cis(head_dim, seq_len, theta=10000.0).to('npu')

# 应用RoPE
q_rotated, k_rotated = apply_rotary_pos_emb(q, k, freqs_cis)
print(f"Q shape: {q.shape} → {q_rotated.shape}")
print(f"K shape: {k.shape} → {k_rotated.shape}")

FlashAttention集成RoPE的三种方式

方式1：RoPE在FlashAttention之前（标准做法）

最常见的方式：先应用RoPE，再把旋转后的Q和K送入FlashAttention。

python 复制代码

def flash_attention_with_rope(q, k, v, freqs_cis, head_num, scale_value):
    """
    FlashAttention + RoPE（方式1：先RoPE后Attention）
    """
    
    # Step 1: 应用RoPE（旋转Q和K）
    q_rotated, k_rotated = apply_rotary_pos_emb(q, k, freqs_cis)
    
    # Step 2: FlashAttention计算（用旋转后的Q和K）
    output = npu_flash_attention(
        q_rotated, k_rotated, v,
        head_num=head_num,
        scale_value=scale_value
    )
    
    return output

# 使用
freqs_cis = precompute_freqs_cis(head_dim, seq_len).to('npu')
output = flash_attention_with_rope(q, k, v, freqs_cis, head_num=32, scale_value=1.0/(head_dim**0.5))

适用场景：推理和训练都适用，最通用的集成方式。

方式2：RoPE融合进FlashAttention（融合kernel）

把RoPE和Attention计算融合成一个kernel，减少HBM读写。

python 复制代码

class FuseRoPEFlashAttention(torch.nn.Module):
    """融合RoPE和FlashAttention的算子"""
    
    def __init__(self, head_dim=128, max_seq_len=4096, theta=10000.0):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        # 预计算旋转角度（避免每次都算）
        self.register_buffer(
            "freqs_cis",
            precompute_freqs_cis(head_dim, max_seq_len, theta)
        )
    
    def forward(self, q, k, v, seq_len=None, head_num=32):
        # 截取需要的seq_len
        if seq_len is None:
            seq_len = q.shape[2]
        
        freqs = self.freqs_cis[:seq_len]
        
        # 融合RoPE + FlashAttention（昇腾NPU支持融合kernel）
        output = fused_rope_flash_attention(
            q, k, v, freqs,
            head_num=head_num,
            block_size=128,
            scale_value=1.0 / (self.head_dim ** 0.5)
        )
        
        return output

⚠️ 踩坑预警：融合kernel的RoPE部分需要正确处理旋转角度。如果旋转角度的预计算精度不够（比如用了float16而不是float32），长期累积的旋转误差会导致位置编码精度退化。

方式3：RoPE应用到子空间（GLM/LLaMA的做法）

LLaMA把RoPE应用到Q和K的一部分维度上（通常是后半部分），前半部分保持不变。

python 复制代码

def apply_partial_rope(q, k, freqs_cis, rotary_dim=None):
    """
    部分RoPE：只旋转前rotary_dim个维度
    
    LLaMA的rotary_dim = head_dim（即全部旋转）
    ChatGLM的rotary_dim = head_dim // 2（只旋转后半部分的一半）
    """
    
    if rotary_dim is None:
        rotary_dim = q.shape[-1]
    
    # 截取要旋转的维度
    q_rot = q[..., :rotary_dim]
    q_pass = q[..., rotary_dim:]
    
    k_rot = k[..., :rotary_dim]
    k_pass = k[..., rotary_dim:]
    
    # 应用RoPE
    q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, freqs_cis)
    
    # 拼接回来
    q_out = torch.cat([q_rot, q_pass], dim=-1)
    k_out = torch.cat([k_rot, k_pass], dim=-1)
    
    return q_out, k_out

常见问题排查

问题1：RoPE跟FlashAttention的seq_len不匹配

python 复制代码

def check_rope_seq_len(q, k, freqs_cis):
    """检查RoPE seq_len是否匹配"""
    
    q_seq_len = q.shape[2]
    k_seq_len = k.shape[2]
    rope_seq_len = freqs_cis.shape[0]
    
    print(f"Q seq_len: {q_seq_len}")
    print(f"K seq_len: {k_seq_len}")
    print(f"RoPE freq_cis seq_len: {rope_seq_len}")
    
    if q_seq_len != rope_seq_len:
        print("⚠️ Q seq_len与RoPE不匹配！")
        print(f"  解决方案：重新生成freqs_cis，seq_len={max(q_seq_len, k_seq_len)}")
        return False
    
    if k_seq_len != rope_seq_len:
        print("⚠️ K seq_len与RoPE不匹配！")
        return False
    
    print("✅ RoPE seq_len匹配")
    return True

# 正确用法：生成足够的RoPE长度
max_seq_len = 8192  # 支持的最大长度
freqs_cis = precompute_freqs_cis(head_dim, max_seq_len).to('npu')

# 推理时传入实际的seq_len
actual_seq_len = current_input_ids.shape[1]
q = model.q_proj(hidden_states)
k = model.k_proj(hidden_states)
q, k = apply_rotary_pos_emb(q, k, freqs_cis[:actual_seq_len])

问题2：RoPE和FlashAttention的梯度不同步

训练时，RoPE的反向传播需要跟FlashAttention的反向传播协调。

python 复制代码

class RoPEFlashAttentionFunction(torch.autograd.Function):
    """带RoPE的FlashAttention反向传播"""
    
    @staticmethod
    def forward(ctx, q, k, v, freqs_cis, head_num, scale_value):
        
        # 前向：先RoPE
        q_rot, k_rot = apply_rotary_pos_emb(q, k, freqs_cis)
        
        # FlashAttention前向
        output = npu_flash_attention(q_rot, k_rot, v, head_num=head_num, scale_value=scale_value)
        
        # 保存需要反向传播的中间结果
        ctx.save_for_backward(q, k, v, freqs_cis, q_rot, k_rot)
        ctx.head_num = head_num
        ctx.scale_value = scale_value
        
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        q, k, v, freqs_cis, q_rot, k_rot = ctx.saved_tensors
        
        # FlashAttention反向
        grad_q_rot, grad_k_rot, grad_v = npu_flash_attention_backward(
            grad_output, q_rot, k_rot, v, output,
            head_num=ctx.head_num,
            scale_value=ctx.scale_value
        )
        
        # RoPE反向（旋转的梯度）
        grad_q, grad_k = apply_rotary_pos_emb_backward(grad_q_rot, grad_k_rot, freqs_cis)
        
        return grad_q, grad_k, grad_v, None, None, None

总结：集成清单

FlashAttention + RoPE集成，按这个清单检查：

检查项	操作	判断标准
seq_len匹配	检查freqs_cis长度≥实际输入长度	长度不足→重新生成
dtype一致	检查freqs_cis是float32	float16→改用float32
rotary_dim	确认模型的rotary_dim设置	跟模型配置一致
梯度同步	检查反向传播中RoPE的梯度	grad_q和grad_k不为空
融合kernel	确认昇腾NPU支持融合版本	昇腾910B+才支持

代码和文档：

https://atomgit.com/cann/ops-transformer