企业官网建设流程全解析

我大概两个多月前想把一个7B模型接到32k长上下文场景里做RAG问答，结果第一次prefill测试直接傻眼——单条带检索片段的长输入，attention部分耗时占整个forward的一多半，KV Cache吃掉十几GB显存，推理服务还没上线就开始骂人。后来读到SHARP这篇工作我才缓过劲来，它的核心思路是把“稀疏注意力”从口号变成可落地的三层管线：先搜每层每头的稀疏率，再用阈值锐化替代硬top-k，最后做带Hessian信息的头部剪枝，配合自定义kernel把稀疏收益真正兑现成延迟下降。

源码分析这个东西，不同领域的打开方式差别很大，有的项目得从协议栈追起，有的则要先理解论文里的公式再进代码。SHARP属于后者，不把论文的动机吃透，直接看源码很容易迷路。所以这篇内容我按论文加源码两条线来组织，把SHARP的动机、方法、关键代码实现、复现结果和踩坑经验完整过一遍。适合正在做LLM长上下文推理优化、模型压缩，或者准备复现稀疏注意力论文的工程师，读代码时建议把官方仓库clone下来对着看，理解会快很多。

1. 长上下文推理的三座大山：复杂度、显存与稀疏困局

1.1 QK^T的二次复杂度怎么把prefill拖慢的

先算笔账。标准注意力是Attention(Q,K,V)=softmax(QK^T/√d_k)V，prefill阶段要一次性处理L个token，QK^T这一步产生L×L的分数矩阵，后面softmax和AV还要再乘出来一遍，所以单头单层的计算量大约在4·L²·d这个量级。d是固定值，L一旦拉长，复杂度是平方级上涨。L从4k涨到32k，计算量直接放大64倍，这种增长不是优化访存能救回来的。

以Llama-2-7B为例，32层、32头、128维head_dim，输入32k token时QK^T单层单头要算10亿个分数对，整层1024个注意力头加起来，光attention这个算子的FLOPs就在0.5 PFLOPs量级。A100的BF16峰值算力312 TFLOPS，理论最快也要快两秒，实际kernel到不了峰值利用率，我在实测里一截32k的prefill，attention部分经常占到整个forward时间的60%到70%。这还不是最致命的，到decode阶段每次虽然只生成一个token，但每个token都要跟前面所有token做注意力，生成2000个token就要做2000次全量注意力扫描，而且每次扫描的序列长度还在不断增加。

有个体感上的类比：prefill像是在超市里一次性把所有商品过一遍收银台，decode像每次只买一件商品但还是要排完整个队伍，队伍还越排越长。所以长上下文推理慢不是某一个阶段的事，两个阶段各有各的痛点，优化方案也得分开对症。

1.2 KV Cache的显存账本

长上下文的第二个大问题是KV Cache。K和V矩阵要缓存每个历史token的键值向量，总量按这个公式算：

KV_Cache_bytes = 2 × num_layers × num_heads × head_dim × seq_len × bytes_per_element

Llama-2-7B在fp16下有32层、32头、128维，套进公式：

• 每个token大约要占 2×32×32×128×2 = 524288 字节，也就是0.5MB
• 8k上下文就要缓存约4GB
• 32k上下文直接飙到16GB
• 如果换成128k，那就是64GB，一张80GB的A100光是KV Cache就快塞满

这还没算模型权重、激活值、优化器状态这些。所以很多号称支持“超长上下文”的部署，实际并行度和batch size全被KV Cache的显存卡死。我见过不少团队为了塞下长上下文，只能把batch压到1，吞吐惨不忍睹，成本还高得吓人。

1.3 稀疏方向喊了很多年，为什么落地这么难

