AI Infra Wiki

WaferLLM Compiler Research Gaps

7分钟阅读

WaferLLM Compiler Research Gaps

TL;DR：WaferLLM 是 PLMR mesh-NoC 加速器上首个完整 LLM 推理系统（OSDI 2025, He et al.）。但论文 §7.5、§8 明确承认 3 个 decode 阶段瓶颈未解。这 3 个瓶颈是编译器视角的 —— 需要 MLIR/TVM 层 pass，而非手写 CSL。这是 Direction 2（compiler-aware decode on mesh-NoC）的研究入口。

三个被作者承认的未解瓶颈

Gap 1: 48KB SRAM 限制 → 5× underutilization

§7.5 原话：“limited local SRAM (48KB) on WSE-2 cores, which hinders efficient tensor parallelism and necessitates pipeline parallelism, causing up to 5× underutilization”

机制：48KB/核装不下 decode 阶段 tensor parallel 的完整 working set，被迫用 pipeline parallelism（而不是 fine-grained tensor parallel），导致 core 利用率只有 1/5。

论文给出的”未来方案”：

§8 “Increasing a core’s local memory by 5-6× could eliminate the need for pipeline parallelism, enabling full tensor parallelism”

但作者没做：有没有不靠硬件升级的编译器手段？

Compiler 视角机会：

  • Persistent kernel —— 把 decode 算子 fusion 化，让一个 kernel 在一个核上常驻，减少 working set 重载
  • Activation recomputation —— 编译器自动识别”用完即可丢弃”的中间 activation，主动 recompute 换 SRAM
  • Operator fission —— 拆小算子，让每个 working set 落在 48KB 内
  • In-place KV cache update —— 编译器保证 KV 更新原地进行，节省 KV 中间 buffer

Gap 2: edge cores underutilization（论文一笔带过）

§7.5 原话：“edge cores are underutilized”（仅一句话提及）

为什么 edge core 利用率低（论文没展开）：

  • mesh 边缘核通信邻居少（2-3 个 vs 内部 4 个）
  • mesh 边缘核可分配的工作负载少（外圈 mesh 长度 < 内圈）
  • 算子 partition 算法默认从 (0,0) 起点放矩形，没考虑 mesh 几何

Compiler 视角机会：

  • Topology-aware placement pass —— 用 ILP/cost model 自动算”哪组核放哪段算子”，让 edge core 不闲置
  • Multi-tile scheduling —— 把 decode 的小算子（如 RMSNorm、Softmax）显式多映射到 edge core，让它们有事可做
  • Geometry-aware NoC 通信预测 —— placement pass 的 cost model 包含 edge 系数

Gap 3: K=2 K-tree allreduce 硬编码

§6.2 原话：“Considering these factors, we have chosen K = 2 for our current implementation”

作者承认（§6.2）：

“a larger K is not always better, as it depends on N and R constraints. Additionally, larger K increases routing complexity and overhead.”

但作者没做：cost model 自动搜 K。

Compiler 视角机会：

  • K-search pass —— 给定 N（核数）+ R（路由路径预算）作为输入，cost model 在 (latency, routing) Pareto 前沿上搜最优 K
  • Adaptive K selection per phase —— prefill 用 K=2，decode 用 K=4（decode 路由预算更紧）
  • Compiler-cost-model-driven autotune —— 类似 TVM Ansor / AutoScheduler 的 cost model + 搜索

进一步：WaferLLM 没提及的 Compiler 视角 Gap

精读论文还发现 3 个作者完全没提的 hook：

Gap 4: Shift-based KV cache 只解决行方向不平衡

Figure 5 描述 shift 是行间老 KV 向上 shift（Y 方向），X 方向 head 维度的不平衡没解

为什么有问题：

  • GQA（Grouped Query Attention）下，多个 Q head 共享 1 个 KV head —— X 方向 head 维度的负载天然不均
  • 长 context 时 X 方向 head 间的 KV 长度差异更大

Compiler 视角机会：

  • 二维 KV shift —— 行 + 列同时做 shift
  • Head-aware shift —— 编译器根据 head 间 KV size 差异，决定 shift 步长
  • Topo-aware shift ordering —— shift 方向根据 mesh 长宽比决定

Gap 5: Routing 资源约束是硬约束但无编译器层 utilization 指标

§3.1: “each core can recognize only message headers with a 5-bit address code, so at most 2^5 = 32 distinct routing paths”

论文假设”软使用 25”（“less than 25 on Cerebras WSE-2”），但编译器层没有任何工具告诉程序员”这个 placement 已用 30/32 paths”。

Compiler 视角机会：

  • Routing-aware placement pass —— placement 输出包含每核已用 routing paths 数 + 余量
  • Routing budget violation diagnostic —— 当 placement 超出 32 时给 warning
  • Multi-pass iterative placement —— 像 register allocation 那样做 spillover

Gap 6: 论文作者承认”未达 7000× 理论上限” —— 这是 honest evaluation 的金矿

§7.5: “WaferLLM does not achieve the theoretical 7,000× improvement. Profiling identifies three contributing factors: (i) cores cannot fully overlap memory and computation; (ii) edge cores are underutilized; (iii) NoC long-range communication overhead persists, despite MeshGEMV mitigating it.”

编译器视角的 3 个对应优化：

  • (i) overlap → compiler-directed prefetching / double buffering
  • (ii) edge cores → 见 Gap 2
  • (iii) NoC long-range → 编译器做通信 schedule 重排，把长程通信隐藏在短程计算中（类 LoopLynx head-wise pipelining）

Direction 2 的具体技术路径

阶段 A（1-2 月）：搭 PLMR-aware MLIR dialect

// 伪代码示意
plmr.mesh_tile {row=0, col=0, sramb_kb=48} : !plmr.tile
plmr.interleave_shift %A, %B {k=2} : (!plmr.tile, !plmr.tile) -> !plmr.tile
plmr.k_tree_allreduce %C {k=2, n_groups=4} : !plmr.tile
plmr.kv_shift %kv {direction="up", stride=1} : !plmr.kv_block

目标：在 MLIR 中表达所有 PLMR-aware 操作

阶段 B（2-3 月）：3 个核心 pass

  1. Routing-Aware Placement Pass —— 输入算子图 + mesh 拓扑 + 路由预算，输出 placement + routing plan
  2. K-Tree Auto-Search Pass —— 输入 N + R，cost model 搜最优 K（覆盖 Gap 3）
  3. 2D KV Cache Shift Pass —— 输入 GQA/MQA head layout + KV 大小，输出二维 shift plan（覆盖 Gap 4）

阶段 C（2-3 月）：仿真评估 + paper

  • 评估：cycle-accurate simulator (SCALE-sim / Timeloop 扩展 + mesh topology)
  • 目标会议：MICRO 2026（11月截）、HPCA 2027（7月截）、MLSys 2026（已过 deadline）
  • 备选：CGO 2027（编译器对口）、ASPLOS 2027

相关页面（已有 wiki 资产）

Citations

[1] WaferLLM_LLM_Inference_at_Wafer_Scale_2025.pdf — He et al. OSDI 2025, §7.5 & §8 [2] waferllm-system.md — 现有 wiki 摘要

关系图谱

反向链接