SGLang Prefill 性能优化：用 GPU Trace 剖析 FlashAttention 内核瓶颈

这是最近折腾 DeepSeek V4.1 长上下文推理时攒下来的一点东西。之前一直在上层调 SGLang 的调度参数什么 chunked prefill 大小、max-running-requests、KV cache 比例调到后面明显感觉天花板不在策略层而在 kernel 本身——尤其是 Prefill 阶段的 FlashAttention 路径慢就是慢在底层那几个 GPU kernel 上。所以这次干脆把 SGLang 的 kernel 源码翻出来配合 Nsight 抓的 GPU Trace把 Prefill 从“看起来慢”追到“具体哪里慢、为什么慢”。这篇先聚焦 Prefill而且是“上篇”主要讲分析思路、FlashAttention 在 SGLang 里的实现选型以及怎么用 Trace 反推源码。整个过程不涉及太玄乎的优化技巧核心是建立一套从源码到性能数据之间来回交叉验证的方法。适合已经在用 SGLang 跑模型、但对 GPU 底层还有好奇心的朋友。如果你只是想快速部署一个服务这篇文章的颗粒度可能太细了但如果你想搞明白 Prefill 的 50ms 到底花在哪那可以接着往下看。1. 为什么先啃 Prefill而不是 decode1.1 Prefill 和 decode 是两种完全不同的性能问题很多人在做推理优化时习惯性地把注意力全放在 decode 上因为 decode 的每个 token 都有显性的延迟用户能直接感受到“一个字一个字往外蹦”。但从 GPU 的角度看decode 其实是个相对“温柔”的阶段每次只生成一个 tokenQ 的序列长度是 1计算量不大瓶颈几乎全部押在内存带宽上——KV cache 的读取、权重矩阵的搬运都是 HBM 带宽的消耗大户。Prefill 完全反过来。它要在一次前向里消化整段 promptQ 的长度可能有几千甚至几万 token矩阵乘法的计算量是 O(N²) 起步FlashAttention 在这里也是 O(N²) 的浮点运算只是把内存访问压到了 O(N) 的量级。这个阶段 GPU 的算力利用率、SM 的占用率、L2 和 HBM 的流量会成为真正的瓶颈。同一个模型decode 阶段看带宽prefill 阶段看算力和 kernel 效率这是两套完全不同的性能分析指标。所以在优化 DeepSeek V4.1 这类长上下文模型时prefill 的耗时往往比 decode 更致命。用户发出一个很长的 prompt如果 prefill 阶段要 3 秒那“首 token 延迟”就是从你接受输入到第一个字出现的全部感知时间。SGLang 的 scheduler 能把 prefill 切块、并发调度但切块只是让多个请求交错执行并没有改变单个 kernel 本身有多快。你要是不知道 FlashAttention 的 kernel 能不能吃满 SM上层再怎么调 chunk size 都是盲调。1.2 SGLang 的 Prefill 调度是怎么绕开 decode 的尾延迟的SGLang 和 vLLM 最大的区别之一就是它引入了 RadixAttention用一棵前缀树来缓存 KV cache。对 prefill 来说如果两个请求共享一段 prompt 前缀第二个请求就可以直接复用第一个请求已经算好的 KV不用重新跑一遍 FlashAttention。这个设计在 DeepSeek 这种有 system prompt、few-shot 模板的场景里收益非常明显因为大量的 token 是重复的。但要注意RadixAttention 解决的是“避免重复计算”的问题它不会让你的 FlashAttention kernel 更快。一旦遇到没有前缀可复用的长 prompt或者前缀命中的部分很短kernel 本身的执行效率就决定了下限。另一个容易被忽略的点是 SGLang 的调度策略prefill 和 decode 可以交错执行而不是像某些框架那样先完整做完 prefill 再做 decode。这里面有个很现实的考量——decode 尾延迟非常敏感如果某个长 prefill 霸占了 GPU 几十毫秒所有正在 decode 的请求都会卡住用户体验立刻崩掉。SGLang 的做法是把 prefill 切成多个 chunk穿插在 decode 的间隙里执行。但这带来一个副作用单个 FlashAttention kernel 的输入长度被切短了kernel 启动的开销占比会上升GPU 的流水线也不容易填满。这个会在后面的 Trace 里清楚地看到。2. FlashAttention 内核在 SGLang 里的三条实现路径2.1 torch SDPA、flash-attn、cuDNN attention 怎么选Prefill 阶段最核心的计算就是 attentionSGLang 在 kernel 层面其实并不自己造轮子它是在几套现成的 attention kernel 实现之间做选择。我在源码里梳理下来主线有三条PyTorch 自带的 SDPAScaled Dot-Product Attention、Dao-AILab 的 flash-attn 库还有 NVIDIA cuDNN 里的 fused attention。先说 SDPA。PyTorch 从 2.0 开始把 attention 统一到torch.nn.functional.scaled_dot_product_attention这个接口后面它会根据输入条件自动选择后端如果是训练且需要反向传播可能走 memory-efficient attention如果是推理fp16/bf16 下会优先选 flash attention 的 CUDA kernel。SGLang 里如果没做特殊配置很多情况下最终跑的就是这条路径。它的好处是兼容性好但问题在于参数不可控你很难针对具体序列长度调整 block 大小。flash-attn 库是更底层的选择。它直接调 FlashAttention 官方的 CUDA kernel支持 fp16、bf16、fp8对不同 head 维度和序列长度做了专门优化。SGLang 的AttentionBackend里对 flash-attn 的支持非常成熟DeepSeek 这种 GQAGrouped Query Attention结构的模型用 flash-attn 的flash_attn_varlen_func会很方便因为它原生支持不同序列长度打包成一个 batch 计算这对 SGLang 的连续 KV cache 特别合适。cuDNN attention 是最近几个版本里冒出来的新选项。NVIDIA 把 fused attention 也塞进了 cuDNN优势是它针对 Hopper 和 Blackwell 架构做了非常激进的调优某些 shape 下比手写 flash-attn 还快。但它的一个问题是 API 比较重workspace 规划、plan 的创建都要显式管理在容器里跑的时候还得小心 cuDNN 版本跟 SGLang 编译时的版本是否匹配。我的建议是如果只是部署服务优先 trust flash-attn 后端如果你想做精细的性能调优最好把三条路径都跑一遍 benchmark不要默认某一个一定最快。2.2 Triton Kernel 源码关键参数解读SGLang 内部还带了一套基于 Triton 手写的 attention kernel这在某些场景下会启用比如 flash-attn 库不适配的 shape或者你想在自定义模型结构里快速迭代。Triton 的好处是语言门槛低写起来比 CUDA 快而且它在 Python 层就能控制 grid、block、num_warps 这些关键参数。看 SGLang 的 Triton attention 实现时我最关注的几个参数是BLOCK_M、BLOCK_N、num_warps和num_stages。BLOCK_M是一次性加载多少个 query token 到 SRAMBLOCK_N是每次内层循环加载多少个 key/value token。M 不能太大因为 shared memory 是有限的你加载的 Q 块越大留给 K、V 块和中间注意力分数矩阵的空间就越小。我实测下来H100 上BLOCK_M取 64 或 128 比较合适A100 上 64 更稳再大并不会线性变快反而可能因为 shared memory 溢出导致 occupancy 下降。如果只看到这里你会觉得 Triton 不过是个“高级模板”。但真正值得琢磨的是它内部的 online softmax 循环。# 伪代码体会结构不是完整实现 m_i -inf # running max l_i 0 # running sum of exp acc 0 # 累积输出 for j in range(0, N, BLOCK_N): k tl.load(K_ptr offs_k) v tl.load(V_ptr offs_v) # 当前块分数 s tl.dot(q, tl.trans(k)) * scale # 更新 running max m_new tl.maximum(m_i, tl.max(s, axis1)) alpha tl.exp(m_i - m_new) p tl.exp(s - m_new[:, None]) # 重新归一化之前的累积 acc acc * alpha[:, None] l_i l_i * alpha tl.sum(p, axis1) acc tl.dot(p, v) m_i m_new这段代码的精髓是你不需要等到整个序列的分数全部算出来再做 softmax而是每处理一个 K/V 块就更新一次 running max 和 running sum。想象你在做一场打分不需要等所有评委到齐才知道最终权重每来一个评委你都动态地把之前的分数据按新冠军重新缩放一下。这就是 FlashAttention 能在 O(1) 额外内存下算完 attention 的原因。理解了这种“分块在线归一化”的思路你才能看懂 GPU Trace 里的计算依赖链每个tl.exp、tl.max、tl.dot都对应一串 SASS 指令任何一个环节的 latency 变大都会直观反映在 kernel 的 stall 原因里。2.3 序列长度、精度、共享内存之间的三角关系Prefill 里最麻烦的一点是输入序列长度是不固定的而 kernel 的形状完全由序列长度和 block 大小决定。当序列长度很短比如小于 1024FlashAttention 的优势其实不明显因为整个 K/V 可能都塞得进 L2 cache分块计算的收益被 kernel launch 开销抵消了。我在 Trace 里经常看到短序列 prefill 阶段的 kernel 占用率很低SM 根本没吃满。这时候反而该考虑用sdpa fused kernel这种更轻量的实现或者干脆走标准的 CUDA 矩阵乘路径。当序列长度冲到 8K、16K 以上的时候事情就不一样了。K/V 矩阵远超 L2 容量必须反复从 HBM 搬运。每个 K/V block 被加载到 SRAM 后会被当前 block 里的所有 query token 复用所以在长序列下BLOCK_M越大K/V 数据被重用的次数就越多HBM 流量越低性能越好。这就是为什么长上下文 prefill 的优化空间远大于短序列。精度选择也很关键。DeepSeek 这类大模型在推理时通常用 bf16但在某些 head 维度比较窄的场景fp8 的 FlashAttention kernel 能带来接近翻倍的吞吐提升。代价是精度损失以及 kernel 里需要额外的 scale 处理。我建议先跑一遍 fp8 的离线验证用实际任务的困惑度或者下游指标判断是否可以接受而不是一上来就全量切。3. 用 GPU Trace 把 Prefill 瓶颈钉死在代码上3.1 一次完整的 Nsight 采集流程源码读得再多最后还是得上 GPU Trace 验证。我这边环境是 CUDA 12.4跑的是 SGLang 比较新的 release 版本采集工具用 Nsight Systems 和 Nsight Compute 配合。Nsight Systems 负责看全局时间线回答“时间都去哪了”的问题nsys profile -t cuda,osrt --output deepseek_prefill -w true \ python bench_prefill.py它会输出每个 CUDA kernel 的开始时间、结束时间、耗时、所在 stream。我一般先用它看 prefill 阶段的整体耗时构成是 attention kernel 占大头还是 MLP 的 GEMM 占大头又或者 kernel 之间的 gaplaunch 延迟、同步等待占了大头。如果 gap 很大说明调度器的显式同步太多或者 CUDA graph 没有正确捕获这就不是 kernel 本身的问题了。Nsight Compute 则负责看单个 kernel 的内部细节回答“kernel 为什么这么慢”ncu --set full --kernel-name regex:flash.* --launch-count 5 \ --export ncu_prefill python bench_prefill.py--kernel-name用正则去匹配 FlashAttention 相关的 kernel--launch-count采集多次取平均值避免第一次启动的 warmup 干扰。这里有个实用的经验ncu的 profile 模式会显著降低 GPU 频率以保证数据稳定性所以别直接用ncu采出来的时间作为线上性能依据那只能看相对关系和 bottleneck 类型真实延迟还是得回到nsys的时间线上看。3.2 一份真实 Trace 里最该看的四组数字拿到ncu报告之后不要被满屏的 counter 淹没我每次只看四组数字。第一是 SM 吞吐率SM Throughput。如果某个 kernel 的 SM 吞吐率在 80% 以上说明它已经把 GPU 的计算资源用得很透了这时候你再优化算法本身的意义也不大考虑的是能否减少计算量或者用更低的精度。如果只有 30% 上下说明 kernel 没有吃满 SM大概率是 memory bound 或 latency bound。第二是内存吞吐Memory Throughput。这里要区分 HBM 和 L2。FlashAttention 的设计目标就是压低 HBM 访问如果你在 Trace 里看到 HBM 吞吐接近峰值说明 K/V 块被重复加载太多次BLOCK_M可能有调整空间或者缓存策略出了问题。第三是 warp stall 原因Warp State Statistics。这是最有价值的数据。它会告诉你 warp 在等什么是等long scoreboard等待 global memory 数据返回还是等fixed latency执行依赖的算术指令又或是等mio throttle共享内存/LSU 端口冲突。我在 FlashAttention kernel 里最常见的 stall 是 long scoreboard这意味着 kernel 在等内存数据落地计算单元其实是空闲的。这时候盲目加大num_warps没用反而可能增加指令发射竞争正确方向是调整 block 大小提高数据复用。第四是 occupancy。但不是越高越好我在 H100 上见过 100% occupancy 的 FlashAttention 反而比 70% occupancy 更慢的情况。原因是高 occupancy 意味着每个线程块分配的 shared memory 更少FlashAttention 的核心优势被打折。所以 occupancy 要结合 stall 原因一起看别单独迷信这个数字。3.3 从 Trace 反推源码一次优化动作的完整验证用一个小例子展示怎么串起来。有次我在 DeepSeek V4.1 的 8K prompt prefill trace 里看到 flash-attn 的 kernel 耗时特别高ncu显示 SM 吞吐只有 45%但 HBM 吞吐到了 85%warp stall 里long scoreboard占了 60% 以上。这套组合拳指向一个结论kernel 是 memory bound且主要瓶颈是 K/V 从 HBM 到 SM 的搬运。因为 8K 序列已经超过 L2 能容纳的范围每个 K/V 块都是直接从 HBM 读取的SM 上的计算单元只能干等着。当时我有两个选择加大BLOCK_M提升 K/V 复用率或者换用 fp8 的 FlashAttention kernel直接减少搬运的字节数。我先试了前者BLOCK_M从 64 改成 128HBM 吞吐从 85% 降到 71%kernel 耗时降了 18%。这个结果符合理论预期K/V 复用率翻倍HBM 流量大约减半但因为 L2 和 TLB 的开销还在不是严格的线性下降。如果继续把BLOCK_M加到 256反而变慢了。原因是 shared memory 不够一个 block 内的寄存器也被大量占用occupancy 掉到 40% 以下没办法再用更多 warp 去掩盖内存延迟。这就是参数调整里最常见的“先升后降”现象。最后我把BLOCK_M固定在 128同时保留 bf16没有切 fp8因为那个场景下游业务对精度更敏感。这个过程本身没有什么魔法关键是你要养成交叉验证的习惯从源码的 block 参数出发预判某个改动会影响哪些 GPU 指标再用 Trace 去验证这个预判而不是漫无目的地试参数。4. 排查实录那些看着像 Kernel 问题的坑4.1 区分“GPU Kernel 报错”和“系统层 Kernel 报错”这个系列的标题提到 kernel但很多人搜到这里的目的是排查各种奇怪的报错。有两条路完全不一样一是 CUDA 层面的 kernelGPU kernel二是操作系统内核OS kernel它们只是中文翻译撞车了。比如有人搜“kernel data inpage error 蓝屏”这个跟 GPU 推理一毛钱关系都没有。它通常是 Windows 内存映射文件读取失败要么是磁盘坏道/掉盘要么是内存条不稳对应的是操作系统内核的 I/O 错误不是 CUDA 的 kernel 报错。你要是刷到了类似的蓝屏先去查硬盘健康度和内存别在 SGLang 的日志里浪费时间。反过来SGLang 在跑的时候如果报 CUDA kernel error描述里一般带CUDA error、an illegal memory access was encountered之类的话这才是 GPU kernel 层面的问题。我的经验是 80% 的情况是 KV cache 越界或者显存不足导致的内存踩踏剩下 20% 是某个算子在大 batch 下触发了 CUDA 的内部 bug。4.2 CUDA kernel errors 和 NVIDIA kernel module 的常见处置跑 SGLang 久了你会遇到两类高频问题。第一类是CUDA kernel errors might be asynchronously reported at some other API call这种报错最烦人因为真正的错误发生点已经过去了你看到的是“延迟上报”。出现这个第一件事不是改代码而是用cuda-memcheck或者compute-sanitizer去定位内存访问违规然后回退请求复现。很多情况下这是某个极端长的 input 触发了一个 kernel 内部的 shared memory 越界跟你的代码逻辑无关但你要能稳定复现才能确认是 SGLang 的 bug 还是你显存不够导致 OOM 后状态错乱。第二类是启动时的 NVIDIA 驱动模块问题The NVIDIA kernel module was not created一般是驱动版本和 CUDA runtime 不匹配或者内核模块没有正确加载。这种问题别在应用层折腾去查nvidia-smi是否正常输出然后重新安装匹配的驱动。我做容器部署时经常遇到容器内 CUDA 版本和宿主机驱动不匹配的问题SGLang 镜像拉下来之后第一件事就是对着nvidia-smi的 driver version 和容器内的libcuda.so版本做检查。4.3 Trace 数据自己就带坑时钟、温度、重放偏差最后分享几个 Trace 数据本身的坑。GPU 动态频率是最大的干扰源。同一段代码冷启动时 GPU 可能跑在 1.0GHz跑热了之后 boost 到 1.8GHz两次采集的 kernel 耗时能差 40%。ncu默认会锁频nsys不会。所以nsys的时间线只能用于定位相对占比不能跨 session 直接对比绝对耗时。做参数调优时要么全程用ncu锁频后的数据要么在脚本里用nvidia-smi -lgc手动锁定 GPU 频率。温度也很重要。GPU 温度超过某个阈值后boost 频率会阶梯式下降。我在做长序列压测时前 10 分钟和最后 10 分钟的数据经常不能直接对比除非你盯着温度曲线做归一化。对于 prefill 这种一次性 burst 计算温度影响相对小但如果你的 benchmark 脚本连续跑几十分钟后面采到的 kernel 时间普遍会变慢。还有ncu的 profile 模式本身会改变程序行为。它默认会禁用 caching allocator 的某些优化还会强制所有 kernel 串行化执行所以你从ncu里看到的 gap 不代表真实部署时的并发情况。正确的做法是先用nsys看真实并发的全局时间线再用ncu深入单个 kernel 的内部机制两者各管一段不要混着解读。写到这里差不多把 Prefill 阶段从源码到 Trace 的主线走通了。我个人习惯是每次调优都先拉一条基线 trace记录 SM 吞吐、内存吞吐、stall 占比三组数字改完参数再看这三组的相对变化不看绝对耗时。这套方法在后面对比 decode 阶段时也会沿用等我把 decode 的 trace 整理出来正好可以跟这篇做一个完整的对照——毕竟 prefill 和 decode 的优化方向在很多地方是反着来的。