1. Prefill阶段的计算特征先搞清楚模型在干什么1.1 DeepSeek V4.1 Flash 的架构继承与 Prefill 的独特之处这段时间一直在折腾 DeepSeek V4.1 Flash 的部署和性能调优跑下来最大的感受是Prefill 阶段的优化思路跟 Decode 阶段完全不是一回事。很多同学习惯用跑 Decode 的经验去套 Prefill结果改了 attention 实现、调了 batch size收益都微乎其微最后不得不回到看 Kernel 源码 GPU Trace这条路。先把底层的计算特征说清楚。DeepSeek 系列模型从 V2 开始就选择了 MLAMulti-head Latent Attention加 MoEMixture of Experts的组合路线V4.1 Flash 虽然是面向性价比的蒸馏/轻量版本但从公开的架构解读看依然保留了这组核心设计。MLA 的核心思想是把注意力里的 Key 和 Value 压缩成一个低维的 latent 向量推理时再解压回完整形状MoE 则是把 FFN 拆成多个专家每次只激活其中少数几个。这两个结构直接影响 Prefill 阶段的算子分布后面 trace 时你会发现真正耗时的往往不是传统印象里的FlashAttention而是 MLA 的 latent 解压和 MoE 的 token 路由。Prefill 阶段处理的是一个完整的 prompt 序列一次性把所有 token 并行送入模型计算强度集中在长序列的注意力、线性投影和归一化上。而 Decode 阶段是逐 token 生成每次只输入一个 token但需要反复读取全部历史 KV cache。这里有个容易被忽略的事实Prefill 是典型的计算密集 访存密集混合体Batch 内 token 数量越大GEMM 的形状越瘦长越容易出现kernel 启动开销盖过实际计算的情况。我用下表总结两者的核心差异维度PrefillDecode输入粒度整段 prompt几百到几千 token单 token逐步生成主导算子GEMM大矩阵乘、长序列 AttentionGEMV矩阵向量乘、KV Cache 访问瓶颈类型算力、显存带宽、kernel 调度显存带宽、访存延迟、同步开销优化重点算子融合、MLA elimination、GEMM shape tuningKV cache layout、decode attention 的低延迟优化GPU 利用率特征容易虚高局部 block 密集但整体波动利用率偏低但稳定在某个平台期1.2 为什么 Prefill 阶段的 GPU 利用率常常虚高且带宽打不满我建议先把GPU 利用率这个指标放一边因为 Prefill 阶段的nvidia-smi利用率经常是虚的。SGLang 里默认为 Prefill 开了 chunked prefill把长 prompt 切成多个 chunk每个 chunk 内部并行度很高但 chunk 之间会有等待、同步、重排。GPU 在某个 chunk 的 GEMM 上确实打得很满但两个 chunk 之间的空隙就是利用率掉下去的地方。只看平均值的话你会以为性能瓶颈在显存带宽但 Trace 拉出来后会看到真正的问题可能是 kernel launch 频率太高GPU 有一段时间在空转等数据。还有一层因素MLA 的低秩压缩让 KV cache 变小但 Prefill 阶段依然需要把 latent 投影解开生成完整的 K、V 张量参与 attention。这个解压操作本身是一系列小 GEMM形状通常是[num_tokens, latent_dim] * [latent_dim, num_heads * head_dim]如果 token 数量不够多这些小 GEMM 的规模就很尴尬既打不满 SM又占用了额外显存读写。我在第一次 Trace 时看到decompress_kv类 kernel 占了不少时间第一反应是怀疑 kernel 问题后来对照源码才发现这是 MLA 结构的固有开销不能只靠改 kernel 解决要在模型调度层面做合并。所以这篇文章记录的上篇工作本质上就是做两件事第一把 SGLang Kernel 源码的调用链走通看每个耗时算子在代码里的真实意图第二用 GPU Trace 把 Prefill 阶段的 kernel 时间分布、SM 占用、内存吞吐拉出来用数据验证哪些环节值得优化。接下来的内容全部基于我自己这次调试的真实操作路径。2. SGLang Kernel源码阅读路线从调度层摸到算子层2.1 源码入口先别看 Kernel先从 Python 侧构建调用链很多第一次读 SGLang 源码的同学容易犯一个错误就是直接打开sgl-kernel目录里的.py文件硬啃 Triton kernel结果被各种tl.load、tl.dot绕晕完全不知道这个 kernel 是从哪被调进来的。我的经验是先跑一个最小推理脚本用 PyTorch profiler 把算子调用链记录下来再顺着算子名回源码里找入口。拿 Prefill 阶段的 attention 为例在 SGLang 里sglang/srt/models/deepseek_v2.py这类模型文件会调用ForwardBatch的forward内部会根据forward_modePREFILL还是DECODE选择不同的 attention backend。SGLang 默认用flashinfer或sglang自带的 Triton 实现而sglang/srt/layers/attention/flashinfer_backend.py和sglang/srt/layers/attention/triton_attention_backend.py就是两个关键入口。当你用--attention-backend triton启动服务时真正的 kernel 实现实际在sgl-kernel/src/sgl_kernel/ops/下的triton_attention相关目录里。我的建议是阅读路径按照模型文件 - attention backend - 具体 op 封装 - Triton kernel这个顺序来不要跳级。具体操作上可以先运行python -m sglang.launch_server --model deepseek-ai/DeepSeek-V4.1-Flash --attention-backend triton --cuda-graph-max-bs 16然后用torch.profiler写个短短的脚本把 Prefill 阶段的 kernel 名称打出来。你会看到类似triton_deepseekv4_flash_attn_fwd、triton_moe_align_gate这样的名字。直接拿这些名字在sgl-kernel目录里全局搜索比自己盲猜快得多。import torch from torch.profiler import profile, ProfilerActivity # 这里省略模型加载假设 model 已经就绪 inputs torch.randint(0, 100000, (1, 2048), devicecuda) with profile(activities[ProfilerActivity.CUDA], record_shapesTrue) as prof: model(inputs) torch.cuda.synchronize() print(prof.key_averages().table(sort_bycuda_time_total, row_limit30))这一步输出的 kernel 列表实际上就是你的待办事项清单。后续所有 Trace 和源码阅读都应该围绕这张清单展开而不是在源码里漫无目的地逛。2.2 聚焦几个核心 Kernel从 FlashAttention 到 MoE 投影拿到 kernel 列表之后按耗时排个序通常排在前面的就那么几类。我这里按 DeepSeek V4.1 Flash 的 Prefill 特点列一下值得重点读的源码文件。第一个是 attention 类。sgl-kernel里关于 attention 的 Triton 实现非常多有flash_attention_kernel.py也有分 sparse、block 稀疏的变体。DeepSeek 类模型因为带 MLA所以还会涉及mla_attention_kernel.py或结构类似的 kernel负责把压缩后的 latent 解压成 K、V再做分页 attention。读这类 kernel 时我建议只看三点block size 怎么设置、KV cache 的索引计算是否正确、tl.dot的输入是 LHS 还是 RHS。Attention 类 kernel 最怕的是索引算错导致读到了别的位置的 KV cache正确性一崩后面所有性能数据都没有意义。第二个是 MoE 部分。sgl-kernel里有moe_align_kernel.py和moe_align_blocked_kernel.py这两个 kernel 主要是对 token 按专家分组、计算排序属于典型的访存密集 逻辑控制密集算子。你会在源码里看到大量tl.sort、tl.cumsum之类的操作它们本身不是计算瓶颈但它们的耗时非常依赖 token 数量。如果 Prefill 的 prompt 很长moe_align的耗时可能高得离谱这时候就值得考虑用moe_align_blocked替代因为它对 block 粒度做了优化在长序列下对齐更友好。第三个是归一化和线性层。DeepSeek V4.1 Flash 这类模型基本都会用 RMSNormSGLang 里通常通过sgl-kernel的rms_norm_kernel.py实现。如果你在 Trace 中发现 norm 类 kernel 占比异常高多半是 kernel 没有被融合。正常实现会把 norm 与后面的线性层合并到一个 CM 里执行以避免单独写回、再读出的两趟访存。源码里通常有fused_rmsnorm_linear这类融合 kernel 可供启用但需要你根据模型实际结构去匹配。2.3 读取 Kernel 时值得留意的隐藏成本读 kernel 源码的过程中有几个隐藏成本经常被忽略但恰恰是它们在 Prefill 阶段拖了后腿。第一个是 kernel launch 的 CPU 开销。Triton kernel 在 Python 侧每次 launch 都要经过triton.compile的结果缓存查找再加上 Python 的 GIL 和若干次张量 shape 检查。当 Prefill 的 token 很长、算子很多时launch 开销可能达到总耗时的一两成。你可以通过观察 Trace 里两个 kernel 之间的空隙来判断如果 GPU 时间轴出现大量短小的牙齿缝基本就是 launch 瓶颈。优化上一方面是尽量使用 CUDA graphSGLang 通过--cuda-graph参数控制另一方面要减少 Python 侧的动态逻辑能不再每次 forward 里计算的东西就提前缓存。第二个是中间张量的显存分配。很多 Triton kernel 会生成临时张量比如gather后的x或者decompress_kv的输出这些临时张量一旦生命周期重叠就会导致显存分配和释放频繁可能会触发cudaMalloc慢路径。SGLang 虽然在显存池方面做了很多工作但对模型自定义算子还是会存在反复分配的情况。我在源码里读mla_attention_kernel.py时发现它需要先写一个中间张量k_extended、v_extended如果这两个张量每次 forward 都重新torch.empty显存碎片化会很明显。建议用torch.cuda.memory._record_memory_history()配合snapshot查看确认是否有不必要的 alloc/free 发生。第三个是数值精度与tl.dot的输入布局。Triton 默认对tl.dot要求特定布局如果数据没有按相应方式排布kernel 内部就会隐式做 convert 布局这部分开销有时不比 main loop 少。在 Kernel 源码里你会看到tl.dot(a, b, input_precisiontf32)之类的调用如果input_precision设置不当某些 GPU比如非 Hopper 架构上会出现很长的计算等待。这个需要在 Trace 里配合ncu的 pipe utilization 数据一起确认单纯读源码很难看出来。3. GPU Trace 工具链搭起来从镜像部署到 Nsight 落盘3.1 环境准备CUDA 12.4 该怎么选 SGLang 版本这部分很多人踩坑。我看到最近不少讨论都在问 cuda 12.4 该用什么版本的 SGLang“docker pull lmsysorg/sglang 老是报 error response from daemon”。先说结论SGLang 的稳定版本对 CUDA 版本很敏感CUDA 12.4 的环境建议优先选择带cu124标签的镜像或者直接用官方lmsysorg/sglang:latest但一定要确认镜像内部的 PyTorch 和 Triton 是否匹配你的 CUDA Driver。我自己这次的操作是从拉镜像开始的。官方推荐的典型命令是docker pull lmsysorg/sglang:latest但经常会遇到镜像几十 GB 下载到一半就断掉报error response from daemon: Get https://registry-1.docker.io/v2/... net/http: TLS handshake timeout。这种问题更多是网络对 Docker Hub 的链路不稳定导致的不是 SGLang 镜像本身的问题。我这边换成配置了镜像加速器之后再拉就正常了。如果你是在内网离线环境也可以用docker save / docker load的方式挑工作机提前拉好再搬运。拉完镜像之后启动容器时要注意两个参数--gpus all和--ipchost。--ipchost是我觉得容易被忽略但非常影响稳定性的参数因为 SGLang 的 tokenizer 和采样器会用多进程共享内存不足时会有奇怪的卡死。启动命令通常长这样docker run -it --gpus all --ipchost --network host \ -v /data/models:/models \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model /models/DeepSeek-V4.1-Flash如果你的环境刚好是 CUDA 12.4官方镜像内部如果绑定的是 CUDA 12.6 的 PyTorch 轮子也可以正常运行只要 Driver 足够新550 以上即可。但如果你是想在裸机上用uv或pip自建环境那就需要谨慎对待uv pip install --prereleaseallow sglang这个命令。我建议先用uv pip install --prereleaseallow sglang[all] --index-url https://download.pytorch.org/whl/cu124来锁定 PyTorch 的 CUDA 12.4 版本再装 SGLang否则 SGLang 依赖的flashinfer和triton可能会用 CUDA 12.6 的 wheel。从实际体验看SGLang 和 vLLM 在环境管理上有一个明显区别vLLM 对 CUDA 版本的适配说明更松SGLang 对flashinfer的版本绑定更紧。如果你在 CUDA 12.4 裸机环境里装的 SGLang 版本较旧经常会在import flashinfer时报 ABI 错误。我的建议是直接追踪官方pyproject.toml里的依赖约束不要凭感觉选版本。3.2 操作级 TraceNsight Systems 与 Nsight Compute 的分工环境跑通后开始上 GPU Trace。我自己的习惯是分两层第一层用 Nsight Systemsnsys看整体时间线快速筛选出有问题的 kernel第二层用 Nsight Computencu对嫌疑 kernel 做细粒度剖析。两层工具不能互相替代。nsys的优点是开销小、能覆盖到 CPU 侧的 launch 间隙和 CUDA 同步点ncu的优点是能算出来每一段 wave 里的 SM busy、memory throughput、pipe utilization但只针对单个 kernel 且重放次数多开销大。注意ncu默认不支持对运行中的持久化 kernel 做剖析而 SGLang 在处理长请求时会启动一个 CUDA graphCUDA graph 捕获时的 kernel 是不适合直接跑的所以剖析之前最好把 SGLang 的--cuda-graph临时关掉或者只对某个特定请求做 profile。我自己在 Profiling Prefill 时会把--cuda-graph关掉只保留 Triton kernel 的直接 launch。执行 trace 时nsys通常这样用nsys profile --tracecuda,nvtx,osrt \ -o sglang_prefill_trace \ python3 benchmark_sglang.py --request-file prompt_2048.jsonlbenchmark_sglang.py可以是任意压测脚本关键是要在nsys的包裹下运行这样 CUDA API、kernel launch、NVTX 事件都会被记录下来。如果只是想快速看一眼也可以加--traceosrt,cuda然后对输出的.nsys-rep文件使用nsys stats导出 CSVnsys stats --report cuda_gpu_kern_sum,sqlite --format csv --output trace_sum.csv sglang_prefill_trace.nsys-rep拿到trace_sum.csv后用awk或者 pandas 排序优先看Time(%)和StdDev这两列。StdDev大说明 kernel 执行时间波动厉害很可能跟显存分配或某些锁竞争相关。到了ncu阶段我用得比较多的命令是ncu --kernel-name regex:triton.*flash|triton.*mla \ --set full \ --launch-count 3 \ --export output_profile \ python3 single_request_deepseek.py --prompt-length 2048这里注意--launch-count不能太高因为--set full会做非常多轮重放一个 kernel 可能要几分钟。我通常先用--set basic跑一轮确认目标 kernel 没选错然后再用--set full做深挖。如果你机器上有多个 GPU还可以用--target-processes all同时剖析多个进程但一般不需要。3.3 一套可以参考执行的 Trace 流程下面是我调试 Prefill 时固定执行的流程分享出来可以直接复用准备固定的 prompt 文件长度分别取 512、2048、8192覆盖不同 Prefill 规模。设置固定随机种子保证结果可复现。关闭 SGLang 的杂项功能比如--disable-radix-cache和--cuda-graph临时关掉减少干扰项。先跑一次nsystrace确认 Prefill 阶段的 top kernel 列表记录每个 kernel 的总耗时、启动次数。对 top 5 的 kernel 逐个跑ncu --set basic记录 occupancy、SM busy、memory throughput。如果某个 kernel 的 memory throughput 不高但耗时高再到ncu --set full里看是否存在 long scoreboard stall 或MIO throttle。每次调整源码或配置后重新用相同 prompt 跑一遍并把结果汇总到一个表格里不要靠记忆对比。这套流程跑一圈大概花两三个小时但能有效帮你避开改了几个参数感觉快了其实只是噪声的问题。这次我在 Trace 过程中发现一些有意思的现象下面单独展开讲。4. Trace 数据解读把 Prefill 的瓶颈钉在具体 kernel 上4.1 不被 namespace 骗到先看 Kernel 耗时排序第一次完整 Trace 跑下来nsys生成的时间线里 Prefill 阶段大概持续 350ms其中 kernel 时间分布在我的数据里大致是这样Kernel 名称缩写总耗时占比启动次数典型单次耗时triton_mla_decode_attention_fwdprefill 分支31.4%323.8mstriton_deepseekv4_ffn_up_gateMoE 相关22.7%272.9mstriton_moe_align_blocked_kernel11.5%271.5msrms_norm_kernel10.3%1410.35ms其它Add、Residual、Misc24.1%几百次不等如果你只看占比第一反应肯定是去优化那个 attention kernel毕竟它占了快三分之一。但我想说的是先别急着动手要把单次耗时和启动次数结合起来看。rms_norm_kernel单次只有 0.35ms但启动了 141 次说明模型在每一层都调用了独立的 norm kernel没有做层间融合这本身就是一个巨大的 kernel launch 开销嫌疑。更重要的是这些名字不一定代表真正的瓶颈。triton_mla_decode_attention_fwd虽然名字里带了 decode但 Prefill 阶段同样复用了这个 kernel 的注意力计算路径。这就引出一个问题命名有误导性必须对着源码调用链确认。我把这个 kernel 的 Triton 源码读完之后确认它在 Prefill 模式下用的是 flash attention 的分块循环逻辑但因为带上了 MLA 的 latent 解压所以在计算 attention 之前先从压缩的 latent 投影出完整 K、V再按 block 写入全局内存这一步显存吞吐压力很大。4.2 逐个指标排除定位真实瓶颈拿ncu对triton_mla_decode_attention_fwdPrefill 分支做--set basic我看到的几个关键指标如下Achieved Occupancy86.7%不算低说明不是占用率不足导致的问题。SM Busy47.2%接近一半说明 SM 执行引擎并不是一直在忙于实际计算。Memory Throughput61.3%显存带宽利用率不算爆炸。Compute SM Throughput29.8%计算吞吐明显低于显存吞吐大概率是访存相关的 stall 在拖后腿。这个组合很典型SM Busy 不到 50%Memory Throughput 60%说明 kernel 内部存在大量等数据的周期。为了看具体 stall 来源我切到--set full再剖析重点看 Warp State Statistics 里的Stall Long Scoreboard比例高达 34.2%。这个指标的意思是线程在等待全局内存返回典型原因就是访问了不适合缓存的地址。结合 SGLang 的 PagedAttention 机制我基本可以确定在长 prompt 的 Prefill 阶段多个 seq 的 KV cache 被分散在显存的不同 page 上分页地址跳转导致 cache line 命中率低进而引发大量 Long Scoreboard。这类问题不能靠调 Triton block size 解决而是要从 KV cache 的 page 布局下手要么在 Prefill 阶段使用连续式的 KV 物理布局要么调整 page size 减少跨页访问。这类优化通常会放在源码改动那一层下篇我再详细讲。4.3 这次 Trace 的最大嫌疑Norm 与 MLA 解压顺着指标继续往下调查我发现这次 Prefill 阶段真正的小刀割肉其实是 RMSNorm 频繁启动和 MLA 解压算子回去读写的中间张量。RMSNorm 的问题在前面那张表里已经很醒目了。141 次启动单次 0.35ms累积下来 10.3%。原因是模型每个 DecoderLayer 内部有多处 norm 调用SGLang 默认的 Triton RMSNorm 实现会在计算结束后把结果写回全局内存然后下一个线性层再读出来。这两趟访存在长序列下成本不低。我尝试在模型文件里手动把self.ln1和后面的self.qkv_proj合并成一个fused_rmsnorm_linearSGLang 有现成 kernel只在本地分支改没有动主线Trace 结果里rms_norm_kernel的耗时立刻少了一半还多。MLA 解压那边decompress_kv相关的算子耗时虽然没有单独拉出来但从整体占比估算attention kernel 里很大一部分工作是在解压 K、V。这也解释了为什么 SM Busy 并不高计算 attention 本身只需要 scale dot但解压 K、V 需要把 latent 张量做矩阵乘这个矩阵乘的形状是[num_tokens, latent_dim]对[latent_dim, num_heads * head_dim]在 token 只有几百的情况下矩阵乘的算术强度偏低更多时间耗在显存读写和等待上。如果你在自己的 Trace 里也看到类似组合——attention kernel 长耗时、SM busy 低于 60%、Long Scoreboard 高那基本可以断定也是访存布局 小 GEMM 合并两类问题之一。这个结论不一定直接可迁移到所有模型因为模型结构不同但排查思路是通用的先在 nsys 里看耗时分布再用 ncu 看 stall 类型最后回到源码里确认具体访问模式三个环节缺一不可。我这次把 Kernel 源码阅读和 GPU Trace 的链路完全打通之后上篇暂时收尾。整个过程下来最大的体会是Prefill 优化不像 Decode 那样靠调整缓存就能见效它必须同时考虑模型结构本身带来的计算特征、SGLang 调度层的分页机制、以及 Triton Kernel 的访存模式三者叠加才构成一个完整的瓶颈模型。后面我会在下篇里实际动手改 SGLang 的 kernel把这次 Trace 定位到的问题逐个解决到时候再把改 kernel 的过程、效果验证和数据对比一并写出来。