WaferLLM Co-Design

导言

WaferLLM 是一篇很适合学习软硬协同设计的系统论文。它没有把晶圆级芯片当作“更大的 GPU”，也没有从某个孤立 Kernel 开始优化；论文先用 PLMR 模型提炼硬件的不可回避约束，再依次重做 LLM 的并行策略、数据布局、KV Cache、GEMM/GEMV 和推理运行时。

本文是系列第一篇，只回答两个问题：WaferLLM 的总体设计链是什么，以及这条链上有哪些值得单独深挖的方向。MeshGEMM 的对齐、循环移位、Interleave、缓冲区生命周期与完整伪代码会放在后续专文中。

不是更大的 GPU

LLM 推理分成 Prefill 和 Decode 两个阶段。Prefill 一次处理多 Token，主要由矩阵乘矩阵的 GEMM 主导；Decode 每次生成一个 Token，许多线性层退化为矩阵乘向量的 GEMV，反复读取模型权重，因此更容易受内存带宽限制。Wafer-scale 设备试图用数十 GB 的分布式片上 SRAM 和 PB/s 级聚合片上带宽改变这个瓶颈。[^paper]

但晶圆面积变大以后，硬件不能继续维持“所有计算核心低成本访问一块统一内存”的抽象。论文使用的 Cerebras WSE-2 有 850,000 个核心，每核只有 48 KB SRAM，总计约 40 GB；核心通过二维 Mesh Network-on-Chip（NoC，片上网络）连接。存储容量很大，但它被切成了几十万块；片上带宽很高，但远距离访问必须穿过很多跳。[^paper-setup]

这使程序的基本问题发生改变：权重放在哪个 Processing Element（PE，处理单元）、一个中间张量跨几跳移动、每核要保留多少通信缓冲区、路由器能同时识别多少路径，都不再是编译器可以完全隐藏的实现细节，而会直接决定算法能否扩展。

作者在知乎说明中把这个变化概括为“分布式与片上的融合”：过去在集群上研究的计算、通信和访存调度，被重新压回一块晶圆之内。这个视角比“WSE 的带宽比 GPU 高多少”更重要，因为它解释了为什么原有软件不能自动继承硬件红利。

{ width=94% }

软硬协同不是把现成软件铺在新硬件上，而是用 P、L、M、R 四类硬件约束重新裁剪软件。图为 Ian 小黑风格自绘认知示意。

约束建模：PLMR

WaferLLM 最关键的设计动作，是先提出一个足够小、但能驱动后续决策的 PLMR 设备模型。PLMR 不是性能预测器，也不是所有加速器都必须满足的定律；它是针对大规模 Mesh、分布式本地内存设备的一组设计约束。[^paper-plmr]

约束	硬件事实	对软件的直接要求	常见失配
P：Massive Parallelism	数十万到百万级细粒度核心	同时切分多个张量维度，让足够多 PE 有工作	只沿一个模型维度切分，平行任务数远小于核心数
L：Non-uniform Latency	Mesh 跳数使远端访问显著慢于本地访问	缩短关键通信路径，避免全局转置和远距离随机访问	把片上分布式内存当成延迟近似均匀的共享内存
M：Constrained Local Memory	每核只有几十 KB 到数 MB 本地内存	Tile 化权重、激活与工作区，控制同时存活对象	AllGather 或复制产生超过本地容量的工作集
R：Constrained Routing	每核路由表项和硬件路径有限	使用可静态规划、数量受控的通信邻居和树结构	为大量远端目标建立独立路径，迫使软件转发

这个模型的价值在于把“硬件特性”变成了可否决软件方案的检查项。例如，一个算法即使计算量正确，只要需要每核维护随 Mesh 宽度增长的路由路径，就会违反 R；一个方案即使本地缓冲区足够小，只要每一步都跨越整条 Mesh，也会违反 L。

更重要的是，P、L、M、R 彼此牵制。复制短向量可以提高 P、减少全局通信，却消耗 M；增加 K-tree 的层级可以缩短归约延迟，却增加根节点需要的路由路径 R；把权重预先排成 Decode 需要的方向，可以消除每 Token 的转置，却在 Prefill 到 Decode 之间增加一次重映射。

因此，WaferLLM 的主线不是“逐项满足四个字母”，而是：先找到当前阶段的主瓶颈，再在四类约束下选择可以成立的物理动作。 下面的五视图把这条关系同时投影为前后对比、因果链、运行流程、对象时序和张量数据流。

{ width=100% }