既然注意力天然稀疏，很多人第一反应就是“只算高分的部分不就行了”。但这里有几层现实问题。第一，非结构化稀疏在GPU上收益很玄学，A100这类GPU的tensor core是为稠密矩阵乘设计的，你告诉它“每行只保留几个非零点”，它要么走gather/scatter路径被访存卡死，要么把稀疏转成密集计算，收益直接蒸发。第二，固定窗口、滑窗这种方案虽然好实现，但长距离依赖确实会丢，测试集上PPL没怎么涨，实际业务里跨段落的指代、摘要、多跳问答全崩。第三，即便有人把稀疏attention做出来了，每个头该保留多少、哪些头根本不重要，这两件事如果拍脑袋定，模型智力损伤得很厉害。

SHARP的出发点正是这三个痛点。它不是简单地把attention换成一个近似算子，而是把“哪些位置值得保留”设计成一套可搜索、可剪枝、可融合kernel的完整流程。这个方法拆开看其实不复杂，但每一层都踩在了工程落地的关键点上。接下来我把论文里那套流程逐个拆开。

2. SHARP的方法拆解：从“观察”到“行动”的三层管线

2.1 论文最关键的观察：注意力头不是等价的

作者对训练好的模型做了一件事：跑一批校准数据，把每一层每一个头的注意力矩阵都dump出来，然后统计能量分布。结论非常有意思——注意力头并不是等价的，大体能分成两类：一类是“局部头”，注意力主要落在附近的token上，适合用固定窗口表达；另一类是“全局头”，少数几个token承担了绝大部分注意力权重，分布是又尖又稀疏的。

更重要的是，几乎所有头都存在一个共性：注意力分数矩阵里大量位置是低值噪声。把每行的分数按从大到小排序，前5%到30%的位置通常就贡献了95%以上的attention能量。这个观察为什么重要？因为它意味着每行保留top-k个分数，理论上对模型输出影响很小，而且这个k是可以通过校准集量化出来的。换句话说，稀疏率不是一个需要手调的玄学超参，而是一个可以从模型自身分布里算出来的量。

另外论文还发现，不同层之间需要的稀疏度差异很大。靠近输入层的头往往更local，中间层开始出现大量稀疏的全局头，靠近输出层又会有部分头重新密集化。如果全局用同一个稀疏率，要么剪少了浪费，要么剪多了伤模型。这就是第一层管线存在的理由——per-head的精细化稀疏，而不是一刀切。

2.2 第一步：per-head稀疏率搜索

SHARP第一步是给每个头算出一个稀疏率。做法是用校准集前向一遍，记录每个头在每层的注意力矩阵，然后对每行做排序，统计“能量保留比例”。给定一个能量阈值（比如保留95%的attention能量），反推需要保留多少个top-k位置，最后取所有行的分位数（比如95分位）作为这个头的稀疏率。

这里有个值得注意的细节：按行分位数而不是平均值来定稀疏率。因为注意力矩阵里存在少量“集中行”，这些行可能只要3%的位置就能保住95%能量，但同一头里也可能有大量注意力本来就分散的行，如果按平均值定稀疏率，后者会被剪成残废。取分位数是在“保能量”和“留余量”之间取得平衡。源码里这一步的结果会存成一个per-layer-per-head的稀疏率表，后续的kernel会用这张表决定每个head走稠密路径还是稀疏路径，以及稀疏路径需要分配多大的top-k空间。

2.3 第二步：阈值锐化替代硬top-k

拿到稀疏率表之后，下一个问题是“怎么在推理时执行稀疏”。最朴素的实现是每行torch.topk取前k个，把其余置0。但这会带来两个工程问题：一是在CUDA上做严格top-k有额外排序开销，k一旦变化kernel分支也不好写；二是硬top-k对注意力分数分布很敏感，某些行分数整体都高，top-k硬截断后剩余分数的重归一化会引入明显误差。

SHARP采用的是“锐化+阈值”的组合。在推理时并不做精确top-k，而是先对注意力logits做一个阈值收缩，把小于等于某个百分位阈值的分数直接置为负无穷，然后进softmax。这样做的效果等于给注意力分布加了一个可导的尖锐化操作，同时天然完成重归一化。阈值来自离线统计阶段算出来的每头百分位，线上就是一次compare加select，比top-k排序便宜得多。

