如果你最近在折腾大模型推理服务化SGLang 这个名字应该不陌生。它靠着 RadixAttention 前缀复用、连续批处理和 chunked prefill 这套组合拳在同类框架里把吞吐做得很能打。但我们团队接到的任务更拧巴不只是在主流加速卡上跑通还要把 SGLang 迁移到昆仑芯XPU平台上并且要求支持多芯并行。绕不开的第一道坎就是多芯插件机制——一层让上层框架不关心底层是 CUDA、ROCm 还是 XPU 的隔离层。这篇文章不写 PPT 式的总结只讲我们实际落地 SGLang-Kunlun 时怎么切接口、怎么管显存、怎么调吞吐以及最后沉淀下来的工程化最佳实践。1. 我们为什么要折腾一个“多芯插件机制”1.1 从一处 if 判断说起后端硬编码带来的连锁问题绝大多数推理框架的第一版都是从单卡单硬件跑通开始的。SGLang 也不例外早期路径上 CUDA 的烙印相当明显显存分配直接走torch.cudakernel 选择写成if device cuda集合通信默认调 NCCL。这套写法在“只有一种卡”的封闭环境里没什么问题但一旦要支持昆仑芯这类 XPU 后端口麻烦就全冒出来了。我们最初试过最朴素的做法在关键函数里加if device.type xpu的分支把显存分配换成torch.xpu对应的接口。前两三个点改起来还挺爽改到第七八个点的时候开始难受改到调度器、显存池和 kernel 选择层时基本就失控了。原因很简单硬件相关的代码散落在框架各处有的是显式依赖有的是隐式假设比如默认device_ids从 0 连续编号。每一处都要打补丁等于把一个本来可以独立演进的硬件后端硬生生揉进了核心代码里。这还没完。多芯并行不是“多插几张卡”就完事它还涉及设备间通信、按芯分片的显存规划、负载均衡和故障隔离。如果这些逻辑跟具体硬件绑定每引入一种新加速器都要把调度器重写一遍维护成本是指数级上升的。1.2 SGLang 本身的扩展点在哪里SGLang 的架构相对清爽核心可以拆成三层管理请求和前缀的 Scheduler、负责模型前向的 ModelRunner、以及提供算子与显存原语的 Runtime 层。我们需要的插件机制其实就是把 Runtime 层抽象成一组稳定的接口让 Scheduler 和 ModelRunner 只依赖接口不依赖具体硬件。我在实际切分时参考了社区在 ROCm、IPEX 等后端口上的通用做法把接口收敛成四个维度设备原语、显存管理、算子分派、集合通信。设备原语管 device 枚举和 stream 同步显存管理管块分配和缓存算子分派负责把 Attention、GEMM 这些高频算子路由到 XPU 上对应的实现集合通信则屏蔽不同硬件厂商的通信库差异。提示插件机制的重点不是“多写一层包装”而是把“变化的部分”和“不变的部分”切开。Scheduler 的排队逻辑、RadixAttention 的树结构、连续批处理的策略都跟硬件无关千万不要为了适配后端把它们也改掉。1.3 三种接入方式的对比给 SGLang 加一个新硬件后端常见路径有三条我们用一张表说清楚区别接入方式改动范围维护成本性能上限适用场景散改 if 分支小零散高每次上游更新都可能冲突低很难做深度优化临时验证不推荐覆写关键类中需理解核心类层次中需跟随框架版本同步调整中能做部分算子替换后端差异较小的场景独立后端插件大但集中在接口层低核心框架升级时插件层基本不动高可针对多芯做专门调度长期维护、多芯片并存的平台我们最终选了第三种。前期投入确实多但后面每次 SGLang 上游发版我们只需要回归插件接口层不用再追着每个 commit 打补丁。对多芯并行的支持也是基于这套插件机制才真正做起来的。2. SGLang-Kunlun 适配的分层架构接口怎么切才合理2.1 四层切分DeviceOps、Memory、Kernel、Communication把后端抽成插件最难的是接口粒度。切得太粗底层差异还是会漏上来切得太细每个后端都要实现一大堆没用的方法接入成本陡增。我们最终定下来的切法以 XPU 适配的实际经验为主参考了社区常见的后端抽象方式分成四个接口层第一层是 DeviceOps管最基础的东西设备数量枚举、当前设备切换、stream 创建与同步、事件记录。这些看着不起眼但多芯场景下全是坑。比如昆仑芯这类 XPU 设备stream 的语义和 CUDA 不一定完全一致有的型号不支持细粒度事件有的需要显式xpu.synchronize()才能保证跨芯可见性。插件层必须把“同步”这个语义统一成上层能理解的状态。第二层是 Memory负责显存分配、释放、拷贝和置零。这里我强烈建议直接做成“基于块的内存池”不要暴露裸的malloc/free。原因后面细说简单讲就是多芯并行下裸分配会导致显存碎片和跨芯内存互相干扰调度器完全没法控制。第三层是 Kernel做算子路由。Attention、RoPE、GEMM 这类重算子在 XPU 上通常有专门优化的实现一些冷门算子则可能要走 fallback。插件层提供dispatch方法由后端自己决定某个算子用原生 kernel 还是降级路径。第四层是 Communication抽象集合通信。对多芯来说这一层决定了 tensor parallel 能不能真正跑起来。CUDA 后端用 NCCLXPU 后端通常对接厂商自己的通信库接口必须设计成 “AllReduce、AllGather、ReduceScatter、Send/Recv” 这几个原语上层只描述“我要做什么通信”不关心底层是谁在传。class XpubBackend(Backend): def device_type(self) - str: return xpu def get_device(self, local_rank: int) - Device: return XPUDevice(local_rank) def create_memory_pool(self, device: Device, capacity: int) - MemoryPool: return XPUMemoryPool(device, capacity) def dispatch(self, op_name: str, *args, **kwargs): if op_name in XPU_KERNEL_TABLE: return XPU_KERNEL_TABLE[op_name](*args, **kwargs) return self.fallback(op_name, *args, **kwargs) def create_communicator(self, world_size: int, rank: int): return XPUCommunicator(world_size, rank)这段代码是简化示意但接口的“味道”是准的上层拿到的永远是统一抽象的Device、MemoryPool、Communicator真正干活的是插件内部的对象。SGLang-Kunlun 的适配本质上就是把这个XPUBackend的实现填满、填对。2.2 插件注册与生命周期管理接口定好之后下一步是让框架在启动时能发现并加载插件。我们在实践中采用注册表机制框架启动时读环境变量或启动参数里的--backend到注册表里找到对应的插件类然后依次调用on_init、on_device_acquire、on_model_load这几个生命周期钩子。有几个细节值得一提。第一插件加载必须支持“延迟初始化”。就是说on_init阶段只做配置解析和参数校验不要急着占显存、建通信域。否则多芯场景下如果某个芯初始化失败整个进程就得重启没有回旋余地。第二每个生命周期钩子都要有降级行为比如on_device_acquire失败时插件可以选择用较少的芯继续跑而不是直接崩掉。这在工程上很有价值尤其在生产环境里一台机器 8 芯有一芯故障能降级成 7 芯继续服务比直接挂掉好太多了。注意生命周期方法不要设计成纯虚类让每个后端硬实现最好给一个默认实现。新的后端接进来时可以先只改最必要的三五个方法跑通流程后面再逐步优化。3. 多芯并行真正难的地方显存、调度与集合通信3.1 多芯显存池从连续裸显存到 BlockAllocator单卡时显存管理已经够烦了多芯并行时问题直接放大好几倍。我们最初在昆仑芯上跑 7B 模型两张芯各分一半参数但 KV cache 的分配非常随意torch.empty一次性申请连续显存然后按请求串行切。结果跑了一下午显存碎片化严重有的芯剩 2GB有的芯只剩 200MB调度器明明看到总量充足却分配不出一个完整的 block。后来我们做了一个关键决定不再用裸显存而是引入 BlockAllocator。做法是把每张芯的显存切成等长 blockblock 大小一般取单个 KV cache 单元对齐到 512 字节或 1KB 的整数倍。调度器申请时按 block 粒度申请释放时按 block 粒度归还块之间允许不连续。这么做有三个直接收益。一是碎片问题基本消失因为最小分配单元固定了不会出现“大块被拆成小块小块又拼不回去”的尴尬。二是方便做前缀复用RadixAttention 的树节点天然可以和 block 一一对应缓存命中时直接把 block 引用接上不需要搬数据。三是可以按芯做配额控制每张芯有多少可用 block 一目了然调度器可以据此决定把一个请求放到哪张芯上。class XPUMemoryPool: def __init__(self, device, num_blocks, block_size): self.blocks [None] * num_blocks self.free_blocks set(range(num_blocks)) self.block_size block_size def alloc(self, n: int) - list[int]: if len(self.free_blocks) n: return None blocks list(self.free_blocks)[:n] self.free_blocks.difference_update(blocks) return blocks def free(self, blocks: list[int]): self.free_blocks.update(blocks)注意这里 block 是逻辑块不是物理地址。真正把逻辑块映射到显存地址发生在 kernel 启动前的bind阶段。这个间接映射层是多芯显存管理的关键没有它前缀复用和跨芯迁移都做不了。3.2 调度器怎么感知“多芯”多芯并行下调度器面对的不再是一个显存池而是多个。请求到了之后要决定放哪张芯、占多少块、要不要跨芯做张量并行。我们的做法是把调度分成两级全局调度和芯内调度。全局调度器维护一个“芯负载表”每张芯记录当前剩余 block 数、正在执行的请求数、算力占用估计值。新请求进来时先按负载选择候选芯集合再在集合内部用 SGLang 原有的调度策略如最短前缀优先决定优先级。这样做的好处是前缀复用逻辑仍然保留在原来的调度器里多芯只是外部加了一层路由。这里有一个很微妙的问题张量并行请求必须同时占用多张芯的显存而且要求这多张芯的空闲 block 数量都够。如果只按单芯负载路由请求可能卡在“第一张芯够、第二张芯不够”的状态。我们的经验是在全局调度表里预先把“芯组”作为调度单元比如 2 芯组、4 芯组显存按组维度统计。宁可牺牲一点灵活性也要保证请求不会半路被显存卡死。3.3 集合通信的坑多芯 tensor parallel 离不开集合通信但硬件厂商的通信库往往没有 NCCL 那么成熟。我们在昆仑芯上调试时遇到过一个很隐蔽的问题AllReduce 默认实现是同步阻塞的但 SGLang 的流水线设计里前向计算和通信本来是可以重叠的。通信库一旦阻塞GPU 流水线就空转吞吐直接掉了一半。解决思路有两步。第一步是在抽象层把通信方法改成“提交即返回”通过event或future做异步等待这样调度端可以继续安排计算任务。第二步是在通信库内部做优化比如把多个小包的 AllReduce 合并成一个大包减少同步点。这一点对 XPU 尤为重要因为小包通信的启动开销在不少硬件上比 CUDA 还大合并批量收益非常明显。经验跑多芯并行第一步先别急着调 kernel先把通信和计算的重叠结构画出来。用 profiler 看两个关键指标通信等待占比、kernel 启动间隙占比。这两个指标优化到位性能基本就稳了。4. 从能跑到跑得快SGLang-Kunlun 的工程化调优4.1 模型并行切分策略模型并行切分Tensor Parallel在多芯场景里是必选项但切多少芯得算清楚。切的芯数越多单芯显存压力越小但通信开销越大切的芯数太少单芯显存可能放不下大模型。我们在 7B、14B、72B 几档模型上都做过对比得到一个粗略经验模型参数量 PGB与单芯可用显存 MGB之间当P / M 2时再考虑切 4 芯以上否则 2 芯性价比最高。切分时要特别留意 Attention 层的切分方式。多头注意力按头切分天然适合多芯并行但 GQAGrouped Query Attention结构下KV head 数量远小于 Q head如果照搬 Q 的切分方式KV cache 的通信会非常浪费。我们最终的做法是只切 Q 和 O 的投影K 和 V 的投影每芯保留完整副本代价是显存多占一点换来的是通信量大幅下降。对 XPU 这种集合通信能力相对保守的平台这买卖很划算。4.2 Chunked Prefill 和 KV Cache 策略SGLang 的 chunked prefill 是性能利器但在多芯插件下需要重新校准 chunk size。chunk size 太小kernel 启动次数变多XPU 的启动开销会让你很难受chunk size 太大长请求会长时间独占芯组短请求被饿死。我们的调法是用固定 batch 压测把 chunk size 从 512 调到 2048画吞吐曲线取峰值对应的值再适当留 20% 余量。KV cache 方面除了前面说的 BlockAllocator还有一个容易被忽略的点cache 的布局要按芯组对齐。比如 2 芯组KV cache 的 block 最好按偶数芯对齐分配避免一个请求的 KV 块在这张芯上占 3 块、在另一张芯上占 4 块的情况。不对齐的后果是每次 Attention 计算都要做一次跨芯的 gather性能损耗肉眼可见。4.3 数据流与重叠工程化调优最后拼的是流水线。SGLang 本身已经做了 prefill 和 decode 的 interleaving但多芯插件接入后我们经常看到现象是某张芯在闲等其他芯忙得要死。排查来排查去根因往往是 Host 端的数据拷贝和预处理没有并行。我的建议是把 tokenizer、detokenizer 和样本预处理放到独立的线程池里不要让它们占住芯组调度的时间片。我们在一台双芯机器上做过对比把预处理线程从 1 加到 4端到端吞吐提升约 18%而显存占用几乎没变。这个优化方向很像流水线工厂里的“瓶颈前移”成本低收益直接。调节参数调大后的效果主要风险我们的推荐区间chunked prefill size计算效率提升长请求占芯组、短请求饥饿512 ~ 2048最大 batch size吞吐提升显存压力、单请求延迟变大按显存测试后留 20% 余量预处理线程数减少 Host 侧瓶颈线程切换开销等于物理核数的一半以内tensor parallel 芯数单芯显存压力降低通信开销上升按 P/M 比值判断5. 实测中的问题排查与反向验证5.1 显存泄漏从一个监控告警说起多芯插件跑稳定之后第一个真正的线上问题是显存泄漏。表现很典型服务运行 3 到 4 小时某一张芯的显存占用缓慢爬升直到 OOM。一开始怀疑是 KV cache 没释放查了一圈发现 block 都在正常归还。后来用 profiler 盯着显存分配调用栈才定位到问题RoPE 算子在 XPU 上有动态 shape 的临时缓存每次请求 shape 不同缓存重新分配旧的没有释放。这类问题在 CUDA 上也可能存在但 XPU 的显存分配路径更敏感一旦中间层缓存没有正确清理内存就直接留在设备侧。我们的修复方式是给插件的临时缓存加上“生命周期标记”每次请求结束时统一回收而不是依赖 operator 内部自己清理。5.2 性能不符合预期的定位链路很多时候多芯并行跑起来性能就是上不去。我们在调 SGLang-Kunlun 时总结了一条定位链路从宏观到微观依次排查看端到端吞吐是不是线性加速如果 2 芯只比 1 芯快 1.2 倍优先怀疑通信锁。看芯组空闲率如果某张芯的空闲率超过 20%优先怀疑调度路由不均衡或显存配额不足。看 kernel 层面的耗时分解用厂商 profiler 找最大的热点函数。最后才去调 kernel 本身的实现和 launch 参数。这条链路帮助我们避开了很多“无头苍蝇式”调优。比如有一次 2 芯加速比只有 1.15我们一开始疯狂调 kernel后来才发现是全局调度器把请求总是路由到同一张芯上另一张芯大部分时间在摸鱼。改一下负载权重就好跟 kernel 一点关系都没有。5.3 回归测试与多芯验证方案插件机制改多了很容易出现“这周修了那周坏”的情况。我们的做法是把三部分测试固化到 CI 里。第一部分是接口契约测试确保每个后端实现的接口都能被标准用例驱动第二部分是数值一致性测试用同一份模型权重在单芯和双芯上跑同一批输入比对 logits 和采样分布第三部分是压力测试跑 24 小时长稳监控显存曲线和延迟分布。数值一致性测试在多芯场景下特别重要因为 tensor parallel 的通信顺序差异可能导致小数点后几位的不一致。我们定的标准是logits 的 cosine 相似度大于 0.9999且生成文本完全相同即可视为通过。这个标准不苛刻但能有效拦截大多数切分和通信实现上的低级错误。写在最后的一点体会这套多芯插件机制从设计到落地我们踩过的坑比预想多得多但收获也很直接SGLang 上游每次更新我们只需要花很少的时间做插件层适配真正把精力放在模型性能和调度策略上。如果你也在折腾类似的异构后端接入我的建议是——第一版插件接口宁可多留扩展点也不要贪图一时省事写死显存块管理和通信抽象一定要一开始就做对因为这两块是后面多芯并行所有优化的地基。最后再分享一个小技巧插件层里加一个--debug-backend开关把所有后端原语调用都打日志线上出问题时对照日志排查比拿 profiler 猜快得多。