Inference MegaKernel

导言

小 batch LLM 推理并不总是被矩阵乘法本身限制。数十到数百个 kernel 的启动边界、kernel 尾部气泡、中间张量写回 HBM,以及 MoE 的 Dispatch/Combine 通信,都可能让计算单元等数据、让内存等指令、让互联等计算。MegaKernel 的核心不是“把代码写得更大”,而是把 host 侧的算子调度下沉为 GPU 内部的任务系统。

本文沿三条路线展开:Hazy Research 的手写整模型 MegaKernel、MPK 的编译器与常驻运行时、DeepSeek MegaMoE 的 expert-wave 通信计算流水。结论是:这类优化用更强的静态专用化和片上资源约束,换取更少的启动边界与更细的就绪事件;收益必须与寄存器、共享内存、HBM、互联、功耗和动态路由一起核算。

为什么常规推理仍有气泡

一次 Transformer 推理通常被拆成 Norm、QKV GEMM、RoPE、Attention、残差、MLP、通信和采样等 kernel。传统执行依赖 host/driver 按 kernel 边界推进:

1
2
3
4
host launch
→ kernel 执行
→ 全局完成/可见
→ 下一个 launch

即使每个 kernel 内部已经高度优化,端到端时间仍可近似写成:

$$
T_{\text{e2e}}

\sum_i T_{\text{kernel},i}

  • \sum_i T_{\text{launch},i}
  • \sum_i T_{\text{tail},i}
  • \sum_i T_{\text{materialize},i}
  • T_{\text{communication}}.
    $$

其中,tail 是一个 kernel 尾部只有少数 SM 仍在工作造成的气泡,materialize 是跨 kernel 依赖迫使中间结果落到全局内存的成本。Hazy Research 在 batch size 1 的 Llama-3.2-1B 案例中观察到约 100 个 kernel,传统实现的 DRAM 带宽利用率最高约 50%;这正是小 batch、memory-bound decode 容易暴露的边界成本。Hazy Research