源码层面这一步被做成了一个fused kernel：先读QK^T的结果，把低于阈值的元素mask掉，再就地做max-subtract和exp-sum，最后和V做矩阵乘。整个过程只读写一遍分数矩阵，比“先算完整softmax再做稀疏”省掉两三次全局访存。我读源码时觉得这里最有工程含量，后面走读kernel再细说。

2.4 第三步：带Hessian信息的头部剪枝

稀疏化能把prefill的计算量降下来，但KV Cache和decode阶段的带宽问题还得靠剪头来解决。头部剪枝不是新东西，难点在于怎么判断哪些头该剪。简单用L2范数或者平均注意力分数做重要性，剪完以后PPL看着还行，下游任务效果会悄悄掉，因为某些头虽然平均权重不大，但在特定知识型任务上是不可替代的。

SHARP的头部重要性标定借鉴了优化文献里的OBS/SparseGPT思路：用校准集计算每个头对loss的Hessian信息，近似估计“剪掉这个头之后loss会涨多少”。具体公式是围绕二阶近似展开的，但工程实现比较直接——对每个候选头轮流mask掉，跑一个mini-batch校准集，记录loss变化，再结合梯度和曲率信息做修正。剪枝策略也不是简单的“一刀切剪掉最不重要的N%”，而是带层间约束的调度：每层最多剪多少、总参数量预算、剩余KV Cache预算，这几个条件一起送进一个贪心分配器，最后得到每层的剪枝数量。这种做法比固定比例剪枝更符合模型的实际冗余分布。

3. 源码走读：四个关键模块的实现细节

官方仓库的顶层目录大概是这样的结构，读之前建议先把这个架子搭在脑子里：

sharp/ ├── scripts/ │ ├── profile_attention.py │ ├── search_sparsity.py │ ├── prune_heads.py │ └── run_eval.py ├── sharp/ │ ├── kernels/ │ │ ├── topk_attn.cu │ │ ├── fused_threshold_softmax.cu │ │ └── gemv_sparse.cu │ ├── pruning/ │ │ ├── importance.py │ │ ├── schedule.py │ │ └── apply.py │ ├── models/ │ │ ├── llama_patch.py │ │ └── hook_utils.py │ └── utils/ └── configs/

下面我按“数据采集→稀疏化算子→剪枝→推理kernel”四段走读，顺序其实也是论文方法的执行顺序。

3.1 注意力分布采集与稀疏率搜索

profile_attention.py负责把训练好的checkpoint跑一遍，hook住每一层的attention输出，存成numpy数组。关键点是要hook在softmax之后、dropout之前的注意力权重，同时把position_ids设置成包含完整上下文长度的样本，避免短样本统计出来的稀疏率和实际推理对不上。

PyTorch里hook的方式很简单：

attention_maps = {} def hook_fn(name): def hook(module, input, output): # LlamaAttention的output是(attn_output, attn_weights, past_key_value) attention_maps[name] = output[1].detach().cpu().float().numpy() return hook for name, module in model.named_modules(): if isinstance(module, LlamaAttention): module.register_forward_hook(hook_fn(name))

如果你用的Transformer版本比较新，LlamaAttention的输出结构可能变过，更省事的做法是forward时直接传output_attentions=True，把注意力权重带回来，然后再用临时buffer收集。我自己的经验是hook容易受到模型封装影响，直接传参反而更稳。

核心统计逻辑不长，大意如下：

