FlashAttention的Attention Sink现象：为什么模型总是盯着第一个token看？

某团队在昇腾NPU上跑Llama-2-7B的长上下文推理，输入了一篇很长的文章（16384个token），让模型总结文章内容。他们用FlashAttention的注意力可视化工具分析模型的注意力分布，发现了一个奇怪的现象——无论输入多长，模型总是把大量的注意力放在第一个token（通常是""）上，其他token分到的注意力很少。他们怀疑FlashAttention出了问题，或者模型有bug。排查了

徐安安_ye1

224人浏览 · 2026-05-25 16:49:20

徐安安_ye1 · 2026-05-25 16:49:20 发布

FlashAttention的Attention Sink现象：为什么模型总是盯着第一个token看？

某团队在昇腾NPU上跑Llama-2-7B的长上下文推理，输入了一篇很长的文章（16384个token），让模型总结文章内容。他们用FlashAttention的注意力可视化工具分析模型的注意力分布，发现了一个奇怪的现象——无论输入多长，模型总是把大量的注意力放在第一个token（通常是"~~"）上，其他token分到的注意力很少。~~

他们怀疑FlashAttention出了问题，或者模型有bug。排查了一圈之后发现，FlashAttention没问题，这是大模型的正常现象，叫做Attention Sink。

Attention Sink是LLM的一个特性，不是bug。理解这个现象，对于正确使用FlashAttention和解释模型行为很重要。今天把这个机制讲清楚，以及怎么在开发中处理它。

先打个比方：会议室里的"锚点"

想象一个会议室里有100个人在讨论问题。每次讨论一个问题，大家都会先看向会议主持人（第一个到场的人），然后再互相讨论。主持人就像一个"锚"——大家都看他，是为了"对齐"讨论的基准。如果不先看主持人，讨论就容易跑偏。

Attention Sink就是模型的"主持人"：第一个token（通常是<s>或[CLS]）被大量注意力"盯着看"，不是因为它有多重要，而是因为它起到一个"参考基准"的作用。后续token通过"看"第一个token，来"对齐"自己的位置和语义。

Attention Sink是怎么产生的？

数学上的解释

Softmax的归一化特性决定了Attention Sink的存在。

Softmax归一化：
  attention(i) = exp(s_i) / Σ exp(s_j)
  
  当某个token j的score s_j特别大时（大于其他所有score很多），
  exp(s_j)会占主导地位，其他token的attention趋近于0
  
  例如：
    s_1 = 10（第一个token）
    s_2 = s_3 = ... = s_100 = -10
    
    exp(10) = 22026
    exp(-10) ≈ 0.00005
    
    attention(1) = 22026 / (22026 + 99 × 0.00005) ≈ 99.999%
    attention(2..100) ≈ 0.0001%

为什么第一个token会成为Sink？

Transformer的训练过程中，第一个token<s>天然有特殊地位：

位置优势：第一个token的位置编码是基准（所有其他位置都是相对于它定义的）
语义积累：所有后续token的信息都会通过残差连接累加到第一个token上
梯度稳定：模型在训练中发现，把注意力分给第一个token可以让梯度更稳定

def analyze_attention_sink(attention_weights, num_layers=32):
    """
    分析Attention Sink现象
    
    参数：
      attention_weights: [num_layers, num_heads, seq_len, seq_len]
    """
    
    results = []
    
    for layer_idx in range(num_layers):
        for head_idx in range(attention_weights.shape[1]):
            attn = attention_weights[layer_idx, head_idx]  # [S, S]
            
            # 第一个token的平均注意力
            sink_attn = attn[:, 0].mean().item()  # 所有位置对第1个token的注意力
            
            # 计算"头部集中度"（前10%的token占多少注意力）
            attn_per_token = attn.sum(dim=0)  # 每个token分到的总注意力
            top10_pct = attn_per_token.topk(k=int(len(attn_per_token) * 0.1)).values.sum().item()
            
            results.append({
                "layer": layer_idx,
                "head": head_idx,
                "sink_attention": sink_attn,
                "top10_concentration": top10_pct
            })
    
    # 统计
    avg_sink = sum(r["sink_attention"] for r in results) / len(results)
    avg_concentration = sum(r["top10_concentration"] for r in results) / len(results)
    
    print(f"\n=== Attention Sink分析 ===")
    print(f"平均Sink注意力（第一个token）: {avg_sink:.2%}")
    print(f"平均Top10%集中度: {avg_concentration:.2%}")
    print(f"说明: {avg_sink:.0%}的注意力集中在第一个token上")
    
    return results

# 可视化示例
# 正常模型的Attention Sink分布
print("正常模型（Llama-2-7B，seq_len=4096）:")
sink_analysis = analyze_attention_sink(fake_attention_weights)

输出示例：

正常模型（Llama-2-7B，seq_len=4096）:
平均Sink注意力（第一个token）: 35.2%
平均Top10%集中度: 72.8%

Attention Sink量化：
  seq_len=512:  Sink≈25%
  seq_len=2048: Sink≈35%
  seq_len=4096: Sink≈40%
  seq_len=8192: Sink≈48%
  
结论：序列越长，Attention Sink越明显
     因为后续token的相对位置信息越来越难编码，Sink锚点更重要

FlashAttention和Attention Sink的关系

FlashAttention的在线Softmax会改变Sink吗？

不会。Attention Sink是模型学习到的特性，跟计算方式无关。

但FlashAttention的在线Softmax实现可能让Sink现象"看起来"更明显——因为在线Softmax在累加过程中，前面的token会"领先"后面的token。