WaferLLM 五视图语义总览。依据论文 §3–6 与开源仓库 fd1c2daa 自绘；它解释设计关系，不是逐周期执行轨迹。可编辑 SVG 保存在仓库资产目录。

读图时分别关注：第一格看硬件抽象如何变化；第二格看 PLMR 如何传导为软件动作；第三格看 Prefill/Decode 的阶段切换；第四格看权重、激活和 KV 的生命周期；第五格看哪些对象留在本地、哪些对象沿 NoC 移动。

阶段化并行：Prefill 与 Decode

很多系统把“模型并行”当成一套贯穿推理的固定切分。WaferLLM 则从算子形状出发，承认 Prefill 与 Decode 需要不同的并行计划。这不是实现偏好，而是矩阵尺寸、算术强度和通信覆盖能力不同导致的结果。

在 Prefill 中，激活含有较长的序列维度，权重也有较大的输入与输出维度。WaferLLM 同时沿 Mesh 的 X、Y 轴切分激活和权重的两个维度，使几十万 PE 都能持有小 Tile 并参与计算。留在本地的是权重 Tile、激活 Tile 和逐步累积的部分和；沿 NoC 移动的是下一轮计算所需的操作数 Tile。

到了 Decode，每次通常只有一个新 Token，继续切分序列维度不能暴露足够并行度。WaferLLM 选择在一个 Mesh 轴上复制短输入向量，在另一个轴上切分权重，让更多 PE 并行计算局部 GEMV。复制的不是完整模型状态，而是本来很短的 Decode 输入；最后只需归约局部部分和。

{ width=96% }

来自论文 Figure 3 与 Figure 4。左侧展示 Prefill 的二维切分，右侧展示 Decode 的细粒度复制与 Dist-GEMV 路径。

这张论文原图证明的是：作者为两个阶段定义了不同的张量放置和算子序列。它不直接证明每一个局部设计带来多少端到端收益，也没有说明任意模型形状都能达到相同利用率；头数、层宽、序列长度、每核容量和 Mesh 尺寸都会改变最佳映射。

用存储换通信

Decode 的复制策略是一个典型协同设计取舍：它增加每核本地驻留，却避免为了一个很短的向量进行更昂贵的全局同步。判断是否值得，必须同时检查复制对象大小、可用 PE 数、归约路径和每核 SRAM 余量。

数据布局：免转置与 KV Shift

共享内存程序里的 transpose、reshape 或 concat 经常只是视图变化，或者由高带宽内存系统承担。在分布式 Mesh 上，若布局变化意味着一个角落的 PE 把数据发到远端角落，它就会变成真实、长距离、需要路由资源的数据搬运。

WaferLLM 对 Prefill 和 Decode 采用两种不同的免转置策略：Prefill 用 transposed distributed GEMM 直接计算需要转置语义的乘法；Decode 则预先按后续 GEMV 的消费方向放置权重。前者改变算子的数据流，后者改变权重的物理布局。两者都没有让转置“免费”，而是把代价移到更容易控制的位置。

KV Cache 的问题更能说明布局为什么属于算法。若把新 K/V 一直拼到逻辑末尾，对应的物理写入和后续计算会集中在少数尾部核心，最忙核心先耗尽本地 SRAM，其余核心却仍有空闲。WaferLLM 让相邻行在写入新 K/V 时逐步向上移动较老的缓存，使每行持有接近均衡的历史片段。

这里要准确描述内存效果：KV Shift 不减少总 KV 元素数，也不把状态移出设备。它降低的是每核驻留的偏斜峰值，并让最大可用上下文由更多 PE 的本地容量共同决定。 新代价是每个 Token 更新时的相邻 NoC 传输、同步和布局管理。论文 Table 5 在其 LLaMA3-8B 与 LLaMA2-13B 配置下报告了 360 倍与 385 倍的最大输出 Token 容量提升，但这个倍率来自特定基线映射，不应外推为所有 KV 管理方案的通用比例。[^paper-kv]

从生命周期看，权重跨多个 Token 持久存在，Prefill 激活只在阶段内存活，Decode 临时向量通常只活一个 Token，而 KV Cache 必须跨 Token 保存。WaferLLM 的布局设计本质上是在安排这些对象在哪里活、活多久、最后由谁消费。

通信算子：MeshGEMM 与 MeshGEMV

并行计划决定“哪些 PE 应该参与”，通信算子决定“它们如何交换操作数和部分结果”。WaferLLM 没有把 GPU/TPU 常见的 AllGather GEMM 或 Ring AllReduce 原样搬到 Mesh，而是分别为 Prefill 和 Decode 的关键算子设计 PLMR 兼容路径。