def search_sparsity(attn_maps, energy=0.95, quantile=0.95): results = {} for layer, heads in attn_maps.items(): for head, attn in heads.items(): sorted_weights = np.sort(attn, axis=-1)[:, ::-1] # 降序 cumsum = np.cumsum(sorted_weights, axis=-1) total = cumsum[:, -1:] # 每条query行需要多少个位置才能累计到energy阈值 k_per_row = (cumsum >= total * energy).argmax(axis=-1) + 1 # 对行取分位数，得到这个head的稀疏率 k = int(np.percentile(k_per_row, quantile * 100)) total_len = attn.shape[-1] results[(layer, head)] = k / total_len return results

这个方法比想象中简单，但效果的关键在于数据。校准集必须覆盖推理时会出现的位置模式——既有长距离依赖，也有局部密集区域。我在复现时拿纯长新闻去统计，结果稀疏率整体偏高，因为长新闻里有很多重复实体和局部窗口模式；后来混入代码、多轮对话、结构化文档，得到的稀疏率表才更贴近真实业务。这是官方文档不会写、但直接影响效果的细节。

3.2 阈值锐化算子的前向实现

论文里“锐化”对应的算子，在代码里不是严格top-k排序，而是用预计算的百分比阈值做mask。PyTorch里一个能跑通的朴素版本可以这样理解：

def sharpened_attention(q, k, v, threshold_pos, scale): # q, k, v: [batch, heads, seq_len, head_dim] scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.size(-1)) # threshold_pos: 每个head的阈值位置，来自3.1的统计结果 kth_scores = torch.kthvalue(scores, k=threshold_pos, dim=-1).values mask = scores < kth_scores.unsqueeze(-1) scores = scores.masked_fill(mask, float("-inf")) scores = scores * scale # 锐化系数 probs = torch.softmax(scores, dim=-1) out = torch.matmul(probs, v) return out

注意这里有两个可以调的东西：一是threshold_pos是每个head的静态值，二是scale这个锐化系数。scale大于1会让softmax之后的分布更尖，等于把“保留位置”里的分数差异进一步拉大。论文实验显示scale在1.0到1.2之间效果较好，过大会导致输出方差变大，长文本生成时出现重复。

实际线上用的CUDA kernel没有走topk，因为torch.kthvalue开销太大。kernel直接以threshold_pos为参数，在计算QK^T的过程中顺便做一个block reduce求第k小的值，然后生成mask。这样把“求阈值+mask+softmax”合并成一次遍历，显存占用少一次整矩阵的读写。源码里这个kernel的block大小取64或128，和后面要说的稀疏块结构正好对齐。

3.3 剪枝调度与模型重写

importance.py里算头部重要性，schedule.py做分配，apply.py把剪枝结果写回模型。剪枝后不是简单给权重矩阵打成零块，而是真的把num_heads改掉重写模型结构，这样后续推理框架才能省掉对应的KV Cache空间。

朴素版本的头部重要性测量是这样：

def measure_head_importance(model, calib_loader, criterion): importance = {} model.eval() base_loss = evaluate_loss(model, calib_loader, criterion) for name, module in model.named_modules(): if not hasattr(module, "num_heads"): continue for head_idx in range(module.num_heads): mask_head(module, head_idx) loss = evaluate_loss(model, calib_loader, criterion) importance[(name, head_idx)] = loss - base_loss unmask_head(module, head_idx) return importance

这个朴素版本计算量很大，7B模型每个head跑一遍完整校准集，几百个头要跑一晚上。源码里做了两个优化：一是只在最后一层输出的loss上做反传，通过梯度估算重要性，不用每个head都重新前向；二是对权重做一阶泰勒展开近似，重要性分数就等于|gradient × weight|在注意力头维度上的均值。这两种近似在大多数模型上已经够用，除非你想精度优先，才用逐个mask的完整版。

层间分配是一个带约束的贪心过程。预算可以是总头部数、总KV Cache容量或总FLOPs，分配器按“单位代价剪掉的影响力损失”排序，优先剪那些“省得多且伤得少”的头。这个思路在源码里实现很朴素，但比固定比例剪枝好很多，建议二次开发时保留。剪枝完成后，apply.py会生成一个新的模型配置，把每层保留的head index写死，实际部署的是这个精简后的模型。