def check_flash_attention_sink(q, k, v, head_num):
    """
    验证FlashAttention的Sink分布是否与标准Attention一致
    """
    
    # 标准Attention（ground truth）
    scale = 1.0 / (q.shape[-1] ** 0.5)
    scores_std = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
    attn_std = F.softmax(scores_std, dim=-1)
    
    # FlashAttention
    attn_flash = npu_flash_attention(q, k, v, head_num=head_num, return_attention=True)
    
    # 对比第一个token的注意力
    sink_std = attn_std[:, :, :, 0].mean().item()  # 标准Attention的Sink
    sink_flash = attn_flash[:, :, :, 0].mean().item()  # FlashAttention的Sink
    
    print(f"\n=== Attention Sink对比 ===")
    print(f"标准Attention Sink: {sink_std:.2%}")
    print(f"FlashAttention Sink: {sink_flash:.2%}")
    print(f"差异: {abs(sink_std - sink_flash):.4f}")
    
    if abs(sink_std - sink_flash) < 0.05:
        print("✅ FlashAttention的Sink分布与标准Attention一致")
    else:
        print("⚠️ FlashAttention的Sink分布与标准Attention有显著差异！")
        print("建议：检查scale_value和在线Softmax实现")

Sink对FlashAttention性能的影响

Attention Sink意味着：无论seq_len多长，第一个token始终是一个"热点"。FlashAttention的SRAM tile策略会优先处理Sink相关的分块。

Sink对性能的影响：
  正面：Sink位置可以被缓存，减少重复计算
  负面：Sink位置容易成为瓶颈（HBM读写集中在这个区域）
  
实际影响：很小，Sink只占Attention计算的一小部分

怎么利用Attention Sink？

应用1：高效的Prefix Caching

Attention Sink意味着第一个token的信息是"共享"的——所有后续token都会参考它。可以利用这个特性做Prefix Caching。

class SinkAwarePrefixCache:
    """
    利用Attention Sink的Prefix Caching
    
    思路：
      如果多个请求共享相同的前缀（prompt），
      只需要计算一次前缀的KV Cache，
      后续请求复用这个KV Cache即可
    """
    
    def __init__(self, max_prefix_len=1024):
        self.max_prefix_len = max_prefix_len
        self.prefix_cache = {}  # {prompt_hash: kv_cache}
    
    def compute_prefix_kv(self, model, prompt_ids):
        """
        计算前缀的KV Cache
        """
        prompt_hash = hash(prompt_ids.tolist())
        
        if prompt_hash in self.prefix_cache:
            print(f"命中Prefix Cache: {prompt_hash}")
            return self.prefix_cache[prompt_hash]
        
        # 第一次计算
        prompt = prompt_ids[:self.max_prefix_len]
        with torch.no_grad():
            outputs = model(torch.unsqueeze(prompt, 0), use_cache=True)
        
        kv_cache = outputs.past_key_values  # 保存KV Cache
        self.prefix_cache[prompt_hash] = kv_cache
        
        print(f"新增Prefix Cache: {prompt_hash}")
        return kv_cache
    
    def generate_with_prefix(self, model, prompt_ids, prefix_ids, prefix_kv):
        """
        用Prefix Cache生成
        """
        # 拼接
        full_ids = torch.cat([prefix_ids, prompt_ids])
        
        # FlashAttention前向（前半部分用Cache，后半部分正常计算）
        outputs = model(
            input_ids=prompt_ids.unsqueeze(0),
            past_key_values=prefix_kv,  # 复用前缀的KV Cache
            use_cache=True
        )
        
        return outputs

应用2：Sink感知的Prompt压缩

Attention Sink的强度跟序列长度正相关——序列越长，Sink越强。可以利用这个特性做Prompt压缩：只保留Sink周围的token，丢弃其他。

class SinkAwarePromptCompressor:
    """
    基于Attention Sink的Prompt压缩
    
    思路：
      保留前N个token（Sink锚点）+ 最后M个token（近期上下文）
      中间部分（远离Sink的token）压缩掉
    """
    
    def __init__(self, prefix_keep=128, suffix_keep=512):
        self.prefix_keep = prefix_keep  # 保留前缀（Sink锚点）
        self.suffix_keep = suffix_keep  # 保留后缀（近期上下文）
    
    def compress(self, input_ids, original_seq_len):
        """
        压缩Prompt
        """
        if original_seq_len <= self.prefix_keep + self.suffix_keep:
            return input_ids  # 太短，不需要压缩
        
        # 保留前缀（Sink锚点）
        prefix = input_ids[:self.prefix_keep]
        
        # 保留后缀（近期上下文）
        suffix = input_ids[-self.suffix_keep:]
        
        # 中间部分压缩掉
        # 可以用 summarization 或 直接截断
        compressed = torch.cat([prefix, suffix])
        
        print(f"Prompt压缩: {original_seq_len} → {len(compressed)} "
              f"(保留{self.prefix_keep}前缀 + {self.suffix_keep}后缀)")
        
        return compressed

总结：Attention Sink理解清单

FlashAttention中的Attention Sink现象，按这个清单理解：

问题	答案
Attention Sink是什么	第一个token聚集了大量注意力的现象
为什么会有Sink	Softmax的归一化特性 + 训练过程
跟FlashAttention有关吗	无关，是模型特性
影响性能吗	几乎不影响
可以利用吗	可以做Prefix Caching和Prompt压缩
需要修复吗	不需要，是正常现象