算子	主要阶段	本地对象	需要移动的对象	设计抓手	主要代价
MeshGEMM	Prefill	A_ij、B_ij、累积 C_ij	每轮下一块 A/B Tile	循环移位保证覆盖；Interleave 将关键移动限制在近邻两跳	N 轮计算-移位、对齐状态、专用布局
MeshGEMV	Decode	权重 Tile、向量副本、局部部分和	归约中的部分结果	K-tree AllReduce 缩短串行归约路径，用 K 调节延迟与路由资源	根节点路径数、阶段同步、广播需求

MeshGEMM 的起点是 Cannon 式循环移位：每个 PE 只保留本地子矩阵，通过多轮移动最终覆盖乘法需要的 Tile 组合。这解决正确性、工作集和固定邻居问题；Interleave 再重排逻辑环到物理位置的映射，把原本可能跨整行的首尾通信压缩为两跳。论文 Figure 6–7 给出了设计与可扩展性分析。[^paper-gemm]

MeshGEMV 的本地乘法很短，通信更难被计算掩盖，因此重点变成归约关键路径。K-tree 把顺序经过大量 PE 的 Reduce 拆成分阶段并行归约；K 越大，路径可以更短，但根核心需要更多路由资源，论文实现选择 K=2。[^paper-gemv]

本文的边界

本文只解释两类算子在系统中的职责，不把概述伪装成实现教程。MeshGEMM 的对齐、Interleave 索引、奇偶 Mesh 边界、非方形 Mesh、异步收发缓冲和完整计算-移位伪代码，会在下一篇单独建立对象账本并对照源码讲解。

论文微基准报告 MeshGEMM 相比 SUMMA 和 Cannon 快约 2–3 倍，MeshGEMV 相比 Cerebras 基线快约 4–8 倍。应把这些结果限制在论文矩阵形状、核心规模和 WSE-2 实现内；它们说明 PLMR 兼容通信在目标设备上有效，不证明任意 Mesh、任意矩阵都保持同一倍率。[^paper-eval]

系统闭环：运行时与硬件反馈

系统闭环包含两个方向：运行时把 PLMR 约束落实为具体执行计划，端到端剖析再把暴露出的新瓶颈反馈给软件、模型和下一代硬件。二者缺一，都只能算局部适配。

运行时控制流

只有并行计划和新算子还不构成完整推理系统。WaferLLM 还要加载模型、选择核心数、组织 Prefill/Decode、在阶段边界重新放置权重与 KV，并将 Transformer 中的 RMSNorm、Softmax、Attention 和 Feed Forward 接到新的通信路径上。

可以把系统控制流压缩成下面这段架构级伪代码。它故意不隐藏阶段切换和对象所有权，但不展开 MeshGEMM/MeshGEMV 的内部机制：

hardware = inspect_mesh(core_count, local_sram, hop_cost, route_budget)
constraints = build_plmr(hardware)

prefill_plan = autotune_prefill(model_shape, prompt_length, constraints)
decode_plan = autotune_decode(model_shape, expected_output_length, constraints)

weights_prefill = place_weight_tiles(checkpoint, prefill_plan)
activations, kv_cache = run_prefill(prompt, weights_prefill, prefill_plan)

weights_decode = remap_weight_tiles(weights_prefill, decode_plan)
kv_cache = remap_kv_cache(kv_cache, decode_plan)

for token_index in range(expected_output_length):
    input_vector = embed_last_token(token_index)
    local_partial_sums = run_local_gemv(input_vector, weights_decode)
    hidden_state = k_tree_allreduce(local_partial_sums, decode_plan)
    new_k, new_v = project_kv(hidden_state, weights_decode)
    kv_cache = adjacent_shift_and_write(kv_cache, new_k, new_v)
    token = sample_next_token(hidden_state)
    if token_is_eos(token):
        break

return generated_tokens

硬件反馈与版本边界

这段流程揭示两类协同：一类是硬件约束向下传导，PLMR 决定映射和算子；另一类是软件瓶颈向上反馈。论文指出 WSE-2 每核 SRAM 太小，迫使系统使用会产生空泡的流水线并行；如果每核内存增加约 5–6 倍，就可能改用更充分的张量并行。作者也观察到 RMSNorm 等本地算子的软件实现会成为新的端到端瓶颈。[^paper-future]