3.4 推理侧的fused kernel与显存处理

线上推理时，稀疏attention最怕的是把稀疏矩阵转成稀疏格式后，操作开销比省下的计算还大。SHARP的做法是块稀疏而不是纯点稀疏：把沿着序列维的K/V分成固定大小（比如64）的块，离线统计时按块内最高分数决定哪些块需要计算。这样kernel可以做“跳过整块”的矩阵乘，命中tensor core的稠密小块计算。块大小64在A100上基本能吃到比较高的计算效率，块太小访存碎片化，块太大稀疏率上不去。

显存方面，被剪掉的头部在KV Cache初始化时就不分配空间，所以KV Cache省下来的量和头部剪枝比例基本线性。稀疏化本身不省KV Cache，但能省掉中间结果矩阵的显存占用——因为mask+softmax的中间分数矩阵不再需要完整落盘，kernel内部一块一块处理完就释放。在32k上下文、batch为4的实验里，峰值显存能比原版eager模式降一小半，主力来自KV Cache剪枝后的减少。

4. 复现笔记：环境、数据与实测结果

4.1 复现环境与配置

我复现时的环境如下：

• 硬件：单张A100 80G
• 模型：Llama-2-7B，base和chat都跑了一遍
• 框架：PyTorch 2.1 + CUDA 12.1 + Transformers 4.36
• 校准集：从LongBench的多个子任务里抽了约512条样本，截断到16k

这里遇到一个环境上的硬约束：要跑SHARP自己的kernel，必须用eager attention而不是flash attention。Transformers里默认的attn_implementation可能是sdpa或flash_attention_2，会把attention计算截胡。需要在加载模型时显式设置：

model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16, attn_implementation="eager", )

如果不开eager模式，后面hook注意力权重和替换kernel都会失效，而且flash attention的中间结果根本不给你看。这个点建议所有复现稀疏注意力论文的人都先检查一遍，能省一晚上的排查时间。

4.2 在7B模型上的实测数据

以原始Llama-2-7B为baseline，在32k上下文、batch=2的配置下，我复现的核心数据如下：

这是我自己机器上的复现数据，跟官方卡型和数值可能略有出入，但趋势是一致的：prefill能获得2到3倍加速，decode提升在50%以上，KV Cache省下来的显存和剪枝比例基本线性。PPL损失在25%剪枝时很小，但LongBench均分已经掉了1.7个点，说明纯PPL指标对“智能损伤”不敏感。另外我也顺手对比了几个常见的稀疏baseline，在LongBench上SHARP明显比StreamingLLM和H2O这两个纯推理期方案稳，主要原因是后两者对历史token的取舍策略是固定的，不会像SHARP这样根据每个头单独定稀疏模式。

4.3 我踩到的三个和论文不一致的细节

第一个意外是短上下文下加速比会变负。模型在4k上下文、batch=8的场景里，加了SHARP的稀疏kernel反而比eager慢。原因很直接：短序列时QK^T计算本身不大，稀疏kernel的block判断、mask、重归一化开销占了主导；加上短文本里注意力不容易出现极稀疏分布，能量保留比例不高，稀疏化等于白做。所以SHARP这类方法面向的是8k以上、16k起步的场景，应用前一定要先看自己的平均序列长度。

第二个意外是某些head被剪掉后PPL几乎不变，但特定任务崩掉。我单独测了一下摘要任务和代码补全，发现在PPL上排名后20%的头里，有一个头对代码缩进的注意力模式很重要，剪掉后代码补全的准确率掉了4%。这说明头部重要性不能只看平均损失，得结合任务覆盖的校准集。后来我是把校准集里混入更多代码样本，让重要性排序在“通用能力”和“关键任务”之间做加权，效果才恢复。

