DeepGEMM/megamoe-research-reports/pr357_green_context_split_megamoe_notes.md

PR357 Green-Context Split MegaMoE 代码review报告

1. 范围

• Remote: git@github.com:RayWang96/DeepGEMM.git
• 分支: split_mega_moe
• 本地 worktree: pr-357
• HEAD: bb837421b
• 审查方式: 代码review，不跑实验

pr-357 分支包含两段独立改动：

范围	Commit	主题	规模
核心 PR357	bb837421b	Add green-context split-kernel MegaMoE	13 files, +3803/-4
旁路改动	41d89ee6c	Add FP16-weights variant of FP8 MQA logits kernel	7 files, +668/-9

核心工作是 SM100 FP8/FP4 MegaMoE split pipeline，不是 SM90/Hopper 实现。它把原本 fused MegaMoE 拆成 dispatch_l1_swiglu、l2_combine、combine_reduce 三个 kernel，并用 CUDA Runtime 13.1 green context 将 K1/K2 放到互斥 SM 分区中并发执行。

---

实现概述

K1: dispatch_l1_swiglu (dispatch + L1 GEMM + SwiGLU)

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/dispatch_l1_swiglu.cuh (1013 lines)

K1 将 fused kernel 的 dispatch + L1 + SwiGLU 各阶段整合到一个 kernel 中。

warp 分工:

Warp Index	数量	Threads	Role	所属 Warpgroup
0..kNumDispatchWarps-1	4 (128 threads)	128	Route-based dispatch: 扫描 topk_idx，写 route count/entry，NVLink pull token+SF	WG0
kNumDispatchWarps	1	32	TMA A (l1_acts) + SFA loader	WG1
kNumDispatchWarps+1	1	32	TMA B (l1_weights) + SFB loader	WG1
kNumDispatchWarps+2	1	32	MMA issue warp (leader CTA only)	WG1
kNumDispatchWarps+3	1	32	Idle (warpgroup 占位)	WG1
剩余	≥4 (≥128 threads)	128–256	Math UMMA + SwiGLU epilogue + FP8 quant + TMA store + L2 arrival mask write	WG2+

寄存器分配 (与 SM100 fused kernel 一致, dispatch=128, TMA=128, epilogue=256):

Role	Reg/thread	Threads	总 Reg
Dispatch	48	128 (4 warps)	6,144
TMA/non-epilogue	40	128 (4 warps)	5,120
Math/epilogue	208	256 (8 warps)	53,248
总计		512 (16 warps)	64,512/64K

关键约束:

• kNumDispatchThreads % 128 == 0 → dispatch 独占整数个 warpgroup
• kNumNonEpilogueThreads == 128 → TMA 恰好一个 warpgroup
• Route-based dispatch: K1 不物化完整 routed token pool，而是在 dispatch warp 中通过 route metadata 边拉取边写 L1 arrival count
• K1 epilogue 写 l2_arrival_mask（release-or），通知 K2 可以开始消费

K2: l2_combine (L2 GEMM + Cross-Rank Combine Scatter)

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/l2_combine.cuh (660 lines)

K2 等待 K1 的 L1 output 就绪，执行 L2 GEMM 并将结果 scatter 到目标 rank 的 combine buffer。K2 无 dispatch warp，线程布局从 K1 config 继承 non_epilogue + epilogue。

warp 分工 (non_epilogue=128, epilogue=256):

Warp Index	数量	Threads	Role	所属 Warpgroup
0	1 warp	32	TMA A (l2_acts) + SFA loader	WG0
1	1 warp	32	TMA B (l2_weights) + SFB loader	WG0
2	1 warp	32	MMA issue warp	WG0
3	1 warp	32	Idle (warpgroup 占位)	WG0
4–11	8 warps	256	Math UMMA + BF16 NVLink scatter epilogue	WG1, WG2

精确 warp 统计: 12 warps = 384 threads = 3 warpgroups

寄存器分配:

Role	Reg/thread	Threads	总 Reg
TMA/non-epilogue	40	128 (4 warps)	5,120
Math/epilogue	208	256 (8 warps)	53,248
总计		384 (12 warps)	58,368/64K

关键设计:

• 使用独立的 Kernel2L2Scheduler，等待 expert_recv_count_sum 高位达到 kKernel1SMs * kNumRanks 才开始调度
• 每个 pool block 等待 l2_arrival_mask == expected_mask（由 K1 epilogue 写入）
• K2 结束前执行 NVLink barrier，保证 K3 reduce 时 combine buffer 跨 rank 可见

K3: combine_reduce (最终 Top-K Reduce)

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/combine_reduce.cuh (142 lines)

