如果你最近在折腾大模型推理服务化，SGLang 这个名字应该不陌生。它靠着 RadixAttention 前缀复用、连续批处理和 chunked prefill 这套组合拳，在同类框架里把吞吐做得很能打。但我们团队接到的任务更拧巴：不只是在主流加速卡上跑通，还要把 SGLang 迁移到昆仑芯（XPU）平台上，并且要求支持多芯并行。绕不开的第一道坎，就是多芯插件机制——一层让上层框架不关心底层是 CUDA、ROCm 还是 XPU 的隔离层。这篇文章不写 PPT 式的总结，只讲我们实际落地 SGLang-Kunlun 时怎么切接口、怎么管显存、怎么调吞吐，以及最后沉淀下来的工程化最佳实践。
1. 我们为什么要折腾一个“多芯插件机制”
1.1 从一处 if 判断说起：后端硬编码带来的连锁问题
绝大多数推理框架的第一版，都是从单卡单硬件跑通开始的。SGLang 也不例外，早期路径上 CUDA 的烙印相当明显：显存分配直接走torch.cuda，kernel 选择写成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 后端用 NCCL，XPU 后端通常对接厂商自己的通信库，接口必须设计成 “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。做法是把每张芯的显存切成等长 block，block 大小一般取单个 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 几档模型上都做过对比，得到一个粗略经验：模型参数量 P（GB）与单芯可用显存 M（GB）之间，当P / M > 2时再考虑切 4 芯以上，否则 2 芯性价比最高。
切分时要特别留意 Attention 层的切分方式。多头注意力按头切分天然适合多芯并行，但 GQA（Grouped 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 猜快得多。