第三个意外是INT8量化叠加稀疏会放大误差。模型先做GPTQ 4bit量化，再接SHARP剪枝时，LongBench分数掉得比单独做任一操作都严重。原因也不难理解：剪枝和量化都是在损失信息，两种近似的误差在深层网络里会叠加而不是抵消。如果想两个都要，得把量化放进校准流程里一起考虑，而不是流水线式地先量化再剪枝。

5. 从复现到落地：调参、兼容与工程化经验

5.1 稀疏率怎么调才不伤模型能力

官方默认能量阈值0.95、分位数0.95这套参数在通用语料上是不错的起点，但真正调到业务场景还是要多跑几组。我的经验是：先用一组很小的校准集（128条）快速跑几组阈值组合，画出PPL-稀疏率曲线，看拐点在哪。通常在能量阈值低于0.9之后PPL开始明显上翘，0.85以下基本不可接受。但不同层的情况不一样，靠近输出层的最后几层对稀疏化非常敏感，搜索出来的k往往偏大，如果为了让整体稀疏率好看而压缩这几层的k，输出分布会被破坏。实际使用时可以把“关照层”配一个单独的更低稀疏率。

对下游任务，我强烈建议不要只看PPL，用两个有代表性的任务做探针，一个偏知识问答，一个偏代码或结构化文本。因为PPL对局部流畅度敏感，但对跨句推理、符号操作不敏感，稀疏化把long-range head剪掉后，PPL损失可能很小，但任务效果掉落很明显。调参的时候把这两个探针任务的分数一并打出来，比看PPL可靠得多。

5.2 与推理框架集成的兼容性坑

把SHARP的稀疏attention塞进vLLM或TensorRT-LLM这类工程框架，比在PyTorch里复现麻烦得多。原因是这些框架的算子已经跟自身的显存管理深度绑定。vLLM用PagedAttention管理KV Cache块，剪枝头可以通过改模型结构省空间，但稀疏attention就没法直接套用——PagedAttention假定每个block内的K/V都是稠密有效token，你改成稀疏后，block的分配和回收逻辑全部要跟着变。

我的建议是分两步走：第一步先在当前模型上把“头部剪枝”部分落地，这一步改的是模型结构，对现有推理框架最友好；第二步再考虑稀疏kernel，最好作为独立推理路径而不是试图改框架内置算子。如果一定要在vLLM里上稀疏，可以退而求其次，把“稀疏化”作为调度层策略——对历史token做窗口划分，某些head只看局部窗口，某些head才做full attention，这种静态模式可以在不改造kernel核心的情况下塞进现有框架。

5.3 值得继续做的方向

SHARP这种“先统计后剪枝”的范式还有很多可以扩展的空间。一是把稀疏率搜索做成在线动态的，模型在长文档内根据局部复杂度切换稀疏模式；二是把稀疏化和KV Cache的量化压缩结合起来，因为被剪掉的头腾出了显存，可以拿这部分预算给剩余头部做更精细的KV量化；三是将头部重要性标定扩展到多任务场景，一个模型部署在多业务上时，重要性应该按流量加权而不是均匀混合。这些方向我现在也在继续试。SHARP的定位更像一个框架，把论文里那套“观测注意力→搜索稀疏模式→结构化剪枝→定制kernel”的方法固化下来，留了很多接口给后续研究者应用。

最后分享一个读代码阶段的体会：不要一上来就看CUDA kernel，先把profile_attention.py和search_sparsity.py跑通，得到一张属于自己模型的稀疏率表，再回头看kernel就清楚它为什么这么设计了。代码里最值钱的不是某个trick，而是整个流程的顺序——先观察、再剪枝、最后优化计算，这个顺序反过来做，往往会白费很多功夫。我最初拿到项目时先试着魔改kernel，发现怎么调收益都很小，后来老老实实把能量分布统计出来，才知道瓶颈根本不在kernel效率，而在没有按真实分布设计稀疏模式。这个教训应该对很多做模型优化的人都有参考价值。