warp 分工 (kNumThreads=512, JIT 模板实例化时设置，覆盖默认值 256):

Warp Index	数量	Threads	Role
0–15	16 warps	512	Top-K reduce: 读取 top-k combine partials (BF16)，FP32 accumulate，BF16 store

精确 warp 统计: 16 warps = 512 threads

寄存器分配 (轻量 kernel, register 压力低):

Role	Reg/thread	Threads	总 Reg
Reduce	~32	512 (16 warps)	~16,384

K3 不使用 green context，运行在 primary context。它依赖 graph node dependency（等 K1/K2 graph node 结束）和 K2 的 NVLink barrier 保证数据可见。

Green Context 资源划分

Context	Kernel	SM 分配 (默认)	Cluster Dim
Green Context 0	K1 dispatch_l1_swiglu	kernel1_sms = 96	2
Green Context 1	K2 l2_combine	kernel2_sms = 52	2
Primary Context	K3 combine_reduce	无限制	1

---

2. 代码差异总览

2.1 核心 split MegaMoE 文件

文件	作用
deep_gemm/mega/__init__.py	Python API，新增 split buffer 和 SM100FP8FP4MegaMoESplitGraph 包装类
csrc/apis/mega.hpp	C++ binding，新增 split buffer size API 和 pybind graph class
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_split.hpp	JIT runtime、配置推导、green context 创建、CUDA graph 构建
deep_gemm/include/deep_gemm/layout/mega_moe_split.cuh	新增 SplitWorkspace，扩展 route-based dispatch metadata
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/dispatch_l1_swiglu.cuh	K1，路由 dispatch、L1 GEMM、SwiGLU、FP8 quant、L2 arrival mask
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/l2_combine.cuh	K2，等待 K1 output、L2 GEMM、跨 rank combine scatter
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/combine_reduce.cuh	K3，按 top-k 对 combine partial 做最终 reduce
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/common.cuh	split pipeline 的 state tensor offset 定义
csrc/jit/handle.hpp	lazy-load CUDA 13.1 green context 和 graph node API
deep_gemm/include/deep_gemm/common/utils.cuh	新增 rank/token peel helper，用于 route-based pool slot 映射
deep_gemm/include/deep_gemm/ptx/ld_st.cuh	新增 CUDA 13 longlong4_32a 兼容和部分 load/store helper
tests/test_mega_moe_split.py	correctness 和 perf 测试，对 fused kernel 做 bitwise 对比

2.2 旁路 MQA 改动

文件	作用
csrc/apis/attention.hpp	weights.dtype == fp16 时选择新 SM100 FP8 MQA logits kernel
csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp	新增 SM100FP8MQALogitsF16WeightsRuntime
deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits_f16_weights.cuh	FP8 input + FP16 weights 的 SM100 2-CTA logits kernel
tests/test_attention.py	增加 FP16 weights 路径覆盖

---

3. Public API 变化

3.1 Split symmetric buffer

deep_gemm/mega/__init__.py 中 SymmBuffer 增加 split: bool = False。当 split=True 时，buffer sizing 使用 _C.get_symm_buffer_size_for_mega_moe_split。

新增 API：

deep_gemm.get_symm_buffer_for_mega_moe_split(...)

与 fused buffer 暴露相同的 tensor view：

x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf

区别在于 split buffer 的 workspace 头部更大，因为需要 route-based dispatch metadata。输入、pool、combine 区域的语义与 fused 对齐，方便 correctness 对比。

3.2 Weight transform

transform_weights_for_mega_moe 负责把权重变成 SM100 MegaMoE 需要的 layout：

权重	变换
L1 weight	gate/up interleave
L1 scale factor	gate/up interleave，再做 UTCCP 需要的 transpose
L2 weight	不 interleave
L2 scale factor	做 UTCCP transpose

3.3 Split graph object

新增 Python 包装类：

deep_gemm.SM100FP8FP4MegaMoESplitGraph(...)
graph.replay()
graph.get_green_context_ids()

构造函数接受 states、ys、sym_buffers、L1/L2 weights、stats，以及 kernel1_sms、kernel2_sms、reduce_sms 等参数。实际 C++ 构造函数会 JIT 编译三个 kernel，创建两个 green context，并构建一个 CUDA graph。

---

4. SplitWorkspace 布局

layout::SplitWorkspace 继承原 fused Workspace，保留 barrier、expert send/recv count、expert recv count sum、L1 arrival count、L2 arrival mask 等公共区域，同时新增 route-based dispatch 区域。

新增区域包括：

