最近在调一个大模型推理服务，压测时发现一个很典型的现象：模型参数明明不大，但一上长上下文，吞吐就垮掉。手里的目标模型是MoE架构、体量对标V3/V3.2这一代开源模型，短文本还能靠batch堆起来，一旦上下文长度跑到64K以上，显存先扛不住，接着通信也开始拖后腿。后来我把Ulysses序列并行从训练场景挪到推理链路里，重新设计了Attention的切分和KV cache组织方式，端到端吞吐提升了3倍左右，显存占用反而还降了。这篇文章就把这套优化从原理到踩坑完完整整拆开说，给正在搞长文本推理的人一个可直接参考的落地方案。1. 项目背景与性能瓶颈分析1.1 为什么长序列推理卡在“算不完”而不是“算不快”很多人在优化大模型推理时，第一反应是看GPU利用率高不高。短序列场景下，只要batch足够大，计算就能把硬件喂饱，问题确实是“算不快”。但长序列场景完全是另一回事：每个请求可能一次性吃掉2万、8万甚至12万token的KV cache，显存很快见底，能塞进去的batch数会被强制压到很小。结果就是GPU的算力大量闲置，卡在“没活可干”的状态。以目标模型为例，单卡80GB显存，32K上下文下最多只能装4个并发请求，而且其中超过一半显存都贡献给了KV cache。这时真正的瓶颈不是FLOPS，而是“显存容量决定batch上限，batch上限决定算力利用率”。换句话说，长序列推理通常卡在资源总量不足，而不是单次计算太慢。要打破这个局面，必须把显存和计算负载同时摊到多张卡上，而不是只在单卡上做文章。1.2 常规并行方案到了长上下文就失灵常见的并行手段无非四种：数据并行、张量并行、流水线并行和专家并行。数据并行最简单，但每个GPU都持有完整模型参数和完整KV cache，显存瓶颈完全没解决。张量并行把模型参数切开，能解决单卡放不下的问题，可是每一层的矩阵乘法和Attention都要求各卡之间做同步通信，序列越长，通信量跟着线性涨。流水线并行把模型按层切段，但切得太细会引入大量气泡，切得不够细又解决不了单层显存爆炸的问题。专家并行对MoE模型有效，可它只负责FFN的稀疏路由部分，对Attention和KV cache帮不上忙。这套组合拳打下来，长上下文的Attention部分依然是全局依赖：每个token都要看所有token，序列维度完全无法并行化，导致Attention成为长文本推理里“木桶最短的那块板”。1.3 Ulysses并行的切入角度Ulysses给我的启发是：既然Attention的瓶颈在“序列太长”，那就直接把序列切了。它不切模型权重，也不切层，而是把输入token按位置分成多段，分别放到不同GPU上。每个GPU只对自己的那一段序列做局部运算，到了Attention那里，再通过一次全交换通信把“序列分片”转成“头分片”，让每个GPU拿到完整序列但只负责一部分注意力头。这种思路巧妙在把Attention的二维计算做了分解：序列维度和头维度可以在不同并行方式间切换。计算量被摊开到多卡，同时KV cache也按头分散存储，单卡显存压力大幅下降。关键的是，Ulysses的通信量只跟序列长度和头数相关，不跟隐藏维度强耦合，在长序列场景下反而比张量并行里的全归约更划算。这就是它能带来几倍提升的根本原因。2. Ulysses并行的核心原理拆解2.1 一句话版原理：把“序列”而不是“权重”切开传统张量并行可以理解成“把矩阵按列劈开，一人算一块，最后拼起来”。Ulysses则是“把句子按词切开，一人看一段，但要互相交换看懂全局”。在Attention之前，每个GPU只用自己手里的那一段token去计算Query、Key、Value；如果直接在这个状态上做局部Attention，那每个token只能看到自己这一段，语义就废了。所以Ulysses在Attention之前做一次all-to-all，把所有GPU上同一批头的Q、K、V汇总到同一个GPU，让每个GPU拥有“完整序列、部分头”。这样每个GPU可以独立计算给指定头分配的Attention结果，不需要再依赖其他GPU的序列片段。计算完之后再做一次反向all-to-all，把输出恢复到“序列分片、全量头”的状态，再继续后面的MLP或MoE层。2.2 多头注意力下的all-to-all重排机制具体到实现层面，假设并行度是4，序列切成4段，每段有16个token，模型一共8个注意力头。切分前，4个GPU各持有一份[16, 8, head_dim]的Q/K/V。第一次all-to-all不是简单地把数据合并，而是按照“头”进行重排：GPU0拿到完整序列64个token的第0、1号头，GPU1拿到第2、3号头，依此类推。下面是我在改造时用的核心伪代码结构：# 序列并行Attention的伪代码 def attention_with_ulysses(qkv_local, sp_size, head_dim): # qkv_local: [local_seq, all_heads, head_dim] # 第一步: 序列分片 -> 头分片 q = all_to_all(qkv_local[:, :, 0], dim=0) k = all_to_all(qkv_local[:, :, 1], dim=0) v = all_to_all(qkv_local[:, :, 2], dim=0) # 现在q/k/v都是 [global_seq, local_heads, head_dim] attn_out = flash_attention(q, k, v) # 第二步: 头分片 -> 序列分片 return all_to_all(attn_out, dim=1)这里两个方向的all-to-all缺一不可。第一次是把“序列分片”变换成“头分片”，第二次是把“头分片”变换回“序列分片”。如果你的框架里有现成的all-to-all接口，要注意接口的dim参数到底指的是序列维还是头维，方向搞反了轻则结果错乱，重则直接死锁。2.3 推理场景和训练场景的差异：KV cache怎么跟着切训练场景下Ulysses推广得很好，因为不需要额外考虑推理状态。但推理时KV cache是持续存活的，不能像训练那样算完就丢。我一开始想当然地认为KV cache也应该按序列分片存，但仔细推了一遍发现不行：Ulysses的all-to-all要求每个GPU手里有完整序列的部分头，因此KV cache必须在“头分片”维度上存储，而不是“序列分片”。也就是说，如果模型是8个头，SP=4，那么每个GPU永久保存2个头的完整序列KV。这样做的好处是显存被平均摊到了4张卡，坏处是查询时每个请求都要访问多张卡的KV cache，这对缓存命中率和局部性是个挑战。我在实践里把KV cache按“头”分块放到独立显存池，并给每块分配一个global head index，这样第二次all-to-all时能精确恢复顺序。2.4 为什么收益不是线性而是超线性很多人看到3倍提升，第一反应是“不就是4卡并行，顶多4倍，怎么才3倍？”实际上这个3倍不是并行度直接换来的。收益来自三个效应的乘法叠加。第一，batch容量翻倍。KV cache分散后，单卡显存占用大幅下降，同一个请求时长上下文能装的并发batch数明显增加。第二，Attention的QKV投影和中间激活不再需要保留完整序列，每卡的临时显存减少，计算局部性变好。第三，all-to-all通信可以和Attention计算重叠，通信时间被藏到计算背后。三者叠加出来的效果是：batch提升约1.8倍，单请求attention延迟下降约1.6倍，通信隐藏再带来10%-20%的余量，最后端到端吞吐就奔着3倍去了。3. 实操：在V3/V3.2级MoE模型上落地Ulysses推理优化3.1 环境与基线配置这次实验用的是一台8卡GPU服务器，每卡显存80GB。模型是MoE架构，总参数量百B级，激活参数量几十B，注意力结构是带KV压缩的低秩形式，本质上是GQA的变种。推理框架基于PyTorch自研流程，对外提供标准的预填充和解码接口。基线配置是“TP=8、DP=2、无序列并行”，即所有序列长度都由8张卡分担隐藏维，张量并行度拉满。预填充阶段batch=8，上下文长度设了32K、64K、128K三档。测下来32K还能跑，64K开始显存告急，128K直接OOM。这就是最原始的痛点：长上下文请求没法正常服务。优化后的配置调整为“TP=4、SP=4、DP=2”，设备总数还是8。其中TP管隐藏维度切分，SP管序列维度切分。Attention层走Ulysses路径，MLP和MoE层维持TP和专家并行。这样既保留了MoE的路由收益，又把Attention的最大痛点独立拎出来解决。3.2 Attention模块改造改造的核心是替换原来的多头注意力实现。原实现是标准的TP式Attention：每个GPU计算部分头，但输入序列是全长的，QKV投影也是在全序列上运算。改成Ulysses后，先对输入做序列切分，每卡只处理local_seq = global_seq / SP的token，再走一次投影。这里有个容易被忽略的细节：Rope位置编码。因为每卡的local sequence只是全局序列的一小段，不能直接用本地下标去做Rope，否则位置信息就乱了。我维护了一张“全局token索引表”，每个local token绑定它在原始序列里的绝对位置，Rope也按绝对位置生成。否则Attention分数计算出来是乱的，位置一错位，长距离依赖全丢。另外，GQA的KV头数往往比Query头少。以目标模型为例，Query头是64，KV头只有8。SP=4时，每个GPU分配2个KV头，这个还够用。如果KV头数少于SP，就需要做“KV头复制”或改用MQA，这个后面在常见问题里细说。3.3 MoE expert并行的协同目标模型是MoE，每个Expert本身不大，但数量多，如果全部复制到每张卡上，显存直接爆炸。所以MoE层的专家并行和Attention的Ulysses并行必须协同工作。我的做法是：Attention层输出先恢复到“序列分片，全量头”状态，然后进入MoE层时，每个token独立走路由。这里SP切分和expert切分是正交的：SP把序列分成4段，expert把FFN权重分成若干份。Token路由到哪个expert，就把token发送到对应expert所在的设备。两次通信可以合并到同一个通信域里，减少Kernel启动次数。最麻烦的是负载均衡。MoE专家的热门程度往往不均，某个expert可能收到其他设备转发来的大量token，造成该设备过载。我用了两层Buffer：第一层按序列分片，第二层按expert分片，每次通信前先做一次全局token计数，动态调整发送阈值。这个细节对端到端性能影响很大，做完之后MoE部分的抖动明显下降。3.4 通信与计算重叠Ulysses能不能发挥出效果，很大程度上取决于all-to-all通信有没有被藏起来。我的做法是把Attention拆成“Q/K/V投影”和“Attention计算”两个阶段，投影阶段发起的all-to-all使用异步通道，计算阶段只等真正需要的那部分数据到位。更具体的做法是引入“双缓冲”机制：对连续的两个序列分片，先预取下一个分片的Q/K/V，当前分片在算Attention时，下一个分片的通信已经在后台进行。这样通信延迟不会直接加在关键路径上。实测下来，纯all-to-all裸跑要占整个Attention层时间的35%，但叠加双缓冲后，真实暴露的通信开销只有12%左右。代码层面我封装了一个带事件同步的all-to-all管理器：# 伪代码：双缓冲重叠 def overlapped_ulysses(query_list, key_list, value_list): # 第一个分片先发起通信 q_comm = async_all_to_all(query_list[0]) k_comm = async_all_to_all(key_list[0]) v_comm = async_all_to_all(value_list[0]) for i in range(1, len(query_list)): wait(q_comm, k_comm, v_comm) attn = flash_attention(q_comm.result(), k_comm.result(), v_comm.result()) # 提前发起下一个分片的通信 q_comm = async_all_to_all(query_list[i]) k_comm = async_all_to_all(key_list[i]) v_comm = async_all_to_all(value_list[i]) wait_and_handle_last(attn)注意异步通信的结果和后续算子要在同一个CUDA stream上做事件同步，不能用Python层面的sleep。很多同学在这里踩坑：异步接口返回时数据根本没到位，直接计算就数值错乱；而强行同步又让本来的重叠效果归零。3.5 验证正确性的几个招序列并行最怕“看起来能跑，实际结果错”。我总结了三个必备校验步骤。第一步，数值对齐测试。固定所有随机种子，同一组输入分别跑“单卡普通Attention”和“多卡Ulysses Attention”，比较Attention输出的logits，绝对误差阈值设1e-3。第二步，位置平移测试。构造一个只有前一半token有意义的句子，让模型预测最后一个token，检查不同SP配置下输出是否一致。这个测试能暴露出Rope和位置索引的错误。第三步，状态持久化测试。多次调用同一份KV cache，改掉batch顺序，确认输出不受历史batch残留影响。因为KV cache按头分片后，如果跨batch复用没做清理，推理结果会越跑越脏。这三步全过了再走性能压测，能省掉后面一半的排查时间。4. 实测数据与调优记录4.1 测试矩阵为了把收益看透，我做了几组对比测试，每个测试固定输入长度，变化并发batch数，统计“端到端吞吐”和“单卡显存峰值”。端到端吞吐的定义是：预填充阶段每秒处理的输入token数加上解码阶段每秒生成的token数，按整段请求归一化。测试长度选了32K、64K、128K三档。32K代表中等场景，64K是长文本临界点，128K是压力测试。每档各跑50次请求取中位数，排除冷启动和波动影响。并行配置就是前面说的两种：基线TP=8 vs 优化后TP=4+SP=4。4.2 实测结果表格上下文长度配置最大batch吞吐(tokens/s)单卡显存峰值相对加速32K基线 TP=88278068GB1.0x32K优化 TP=4+SP=414712052GB2.56x64K基线 TP=8498078GB1.0x64K优化 TP=4+SP=48296061GB3.02x128K基线 TP=8124082GB1.0x128K优化 TP=4+SP=4382068GB3.42x可以看到，序列越长，收益越明显。32K时因为通信占比相对高，加速只有2.5倍；到128K时单卡显存省下14GB，batch从1提升到3，Attention计算又因分片更充分，整体加速达到3.4倍。标题说3倍，不是挑好听的说，是抓了一个保守的中间值。4.3 关键参数怎么选SP大小是第一个参数。不是越大越好，因为SP越大，每卡要交换的QKV片段就越多。在8卡环境里，TP=4+SP=4是最稳的组合；TP=2+SP=8虽然Attention头分得更散，但all-to-all通信矩阵变大，NVLink带宽接近饱和，端到端反而降到2.4倍。所以SP的上限取决于模型头部数量和通信带宽。micro-batch是第二个参数。Ulysses的all-to-all本质上是多个请求一起做通信效率才高。我把连续8个请求打包成一组，先做QKV投影，再一起进入attention的all-to-all，比逐请求通信吞吐高了30%。代价是首延迟稍微变差，但整体吞吐优先的场景完全值得。通信后端的选择也影响显著。如果设备间是NVLink全互联，all-to-all可以直接映射到硬件全交换；如果是普通以太网，效果就大打折扣。我们的测试机上NVLink全互联是标配，所以all-to-all的带宽发挥到了接近峰值。4.4 3x到底从哪里来把128K场景拆一下账。基线上batch只能放1个请求，单请求预填充大约240 tokens/s，GPU利用率不到一半。优化后batch放到3，单请求的Attention计算时间还因为TP从8降到4而略微下降？这里我需要解释清楚：我们保留TP=4负责隐藏维切分，SP=4负责序列切分，Attention变成“每个设备只计算一部分序列状态下的一部分头”，实际单设备处理的有效token数大幅下降，所以预填充的墙钟时间缩短，同时batch也从1涨到3。吞吐 = 并发数 × 单请求速率，两个因子叠起来就是3倍以上。另外还有一部分收益来自显存。128K场景优化前KV cache每个请求接近70GB，单卡根本放不下完整一套，只能靠TP把KV分散到8卡，但调度开销巨大。优化后KV cache按头分到SP=4的卡上，每卡只存四分之一，显存一旦有富余，就能多塞并发请求。显存是长文本推理的硬上限，这个账算清楚了，优化方向就自然清晰了。5. 常见问题与排查技巧实录5.1 all-to-all通信风暴怎么缓解症状：优化后吞吐不但没升，反而比基线还低。查一下各卡网络带宽利用率，如果长时间100%，说明all-to-all成了新瓶颈。原因通常是通信粒度过细，频繁小包通信导致握手开销远大于数据传输。解法有两个。第一个是把多次all-to-all合并成一次超大通信，比如把Q、K、V三个张量拼接成一个维度再去交换，减少Kernel启动数量。第二个是分片限流，不要所有卡同时发起全交换，而是按卡号错峰启动；等前一组头的通信完成一半，再启动下一组。我在实际项目里给all-to-all加了“令牌桶”调度器，峰值带宽占用下降了40%，通信总时间反而缩短了。还有一个很容易忽略的点：all-to-all用的是全局通信域，不要让不相关的集合通信插进来抢占带宽。MoE的token分发如果和Attention的all-to-all共用同一张网卡队列，会互相踩踏。我在框架里把两类通信放到不同优先级流上，问题立刻缓解。5.2 头数不够分怎么办如果你遇到模型KV头只有2个，但SP想设4，直接分是分不出来的。Ulysses要求序列切分后每个GPU至少分到1个完整的KV头，否则一个头要被多个GPU重复持有，就失去切分意义了。常见做法是“先复制再切”：逻辑上把KV头从2个复制到4个，再按4路切。这样每个GPU持有一个复制的KV头，但复制的KV头是同一个物理头的拷贝，显存反而多了一倍，压缩了收益。更好的方案是改用MQA或降级为TP+SP混合：SP只做2，TP补够并行度。或者干脆放弃头分片，改用Ring Attention那种按序列块交换KV的方案。总之头数是硬约束，设计并行方案前先数清模型有几组KV头。5.3 mask和位置编码错位我踩过最深的一个坑是attention_mask。序列并行下，每个GPU的局部序列没有完整的因果依赖信息。比如GPU0处理第0-15个token，GPU1处理第16-31个token，如果GPU0要看到第20个token，本地mask根本不知道有这个位置。解决办法有两个。一个是采用全局mask：在第一次all-to-all之后，每个GPU虽然只有部分头，但它拥有全部序列的Q/K/V，此时mask必然是全量的[full_seq, full_seq]，按设备号裁剪到对应头即可。另一个是分块mask拼接：在局部Attention前，把其他设备的K/V长度补进mask，构造[local_q, global_kv]的矩阵。第二种省显存，但写起来容易错位。我更推荐第一种，正确性优先。Rope的坑和mask类似。每个设备的local token不能直接用0、1、2这样的本地序号，而是要用它在完整序列里的绝对位置。错误的现象是：短文本没问题，长文本精度崩，而且崩溃点集中在上下文超过SP段长度的地方。5.4 和投机解码/量化组合的坑在优化完Attention后，我又尝试叠加量化和投机解码，结果踩了个新坑。量化后的KV cache和Ulysses的all-to-all不太兼容，因为all-to-all要求交换的张量具备连续内存布局；而量化常用分组位宽，不同头之间可能有不连续对齐。我调试半天发现，先反量化再用FP16做all-to-all，比直接发量化张量还要快，因为省掉了目标端的解包对齐操作。投机解码也会受影响。投机模型抽出的草稿token和真实token长度不一致，导致序列并行下各设备的local sequence长度不齐。Ulysses要求各设备序列分片等长才能高效通信。我最后做了一个“padding对齐”方案：给短的分片补上特殊mask token，计算时不影响loss，但在通信前统一长度。注意这里是推理，不涉及反向传播，所以padding不会有梯度污染。5.5 一个“照抄作业”的自检清单如果你正准备在自己模型上复制这套优化，建议按这个顺序自查：数清Query头数和KV头数，确认SP最大值。确认Rope的绝对位置索引是否在切分后维持正确。确认all-to-all两个方向的dim参数没有搞反。确认Attention mask是全局长度的，而不是local长度。确认KV cache按头分片存储，而非按序列分片。确认MoE层的expert分发与序列并行没有共用通信队列。确认异步通信有事件同步，没有数据竞争。先过数值对齐，再做batch=1的稳定性测试，最后才调batch冲刺性能。按这个表逐项对照，至少能避开我花了一周才趟完的坑。最后说一点个人体会。Ulysses这套思路不是银弹，它对超长上下文、需要多卡并行、且Attention头数足够多的场景非常有效；但如果序列长度只有8K以内，或者模型本身是单头注意力，耗费那么多通信figure来切分显然不划算。我这次优化做完最大的收获不是那3倍性能，而是重新理解了“并行”的真正含义：并行的对象不一定非是权重和层，只要能把计算依赖解开，序列本身也是一种极好的并行维度。
最近在调一个大模型推理服务，压测时发现一个很典型的现象：模型参数明明不大，但一上长上下文，吞吐就垮掉。手里的目标模型是MoE架构、体量对标V3/V3.2这一代开源模型，短文本还能靠batch堆起来，一旦上下文长度跑到64K以上，显存先扛不住，接着通信也开始拖后腿。后来我把Ulysses序列并行从训练场景挪到推理链路里，重新设计了Attention的切分和KV cache组织方式，端到端吞吐提升了3倍左右，显存占用反而还降了。这篇文章就把这套优化从原理到踩坑完完整整拆开说，给正在搞长文本推理的人一个可直接参考的落地方案。1. 项目背景与性能瓶颈分析1.1 为什么长序列推理卡在“算不完”而不是“算不快”很多人在优化大模型推理时，第一反应是看GPU利用率高不高。短序列场景下，只要batch足够大，计算就能把硬件喂饱，问题确实是“算不快”。但长序列场景完全是另一回事：每个请求可能一次性吃掉2万、8万甚至12万token的KV cache，显存很快见底，能塞进去的batch数会被强制压到很小。结果就是GPU的算力大量闲置，卡在“没活可干”的状态。以目标模型为例，单卡80GB显存，32K上下文下最多只能装4个并发请求，而且其中超过一半显存都贡献给了KV cache。这时真正的瓶颈不是FLOPS，而是“显存容量决定batch上限，batch上限决定算力利用率”。换句话说，长序列推理通常卡在资源总量不足，而不是单次计算太慢。要打破这个局面，必须把显存和计算负载同时摊到多张卡上，而不是只在单卡上做文章。1.2 常规并行方案到了长上下文就失灵常见的并行手段无非四种：数据并行、张量并行、流水线并行和专家并行。数据并行最简单，但每个GPU都持有完整模型参数和完整KV cache，显存瓶颈完全没解决。张量并行把模型参数切开，能解决单卡放不下的问题，可是每一层的矩阵乘法和Attention都要求各卡之间做同步通信，序列越长，通信量跟着线性涨。流水线并行把模型按层切段，但切得太细会引入大量气泡，切得不够细又解决不了单层显存爆炸的问题。专家并行对MoE模型有效，可它只负责FFN的稀疏路由部分，对Attention和KV cache帮不上忙。这套组合拳打下来，长上下文的Attention部分依然是全局依赖：每个token都要看所有token，序列维度完全无法并行化，导致Attention成为长文本推理里“木桶最短的那块板”。1.3 Ulysses并行的切入角度Ulysses给我的启发是：既然Attention的瓶颈在“序列太长”，那就直接把序列切了。它不切模型权重，也不切层，而是把输入token按位置分成多段，分别放到不同GPU上。每个GPU只对自己的那一段序列做局部运算，到了Attention那里，再通过一次全交换通信把“序列分片”转成“头分片”，让每个GPU拿到完整序列但只负责一部分注意力头。这种思路巧妙在把Attention的二维计算做了分解：序列维度和头维度可以在不同并行方式间切换。计算量被摊开到多卡，同时KV cache也按头分散存储，单卡显存压力大幅下降。关键的是，Ulysses的通信量只跟序列长度和头数相关，不跟隐藏维度强耦合，在长序列场景下反而比张量并行里的全归约更划算。这就是它能带来几倍提升的根本原因。2. Ulysses并行的核心原理拆解2.1 一句话版原理：把“序列”而不是“权重”切开传统张量并行可以理解成“把矩阵按列劈开，一人算一块，最后拼起来”。Ulysses则是“把句子按词切开，一人看一段，但要互相交换看懂全局”。在Attention之前，每个GPU只用自己手里的那一段token去计算Query、Key、Value；如果直接在这个状态上做局部Attention，那每个token只能看到自己这一段，语义就废了。所以Ulysses在Attention之前做一次all-to-all，把所有GPU上同一批头的Q、K、V汇总到同一个GPU，让每个GPU拥有“完整序列、部分头”。这样每个GPU可以独立计算给指定头分配的Attention结果，不需要再依赖其他GPU的序列片段。计算完之后再做一次反向all-to-all，把输出恢复到“序列分片、全量头”的状态，再继续后面的MLP或MoE层。2.2 多头注意力下的all-to-all重排机制具体到实现层面，假设并行度是4，序列切成4段，每段有16个token，模型一共8个注意力头。切分前，4个GPU各持有一份[16, 8, head_dim]的Q/K/V。第一次all-to-all不是简单地把数据合并，而是按照“头”进行重排：GPU0拿到完整序列64个token的第0、1号头，GPU1拿到第2、3号头，依此类推。下面是我在改造时用的核心伪代码结构：# 序列并行Attention的伪代码 def attention_with_ulysses(qkv_local, sp_size, head_dim): # qkv_local: [local_seq, all_heads, head_dim] # 第一步: 序列分片 -> 头分片 q = all_to_all(qkv_local[:, :, 0], dim=0) k = all_to_all(qkv_local[:, :, 1], dim=0) v = all_to_all(qkv_local[:, :, 2], dim=0) # 现在q/k/v都是 [global_seq, local_heads, head_dim] attn_out = flash_attention(q, k, v) # 第二步: 头分片 -> 序列分片 return all_to_all(attn_out, dim=1)这里两个方向的all-to-all缺一不可。第一次是把“序列分片”变换成“头分片”，第二次是把“头分片”变换回“序列分片”。如果你的框架里有现成的all-to-all接口，要注意接口的dim参数到底指的是序列维还是头维，方向搞反了轻则结果错乱，重则直接死锁。2.3 推理场景和训练场景的差异：KV cache怎么跟着切训练场景下Ulysses推广得很好，因为不需要额外考虑推理状态。但推理时KV cache是持续存活的，不能像训练那样算完就丢。我一开始想当然地认为KV cache也应该按序列分片存，但仔细推了一遍发现不行：Ulysses的all-to-all要求每个GPU手里有完整序列的部分头，因此KV cache必须在“头分片”维度上存储，而不是“序列分片”。也就是说，如果模型是8个头，SP=4，那么每个GPU永久保存2个头的完整序列KV。这样做的好处是显存被平均摊到了4张卡，坏处是查询时每个请求都要访问多张卡的KV cache，这对缓存命中率和局部性是个挑战。我在实践里把KV cache按“头”分块放到独立显存池，并给每块分配一个global head index，这样第二次all-to-all时能精确恢复顺序。2.4 为什么收益不是线性而是超线性很多人看到3倍提升，第一反应是“不就是4卡并行，顶多4倍，怎么才3倍？”实际上这个3倍不是并行度直接换来的。收益来自三个效应的乘法叠加。第一，batch容量翻倍。KV cache分散后，单卡显存占用大幅下降，同一个请求时长上下文能装的并发batch数明显增加。第二，Attention的QKV投影和中间激活不再需要保留完整序列，每卡的临时显存减少，计算局部性变好。第三，all-to-all通信可以和Attention计算重叠，通信时间被藏到计算背后。三者叠加出来的效果是：batch提升约1.8倍，单请求attention延迟下降约1.6倍，通信隐藏再带来10%-20%的余量，最后端到端吞吐就奔着3倍去了。3. 实操：在V3/V3.2级MoE模型上落地Ulysses推理优化3.1 环境与基线配置这次实验用的是一台8卡GPU服务器，每卡显存80GB。模型是MoE架构，总参数量百B级，激活参数量几十B，注意力结构是带KV压缩的低秩形式，本质上是GQA的变种。推理框架基于PyTorch自研流程，对外提供标准的预填充和解码接口。基线配置是“TP=8、DP=2、无序列并行”，即所有序列长度都由8张卡分担隐藏维，张量并行度拉满。预填充阶段batch=8，上下文长度设了32K、64K、128K三档。测下来32K还能跑，64K开始显存告急，128K直接OOM。这就是最原始的痛点：长上下文请求没法正常服务。优化后的配置调整为“TP=4、SP=4、DP=2”，设备总数还是8。其中TP管隐藏维度切分，SP管序列维度切分。Attention层走Ulysses路径，MLP和MoE层维持TP和专家并行。这样既保留了MoE的路由收益，又把Attention的最大痛点独立拎出来解决。3.2 Attention模块改造改造的核心是替换原来的多头注意力实现。原实现是标准的TP式Attention：每个GPU计算部分头，但输入序列是全长的，QKV投影也是在全序列上运算。改成Ulysses后，先对输入做序列切分，每卡只处理local_seq = global_seq / SP的token，再走一次投影。这里有个容易被忽略的细节：Rope位置编码。因为每卡的local sequence只是全局序列的一小段，不能直接用本地下标去做Rope，否则位置信息就乱了。我维护了一张“全局token索引表”，每个local token绑定它在原始序列里的绝对位置，Rope也按绝对位置生成。否则Attention分数计算出来是乱的，位置一错位，长距离依赖全丢。另外，GQA的KV头数往往比Query头少。以目标模型为例，Query头是64，KV头只有8。SP=4时，每个GPU分配2个KV头，这个还够用。如果KV头数少于SP，就需要做“KV头复制”或改用MQA，这个后面在常见问题里细说。3.3 MoE expert并行的协同目标模型是MoE，每个Expert本身不大，但数量多，如果全部复制到每张卡上，显存直接爆炸。所以MoE层的专家并行和Attention的Ulysses并行必须协同工作。我的做法是：Attention层输出先恢复到“序列分片，全量头”状态，然后进入MoE层时，每个token独立走路由。这里SP切分和expert切分是正交的：SP把序列分成4段，expert把FFN权重分成若干份。Token路由到哪个expert，就把token发送到对应expert所在的设备。两次通信可以合并到同一个通信域里，减少Kernel启动次数。最麻烦的是负载均衡。MoE专家的热门程度往往不均，某个expert可能收到其他设备转发来的大量token，造成该设备过载。我用了两层Buffer：第一层按序列分片，第二层按expert分片，每次通信前先做一次全局token计数，动态调整发送阈值。这个细节对端到端性能影响很大，做完之后MoE部分的抖动明显下降。3.4 通信与计算重叠Ulysses能不能发挥出效果，很大程度上取决于all-to-all通信有没有被藏起来。我的做法是把Attention拆成“Q/K/V投影”和“Attention计算”两个阶段，投影阶段发起的all-to-all使用异步通道，计算阶段只等真正需要的那部分数据到位。更具体的做法是引入“双缓冲”机制：对连续的两个序列分片，先预取下一个分片的Q/K/V，当前分片在算Attention时，下一个分片的通信已经在后台进行。这样通信延迟不会直接加在关键路径上。实测下来，纯all-to-all裸跑要占整个Attention层时间的35%，但叠加双缓冲后，真实暴露的通信开销只有12%左右。代码层面我封装了一个带事件同步的all-to-all管理器：# 伪代码：双缓冲重叠 def overlapped_ulysses(query_list, key_list, value_list): # 第一个分片先发起通信 q_comm = async_all_to_all(query_list[0]) k_comm = async_all_to_all(key_list[0]) v_comm = async_all_to_all(value_list[0]) for i in range(1, len(query_list)): wait(q_comm, k_comm, v_comm) attn = flash_attention(q_comm.result(), k_comm.result(), v_comm.result()) # 提前发起下一个分片的通信 q_comm = async_all_to_all(query_list[i]) k_comm = async_all_to_all(key_list[i]) v_comm = async_all_to_all(value_list[i]) wait_and_handle_last(attn)注意异步通信的结果和后续算子要在同一个CUDA stream上做事件同步，不能用Python层面的sleep。很多同学在这里踩坑：异步接口返回时数据根本没到位，直接计算就数值错乱；而强行同步又让本来的重叠效果归零。3.5 验证正确性的几个招序列并行最怕“看起来能跑，实际结果错”。我总结了三个必备校验步骤。第一步，数值对齐测试。固定所有随机种子，同一组输入分别跑“单卡普通Attention”和“多卡Ulysses Attention”，比较Attention输出的logits，绝对误差阈值设1e-3。第二步，位置平移测试。构造一个只有前一半token有意义的句子，让模型预测最后一个token，检查不同SP配置下输出是否一致。这个测试能暴露出Rope和位置索引的错误。第三步，状态持久化测试。多次调用同一份KV cache，改掉batch顺序，确认输出不受历史batch残留影响。因为KV cache按头分片后，如果跨batch复用没做清理，推理结果会越跑越脏。这三步全过了再走性能压测，能省掉后面一半的排查时间。4. 实测数据与调优记录4.1 测试矩阵为了把收益看透，我做了几组对比测试，每个测试固定输入长度，变化并发batch数，统计“端到端吞吐”和“单卡显存峰值”。端到端吞吐的定义是：预填充阶段每秒处理的输入token数加上解码阶段每秒生成的token数，按整段请求归一化。测试长度选了32K、64K、128K三档。32K代表中等场景，64K是长文本临界点，128K是压力测试。每档各跑50次请求取中位数，排除冷启动和波动影响。并行配置就是前面说的两种：基线TP=8 vs 优化后TP=4+SP=4。4.2 实测结果表格上下文长度配置最大batch吞吐(tokens/s)单卡显存峰值相对加速32K基线 TP=88278068GB1.0x32K优化 TP=4+SP=414712052GB2.56x64K基线 TP=8498078GB1.0x64K优化 TP=4+SP=48296061GB3.02x128K基线 TP=8124082GB1.0x128K优化 TP=4+SP=4382068GB3.42x可以看到，序列越长，收益越明显。32K时因为通信占比相对高，加速只有2.5倍；到128K时单卡显存省下14GB，batch从1提升到3，Attention计算又因分片更充分，整体加速达到3.4倍。标题说3倍，不是挑好听的说，是抓了一个保守的中间值。4.3 关键参数怎么选SP大小是第一个参数。不是越大越好，因为SP越大，每卡要交换的QKV片段就越多。在8卡环境里，TP=4+SP=4是最稳的组合；TP=2+SP=8虽然Attention头分得更散，但all-to-all通信矩阵变大，NVLink带宽接近饱和，端到端反而降到2.4倍。所以SP的上限取决于模型头部数量和通信带宽。micro-batch是第二个参数。Ulysses的all-to-all本质上是多个请求一起做通信效率才高。我把连续8个请求打包成一组，先做QKV投影，再一起进入attention的all-to-all，比逐请求通信吞吐高了30%。代价是首延迟稍微变差，但整体吞吐优先的场景完全值得。通信后端的选择也影响显著。如果设备间是NVLink全互联，all-to-all可以直接映射到硬件全交换；如果是普通以太网，效果就大打折扣。我们的测试机上NVLink全互联是标配，所以all-to-all的带宽发挥到了接近峰值。4.4 3x到底从哪里来把128K场景拆一下账。基线上batch只能放1个请求，单请求预填充大约240 tokens/s，GPU利用率不到一半。优化后batch放到3，单请求的Attention计算时间还因为TP从8降到4而略微下降？这里我需要解释清楚：我们保留TP=4负责隐藏维切分，SP=4负责序列切分，Attention变成“每个设备只计算一部分序列状态下的一部分头”，实际单设备处理的有效token数大幅下降，所以预填充的墙钟时间缩短，同时batch也从1涨到3。吞吐 = 并发数 × 单请求速率，两个因子叠起来就是3倍以上。另外还有一部分收益来自显存。128K场景优化前KV cache每个请求接近70GB，单卡根本放不下完整一套，只能靠TP把KV分散到8卡，但调度开销巨大。优化后KV cache按头分到SP=4的卡上，每卡只存四分之一，显存一旦有富余，就能多塞并发请求。显存是长文本推理的硬上限，这个账算清楚了，优化方向就自然清晰了。5. 常见问题与排查技巧实录5.1 all-to-all通信风暴怎么缓解症状：优化后吞吐不但没升，反而比基线还低。查一下各卡网络带宽利用率，如果长时间100%，说明all-to-all成了新瓶颈。原因通常是通信粒度过细，频繁小包通信导致握手开销远大于数据传输。解法有两个。第一个是把多次all-to-all合并成一次超大通信，比如把Q、K、V三个张量拼接成一个维度再去交换，减少Kernel启动数量。第二个是分片限流，不要所有卡同时发起全交换，而是按卡号错峰启动；等前一组头的通信完成一半，再启动下一组。我在实际项目里给all-to-all加了“令牌桶”调度器，峰值带宽占用下降了40%，通信总时间反而缩短了。还有一个很容易忽略的点：all-to-all用的是全局通信域，不要让不相关的集合通信插进来抢占带宽。MoE的token分发如果和Attention的all-to-all共用同一张网卡队列，会互相踩踏。我在框架里把两类通信放到不同优先级流上，问题立刻缓解。5.2 头数不够分怎么办如果你遇到模型KV头只有2个，但SP想设4，直接分是分不出来的。Ulysses要求序列切分后每个GPU至少分到1个完整的KV头，否则一个头要被多个GPU重复持有，就失去切分意义了。常见做法是“先复制再切”：逻辑上把KV头从2个复制到4个，再按4路切。这样每个GPU持有一个复制的KV头，但复制的KV头是同一个物理头的拷贝，显存反而多了一倍，压缩了收益。更好的方案是改用MQA或降级为TP+SP混合：SP只做2，TP补够并行度。或者干脆放弃头分片，改用Ring Attention那种按序列块交换KV的方案。总之头数是硬约束，设计并行方案前先数清模型有几组KV头。5.3 mask和位置编码错位我踩过最深的一个坑是attention_mask。序列并行下，每个GPU的局部序列没有完整的因果依赖信息。比如GPU0处理第0-15个token，GPU1处理第16-31个token，如果GPU0要看到第20个token，本地mask根本不知道有这个位置。解决办法有两个。一个是采用全局mask：在第一次all-to-all之后，每个GPU虽然只有部分头，但它拥有全部序列的Q/K/V，此时mask必然是全量的[full_seq, full_seq]，按设备号裁剪到对应头即可。另一个是分块mask拼接：在局部Attention前，把其他设备的K/V长度补进mask，构造[local_q, global_kv]的矩阵。第二种省显存，但写起来容易错位。我更推荐第一种，正确性优先。Rope的坑和mask类似。每个设备的local token不能直接用0、1、2这样的本地序号，而是要用它在完整序列里的绝对位置。错误的现象是：短文本没问题，长文本精度崩，而且崩溃点集中在上下文超过SP段长度的地方。5.4 和投机解码/量化组合的坑在优化完Attention后，我又尝试叠加量化和投机解码，结果踩了个新坑。量化后的KV cache和Ulysses的all-to-all不太兼容，因为all-to-all要求交换的张量具备连续内存布局；而量化常用分组位宽，不同头之间可能有不连续对齐。我调试半天发现，先反量化再用FP16做all-to-all，比直接发量化张量还要快，因为省掉了目标端的解包对齐操作。投机解码也会受影响。投机模型抽出的草稿token和真实token长度不一致，导致序列并行下各设备的local sequence长度不齐。Ulysses要求各设备序列分片等长才能高效通信。我最后做了一个“padding对齐”方案：给短的分片补上特殊mask token，计算时不影响loss，但在通信前统一长度。注意这里是推理，不涉及反向传播，所以padding不会有梯度污染。5.5 一个“照抄作业”的自检清单如果你正准备在自己模型上复制这套优化，建议按这个顺序自查：数清Query头数和KV头数，确认SP最大值。确认Rope的绝对位置索引是否在切分后维持正确。确认all-to-all两个方向的dim参数没有搞反。确认Attention mask是全局长度的，而不是local长度。确认KV cache按头分片存储，而非按序列分片。确认MoE层的expert分发与序列并行没有共用通信队列。确认异步通信有事件同步，没有数据竞争。先过数值对齐，再做batch=1的稳定性测试，最后才调batch冲刺性能。按这个表逐项对照，至少能避开我花了一周才趟完的坑。最后说一点个人体会。Ulysses这套思路不是银弹，它对超长上下文、需要多卡并行、且Attention头数足够多的场景非常有效；但如果序列长度只有8K以内，或者模型本身是单头注意力，耗费那么多通信figure来切分显然不划算。我这次优化做完最大的收获不是那3倍性能，而是重新理解了“并行”的真正含义：并行的对象不一定非是权重和层，只要能把计算依赖解开，序列本身也是一种极好的并行维度。