跳转至

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 边界推进:

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

认知锚点:优化目标不是单独跑满某个部件,而是让 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

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

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

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 五视图

手写整模型 MegaKernel 五视图。阅读重点:A 的 kernel boundary 变化,D 的常驻对象生命周期,以及 E 中逻辑张量与物理 materialization 的区别。

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

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

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 内部边界

Hazy 官方逻辑图:上方是传统 kernel 边界,下方把工作切成可跨算子交错的 tile。来源:Hazy Research。

Hazy MegaKernel 官方性能图

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 编译器与运行时五视图

MPK 五视图。阅读重点:A 的 kernel graph→tGraph,D 的 worker/scheduler SM 所有权,E 的“先写数据、后发事件”正确性链。

MPK Figure 4:从算子图到 SM-level tGraph

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 中。

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:端到端相对性能

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 的物理阶段通常是:

Dispatch All-to-All
→ Linear-1
→ SwiGLU
→ Linear-2
→ Combine All-to-All

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

下一 wave 的 Dispatch/Pull
+ 当前 wave 的 L1/Act/L2
+ 上一 wave 的 Combine

MegaMoE 五视图

MegaMoE 五视图。阅读重点:A 的三段稳态流水,D 的 EP rank 与 bounded ring,E 的远端 pull 和 combine slot 回写。

DeepSeek-V4 Figure 5:MegaMoE wave 流水

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

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

按 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.

评论