区域	用途
src_token_topk_idx[expert][rank][token]	记录每个 local expert 收到的来自哪个 rank/token/topk slot
src_route_count[rank][token]	记录一个 token 在目标 rank 上命中了几条 top-k route
src_route_entry[rank][token][topk]	多 route 时记录 packed route entry
token_src_metadata[pool_token]	K2 scatter 时恢复目标 rank、原 token、topk slot

这个设计把 dispatch 从“先复制完整 routed token pool”改成“写路由元数据，按 expert pool 顺序拉取 token”，使 K1 可以边 dispatch、边 L1、边产生 L2 input。

---

5. CUDA Graph 与 Green Context

5.1 CUDA 13.1 依赖

sm100_fp8_fp4_mega_moe_split.hpp 只在 CUDART_VERSION >= 13010 下启用真实实现。低于 CUDA Runtime 13.1 时，SM100FP8FP4MegaMoESplitGraph 是 stub，构造和 replay 都直接报错。

csrc/jit/handle.hpp lazy-load 了以下 runtime symbol：

API	用途
cudaDeviceGetExecutionCtx	获取 primary execution context
cudaDeviceGetDevResource	查询 SM/workqueue resource
cudaDevSmResourceSplit	把 SM 资源切成两个 group
cudaGreenCtxCreate	创建 green context
cudaGraphAddNode	添加带 context 的 graph node
cudaGraphKernelNodeSetAttribute	设置 cluster dim
cudaGraphLaunch	replay graph

5.2 SM 资源划分

create_green_contexts() 读取 device SM resource，按 kernel1_sms 和 kernel2_sms 切成两个 SM group：

Context	Kernel	资源
green_contexts_[0]	K1 dispatch_l1_swiglu	kernel1_sms
green_contexts_[1]	K2 l2_combine	kernel2_sms
primary_context_	K3 combine_reduce	未 green-context 限制

K1/K2 都用 cluster_dim=2，并设置 coscheduledSmCount=2。workqueue resource 使用 cudaDevWorkqueueConfigScopeGreenCtxBalanced。

5.3 Graph node dependency

build_graph() 的 dependency 关系：

Node chain	Dependency
K1 nodes	多 buffer 时 K1 串行依赖前一个 K1
K2 nodes	多 buffer 时 K2 串行依赖前一个 K2
第一个 K3 node	依赖最后一个 K1 和最后一个 K2
后续 K3 nodes	串行依赖前一个 K3

单 buffer 下语义直接：K1 和 K2 同时启动，K2 在 kernel 内 busy-wait K1 的 arrival mask，K3 等 K1/K2 graph node 都结束后 reduce。

多 buffer 下 K1 chain 和 K2 chain 可以重叠，但 K3 等所有 K1/K2 chain 结束后再开始，不是 per-buffer K3 逐个接在各自 K1/K2 后面。

---

6. 三段 Kernel 实现

6.1 K1: dispatch_l1_swiglu

入口：

sm100_fp8_fp4_mega_moe_split_dispatch_l1_swiglu_impl(...)

K1 同时承担三件事：

阶段	行为
Dispatch metadata	扫描 topk_idx，写 expert send count、route count、route entry、remote src_token_topk_idx
L1 GEMM	按 expert pool 顺序拉取输入 token 和 SF，做 L1 FP8/FP4 UMMA
SwiGLU/quant	epilogue 中做 clamp、SwiGLU、乘 top-k weight、量化成 FP8，写 l2_acts 和 l2_acts_sf

关键机制：

机制	说明
route-based pull	dispatch warp 通过 route metadata 拉取 remote input token，不提前物化全部 token pool
duplicate route 去重	如果一个 token 对同一目标 rank 有多个 top-k route，用 route_count 和 packed route entry 复用一次 remote pull
expert count 同步	本地写 send count，跨 rank 用 atomic_add_sys 汇总 recv count
NVLink barrier	dispatch metadata 完成后，通过 comm::nvlink_barrier 保证各 rank 可见
L1 arrival count	每个 pool block 的 token 写入后，l1_arrival_count release-add，L1 GEMM warp 等待完整 block
L2 arrival mask	K1 epilogue 完成一个 L1 output block 后，对 l2_arrival_mask 做 release-or，通知 K2 可消费

K1 的 shared memory 与 fused kernel 不同。JIT 侧新增 get_mega_moe_split_kernel1_pipeline()，把 dispatch send buffer、expert count、CD staging、barrier、TMA stage 全部纳入 SMEM sizing，避免复用 fused heuristic 导致 SMEM 不足。

6.2 K2: l2_combine

入口：

sm100_fp8_fp4_mega_moe_split_l2_combine_impl(...)

K2 只做 L2 和 combine scatter：