这正是软硬协同区别于“硬件适配”的地方：软件不只服从当前硬件，还要用端到端剖析告诉下一代硬件，增加哪类资源能让整个系统而非单个 Kernel 受益。

开源代码的版本边界也必须说明。本文研究时固定仓库提交为 fd1c2daae37cd68706c03fc8009887ecee9900f8；其 README 已声明当前代码面向 WSE-3 和 Cerebras SDK 2.10.0，而 OSDI 论文结果来自 WSE-2 与 SDK 1.x。当前代码能证明项目持续维护和具体模块归属，但不能直接充当论文数值的原样复现实验。[^repo]

如何读实验结果

论文使用 Throughput per Request（TPR，单请求吞吐）比较 WSE-2 上的 WaferLLM、T10、Ladder，以及 A100 上的 SGLang。Table 2 同时改变了硬件、编程模型、并行拓扑和软件栈，因此它是完整系统比较，不是只隔离 MeshGEMM 或片上带宽的单变量实验。

{ width=82% }

来自论文 Table 2，展示 LLaMA3-8B 与 LLaMA2-13B 在不同输入/输出长度下的端到端 TPR。

这张表支持两个结论：第一，直接把共享内存编译器或面向小规模互连的分布式方案映射到 WSE-2，不能利用大规模 Mesh；第二，WaferLLM 在论文列出的单请求场景中明显超过 A100/SGLang 配置。它不证明 WSE-2 在所有服务吞吐、并发数、P99、模型规模、成本或功耗条件下都优于 GPU 集群。

论文也主动报告了收益缩小：单独 GEMV 的巨大带宽优势没有完整转化为端到端同倍率优势，因为本地 SRAM、流水线空泡、边缘核心利用率、长距离 NoC 和 GPU 导向的模型层宽都会进入系统关键路径。微基准回答“这个机制是否改善目标瓶颈”，端到端实验回答“改善后还有什么成为新瓶颈”。

数字必须带实验边界

引用 WaferLLM 的倍率时，至少同时带上模型、输入/输出长度、WSE-2 或 A100 配置、SGLang/T10/Ladder 基线和 TPR 指标。作者知乎文章适合理解动机，但个别软件标签和头部数字与论文 v3 表述不同；本文的定量结论以论文 v3 和 OSDI 页面为准。

系列拆分

WaferLLM 适合拆成一个从抽象到实现的软硬协同系列。本文建立“硬件事实 → 约束模型 → 软件方向 → 实验反馈”的总地图，后续文章分别保持一个清晰的机制边界：

1. MeshGEMM： 从 Cannon 的长关键路径出发，完整解释 Alignment、Cyclic Shift、Interleave、两跳通信、非方形 Mesh 与源码对象。
2. MeshGEMV： 解释 Pipeline/Ring/K-tree AllReduce 的路径长度、路由预算、K 的选择和 Decode 归约对象。
3. Wafer LLM Parallelism： 单独分析 Prefill 二维切分、Decode 复制、免转置布局与阶段转换。
4. KV Shift： 以对象生命周期和每核峰值为中心，对比逻辑 concat 与物理相邻移位。
5. Co-Design Evaluation： 区分算子微基准、阶段 TPR、端到端 TPR、能效和当前 WSE-3 复现边界。

最值得复用的设计思路不是某个固定通信 Pattern，而是：先把硬件差异写成可以否决方案的约束，再让每一层软件明确说明它改变了什么对象、移动了什么数据、付出了什么新代价，最后用端到端瓶颈反向修正硬件与模型。

参考资料

[^paper]: Congjie He et al., WaferLLM: Large Language Model Inference at Wafer Scale, arXiv v3 / OSDI 2025, Abstract and §1–2.
[^paper-setup]: WaferLLM, §7 Experiment setup; WSE-2 core count, clock, per-core SRAM and aggregate SRAM.
[^paper-plmr]: WaferLLM, §3.1–3.2, PLMR device model and limitations of shared/distributed-memory baselines.
[^paper-kv]: WaferLLM, §4.3 and §7.4, Figure 5 and Table 5.
[^paper-gemm]: WaferLLM, §5.1–5.4, Figure 6–7 and Algorithm 1.
[^paper-gemv]: WaferLLM, §6.1–6.2, Figure 8.
[^paper-eval]: WaferLLM, §7.2–7.3, Figure 9–10.
[^paper-future]: WaferLLM, §8, hardware architecture, memory-to-compute ratio and software limitations.
[^repo]: MeshInfra/WaferLLM, pinned at fd1c2daae37cd68706c03fc8009887ecee9900f8; README reproduction boundary.