先解释一下背景。团队有一条基于 QSAQuantized Sparse Accelerator自研推理引擎的线上服务链路底层算子、显存池和调度器都是围绕自家框架设计的。但最近算法侧一直在用 HiSparse 做长上下文稀疏注意力实验效果不错——在 128K 乃至 256K 长度下分层局部性Hierarchical Locality带来的加速比非常明显。问题来了HiSparse 原型代码和 QSA 的接口完全是两套体系模型想上线就得把 HiSparse 的稀疏注意力核心逻辑完整移植到 QSA 里还得在双 RTX 4090 48GB 上跑通 8×256K 的部署配置。折腾了大概两周半中间踩了不少坑也把一些底层细节重新梳理了一遍。这篇就当是移植过程的完整记录包括分层局部性在代码里到底怎么落、8×256K 意味着多夸张的显存和调度压力、以及哪些地方容易翻车。如果你也在做长上下文推理或者稀疏注意力方向应该能省不少时间。1. 项目背景与方案选型为什么非要把 HiSparse 塞进 QSA动这个移植项目之前团队其实有过一轮讨论为什么不直接在 HiSparse 基础上封装 API非要往 QSA 里搬这个问题直接决定了后续所有技术选型值得先把逻辑讲清楚。1.1 HiSparse 在长上下文里的定位和优势HiSparse 不是一个单纯的稀疏注意力实现它最核心的卖点是分层局部性的建模方式。传统稀疏注意力比如 Longformer、BigBird用的是固定稀疏模式滑窗加随机连接或者全局 token 加窗口。这种做法的优点是简单、可控但缺点是模式固定对输入内容的自适应性很差——某些需要长期依赖的 token 对如果不在预设模式里注意力就会被硬切掉精度损失不可控。HiSparse 的做法是分层来做先在一个很粗的粒度上判断哪些区域值得关注再逐步细算。它把上下文划分成多个层级的块每一层通过轻量的向量相似度计算筛掉大量无关区域最终保留的稀疏模式是随输入动态变化的。这个设计在 128K 以上的超长上下文里优势非常明显因为长上下文里真正重要的信息占比通常极低固定模式要么覆盖不够、要么计算浪费太多。算法侧在几个长文本检索和摘要任务上测过HiSparse 在 256K 输入下的 prefill 延迟比 FlashAttention-2 全量注意力低了一个数量级同时精度几乎没有掉。所以这个库不是玩具是真能扛事的。1.2 QSA 框架的定位和移植的必要性QSA 是我们内部维护的一套推理加速框架主要解决三个问题多模型服务的显存隔离、动态 batch 调度、以及量化算子的统一管理。线上所有对外提供的 LLM 服务都是跑在 QSA 上的它对显存池的管理非常细——算子申请显存不直接走 CUDA API而是统一走 QSA 的显存池Memory Pool按请求生命周期分配和回收。这就带来一个现实问题HiSparse 的原型代码完全是自己管显存CUDA malloc 满天飞注意力 kernel 也是独立的 CUDA 实现根本不走 QSA 的算子注册机制。如果直接拿 HiSparse 上线意味着要单独给它维护一套显存分配逻辑这和 QSA 的统一管理是冲突的线上 OOM 的概率会大幅上升。所以项目目标很明确把 HiSparse 的核心算法逻辑尤其是分层局部性这一块重写为 QSA 框架下的算子纳入统一的显存池、算子调度和量化体系。另一条路——把 QSA 改造成适配 HiSparse 的壳——在架构上完全不可取QSA 的调度器会和外部显存管理冲突。1.3 双 RTX 4090 48GB 平台的考量和约束硬件选择其实是被项目需求推着走的。算法侧要求单条请求支持 256K 上下文线上想要 8 个并发请求同时跑也就是标题里的 8×256K。这个要求用显存算一下就知道了如果完全不做稀疏化256K 长度的 KV Cache 会吃掉几十 GB 显存8 个并发根本不可能塞进任何单卡。为什么要强调 RTX 4090 48GB这里解释一下。标准 RTX 4090 是 24GB 显存但我们这批卡是特定厂商推出的 48GB 版本显存翻倍但核心规格基本一致单卡显存带宽和 24GB 版本是同一个量级都在 1TB/s 左右。双卡就有 96GB 总显存加上 NVLink 并不存在——4090 基本没有 NVLink 互联——所以双卡之间的通信完全依赖 PCIe 5.0 的带宽这其实是一个很重要的约束。选这个平台的原因也很朴素成本敏感。A100/H100 的采购周期和预算都不现实而长上下文推理对显存容量的需求远大于对计算峰值的需求。稀疏注意力已经把计算量降下来了对显存带宽的需求也降了所以 4090 级别的显存带宽是够用的瓶颈反而在容量和通信上。2. 分层局部性8×256K 长上下文加速的核心密码如果说移植地址、接口、显存这些都是体力活那一套能把 8×256K 显存需求压到双 4090 能扛住的分层局部性设计才是这个项目的灵魂。这一节详细拆解整个算法思路以及它怎么和部署硬件匹配起来。2.1 长上下文里的注意力分布规律在做任何稀疏化之前先得搞清楚一个问题256K 长度下注意力权重到底是分布在哪里。以前跑全量注意力的时候其实看不出来因为 256K × 256K 的注意力矩阵太大了根本不可能显式存下来。HiSparse 的做法是先在小规模实验里用 FlashAttention 的统计功能输出每个 query 对历史的注意力分布。观察下来的结果非常有规律一个 query 至少会把 60%-70% 的注意力权重放在离它最近的 1K-2K 个 token 上这符合语言模型的局部性偏好但剩下 30%-40% 的权重会散落在更早的上下文中这些位置因 query 而异——有的 query 在找文档开头的人名有的在找段落标题有的在找某段代码定义。如果只用固定窗口这 30%-40% 的信息就全丢了模型表现为记性变差。而全量注意力在这个场景下的问题也很明显每一个 token 都要和 256K 个历史 token 做内积。在 4090 上实测虽然 FlashAttention 已经做到了接近理论峰值但 prefill 时间依然随长度线性增长到 256K 时已经无法接受。2.2 从细到粗的三级稀疏筛选HiSparse 的分层局部性最终在代码里落地成了三级结构。这种设计很像人看文章的方式先扫一眼标题和段落摘要确定大概范围然后翻到具体段落快速扫读最后才逐句精读关键部分。具体来说第一级是全球粗筛级Global Coarse Layer。整个上下文按 1024 个 token 切成一个块块之间用一个可学习的汇总向量表示本质上是这个块的注意力输出平均池化。对每个 query先和所有历史块做一次轻量内积选出 top-k 个候选块。256K 上下文切成 256 个块这一步只需要 256 次向量内积计算量几乎可以忽略。top-k 我们通常取 32 或 48。第二级是局部精筛级Local Fine Layer。对上一级选出的候选块在块内部再按 64 个 token 切片同样用轻量内积筛出最相关的若干片。这一步的复杂度是 top-k 乘以每个块内的片数也就是 32 × 16 512 次内积依然很低。第三级是窗口注意力级Window Attention。不管前面的粗筛结果如何每个 query 都会直接拿最近的 512 个 token 做全量注意力。这一层用来保证语言建模的基本局部性不丢失也避免远端筛选错误导致完全漏掉关键信息。三级筛选合并之后每个 query 实际需要参与的 KV 范围大概是512 个局部 token 粗筛选出的若干片64 token 精度的并集。实测下来每个 query 平均只需要和 2K-4K 个 KV token 做注意力相比原始 256K 直接少了 64-128 倍。2.3 稀疏筛选在硬件上的映射分层局部性不只是一个算法概念它在部署时要解决一个很具体的问题稀疏筛选出来的 KV 在显存里怎么存放才能保证访存效率。最开始踩过一个坑如果把选出来的 KV 拷贝到一个紧凑的连续内存块里再做注意力虽然计算很快但拷贝本身的耗时很夸张——因为每次都要从不同位置拉数据到同一块缓存随机访问很多。后来改成只在块边界上做稀疏化粗筛时以 1024 长度块为单位决定要不要整体加载这样保证从显存搬运时大部分数据是连续的。只有在确定了要加载的块之后才在块内部做更细的筛选。这一层缓冲区设计在移植过程中反复调了很多次是性能影响最大的部分之一。另外QSA 里本来就有基于 CUDA Graph 的显存预取机制我们把块筛选结果作为图的输入使稀疏模式固定的帧内执行可以完全落进预取流水线里避免了 kernel 间同步等待。这也是移植到 QSA 之后吞吐能明显优于 HiSparse 原型的重要原因——原型的 kernel 之间没有这种整体调度。3. 移植实战接口对齐、内核重写和显存管理算法方案定了剩下的就是体力活。这部分按实际的移植顺序来写包括接口对齐、kernel 重写、双卡显存管理三个大块以及每个环节里我认为值得记录的细节。3.1 算子注册和接口对齐QSA 框架里每个算子都注册成一个 Operator对外暴露统一的 Init、Forward、Free 接口。HiSparse 原型的稀疏参数是直接暴露成 Python 对象的字典移植第一步就是把稀疏模式配置从 Python 侧挪到 C 侧注册成 QSA::SparseAttentionConfig。关键的接口参数有这几个window_size局部窗口大小我们固定为 512block_size_coarse粗筛块大小固定为 1024block_size_fine精筛片大小固定为 64top_k_coarse粗筛保留块数默认 32top_k_fine精筛保留片数默认 16这些参数如果暴露成 Python 字典每个请求进来都要解析一遍字符串解析成本不高但积少成多。注册成 C 结构体之后请求配置只传一个枚举 ID查找表直接拿 ID 索引速度快很多。这块属于典型的看着不起眼、改完真舒服的优化。3.2 稀疏模式计算从 CPU 搬到 GPU最初版本的分层稀疏筛选逻辑是放在 CPU 上算的。原因很简单——HiSparse 原型就是这么写的当时的思路是 CPU 算筛选GPU 只负责执行已经被筛选好的注意力 kernel。256K 上下文的筛选逻辑大概要做 256 次块级内积排序在 CPU 上大约需要 3-5 毫秒。一开始觉得这个时间可以接受后来发现部署场景下完全不是这么回事。因为 8×256K 意味着 8 个并发请求同时到达如果每个请求的筛选都在 CPU 上做那么 CPU 侧单线程跑一轮就需要 24-40 毫秒。在 prefill 阶段这是不可接受的——prefill 本身在高稀疏率下也就几十毫秒CPU 端筛选占到一半以上的时间太浪费了。后来把整个筛选逻辑写成了两个 CUDA kernel第一个 kernel 负责块级向量内积和 top-k 选择。把每个 query 的隐藏状态和所有块的汇总向量做矩阵乘然后对结果做 top-k。这里 top-k 我们没有自己实现直接用 CUB 库的 DeviceSegmentedRadixSort 对分数排序去索引实测在 256K 上下文下这个 kernel 加排序的耗时在 0.3-0.5 毫秒左右。第二个 kernel 负责把选出来的 KV 块从显存池中 gather 到一块临时缓冲区。gather 过程用了向量化加载float4 或 half2确保访存效率接近峰值。这个 kernel 耗时大约 0.2-0.4 毫秒。两项加一起一次筛选从 3-5 毫秒降到了不到 1 毫秒而且彻底释放了 CPU 的调度压力。这一步是移植之后性能提升最大的单项优化强烈建议做。3.3 注意力内核和 FlashAttention 的融合筛选之后的注意力计算我们直接沿用 QSA 里已有的 FlashAttention 风格内核但做了一个关键改动把 block mask 融合进内核里。标准 FlashAttention 的流程是遍历 KV 块、在线 softmax 更新我们改成了遍历被选中的 KV 块跳过未选中的。因为稀疏筛选的结果已经确定了每个 query 需要访问哪些块这个信息被编码成一个 compact 索引数组kernel 启动时直接把它加载进共享内存循环体按索引跳转。这带来一个额外的好处IO 量大幅降低。全量 FlashAttention 每次迭代要把 KV 块从 HBM 读到 SRAM而我们的内核只读取被选中的块256K 上下文下从原来的 256 次迭代减到 30-50 次迭代。这意味着从 HBM 读取的数据量降低了 5-8 倍而注意力计算本身在 4090 上的瓶颈恰恰是 HBM 带宽而不是计算单元。3.4 双卡显存分配和流水线切分双卡 96GB 不是一条简单的把模型放卡 0、把 KV 放卡 1就能完事的。模型本身结构不大以 7B 参数量的 GQA 模型为例权重量化后不到 8GB单卡放得下。真正的显存大头在 KV Cache。我们的切分方案是双卡按层做流水线并行卡 0 负责前一半层卡 1 负责后一半层。每个请求的激活从卡 0 进入前向算到中间层后通过 PCIe 把激活传到卡 1。KV Cache 跟随所在层存放所以卡 0 只缓存前一半层的 KV卡 1 只缓存后一半层。这个切分方案在 8×256K 场景下很合适两块卡各自承担的 KV 压力大约是总量的一半避免了任何一张卡先爆。相比之下张量并行TP方案在 4090 上没有 NVLink 的情况下不可取——每层都要跨卡通信PCIe 的延迟和带宽会成为严重瓶颈。流水线并行只需要每请求传一次中间激活通信频率低得多。实测双卡之间传输一层 7B 模型中间激活约 0.5MB 每 token 每层在 batch8、序列片段为 1K 时耗时约 2-3 毫秒占整体延迟的 8% 左右可接受。4. 8×256K 部署显存预算、调度策略和性能数据算法移植完了配置也调通了但真正达到8×256K 可上线还是经历了很痛苦的调优过程。这一节给出具体的显存计算、batch 调度方案和最终性能数据。4.1 8×256K 的显存预算到底怎么算先把账算清楚。8×256K 有两层含义一是模型并发处理 8 个请求每个请求最长支持 256K 上下文二是显存里需要为每个请求维护 KV Cache。以 7B GQA 模型假设 28 层、8 个 KV head、head dim 128、FP16 存储为例一个 256K 长度请求的全量 KV Cache 是256 × 1024token 数× 8KV heads× 128head dim× 2K 和 V 两份× 2FP16 两个字节按 1K1024 计≈ 2.15GB最后除 1024^3再乘 28 层 ≈ 60GB。这个数字很吓人一个请求就 60GB8 个请求要 480GB双卡 96GB 根本不可能。但使用分层局部性之后KV Cache 不再全量存储所有 token 的 KV。我们实际只存两部分局部窗口每层 512 个 token和粗筛选出的远端块每层 32 个 1024 块。按 256K 输入计算实际存储的 token 数大概是 512 加上粗筛保留的 32 × 1024 33280约 33K token相比 256K 压缩了 8 倍左右。这样每个请求的实际 KV 显存大约是 60GB / 8 7.5GB。8 个请求就是 60GB双卡各承担 30GB加上模型权重和中间激活每卡 50GB 以内的预算就够用了刚好卡在双 4090 48GB 的总空间内。所以 8×256K 能落地核心是稀疏化把显存需求从 480GB 压到了 60GB而不是靠堆硬件。4.2 面向稀疏度的动态 batch 调度8 个并发请求各自在不同的解码阶段它们的稀疏度也是动态变化的。请求刚进来时历史 token 少稀疏筛选用不上随着 context 变长远距离依赖开始增多稀疏模式也在变化。QSA 原有调度器是按请求优先级和显存余量来做 batch 拼接的但移植后我们加入了一个新的调度维度稀疏度感知。每当有新请求要加入 batch 时调度器会预估它下一步的 KV 显存增长量而这个增长量是稀疏模式的函数——同样增加 1024 个 token如果粗筛命中的块多显存增长就快命中的块少增长就慢。这个预估模型其实非常简单把粗筛 top-k 的分数缓存下来预测下一个 chunk 能命中多少块命中率乘以块大小乘以 KV 单位大小就是后续显存增量。预测精度不需要太高够调度器做拒绝或排队判断就行。实际跑下来这个机制避免了很多次隐性 OOM——如果只按全量 KV 的增长来估批大小会过于保守吞吐上不去如果完全忽略稀疏度8 个长请求同时增长时会在某一个时间点集体越界。调度器和稀疏度挂钩之后系统能在显存安全的前提下把 batch 尽量打满。4.3 性能实测数据部署完成后跑了一组对比数据硬件环境是双 RTX 4090 48GBCUDA 12.2模型是 7B GQA 量化版本。这里列三个最有代表性的指标。Prefill 阶段256K 输入单请求HiSparse 原型在单卡上大约是 4.2 秒移植到 QSA 后由于显存预取和 kernel 融合降到 3.1 秒左右。这个提升主要来自 CUDA Graph 优化和避免 CPU-GPU 同步而不是算法本身的变化。解码阶段8 并发每请求 256K 上下文batch8有效吞吐大约 1800 tokens/s单请求平均时延 4.5 秒也就是 8 个请求一起跑平均每个请求每秒生成约 225 token。对比同类硬件的全量注意力实现吞吐大约 300-400 tokens/s提升 4-5 倍。显存占用8 个请求全部达到 256K 上下文时卡 0 显存峰值 43GB卡 1 显存峰值 41GB留有大约 5GB 的头部空间给临时激活和通信缓冲区。注意性能数据强烈依赖稀疏度配置。如果 top_k_coarse 从 32 提高到 64吞吐会下降约 30%但精度会更稳。部署时不要盲目追求最高吞吐建议先在目标任务上做一个精度-吞吐的权衡曲线再定参数。5. 移植与部署中的坑常见问题和排查记录最后这部分是踩坑实录很多问题都不是一次定位到的花了大把时间在 profiler 和日志之间来回。整理成速查表方便后面的人直接对照。5.1 问题一块稀疏模式导致精度波动现象移植初期在 RAG 问答任务上 F1 分数比 HiSparse 原型掉了将近 9 个百分点。最初怀疑是 kernel 里 softmax 数值精度问题排查很久才发现是粗筛 top-k 的选择逻辑差异。原因HiSparse 原型在 CPU 端做筛选时用了完整的浮点精度比较移植到 GPU kernel 时为了省寄存器把块分数转成了半精度浮点存储。256 个块分数在半精度下排序相邻分数的区分度不够导致 top-32 的边界块选错。解决把块分数改成 float32 存储排序比较时用全精度。在排序前不做任何量化或截断。修改后精度回归到与原型一致。这个教训是稀疏筛选的关键路径上不要自作聪明做精度压缩省下几百字节显存却亏掉几个点精度不值。5.2 问题二显存池碎片化和隐性 OOM现象8 个请求跑了一段时间后偶发显存分配失败但看峰值显存明明还有余量。QSA 的显存池是池化的池内块大小不等长请求的 KV 增长需要连续的大块内存池子被小的残留块占满后大请求就会分配失败。解决QSA 显存池本身有大块预留机制但默认没开启。改成给 KV Cache 单独建一个大块专用池池内分配单位是 1MB 的块和模型权重、临时缓冲区分开管理。这样 KV 增长时可以直接从大块池里取不会和 activator 的短生命周期分配互相干扰。另外一个习惯性建议定期查看显存碎片率。可以在 QSA 的监控接口里加一个显存池最大连续空闲块大小的统计这个值如果长期明显低于总空闲量就是碎片化的信号。5.3 问题三双卡通信是 PCIe 5.0但时延仍偶发毛刺现象解码阶段时延曲线不是平滑的每几十秒会冒出一个 200-300ms 的尖峰。用 nvidia-smi 看 GPU 利用率卡 0 在尖峰时利用率掉到接近 0卡 1 利用率正常。原因中间激活从卡 0 传到卡 1 是同步等待的当卡 1 正好在处理上一个 batch 的尾部时卡 0 就要空等不短的时间。后来定位到是双方没有做流水线重叠——卡 0 算完后不预取下一份数据卡 1 还在忙两边就串行了。解决在卡 0 的 kernel 执行完之后不等卡 1 反馈直接把下一份中间激活放到 PCIe 传输队列里用 CUDA Graph 的并行分支配合双流实现。这样卡 1 在处理 A 时卡 0 已经在传输 B 了通信和计算完全重叠。修改之后尖峰基本消失P99 时延从 230ms 降到 45ms。5.4 问题四8×256K 预热启动时间过长现象服务重启后第一次 256K 长请求耗时是稳定状态的 3 倍多大约 9 秒。原因QSA 的显存池和 CUDA Graph 在首轮请求时需要冷启动大量显存分配和 kernel 编译触发。这是老问题在短上下文场景不明显256K 场景下被放大。解决部署脚本里加了一个预热请求服务启动后模拟一次 32K 长度的生成请求让显存池预分配到位、CUDA Graph 编译完成。预热后重启成本平摊到首次请求实际体验提升非常明显。下面整理成一张速查表方便直接对照排查问题现象根因解决方式是否容易忽略精度掉点半精度分数排序不稳定块分数改全精度比较容易忽略建议默认全精度间歇性 OOM显存池碎片化KV Cache 独立大块池需要监控空闲块统计P99 时延尖峰跨卡通信串行双流 pipeline 通信重叠需要 CUDA Graph 支持冷启动长尾显存池和 kernel 冷启动预热请求 CUDA Graph 预编译部署流程必做从 HiSparse 到 QSA 的移植过程比预想中复杂但最终效果对得起投入。我个人最大的体会是稀疏注意力这种算法只有在和推理框架的显存管理、调度策略深度耦合之后才能真正发挥它的硬件效率。算法原型阶段用 Python 拼 C 很容易但上线部署时显存分配、kernel 调度、通信重叠这种脏活累活往往会决定最终性能上限。最后再分享一个小技巧如果后面有人接手这个项目建议优先看稀疏模式生成后的访存局部性不要一上来就调 kernel 计算效率——访存模式顺了计算效率稍微差点整体也能跑得很好。
8×256K长上下文稀疏注意力：从HiSparse到QSA的移植实战
先解释一下背景。团队有一条基于 QSAQuantized Sparse Accelerator自研推理引擎的线上服务链路底层算子、显存池和调度器都是围绕自家框架设计的。但最近算法侧一直在用 HiSparse 做长上下文稀疏注意力实验效果不错——在 128K 乃至 256K 长度下分层局部性Hierarchical Locality带来的加速比非常明显。问题来了HiSparse 原型代码和 QSA 的接口完全是两套体系模型想上线就得把 HiSparse 的稀疏注意力核心逻辑完整移植到 QSA 里还得在双 RTX 4090 48GB 上跑通 8×256K 的部署配置。折腾了大概两周半中间踩了不少坑也把一些底层细节重新梳理了一遍。这篇就当是移植过程的完整记录包括分层局部性在代码里到底怎么落、8×256K 意味着多夸张的显存和调度压力、以及哪些地方容易翻车。如果你也在做长上下文推理或者稀疏注意力方向应该能省不少时间。1. 项目背景与方案选型为什么非要把 HiSparse 塞进 QSA动这个移植项目之前团队其实有过一轮讨论为什么不直接在 HiSparse 基础上封装 API非要往 QSA 里搬这个问题直接决定了后续所有技术选型值得先把逻辑讲清楚。1.1 HiSparse 在长上下文里的定位和优势HiSparse 不是一个单纯的稀疏注意力实现它最核心的卖点是分层局部性的建模方式。传统稀疏注意力比如 Longformer、BigBird用的是固定稀疏模式滑窗加随机连接或者全局 token 加窗口。这种做法的优点是简单、可控但缺点是模式固定对输入内容的自适应性很差——某些需要长期依赖的 token 对如果不在预设模式里注意力就会被硬切掉精度损失不可控。HiSparse 的做法是分层来做先在一个很粗的粒度上判断哪些区域值得关注再逐步细算。它把上下文划分成多个层级的块每一层通过轻量的向量相似度计算筛掉大量无关区域最终保留的稀疏模式是随输入动态变化的。这个设计在 128K 以上的超长上下文里优势非常明显因为长上下文里真正重要的信息占比通常极低固定模式要么覆盖不够、要么计算浪费太多。算法侧在几个长文本检索和摘要任务上测过HiSparse 在 256K 输入下的 prefill 延迟比 FlashAttention-2 全量注意力低了一个数量级同时精度几乎没有掉。所以这个库不是玩具是真能扛事的。1.2 QSA 框架的定位和移植的必要性QSA 是我们内部维护的一套推理加速框架主要解决三个问题多模型服务的显存隔离、动态 batch 调度、以及量化算子的统一管理。线上所有对外提供的 LLM 服务都是跑在 QSA 上的它对显存池的管理非常细——算子申请显存不直接走 CUDA API而是统一走 QSA 的显存池Memory Pool按请求生命周期分配和回收。这就带来一个现实问题HiSparse 的原型代码完全是自己管显存CUDA malloc 满天飞注意力 kernel 也是独立的 CUDA 实现根本不走 QSA 的算子注册机制。如果直接拿 HiSparse 上线意味着要单独给它维护一套显存分配逻辑这和 QSA 的统一管理是冲突的线上 OOM 的概率会大幅上升。所以项目目标很明确把 HiSparse 的核心算法逻辑尤其是分层局部性这一块重写为 QSA 框架下的算子纳入统一的显存池、算子调度和量化体系。另一条路——把 QSA 改造成适配 HiSparse 的壳——在架构上完全不可取QSA 的调度器会和外部显存管理冲突。1.3 双 RTX 4090 48GB 平台的考量和约束硬件选择其实是被项目需求推着走的。算法侧要求单条请求支持 256K 上下文线上想要 8 个并发请求同时跑也就是标题里的 8×256K。这个要求用显存算一下就知道了如果完全不做稀疏化256K 长度的 KV Cache 会吃掉几十 GB 显存8 个并发根本不可能塞进任何单卡。为什么要强调 RTX 4090 48GB这里解释一下。标准 RTX 4090 是 24GB 显存但我们这批卡是特定厂商推出的 48GB 版本显存翻倍但核心规格基本一致单卡显存带宽和 24GB 版本是同一个量级都在 1TB/s 左右。双卡就有 96GB 总显存加上 NVLink 并不存在——4090 基本没有 NVLink 互联——所以双卡之间的通信完全依赖 PCIe 5.0 的带宽这其实是一个很重要的约束。选这个平台的原因也很朴素成本敏感。A100/H100 的采购周期和预算都不现实而长上下文推理对显存容量的需求远大于对计算峰值的需求。稀疏注意力已经把计算量降下来了对显存带宽的需求也降了所以 4090 级别的显存带宽是够用的瓶颈反而在容量和通信上。2. 分层局部性8×256K 长上下文加速的核心密码如果说移植地址、接口、显存这些都是体力活那一套能把 8×256K 显存需求压到双 4090 能扛住的分层局部性设计才是这个项目的灵魂。这一节详细拆解整个算法思路以及它怎么和部署硬件匹配起来。2.1 长上下文里的注意力分布规律在做任何稀疏化之前先得搞清楚一个问题256K 长度下注意力权重到底是分布在哪里。以前跑全量注意力的时候其实看不出来因为 256K × 256K 的注意力矩阵太大了根本不可能显式存下来。HiSparse 的做法是先在小规模实验里用 FlashAttention 的统计功能输出每个 query 对历史的注意力分布。观察下来的结果非常有规律一个 query 至少会把 60%-70% 的注意力权重放在离它最近的 1K-2K 个 token 上这符合语言模型的局部性偏好但剩下 30%-40% 的权重会散落在更早的上下文中这些位置因 query 而异——有的 query 在找文档开头的人名有的在找段落标题有的在找某段代码定义。如果只用固定窗口这 30%-40% 的信息就全丢了模型表现为记性变差。而全量注意力在这个场景下的问题也很明显每一个 token 都要和 256K 个历史 token 做内积。在 4090 上实测虽然 FlashAttention 已经做到了接近理论峰值但 prefill 时间依然随长度线性增长到 256K 时已经无法接受。2.2 从细到粗的三级稀疏筛选HiSparse 的分层局部性最终在代码里落地成了三级结构。这种设计很像人看文章的方式先扫一眼标题和段落摘要确定大概范围然后翻到具体段落快速扫读最后才逐句精读关键部分。具体来说第一级是全球粗筛级Global Coarse Layer。整个上下文按 1024 个 token 切成一个块块之间用一个可学习的汇总向量表示本质上是这个块的注意力输出平均池化。对每个 query先和所有历史块做一次轻量内积选出 top-k 个候选块。256K 上下文切成 256 个块这一步只需要 256 次向量内积计算量几乎可以忽略。top-k 我们通常取 32 或 48。第二级是局部精筛级Local Fine Layer。对上一级选出的候选块在块内部再按 64 个 token 切片同样用轻量内积筛出最相关的若干片。这一步的复杂度是 top-k 乘以每个块内的片数也就是 32 × 16 512 次内积依然很低。第三级是窗口注意力级Window Attention。不管前面的粗筛结果如何每个 query 都会直接拿最近的 512 个 token 做全量注意力。这一层用来保证语言建模的基本局部性不丢失也避免远端筛选错误导致完全漏掉关键信息。三级筛选合并之后每个 query 实际需要参与的 KV 范围大概是512 个局部 token 粗筛选出的若干片64 token 精度的并集。实测下来每个 query 平均只需要和 2K-4K 个 KV token 做注意力相比原始 256K 直接少了 64-128 倍。2.3 稀疏筛选在硬件上的映射分层局部性不只是一个算法概念它在部署时要解决一个很具体的问题稀疏筛选出来的 KV 在显存里怎么存放才能保证访存效率。最开始踩过一个坑如果把选出来的 KV 拷贝到一个紧凑的连续内存块里再做注意力虽然计算很快但拷贝本身的耗时很夸张——因为每次都要从不同位置拉数据到同一块缓存随机访问很多。后来改成只在块边界上做稀疏化粗筛时以 1024 长度块为单位决定要不要整体加载这样保证从显存搬运时大部分数据是连续的。只有在确定了要加载的块之后才在块内部做更细的筛选。这一层缓冲区设计在移植过程中反复调了很多次是性能影响最大的部分之一。另外QSA 里本来就有基于 CUDA Graph 的显存预取机制我们把块筛选结果作为图的输入使稀疏模式固定的帧内执行可以完全落进预取流水线里避免了 kernel 间同步等待。这也是移植到 QSA 之后吞吐能明显优于 HiSparse 原型的重要原因——原型的 kernel 之间没有这种整体调度。3. 移植实战接口对齐、内核重写和显存管理算法方案定了剩下的就是体力活。这部分按实际的移植顺序来写包括接口对齐、kernel 重写、双卡显存管理三个大块以及每个环节里我认为值得记录的细节。3.1 算子注册和接口对齐QSA 框架里每个算子都注册成一个 Operator对外暴露统一的 Init、Forward、Free 接口。HiSparse 原型的稀疏参数是直接暴露成 Python 对象的字典移植第一步就是把稀疏模式配置从 Python 侧挪到 C 侧注册成 QSA::SparseAttentionConfig。关键的接口参数有这几个window_size局部窗口大小我们固定为 512block_size_coarse粗筛块大小固定为 1024block_size_fine精筛片大小固定为 64top_k_coarse粗筛保留块数默认 32top_k_fine精筛保留片数默认 16这些参数如果暴露成 Python 字典每个请求进来都要解析一遍字符串解析成本不高但积少成多。注册成 C 结构体之后请求配置只传一个枚举 ID查找表直接拿 ID 索引速度快很多。这块属于典型的看着不起眼、改完真舒服的优化。3.2 稀疏模式计算从 CPU 搬到 GPU最初版本的分层稀疏筛选逻辑是放在 CPU 上算的。原因很简单——HiSparse 原型就是这么写的当时的思路是 CPU 算筛选GPU 只负责执行已经被筛选好的注意力 kernel。256K 上下文的筛选逻辑大概要做 256 次块级内积排序在 CPU 上大约需要 3-5 毫秒。一开始觉得这个时间可以接受后来发现部署场景下完全不是这么回事。因为 8×256K 意味着 8 个并发请求同时到达如果每个请求的筛选都在 CPU 上做那么 CPU 侧单线程跑一轮就需要 24-40 毫秒。在 prefill 阶段这是不可接受的——prefill 本身在高稀疏率下也就几十毫秒CPU 端筛选占到一半以上的时间太浪费了。后来把整个筛选逻辑写成了两个 CUDA kernel第一个 kernel 负责块级向量内积和 top-k 选择。把每个 query 的隐藏状态和所有块的汇总向量做矩阵乘然后对结果做 top-k。这里 top-k 我们没有自己实现直接用 CUB 库的 DeviceSegmentedRadixSort 对分数排序去索引实测在 256K 上下文下这个 kernel 加排序的耗时在 0.3-0.5 毫秒左右。第二个 kernel 负责把选出来的 KV 块从显存池中 gather 到一块临时缓冲区。gather 过程用了向量化加载float4 或 half2确保访存效率接近峰值。这个 kernel 耗时大约 0.2-0.4 毫秒。两项加一起一次筛选从 3-5 毫秒降到了不到 1 毫秒而且彻底释放了 CPU 的调度压力。这一步是移植之后性能提升最大的单项优化强烈建议做。3.3 注意力内核和 FlashAttention 的融合筛选之后的注意力计算我们直接沿用 QSA 里已有的 FlashAttention 风格内核但做了一个关键改动把 block mask 融合进内核里。标准 FlashAttention 的流程是遍历 KV 块、在线 softmax 更新我们改成了遍历被选中的 KV 块跳过未选中的。因为稀疏筛选的结果已经确定了每个 query 需要访问哪些块这个信息被编码成一个 compact 索引数组kernel 启动时直接把它加载进共享内存循环体按索引跳转。这带来一个额外的好处IO 量大幅降低。全量 FlashAttention 每次迭代要把 KV 块从 HBM 读到 SRAM而我们的内核只读取被选中的块256K 上下文下从原来的 256 次迭代减到 30-50 次迭代。这意味着从 HBM 读取的数据量降低了 5-8 倍而注意力计算本身在 4090 上的瓶颈恰恰是 HBM 带宽而不是计算单元。3.4 双卡显存分配和流水线切分双卡 96GB 不是一条简单的把模型放卡 0、把 KV 放卡 1就能完事的。模型本身结构不大以 7B 参数量的 GQA 模型为例权重量化后不到 8GB单卡放得下。真正的显存大头在 KV Cache。我们的切分方案是双卡按层做流水线并行卡 0 负责前一半层卡 1 负责后一半层。每个请求的激活从卡 0 进入前向算到中间层后通过 PCIe 把激活传到卡 1。KV Cache 跟随所在层存放所以卡 0 只缓存前一半层的 KV卡 1 只缓存后一半层。这个切分方案在 8×256K 场景下很合适两块卡各自承担的 KV 压力大约是总量的一半避免了任何一张卡先爆。相比之下张量并行TP方案在 4090 上没有 NVLink 的情况下不可取——每层都要跨卡通信PCIe 的延迟和带宽会成为严重瓶颈。流水线并行只需要每请求传一次中间激活通信频率低得多。实测双卡之间传输一层 7B 模型中间激活约 0.5MB 每 token 每层在 batch8、序列片段为 1K 时耗时约 2-3 毫秒占整体延迟的 8% 左右可接受。4. 8×256K 部署显存预算、调度策略和性能数据算法移植完了配置也调通了但真正达到8×256K 可上线还是经历了很痛苦的调优过程。这一节给出具体的显存计算、batch 调度方案和最终性能数据。4.1 8×256K 的显存预算到底怎么算先把账算清楚。8×256K 有两层含义一是模型并发处理 8 个请求每个请求最长支持 256K 上下文二是显存里需要为每个请求维护 KV Cache。以 7B GQA 模型假设 28 层、8 个 KV head、head dim 128、FP16 存储为例一个 256K 长度请求的全量 KV Cache 是256 × 1024token 数× 8KV heads× 128head dim× 2K 和 V 两份× 2FP16 两个字节按 1K1024 计≈ 2.15GB最后除 1024^3再乘 28 层 ≈ 60GB。这个数字很吓人一个请求就 60GB8 个请求要 480GB双卡 96GB 根本不可能。但使用分层局部性之后KV Cache 不再全量存储所有 token 的 KV。我们实际只存两部分局部窗口每层 512 个 token和粗筛选出的远端块每层 32 个 1024 块。按 256K 输入计算实际存储的 token 数大概是 512 加上粗筛保留的 32 × 1024 33280约 33K token相比 256K 压缩了 8 倍左右。这样每个请求的实际 KV 显存大约是 60GB / 8 7.5GB。8 个请求就是 60GB双卡各承担 30GB加上模型权重和中间激活每卡 50GB 以内的预算就够用了刚好卡在双 4090 48GB 的总空间内。所以 8×256K 能落地核心是稀疏化把显存需求从 480GB 压到了 60GB而不是靠堆硬件。4.2 面向稀疏度的动态 batch 调度8 个并发请求各自在不同的解码阶段它们的稀疏度也是动态变化的。请求刚进来时历史 token 少稀疏筛选用不上随着 context 变长远距离依赖开始增多稀疏模式也在变化。QSA 原有调度器是按请求优先级和显存余量来做 batch 拼接的但移植后我们加入了一个新的调度维度稀疏度感知。每当有新请求要加入 batch 时调度器会预估它下一步的 KV 显存增长量而这个增长量是稀疏模式的函数——同样增加 1024 个 token如果粗筛命中的块多显存增长就快命中的块少增长就慢。这个预估模型其实非常简单把粗筛 top-k 的分数缓存下来预测下一个 chunk 能命中多少块命中率乘以块大小乘以 KV 单位大小就是后续显存增量。预测精度不需要太高够调度器做拒绝或排队判断就行。实际跑下来这个机制避免了很多次隐性 OOM——如果只按全量 KV 的增长来估批大小会过于保守吞吐上不去如果完全忽略稀疏度8 个长请求同时增长时会在某一个时间点集体越界。调度器和稀疏度挂钩之后系统能在显存安全的前提下把 batch 尽量打满。4.3 性能实测数据部署完成后跑了一组对比数据硬件环境是双 RTX 4090 48GBCUDA 12.2模型是 7B GQA 量化版本。这里列三个最有代表性的指标。Prefill 阶段256K 输入单请求HiSparse 原型在单卡上大约是 4.2 秒移植到 QSA 后由于显存预取和 kernel 融合降到 3.1 秒左右。这个提升主要来自 CUDA Graph 优化和避免 CPU-GPU 同步而不是算法本身的变化。解码阶段8 并发每请求 256K 上下文batch8有效吞吐大约 1800 tokens/s单请求平均时延 4.5 秒也就是 8 个请求一起跑平均每个请求每秒生成约 225 token。对比同类硬件的全量注意力实现吞吐大约 300-400 tokens/s提升 4-5 倍。显存占用8 个请求全部达到 256K 上下文时卡 0 显存峰值 43GB卡 1 显存峰值 41GB留有大约 5GB 的头部空间给临时激活和通信缓冲区。注意性能数据强烈依赖稀疏度配置。如果 top_k_coarse 从 32 提高到 64吞吐会下降约 30%但精度会更稳。部署时不要盲目追求最高吞吐建议先在目标任务上做一个精度-吞吐的权衡曲线再定参数。5. 移植与部署中的坑常见问题和排查记录最后这部分是踩坑实录很多问题都不是一次定位到的花了大把时间在 profiler 和日志之间来回。整理成速查表方便后面的人直接对照。5.1 问题一块稀疏模式导致精度波动现象移植初期在 RAG 问答任务上 F1 分数比 HiSparse 原型掉了将近 9 个百分点。最初怀疑是 kernel 里 softmax 数值精度问题排查很久才发现是粗筛 top-k 的选择逻辑差异。原因HiSparse 原型在 CPU 端做筛选时用了完整的浮点精度比较移植到 GPU kernel 时为了省寄存器把块分数转成了半精度浮点存储。256 个块分数在半精度下排序相邻分数的区分度不够导致 top-32 的边界块选错。解决把块分数改成 float32 存储排序比较时用全精度。在排序前不做任何量化或截断。修改后精度回归到与原型一致。这个教训是稀疏筛选的关键路径上不要自作聪明做精度压缩省下几百字节显存却亏掉几个点精度不值。5.2 问题二显存池碎片化和隐性 OOM现象8 个请求跑了一段时间后偶发显存分配失败但看峰值显存明明还有余量。QSA 的显存池是池化的池内块大小不等长请求的 KV 增长需要连续的大块内存池子被小的残留块占满后大请求就会分配失败。解决QSA 显存池本身有大块预留机制但默认没开启。改成给 KV Cache 单独建一个大块专用池池内分配单位是 1MB 的块和模型权重、临时缓冲区分开管理。这样 KV 增长时可以直接从大块池里取不会和 activator 的短生命周期分配互相干扰。另外一个习惯性建议定期查看显存碎片率。可以在 QSA 的监控接口里加一个显存池最大连续空闲块大小的统计这个值如果长期明显低于总空闲量就是碎片化的信号。5.3 问题三双卡通信是 PCIe 5.0但时延仍偶发毛刺现象解码阶段时延曲线不是平滑的每几十秒会冒出一个 200-300ms 的尖峰。用 nvidia-smi 看 GPU 利用率卡 0 在尖峰时利用率掉到接近 0卡 1 利用率正常。原因中间激活从卡 0 传到卡 1 是同步等待的当卡 1 正好在处理上一个 batch 的尾部时卡 0 就要空等不短的时间。后来定位到是双方没有做流水线重叠——卡 0 算完后不预取下一份数据卡 1 还在忙两边就串行了。解决在卡 0 的 kernel 执行完之后不等卡 1 反馈直接把下一份中间激活放到 PCIe 传输队列里用 CUDA Graph 的并行分支配合双流实现。这样卡 1 在处理 A 时卡 0 已经在传输 B 了通信和计算完全重叠。修改之后尖峰基本消失P99 时延从 230ms 降到 45ms。5.4 问题四8×256K 预热启动时间过长现象服务重启后第一次 256K 长请求耗时是稳定状态的 3 倍多大约 9 秒。原因QSA 的显存池和 CUDA Graph 在首轮请求时需要冷启动大量显存分配和 kernel 编译触发。这是老问题在短上下文场景不明显256K 场景下被放大。解决部署脚本里加了一个预热请求服务启动后模拟一次 32K 长度的生成请求让显存池预分配到位、CUDA Graph 编译完成。预热后重启成本平摊到首次请求实际体验提升非常明显。下面整理成一张速查表方便直接对照排查问题现象根因解决方式是否容易忽略精度掉点半精度分数排序不稳定块分数改全精度比较容易忽略建议默认全精度间歇性 OOM显存池碎片化KV Cache 独立大块池需要监控空闲块统计P99 时延尖峰跨卡通信串行双流 pipeline 通信重叠需要 CUDA Graph 支持冷启动长尾显存池和 kernel 冷启动预热请求 CUDA Graph 预编译部署流程必做从 HiSparse 到 QSA 的移植过程比预想中复杂但最终效果对得起投入。我个人最大的体会是稀疏注意力这种算法只有在和推理框架的显存管理、调度策略深度耦合之后才能真正发挥它的硬件效率。算法原型阶段用 Python 拼 C 很容易但上线部署时显存分配、kernel 调度、通信重叠这种脏活累活往往会决定最终性能上限。最后再分享一个小技巧如果后面有人接手这个项目建议优先看稀疏模式生成后的访存局部性不要一上来就调 kernel 计算效率——访存模式顺了计算效率稍微差点整体也能跑得很好。