阶段	行为
Scheduling	Kernel2L2Scheduler 等待 expert_recv_count_sum 的高位达到 kKernel1SMs * kNumRanks
Readiness wait	每个 pool block 等待 l2_arrival_mask == expected_mask
L2 GEMM	TMA load l2_acts/l2_weights 和 scale factors，执行 SM100 MXF8F6F4 UMMA
Combine scatter	epilogue 将 BF16 L2 output 通过 sym_buffer.map 写回目标 rank 的 combine buffer
Before reduce barrier	K2 epilogue 结束后执行 NVLink barrier，保证 K3 reduce 前 combine buffer 已跨 rank 可见

K2 的 scheduler 使用 block_idx += kKernel2SMs 方式分配 block，因此 K2 的工作分区与其 green-context SM 数绑定。

6.3 K3: combine_reduce

入口：

sm100_fp8_fp4_mega_moe_split_combine_reduce_impl(...)

K3 是最简单的一段：

阶段	行为
Grid	一个 CTA 对应一个 original token
Load top-k	读取该 token 的 topk_idx
Reduce	对 combine_token_buffer[topk_slot][token] 中的 BF16 partial 做 FP32 accumulate
Store	cast 回 BF16，写 y[token]

K3 依赖 graph node dependency 和 K2 内部 NVLink barrier 来保证数据可见。它不使用 green context，也没有被 reduce_sms 限制。

---

7. 同步与状态设计

7.1 Workspace 同步

同步对象	Producer	Consumer	作用
expert_send_count	K1 dispatch	K1/K2 scheduler	每个 expert 的发送 token 数
expert_recv_count_sum	K1 dispatch	K1/K2 scheduler	每个 local expert 的总接收 token 数和到达计数
l1_arrival_count	K1 dispatch pull	K1 L1 GEMM	pool block 输入 token 是否完整
l2_arrival_mask	K1 epilogue	K2 L2 GEMM	L1 output 的 N-block 是否完整
token_src_metadata	K1 dispatch pull	K2 combine scatter	scatter 回原 rank/token/topk slot

7.2 state tensor

common.cuh 定义了一个至少 7 个 int32 的 state tensor：

Offset	名称	当前用途
0	K1ReadyTasks	已定义，当前实现中未见核心调度使用
1	K1DoneBlocks	已定义，当前实现中未见核心调度使用
2	K2ClaimCounter	已定义，当前实现中未见核心调度使用
3	K2DoneTasks	K2 epilogue 完成 task 后 atomicAdd
4	K2DoneBlocks	K2 kernel block 完成后 atomicAdd
5	K3DoneElements	K3 每 token 完成后 atomicAdd
6	K2Checksum	已定义，当前实现中未见核心调度使用

当前实现主要依赖 workspace arrival/count 和 graph dependency，state 更像是调试/进度计数和未来扩展预留。

---

8. Correctness 与测试策略

tests/test_mega_moe_split.py 是主要验证入口，本文未运行。

测试逻辑：

步骤	内容
构造输入	FP8 activation、FP4 L1/L2 weights、相同 top-k routing
分配 buffer	fused buffer 和 split buffer 分开分配
跑 fused	调用 deep_gemm.fp8_fp4_mega_moe 得到 reference
跑 split	构造 SM100FP8FP4MegaMoESplitGraph 后 replay()
correctness	torch.equal(y_split, y_fused)，要求 bitwise identical
perf	对 fused CUDA graph 和 split graph 做 best-of-N wall-clock，对比 TFLOPS/HBM/NVLink 粗略指标

默认参数：

参数	默认值
num_processes	8
num_tokens	8192
hidden	7168
intermediate_hidden	3072
num_experts	384
num_topk	6
kernel1_sms	96
kernel2_sms	52
reduce_sms	148

测试注释声称 split pipeline 通常略快于 fused，但这是代码注释，非本文实验结论。

---

9. 代码review发现与潜在问题

9.1 reduce_sms 和 reduce_work_iters 当前未实际生效

Python docstring 说 K3 在 reduce_sms 上运行，但 C++ build_graph() 中 K3 node 使用：

add_kernel_node(primary_context_, kernel3_graph_kernel_, num_tokens_, 512, 0, 1, ...)

具体情况：

参数	当前情况
reduce_sms	构造时保存并做正数检查，但未用于 K3 grid 或 green context
reduce_work_iters	构造时保存，但未传给 K3，也未影响调度
K3 SM 限制	不使用 green context，运行在 primary context

这不一定是 bug，可能是 API 预留或尚未完成的功能。但文档和实现存在不一致，后续使用不要把 reduce_sms 当成有效的 SM partition 控制。

