# 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：

```python
deep_gemm.get_symm_buffer_for_mega_moe_split(...)
```

与 fused buffer 暴露相同的 tensor view：

```python
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 包装类：

```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`

入口：

```cpp
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`

入口：

```cpp
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`

入口：

```cpp
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 使用：

```cpp
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 前都会：

```python
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：

```cpp
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 主链路 |