![软硬件协同的认知锚点:同时驱动 HBM、Tensor Core 和 NVLink](https://pic.shaojiemike.top/shaojiemike/2026/07/5c068d008a077a769a0141ea0f4ea7da.png){ width=90% }
认知锚点:优化目标不是单独跑满某个部件,而是让 HBM、Tensor Core 与 NVLink 在依赖允许时持续并行工作。该图不作性能证据。

三个容易混淆的概念

  • Operator Fusion 通常融合局部相邻算子,调度单位仍是 kernel。
  • CUDA Graph 降低 host launch 和图重放开销,但依赖主要仍表达在 kernel 节点之间。
  • MegaKernel 把任务、事件、队列和通信进一步放进一个常驻 kernel,调度粒度下沉到 SM、CTA、warp 或 tile。

NVIDIA 的 Programmatic Dependent Launch(PDL)允许依赖 kernel 在前序 kernel 完全结束前提前启动,但官方文档明确说明这种重叠是机会性的,不能依赖实际并发来保证正确性。因此,PDL 能缓解部分边界,却不能替代一个显式管理细粒度依赖的设备内运行时。CUDA PDL

统一框架:把算子图下沉为片上任务系统

三条路线的共同转换可以概括为:

1
2
3
4
5
kernel-level graph
→ tile / task graph
→ per-SM queue + event/counter
→ resident worker
→ register / SMEM / HBM / interconnect co-scheduling

算子和张量的逻辑语义没有消失。 改变的是物理执行边界:某个中间张量是否必须完整写回 HBM,某个消费者是否要等待整个生产者 kernel,某段通信是否能与下一块 Tensor Core 计算重叠。

路线 调度生成 主要粒度 最适合 主要代价
手写整模型 MegaKernel 离线脚本与手工 CUDA 模型专用 tile/instruction 固定模型、固定 shape、batch 1 极低延迟 工程成本高,硬件和模型专用
MPK 编译器生成 tGraph,运行时调度 SM-level task/event 多算子、多模型、单/多 GPU 自动化 scheduler SM、event traffic、编译复杂度
MegaMoE MoE 专用动态 scheduler expert block/wave SM100 上 EP MoE 的通信计算融合 routing、ring capacity、互联和功耗约束

这不是免费的抽象迁移。原来由 driver 和 kernel boundary 隐式提供的正确性,现在必须由设备内协议显式完成:

  1. 数据可见性:生产者写入必须先于完成事件发布。
  2. 进度保证:等待方不能占满所有 resident resource,使生产者无法获得 SM。
  3. 容量保证:queue、ring、combine slot 和 KV page 必须覆盖配置上限。
  4. 资源保证:最大 task body 的寄存器和共享内存需求不能让 kernel 无法启动或 occupancy 过低。

手写整模型 MegaKernel

原理、调度与数据流

Hazy Research 的路线从一个固定 Llama 配置出发,将每个算子切成 tile,构造依赖 DAG,再把 tile 静态分配到每个 SM 的指令队列。运行时启动一个与目标 GPU SM 数匹配的常驻 grid,每个 CTA 读取自己的队列、等待 counter、租用共享内存页、执行对应 operator body,再发布下游事件。

![手写整模型 MegaKernel 五视图](https://pic.shaojiemike.top/shaojiemike/2026/07/27821db6b287ff5d359de80e0c0596db.png){ width=100% }
手写整模型 MegaKernel 五视图。阅读重点:A 的 kernel boundary 变化,D 的常驻对象生命周期,以及 E 中逻辑张量与物理 materialization 的区别。

源码中的 scheduler.py 先建立 DAG、分配 wave 与 memory/compute pool,再把每个 SM 的计划张量化为指令队列;llama.cu 只做一次模型专用 mk 绑定,权重、KV cache、activation 和 counter 都作为这次常驻执行的对象传入。Megakernels repository

下面的规范化伪代码保留了完整控制流:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def run_full_model_megakernel(sm_id, queue, tensors, pages, counters):
program_counter = 0
while True:
instruction = queue[sm_id][program_counter]
if instruction.opcode == "END":
break

for dependency in instruction.dependencies:
while counters[dependency.counter_id] < dependency.required_value:
device_pause()

leased_pages = []
for page_index in range(instruction.required_pages):
leased_pages.append(pages.acquire(sm_id, page_index))

if instruction.opcode == "RMS_NORM":
execute_rms_norm(instruction, tensors, leased_pages)
elif instruction.opcode == "QKV_MATVEC":
execute_qkv_matvec(instruction, tensors, leased_pages)
elif instruction.opcode == "ATTENTION":
execute_attention(instruction, tensors, leased_pages)
elif instruction.opcode == "MLP_UP_GATE":
execute_mlp_up_gate(instruction, tensors, leased_pages)
elif instruction.opcode == "MLP_DOWN":
execute_mlp_down(instruction, tensors, leased_pages)
elif instruction.opcode == "LM_HEAD":
execute_lm_head(instruction, tensors, leased_pages)
else:
raise_invalid_opcode(instruction.opcode)

device_memory_fence()
for completion in instruction.completions:
atomic_add(counters[completion.counter_id], completion.delta)

for leased_page in leased_pages:
pages.release(sm_id, leased_page)

program_counter += 1

共享内存分页解决了异构算子争用同一 CTA 片上空间的问题:算子不再各自声明互不复用的大块静态 SMEM,而是按阶段租用固定页。全局 counter则承担跨 SM 依赖;生产者完成一个 tile 后增加 counter,消费者只等待自己需要的计数值。代价是 counter 的 L2/atomic 流量与更复杂的 memory ordering。

对象 典型形状/类型 位置 生命周期
per-SM 指令队列 [S,Q,32] int32 一类编码 HBM/L2 整个 MegaKernel
权重 模型固定张量 HBM 模型实例
KV cache 分层分页 K/V HBM request
activation tile operator-specific tile register/SMEM,必要时 HBM 依赖区间
SMEM page 固定页集合 片上共享内存 task 租用区间
counter/barrier 标量或数组 SMEM + HBM/L2 kernel 或依赖区间

效果与适用边界

![常规 kernel 边界与 MegaKernel 内部边界](https://pic.shaojiemike.top/shaojiemike/2026/07/0d622056756b4de566447b5c5e0560e8.png){ width=90% }
Hazy 官方逻辑图:上方是传统 kernel 边界,下方把工作切成可跨算子交错的 tile。来源:Hazy Research。
![Hazy MegaKernel 官方性能图](https://pic.shaojiemike.top/shaojiemike/2026/07/74796821083e45d7f7b4fd7f52032ff9.png){ width=92% }
Hazy 官方结果图。博客报告 H100 超过 1.5×,B200 端到端低于 680 μs;这些是其特定 Llama-3.2-1B、batch 1 配置的结果,本轮未独立复现。

Hazy 报告其 H100 实现达到约 78% DRAM 带宽利用率,并在 B200 上进一步利用更强的 Tensor Core 与内存系统。这里的关键不是“B200 自动让融合更快”,而是静态 schedule 必须针对 SM 数、共享内存、Tensor Core 指令和内存系统重新专用化。

手写路线的边界

固定提交中的 Llama shape、H100/B200 SM 数和 task layout 都是显式配置。它证明了模型专用 MegaKernel 可以工作,不等于任意 Hugging Face 模型都能自动转换。动态 batch、变长 request、量化格式变化和新算子都会改变静态队列。

MPK 编译器与运行时

从 kernel graph 到 tGraph

MPK 将手写路线中的“构图、切 tile、插事件、分配 SM”提升为编译器问题。输入是 tensor program 与 inference config,输出是优化后的 SM-level tGraph。节点分为 task 与 event:task 在一个 SM 上执行计算或通信,event 聚合多个 task 的通知并触发下游 task。

![MPK 编译器与运行时五视图](https://pic.shaojiemike.top/shaojiemike/2026/07/2d5cf09657c04a1b249d449305b5bd09.png){ width=100% }
MPK 五视图。阅读重点:A 的 kernel graph→tGraph,D 的 worker/scheduler SM 所有权,E 的“先写数据、后发事件”正确性链。
![MPK Figure 4:从算子图到 SM-level tGraph](https://pic.shaojiemike.top/shaojiemike/2026/07/343d187165ca63611d8b9921e199b4ed.png){ width=92% }
MPK Figure 4:细粒度 event 只连接真正依赖的 MatMul tile 与 AllReduce tile,避免把整个算子串行化。来源:MPK 论文。

编译阶段包含 event fusion、normalization 和 linearization;运行阶段将部分 SM 作为 worker、部分 SM 作为 scheduler。worker 从自己的 task queue 取任务,完成后向 event queue 发布通知;scheduler warp 聚合通知,把新就绪 task 放入 worker queue。多 GPU tensor 与事件可放在 NVSHMEM symmetric heap 中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def compile_mpk(tensor_program, inference_config):
task_graph = lower_operators_to_sm_tasks(tensor_program)
task_graph = infer_precise_tile_dependencies(task_graph)
task_graph = fuse_equivalent_events(task_graph)
task_graph = normalize_task_and_event_edges(task_graph)
task_graph = linearize_ready_task_ranges(task_graph)
binary = jit_or_aot_compile(task_graph, inference_config)
return binary, task_graph


def run_mpk(runtime, task_graph, request):
runtime.bind_request_tensors(request)
runtime.enqueue_root_events(task_graph.root_events)

while not runtime.request_finished(request):
for scheduler in runtime.local_schedulers:
event = scheduler.try_pop_ready_event()
if event is not None:
for task_id in event.ready_tasks:
worker_id = scheduler.select_worker(task_id)
runtime.worker_queues[worker_id].push_release(task_id)

for worker in runtime.workers:
task_id = runtime.worker_queues[worker.id].try_pop_acquire()
if task_id is not None:
task = task_graph.tasks[task_id]
execute_sm_task(task, runtime.tensor_table)
device_memory_fence()
for event_id in task.triggered_events:
runtime.scheduler_queues[event_id.owner].push_release(event_id)

if runtime.online_mode and runtime.no_active_tokens():
if runtime.shutdown_requested():
runtime.publish_termination_events()
break
device_pause()

固定 Mirage 源码中的 online 模式在没有 active token 时不会退出 kernel,而是轮询 shutdown signal;worker queue 与 scheduler queue 使用 acquire/release 风格的 load/store/atomic 顺序维持数据可见性。Mirage repository

对象 作用 所有者 生命周期
tGraph task operator tile 或 communication tile scheduler→worker graph iteration
event dependency count 与 ready state producer + scheduler graph iteration
worker queue TaskId ring worker SM 常驻 kernel
scheduler queue EventId ring scheduler warp 常驻 kernel
tensor tile task-specific shape/dtype worker CTA task
remote tensor/event symmetric allocation GPU rank/PE runtime/request
KV page queue page id 与 cursor online runtime server

自动化收益与新瓶颈

![MPK Figure 9:端到端相对性能](https://pic.shaojiemike.top/shaojiemike/2026/07/0562ab7acda1d5d4dda5fba2cabb1632.png){ width=96% }
MPK Figure 9:A100/H100/B200、五个模型和多个 batch size 下的论文结果,MPK 相对最佳 baseline 的标注约为 1.0–1.7×。所有柱均按 MPK 归一化,本轮未复现。

论文以 Qwen3-8B 为例统计到 293 次 eager launch,并报告 eager launch 约 3.8 μs、CUDA Graph 约 0.8 μs;MPK 内部 scheduler 开销约占 0.28%。这说明 launch 边界足够多时,下沉调度存在明确窗口。MPK paper

但自动化引入了新的优化面:

  • scheduler 占用 SM:用于调度的 SM 不再直接执行主要算术。
  • 最大寄存器体:单一常驻 kernel 内任一重 task 的寄存器需求可能约束所有 CTA。
  • 事件粒度:过粗失去 overlap,过细增加 queue、atomic 和 L2 流量。
  • 编译搜索:task 切分、event fusion、worker/scheduler 比例和通信 task 都与硬件相关。

MegaMoE 专用流水

语义不变,物理阶段被打散

对 token t,top-k MoE 的语义仍可写为:

$$
y_t

\sum_{j=1}^{K}
p_{t,j}
W^{(2)}{e{t,j}}
\operatorname{SwiGLU}
\left(
W^{(1)}{e{t,j}}x_t
\right).
$$

传统 expert parallelism 的物理阶段通常是:

1
2
3
4
5
Dispatch All-to-All
→ Linear-1
→ SwiGLU
→ Linear-2
→ Combine All-to-All

MegaMoE 将 expert 切成多个 wave。一个 wave 的 token 到齐后立即计算,不等待所有 expert;稳态时同时执行:

1
2
3
下一 wave 的 Dispatch/Pull
+ 当前 wave 的 L1/Act/L2
+ 上一 wave 的 Combine
![MegaMoE 五视图](https://pic.shaojiemike.top/shaojiemike/2026/07/167d9a164e7befac4655fd27ecf459eb.png){ width=100% }
MegaMoE 五视图。阅读重点:A 的三段稳态流水,D 的 EP rank 与 bounded ring,E 的远端 pull 和 combine slot 回写。
![DeepSeek-V4 Figure 5:MegaMoE wave 流水](https://pic.shaojiemike.top/shaojiemike/2026/07/2c6572e4f070d6876780131f63d32cb0.png){ width=94% }
DeepSeek-V4 Figure 5。图中的 1.92× 是 V4-Flash 配置下的理论 speedup,不能当作通用实测值。来源:DeepSeek-V4 技术报告。

对称内存、pull dispatch 与 ring

当前 DeepGEMM 固定提交首先通过 torch.distributed._symmetric_memory 分配每个 rank 形状一致的 buffer,执行 rendezvous 获得 peer 指针,再将 input、scale、top-k、workspace、L1/L2 ring 和 combine slot 切成固定布局。PyTorch 官方将 symmetric memory 标为 alpha,并明确要求各 rank 以相同顺序、相同 shape/dtype 建立 rendezvous。PyTorch Symmetric Memory

Dispatch 不是 source rank 把完整 activation push 到目标 expert rank,而是:

  1. source rank 统计 route,向目标 rank 写入 source-token/top-k 元数据与 expert count;
  2. 目标 rank 读取元数据;
  3. 目标 rank 通过 peer pointer pull 远端 token 与 scale 到本地 L1 ring;
  4. L1/L2 task 消费本地 ring;
  5. L2 epilogue 将 BF16 结果写回 source rank 的 combine slot;
  6. source rank 归约 top-k slot 得到输出。

这种 one-sided 访问符合 NVSHMEM/PyTorch symmetric memory 的 PGAS 思路:数据位置由 <symmetric address, rank> 表示,GPU kernel 可发起远端 load/store/RMA。但通信与同步被解耦后,程序必须自己处理 fence、signal 与 buffer reuse。NVSHMEM

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def run_mega_moe(rank, x, x_scale, topk_idx, topk_weight, w1, w2, symm):
symm.copy_input(rank, x, x_scale, topk_idx, topk_weight)
symm.zero_workspace(rank)
symm.rank_barrier()

local_route_count = count_routes_per_expert(topk_idx)
route_offsets = exclusive_prefix_sum(local_route_count)

for token_id in range(x.shape[0]):
for topk_slot in range(topk_idx.shape[1]):
expert_id = topk_idx[token_id, topk_slot]
if expert_id >= 0:
destination_rank = owner_rank(expert_id)
destination_slot = allocate_route_slot(route_offsets, expert_id)
symm.write_remote_route_metadata(
destination_rank,
expert_id,
destination_slot,
rank,
token_id,
topk_slot,
)

symm.publish_expert_counts(rank, local_route_count)
symm.nvlink_barrier()
scheduler = build_wave_scheduler(symm.read_local_expert_counts(rank))

while scheduler.has_task():
task = scheduler.next_task()
if task.phase == "LINEAR1":
source = symm.read_route_metadata(task.expert_id, task.token_block)
token_block = symm.pull_remote_input_to_l1_ring(source)
l1_output = fp8_fp4_or_bf16_linear1(token_block, w1[task.expert_id])
activated = swiglu_and_cast(l1_output)
symm.commit_l2_ring(task.ring_slot, activated)
scheduler.publish_linear1_completion(task)
elif task.phase == "LINEAR2":
scheduler.wait_linear1_completion(task)
activated = symm.read_l2_ring(task.ring_slot)
expert_output = fp8_fp4_or_bf16_linear2(activated, w2[task.expert_id])
symm.write_remote_combine_slot(task.source, expert_output)
scheduler.publish_linear2_completion(task)
else:
raise_invalid_phase(task.phase)

symm.nvlink_barrier()
y = zeros_bf16(x.shape)
for token_id in range(x.shape[0]):
for topk_slot in range(topk_idx.shape[1]):
expert_output = symm.read_local_combine_slot(token_id, topk_slot)
y[token_id] += topk_weight[token_id, topk_slot] * expert_output
symm.rank_barrier()
return y

源码还包含一个容易忽略的进度保证:scheduler 先预热若干 L1 wave,再交错 L1/L2。原因是 L2 依赖 L1 产生 ring 数据;如果过早分配太多 L2 task,可能占住资源并阻塞继续产生 L1 数据,形成 L1→L2 ring deadlock。

对象 逻辑形状 物理位置 生命周期
x [T,H] FP8/BF16 per-rank symmetric input kernel call
topk_idx/weight [T,K] per-rank symmetric input kernel call
route metadata source rank/token/top-k + count workspace/L2 kernel call
L1 ring [B_r,H] destination rank workspace wave
L2 ring [B_r,I] destination rank workspace L1→L2
combine slot [K(+shared),T,H] BF16 source rank workspace 直到归约
W1/W2 expert-local matrices local HBM 模型实例
y [T,H] BF16 source rank output 下游消费

计算通信比与功耗墙

完全隐藏通信的必要条件可写为:

$$
\frac{C}{B}
\le
\frac{V_{\text{comp}}}{V_{\text{comm}}},
$$

左边是通信时间,右边是可与通信重叠的计算窗口。DeepSeek-V4 报告对 V4-Pro 给出每 token-expert 约 6hd FLOP 与 3h byte,得到:

$$
\frac{C}{B} \le 2d = 6144\ \text{FLOP/Byte}.
$$

换句话说,每 1 GB/s 互联带宽最多可被约 6.1 TFLOP/s 的相应计算窗口覆盖。这是模型化上限,真实系统还会受 route skew、短消息 notification、ring tail 和 memory ordering 影响。DeepSeek-V4 report

更深的软硬协同来自功耗。MegaMoE 同时驱动 Tensor Core、HBM 和 NVLink;如果板卡总功耗触顶导致某一域降频,“重叠”可能只是把三个阶段都变慢。报告因此把 power throttling 明确列为关键问题。性能实验必须同时记录 clocks、power cap、温度和 throttling reason,而不应只看 kernel duration。

开源状态与性能证据

DeepGEMM PR #304 首发时描述为融合 Dispatch、Linear-1、SwiGLU、Linear-2 与 Combine,并只支持 FP8×FP4;当前固定提交 559d79f 已出现 bf16_mega_moe 路径。报告使用 MegaMoE2 名称,但固定提交的公开 API 仍是 mega_moe,没有独立同名 public symbol。PR #304

![按 DeepGEMM PR #316 数据重绘的 MegaMoE 延迟与 speedup](https://pic.shaojiemike.top/shaojiemike/2026/07/a0539804d60fd8ccfabf7f9b1df2b3c8.png){ width=100% }
按 DeepGEMM PR #316 数据确定性重绘。EP8、8 rank 平均;柱高是延迟,badge 是官方报告 speedup。该图不是独立复现。

PR #316 报告 V4-Flash/V4-Pro 在不同 token 数下约 1.50–1.96×。但 PR 正文没有披露 GPU 型号、功耗状态、CUDA/PyTorch 构建、对比 baseline 细节和测量方差;评论区也有人追问这些信息。因此,这些数字只能被称为官方报告值,不能外推为 B200、B300 或任意 SM100 系统的性能。

GPU 与 NPU 证据不能混写

DeepSeek-V4 报告称方案在 NVIDIA GPU 与 Huawei Ascend NPU 上验证,并报告一般推理 1.50–1.73×、低延迟场景最高 1.96×。公开 DeepGEMM 固定提交展示的是 CUDA SM100 实现;本轮没有找到等价的公开 Ascend kernel,因此不能声称 NPU 路径已开源可复现。

软硬协同的五本账

调度、片上资源与内存

账本 MegaKernel 改变了什么 主要风险 必测指标
调度 host kernel graph→device task/event queue contention、event 过细 launch 数、scheduler 周期、L2 atomic
寄存器/SMEM 多种 task 共存,SMEM 可分页复用 max-register body、spill、低 occupancy registers/thread、spill、active CTA/SM
HBM 减少中间 materialization,缩短 tile 生命周期 queue/ring/counter 新流量 DRAM BW、L2 hit、bytes/token
互联 communication task 与 compute task 统一调度 notification、ordering、routing skew NVLink BW、wave tail、remote load latency
功耗 HBM、Tensor Core、NVLink 同时活跃 power cap 触发降频 clocks、board power、throttle reason

CUDA 官方文档给出的底层约束非常直接:一个 block 需要的寄存器数与共享内存必须能放进 SM,否则 kernel 无法启动;即使能启动,也会限制 resident block/warp 数。CUDA Programming Model 对常驻 kernel 而言,这不是局部算子的 occupancy 问题,而可能成为整个模型执行体的 occupancy 问题。

跨 CTA 同步同样必须满足进度契约。Cooperative Groups 提供显式 grid 级同步,但所有参与线程必须到达 collective;分支不一致会导致死锁或数据破坏。CUDA Cooperative Groups

峰值显存应按对象而不是按 kernel 数计算:

$$
\begin{aligned}
M_{\text{peak}}
=&
M_{\text{model state}}

  • M_{\text{live activations}}
  • M_{\text{workspace}} \
    &+ M_{\text{communication buffers}}
  • M_{\text{live I/O}}
  • M_{\text{allocator reserved}}
  • M_{\text{runtime margin}}.
    \end{aligned}
    $$

MegaKernel 可能减少全局中间 activation,却增加常驻 queue/event、SMEM page、symmetric buffer、ring capacity 和 combine slot。权重、KV cache 与最终输出不会因为 kernel 数减少而自动缩小。没有同模型、同 shape、同 allocator 的实测,不能给出通用显存节省比例。

选型与验证

什么时候选择哪条路线

场景 首选 理由
固定模型、固定 shape、batch 1、追求最低延迟 手写整模型 MegaKernel 能做最激进的静态布局与 per-SM schedule
多模型或多 GPU,希望自动跨算子/通信优化 MPK 编译器生成 tGraph,运行时处理 task/event
Blackwell SM100、EP MoE、NVLink 通信显著 MegaMoE 专门融合 Dispatch/MLP/Combine 并使用 expert wave
单个大 GEMM/Attention 已长期饱和 先优化算子 launch/tail 占比可能不足以覆盖 MegaKernel 复杂度
shape 高度动态、控制流复杂、多租户隔离严格 CUDA Graph/局部 fusion 起步 常驻静态 schedule 的维护和隔离成本更高

可复现实验清单

  1. 固定环境
    • GPU/NPU 型号、SM/Cube 数、HBM、互联拓扑。
    • driver、CUDA/CANN、PyTorch、编译器与提交。
    • clocks、power cap、MIG/虚拟化、温度。
  2. 固定语义
    • 模型、精度、量化 scale、batch、prompt/decode 长度。
    • MoE 的 E/K/H/I/P、shared expert、capacity 与 routing 分布。
    • 输出精度和 deterministic tolerance。
  3. 拆分时间
    • host launch、device scheduler、compute、HBM、NVLink、tail。
    • prefill 与 decode 分开,warmup 与 steady state 分开。
  4. 审计资源
    • registers/thread、spill、static/dynamic SMEM、occupancy。
    • queue/ring 高水位、atomic contention、L2 hit。
    • 峰值 allocated/reserved、symmetric heap 与 workspace。
  5. 做压力测试
    • routing skew、hot expert、极小 token、capacity 边界。
    • power throttling、跨节点互联、故障/超时与 shutdown。

推荐的落地顺序

先用 profiler 证明 launch + tail + materialization 占比足够大;再做局部 persistent/fusion 原型;随后审计寄存器、SMEM 和 queue 进度;最后才把通信与多 GPU 加入同一 kernel。否则容易用一个更复杂的调度器,替换掉一个并非主要瓶颈的 kernel graph。

总结

MegaKernel 是执行抽象的迁移,不只是 fusion 的扩大。 它把 host/driver 的 kernel 边界改造成设备内 task、event、queue 和 counter,从而让消费者按 tile 就绪、让中间对象更短命、让通信与计算进入同一调度域。

  • 手写整模型 MegaKernel证明了极致静态专用化可显著降低小 batch 推理气泡,但维护成本和硬件绑定最强。
  • MPK尝试把这一过程编译器化,以 SM-level tGraph 统一计算和通信,代价是 on-device scheduler 与事件系统。
  • MegaMoE把相同思想收缩到最值得专用化的 EP MoE 层,用 symmetric memory、远端 pull、bounded ring 和 expert wave 同时驱动 NVLink 与 Tensor Core。

真正的软硬协同判断不是“是否只有一个 kernel”,而是:调度粒度是否足够细、依赖协议是否正确、片上资源是否可驻留、HBM 与互联是否能被计算窗口覆盖,以及三者并发时是否撞上功耗墙。

参考文献

  1. Hazy Research, No Bubbles at All: Megakernels for LLM Inference.
  2. HazyResearch, Megakernels, pinned at 7309cec.
  3. Cheng et al., MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs.
  4. Mirage Project, Mirage / MPK, pinned at 5715c6f.
  5. DeepSeek-AI, DeepSeek-V4 Technical Report.
  6. DeepSeek-AI, DeepGEMM, pinned at 559d79f.
  7. NVIDIA, CUDA Programmatic Dependent Launch.
  8. PyTorch, Symmetric Memory.
  9. NVIDIA, NVSHMEM Programming Model.
Author

Shaojie Tan

Posted on

2026-07-25

Updated on

2026-07-27

Licensed under