9.2 K1/K2 依赖 busy-wait arrival，异常路径可能死锁

K1/K2 之间的真实数据流依赖 workspace 中的 count/mask，而非 graph dependency。K2 会等待 expert_recv_count_sum 和 l2_arrival_mask。如果 K1 提前失败、route metadata 被破坏、或某个 arrival mask 没有写满，K2 可能长时间 spin 无法退出。

9.3 Split buffer 比 fused buffer 更大，且调用者必须显式 reset

SplitWorkspace 增加了 route metadata 区域，buffer size 大于 fused。测试里每次 replay 前都会：

split_buffer.buffer.zero_()
state.zero_()
fill(split_buffer)

Public graph API 本身不自动 reset state/workspace。调用者如果 replay 多次，需要自行清零并重新填输入，否则 arrival/count/mask 会残留。

9.4 Graph construction 固化 shape 和指针

SM100FP8FP4MegaMoESplitGraph 在构造时固化：

固化对象	后果
num_tokens、hidden、intermediate、experts、topk	不适合动态 shape 直接复用
TMA descriptors	weights/buffer pointer 和 layout 固定
green contexts	生命周期绑定 graph object
graph node params	多 buffer 列表长度固定

动态 token 数或不同权重需要重新构造 graph。

9.5 SM100 only

所有新增 split kernels 都要求 __CUDA_ARCH__ >= 1000，且 graph 实现要求 CUDA Runtime 13.1+。这不是 SM90/H100/H200 路径。

---

10. 与已有 MegaMoE split 方向的关系

pr357 的 split 不是简单的"L1 kernel 后接 L2 kernel"。其关键差异是 green-context concurrency：

维度	常规 split L1/L2	pr357 green-context split
K1/K2 launch	两个 kernel 串行或外部并发	一个 CUDA graph 内两个 green context node
SM partition	默认由调度器抢占/共享	K1/K2 分别绑定固定 SM resource group
K2 等待方式	通常等 L1 kernel 完成	K2 kernel 先启动，按 block 等 K1 arrival mask
数据流	L1 全量完成后 L2 消费	K1 产生 pool block，K2 可逐 block 消费
K3	split 后单独 reduce	graph dependency 等 K1/K2 后 reduce

目标是重叠 dispatch/L1 和 L2/combine，并用固定 SM 切分减少 K1/K2 互相抢占。

---

11. FP16-weight MQA logits 旁路改动

这部分不是 split MegaMoE 主线，但在 pr-357 分支内。

11.1 API dispatch

fp8_fp4_mqa_logits 中 weights.scalar_type() == torch::kFloat16 时选择新 kernel：

sm100_fp8_mqa_logits_f16_weights(...)

约束：

约束	原因
arch_major == 10	只支持 SM100
seq_len % 4 == 0	2-CTA kernel 的 query tiling 没有 per-row bound check
不支持 FP4 Q/KV	代码显式 assert not (is_fp4 and weights_is_f16)
FP16 accumulation	QK score 和 per-head weighted sum 都用 FP16 accumulation，速度快但更容易溢出

11.2 Kernel 结构

sm100_fp8_mqa_logits_f16_weights.cuh 是一个 2-CTA cluster kernel：

组件	说明
Q/KV	FP8 E4M3
weights	FP16
KV scale	FP32
SMEM pipeline	Q、weights、KV、KV scale、KV offsets 多 stage
TMEM	使用 SM100 tensor memory accumulator
Launch	num_specialized_threads=128，num_math_threads=256，cluster_dim=2

这是 attention/MQA 路径优化，不影响 split MegaMoE API。

---

12. 结论

pr357 的核心价值是引入了一个面向 SM100/CUDA 13.1 的 green-context split-kernel MegaMoE execution model。它不是单纯把 fused MegaMoE 拆成三个 kernel，而是让 K1 和 K2 在一个 CUDA graph 中用不同 green context 同时运行，借助 workspace arrival mask 实现 producer/consumer 流水线。

最重要的实现要点：

结论	说明
K1/K2 真正并发	green context 将 kernel1_sms 和 kernel2_sms 切成独立 SM partition
数据流按 pool block 流水化	K1 写 l2_arrival_mask，K2 等 mask 后消费
dispatch 变为 route-based pull	SplitWorkspace 保存 route metadata，减少重复 remote token pull
correctness 目标是 bitwise fused	测试直接要求 torch.equal(y_split, y_fused)
API 仍有预留/不一致	reduce_sms 和 reduce_work_iters 当前未控制 K3
分支含旁路 MQA 优化	FP16-weight MQA logits 是独立 sidecar，不属于 MegaMoE split 主链路