3.1 原理#

Symmetric Memory 是 PyTorch 2.x 通过 torch.distributed._symmetric_memory 导出的特性，底层基于 CUDA P2P memory access over NVLink。多个 GPU 将各自的物理 HBM 映射到相同的虚拟地址范围。每个 GPU 可以通过普通的 load/store 指令直接读写其他 GPU 的物理 HBM——无需经过 cudaMemcpy、NCCL 集合通信或 kernel launch。

从 MegaMoE 的 SymBuffer 实现（sym_buffer.cuh）可以看到完整的地址映射逻辑：

1
template <uint32_t kNumRanks>
2
struct SymBuffer {
3
    int64_t base;                    // 本 GPU 在本地虚拟地址空间的基地址
4
    int64_t offsets[kNumMaxRanks];   // offsets[i] = GPU_i.BASE - base
5
    uint32_t rank_idx;
6
};
7

8
template <typename ptr_t>
9
CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {
10
    int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
11
    return *reinterpret_cast<ptr_t*>(&mapped_ptr);
12
}

offsets[i] 存储的是 (GPU_i 本地物理 HBM 的虚拟地址) - (本 GPU 本地物理 HBM 的虚拟地址)。因此 map(local_ptr, dst_rank) 的返回值就是一个直接指向远端 GPU HBM 的有效指针，可以直接 dereference。

在 Python 侧，SymBuffer 的创建通过 symm_mem.empty() 和 symm_mem.rendezvous() 完成（mega/__init__.py）：

1
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
2
self.handle = symm_mem.rendezvous(self.buffer, group=group)

empty() 在所有 GPU 上分配大小相同的 buffer，rendezvous() 交换各 GPU 的地址信息并填充到 SymBuffer 的 offsets 数组中。

3.2 远端访问延迟#

通过 NVLink 进行 P2P load 的单次延迟约 1∼3 μs，远高于本地 HBM 的数百 ns。如果每个 token 的每个元素都通过细粒度 load 从远端获取，延迟将无法接受。

MegaMoE 通过 TMA 批量传输处理这个问题。在 dispatch 阶段拉取 token 数据时，每次 TMA load 传输的是整个 hidden 维度的连续数据（默认 7168 bytes）（sm100_fp8_fp4_mega_moe.cuh#L548）：

1
ptx::tma_load_1d(
2
    pull_buffer.get_base_ptr(),
3
    sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(),
4
                   current_rank_in_expert_idx),
5
    pull_mbarrier, kHidden);

TMA 是 GPU 的硬件 DMA 引擎，异步执行数据搬运。SM 发射 TMA 指令后可以立即切换到其他 warp 继续计算，不需要自旋等待搬运完成。单次 TMA load 的 payload 达到 KB 级别，摊薄了 NVLink 的固定访问延迟。

3.3 Dispatch Pull vs Combine Push#

MegaMoE 在 dispatch 阶段使用 NVLink read（pull），在 combine 阶段使用 NVLink write（push）。

代码注释指出 pull 的优势：read 可以避免 write 的 flag 额外延迟。Pull 模式下，发起方完全控制读取时机——只需要知晓远端 buffer 已写入完成（通过 NVLink barrier 保证），不需要远端显式设置 per-transaction flag。Push 模式下，发送方写入后远端必须通过某种同步机制确认 buffer 可用，增加了一轮协调。

3.4 适用边界#

Symmetric Memory 的底层 NVLink P2P 仅在同一 NVSwitch 域内有效。这意味着 MegaMoE 的通信计算融合限于单机/超节点内部。跨节点的 expert dispatch 需要其他机制（如 RDMA 转发），不在 MegaMoE 的设计范围内。

4.1 线程布局#

Kernel 模板参数定义了三种角色的线程数量（sm100_fp8_fp4_mega_moe.cuh#L21）：

1
template <
2
    uint32_t kNumDispatchThreads,        // Dispatch: 可配置
3
    uint32_t kNumNonEpilogueThreads,     // MMA: 固定 128 (= 4 warps)
4
    uint32_t kNumEpilogueThreads,        // Epilogue: 可配置
5
    ...
6
>

派生常量：

1
kNumDispatchWarps = kNumDispatchThreads / 32;
2
kNumMMANonEpilogueWarps = 128 / 32;  // = 4
3
kNumEpilogueWarps = kNumEpilogueThreads / 32;

三种 warp 的寄存器分配（sm100_fp8_fp4_mega_moe.cuh#L346）：

1
constexpr uint32_t kNumDispatchRegisters = 48;
2
constexpr uint32_t kNumNonEpilogueRegisters = 40;
3
constexpr uint32_t kNumEpilogueRegisters = 208;
4
// 约束：总和 ≤ 64512（Blackwell 单 SM 寄存器文件大小）

Dispatch warp 需要的寄存器最少（48），因为它的主要操作是索引计算和 TMA 发射。MMA warp 只需要 40 个寄存器——矩阵乘法的主要工作状态在 TMEM 中，GPR 仅存储指令描述符和少量控制变量。Epilogue warp 需要 208 个寄存器，因为 SwiGLU 计算、amax reduction、FP8 cast、NVLink 写回等多步操作需要在寄存器中维护大量中间状态。

1
每个 SM 内部的 warp 角色分配：
2

3
┌──────────────────────┬──────────────────┬──────────────────────┐
4
│  Dispatch Warps      │  MMA Warps (N)    │  Epilogue Warps      │
5
│  48 regs/warp        │  40 regs/warp     │  208 regs/warp       │
6
├──────────────────────┼──────────────────┼──────────────────────┤
7
│ Expert token 计数     │ TMA Load Token    │ L1: SwiGLU + FP8     │
8
│ 跨 GPU 写源 token 索引│ TMA Load Weight   │     重量化 + TMA store│
9
│ NVLink Pull Token    │ UMMA Issue        │ L2: BF16 cast +      │
10
│ 信号 l1_arrival       │ 信号 tmem_full     │     NVLink write     │
11
│ Workspace 清理        │                   │ Combine reduce       │
12
└──────────────────────┴──────────────────┴──────────────────────┘

4.2 Dispatch Warps#

Dispatch warp 的入口在 warp_idx < kNumDispatchWarps 分支（sm100_fp8_fp4_mega_moe.cuh#L359）。

Step 1 — 本地计数

Dispatch warp 遍历本 GPU 上的所有 token，对每个 token 的 top-k expert 进行 atomicAdd_block 累加，得到每个 expert 的本地 token 数（L386）。

Step 2 — 跨 GPU 通告 token 数

每个 expert 的 token 数通过 64-bit atomic_add 写入 Symmetric Memory workspace 的 expert_send_count（L393）：

1
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
2
    const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
3
    smem_expert_count[i] = static_cast<uint32_t>(
4
        ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
5
}

64-bit atomic 的高 32 位是「SM 完成计数器」，所有 SM 都写入后其值等于 kNumSMs * kNumRanks，表示所有 GPU 的所有 SM 都已完成该 expert 的 token 数统计。

Step 3 — 写源 token 索引到远端

dispatch warp 再次遍历 token，这次将 (token_idx, topk_idx) 组合写入对应 expert 所在 GPU 的 workspace 区域（L401）：

1
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;

dst_ptr 指向 expert 所在 GPU 的 src_token_topk_idx 数组中的位置，sym_buffer.map() 将其转换为远端可写的指针。

Step 4 — Grid Sync + NVLink Barrier

所有 dispatch warp 完成写入后，执行 grid sync（SM 级同步），然后由 SM 0 的 dispatch warp 通过 red_add_rel_sys 将 per-rank recv count 搬运到 expert 所在 GPU。最后是 NVLink barrier，确保所有 rank 之间的计数信息完全同步。

Step 5 — Pull Token 数据

这是 dispatch 与计算重叠的关键阶段。dispatch warp 遍历本 GPU 上的所有 local expert，对每个 expert 通过 Round-Robin Min-Peeling 算法在所有 src rank 之间均衡选取 token。每选取一个 token，发射一次 TMA load 从远端 pull 该 token 的 hidden 向量到本地 SMEM（L548）：

1
if (cute::elect_one_sync()) {
2
    ptx::tma_load_1d(
3
        pull_buffer.get_base_ptr(),
4
        sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(),
5
                       current_rank_in_expert_idx),
6
        pull_mbarrier, kHidden);
7
}

TMA load 完成后，dispatch warp 将 token 数据从 pull buffer（SMEM）搬到 L1 token buffer（也位于 Symmetric Memory workspace），写入 token 权重和源元数据，然后通过 red_add_rel 递增 l1_arrival_count，通知 MMA warp 该 token block 已就绪（L599）：

1
ptx::red_add_rel(
2
    workspace.get_l1_arrival_count_ptr(
3
        expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);

4.3 MMA Warps — TMA Load 数据 (warp_idx == kNumDispatchWarps)#

该 warp 负责通过 TMA 将 token 激活量（activations）和 scale factor A 从 global memory 搬运到 SMEM，是 MMA 流水线的生产者。通过 scheduler.for_each_block() 遍历所有待计算的 block（L666）。

关键同步逻辑在 block 循环开始处（L682）：

1
if (block_phase == BlockPhase::Linear1) {
2
    const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
3
    const auto expected = scheduler.get_valid_m<false>();
4
    while (ptx::ld_acq(ptr) != expected);
5
}

这里 ld_acq 是 acquire 语义的 load，等待 dispatch warp 的 red_add_rel（release 语义）。expected 是当前 block 中有效 token 的数量（BLOCK_M 或最后一个 block 的余数）。当所有 token 数据到达后，MMA warp 启动 TMA load 将 token 数据搬入 SMEM。

对于 L2 phase，同步逻辑不同——它等待的是 L1 epilogue 的结果（L694）：

1
const uint64_t expected = ((1ull << num_k_blocks) << num_k_blocks) - 1;
2
while (ptx::ld_acq_gpu(ptr) != expected);

因为 BLOCK_K == BLOCK_N，L1 输出的每个 block 被切分成两半（gate 和 up 分别对应 SwiGLU 的两个输入），所以 L2 需要等待 2 * num_k_blocks 个 bit 的 mask 全为 1。((1ull << n) << n) - 1 等价于 (1ull << 2n) - 1，但避免了 n==32 时的未定义行为。

TMA load 完成后，warp 通过 full_barrier.arrive() 通知 MMA issue warp 数据已就绪。

4.4 MMA Warps — TMA Load 权重 (warp_idx == kNumDispatchWarps + 1)#

与上一 warp 对称，负责通过 TMA 将权重（weights）和 scale factor B 从 global memory 搬运到 SMEM（L735）。权重数据不需要等待 dispatch 完成，只需要等待流水线中的 consumer release（empty_barriers[stage_idx]->wait(phase ^ 1)）。

4.5 MMA Warps — Issue (warp_idx == kNumDispatchWarps + 2)#

这是唯一发射 UMMA 指令的 warp，且只有 leader CTA 执行（L778）。2-CTA cluster 中 leader CTA 负责计算，follower CTA 仅提供 TMA load 带宽。

在每个 K-block 迭代中，warp 先等待 TMA load 完成（full_barriers[stage_idx]->wait(phase)），然后通过 UTCCP 将 scale factor 从 SMEM copy 到 TMEM，最后发射 UMMA 指令（L854）：

1
for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) {
2
    const auto runtime_instr_desc =
3
        mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, k, k);
4
    ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma(
5
        b_desc, a_desc, accum_stage_idx * UMMA_N,
6
        k_block_idx > 0 or k > 0,  // clear: 第一个 K-block 需要清零累加器
7
        runtime_instr_desc,
8
        kTmemStartColOfSFB, kTmemStartColOfSFA);
9
}

BLOCK_K 固定为 128，UMMA_K 为 32，所以每个 block 发射 4 次 UMMA 指令。第一条指令的 clear 标志为 1（累加器初始化为 0），后续指令的 clear 标志为 0（累加到已有结果上）。

MMA 发射完成后通过 empty_barrier_arrive() 释放流水线阶段，允许 TMA load warp 继续加载下一批数据。同时，如果当前 block 的所有 K 都已完成，则通过 tmem_full_barriers[accum_stage_idx]->arrive() 和 umma_arrive_multicast_2x1SM 通知 epilogue warp。

4.6 Epilogue Warps#

Epilogue warp 的入口在 warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps（L887）。这是寄存器需求最高的角色（208 regs/warp）。

L1 Epilogue（L941）：

从 TMEM 加载 GEMM 累加结果，通过 SM100_TMEM_LOAD_16dp256b1x 以 16 个 float 为粒度读出，然后执行 SwiGLU（L997）：

1
// silu(x) = x * sigmoid(x), 其中 sigmoid(x) = 1 / (1 + e^(-x))
2
auto gate = __bfloat1622float2(bf16_gate);
3
auto neg_gate_exp = make_float2(
4
    kFastMath ? __expf(-gate.x) : expf(-gate.x),
5
    kFastMath ? __expf(-gate.y) : expf(-gate.y));
6
const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp);
7
gate = {gate.x / denom.x, gate.y / denom.y};
8
swiglu_values[i * 2 + k] = __fmul2_rn(__fmul2_rn(gate, up), weights);

kFastMath 路径使用 __expf()（内联 PTX 的 ex2.approx 指令，吞吐远高于 expf()）和 fast_rcp()（用牛顿迭代逼近倒数），牺牲末位精度换取吞吐。

SwiGLU 之后，epilogue warp 执行：

• amax reduction：跨 warp 求 max-abs，用于计算 FP8 的 scale factor（L1024）
• FP8 cast：将 BF16 结果量化为 FP8 E4M3，写入 SMEM 的 CD buffer（L1055）
• TMA store：将 FP8 结果通过 TMA 写回 Symmetric Memory workspace 中的 L2 输入区域（L1099）
• 通知 L2：通过 red_or_rel_gpu 设置 l2_arrival_mask 的对应 bit（L1118）

L2 Epilogue（L1124）：

与 L1 epilogue 类似，从 TMEM 读出 GEMM 累加结果，转换为 BF16 后写入 SMEM。然后通过 NVLink write 直接将结果写回远端 GPU 的 combine buffer（L1212）：

1
// 读取 dispatch 阶段存储的源元数据
2
const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block);
3
const uint32_t dst_rank_idx = src_metadata.rank_idx;
4
const uint32_t dst_token_idx = src_metadata.token_idx;
5

6
// 直接通过 Symmetric Memory 写入远端 GPU
7
*sym_buffer.map(dst_ptr, dst_rank_idx) = packed;

Combine reduction（L1242）：

所有 L2 epilogue 完成后，epilogue warp 执行最终的 combine——将同一 token 在不同 expert 上的 top-k 结果按路由权重归约。这通过 TMA load 从 combine buffer 将 chunk 数据搬入 SMEM 逐元素累加，再 TMA store 写回最终的 output tensor。

7.1 启发式配置#

启发式配置需要满足三个约束：正确性、共享内存上限、SM 利用率。

BLOCK_M 选择

get_block_config_for_mega_moe 根据每 expert 期望 token 数选择 BLOCK_M，候选值为 {16, 32, 64, 96, 128, 192}。预期 token 数的计算公式：

1
expected_tokens_per_expert = num_tokens × num_ranks × num_topk / num_experts

分子是全局所有 token 乘以 topk 的总路由次数，分母是 expert 总数。实际路由不均匀，因此所有估算中引入了 kImbalanceFactor = 2，把目标工作量放大两倍以吸收热专家的尾延迟。

BLOCK_M 的选择阶梯为：token 很少时用小 BLOCK_M（16 或 32），避免 M 维度大量 padding 浪费计算；token 多时用大 BLOCK_M（128 或 192），提高单 tile 的算术密度。

Pool 容量

pool 容量必须覆盖所有 rank 的所有 token 全部路由到本 rank 的最坏情况（mega_moe.cuh#L17）：

1
num_max_recv_tokens = num_ranks × num_max_tokens_per_rank
2
num_max_experts_per_token = min(num_topk, num_experts_per_rank)
3
pool_size = align(
4
    num_max_recv_tokens × num_max_experts_per_token
5
    + num_experts_per_rank × (kMaxCandidateBlockM - 1),
6
    kLCMCandidateBlockM)

其中 num_max_experts_per_token 是因为一个 token 的 top-k 选择不可能重复落在同一个 expert 上，所以本 rank 最多收到该 token 的 min(num_topk, num_experts_per_rank) 份副本。额外的 num_experts_per_rank × (kMaxCandidateBlockM - 1) 是各 expert 之间 BLOCK_M 对齐填充的上界。最终对齐到 kLCMCandidateBlockM = 384（所有候选 BLOCK_M 的最小公倍数）。

Expert Wave 粒度

get_num_experts_per_wave_for_mega_moe 的核心逻辑是：在保证所有 SM 都有 block 可算的前提下，wave 内 expert 越少越好。原因是 expert 少 → L1 整体完成得快 → L2 启动得早 → pipeline 深度浅。

具体计算：先估计一个 expert 的 L1 block 数（num_m_blocks × num_n_blocks），然后除以 SM 数得到需要的 expert 数下界。用 kImbalanceFactor = 2 放大以吸收路由不均。最后向上取整到 num_experts_per_rank 的因子（scheduler 要求所有 wave 处理相同数量的 expert）。

SMEM 分配与流水线深度

get_pipeline_config_for_mega_moe 估算共享内存的固定开销和单级流水线开销，然后计算 num_stages = (total_smem - fixed) / per_stage。固定开销包括 dispatch 区的 expert 计数器和 send buffer、CD 输出区（L1/L2 epilogue 共享 SMEM 的 CD buffer，取两者 max）、amax reduction 区和所有 mbarrier 对象。单级流水线开销包括 A tile + B tile + SFA + SFB + full_barrier + empty_barrier。

7.2 Wave 划分#

MoE 计算被划分为多个 wave，每个 wave 处理 kNumExpertsPerWave 个 expert，包含两个 phase：L1、L2。

get_num_experts_per_wave_for_mega_moe 的选取原则：在保证所有 SM 都有活可干的前提下，expert 越少越好。原因是：expert 少 → L1 完成得快 → L2 启动得早 → pipeline 深度浅、延迟低。但 expert 太少会导致部分 SM 闲置。因此最优取值是刚好能打满所有 SM 的最小 expert 数。

6.2 块分配#

同一个 expert 内部按先 N 方向（output channel）、后 M 方向（token）遍历所有 block。分配策略是 round-robin：

1
// fetch_next_l1_block 中
2
n_block_idx = block_idx - m_block_idx * kNumL1BlockNs;
3
block_idx += kNumSMs;  // 跳转到下一轮属于本 SM 的 block

所有 SM 从自己的 sm_id 开始，每轮前进 kNumSMs 步。这保证了每个 SM 的工作量均衡。

6.3 状态机#

核心状态机位于 for_each_block 调用的 get_next_block()（L148）：

1
L1 phase: 遍历当前 wave 所有 expert 的 L1 block
2
  → L1 完成 → 切换到 L2 phase
3
L2 phase: 遍历当前 wave 所有 expert 的 L2 block
4
  → L2 完成 → 切换到下一个 wave 的 L1 phase
5
所有 wave 完成 → BlockPhase::None → kernel 结束

状态机确保了一个关键性质：同一 wave 内，所有 SM 先共同完成所有 expert 的 L1，再切换到 L2。这意味着 L2 开始时大部分 L1 结果已就绪，L2 MMA warp 的自旋等待时间被最小化。

7.1 精度配置#

FP4 E2M1 格式：1-bit 符号、2-bit 指数、1-bit 尾数，共 16 个可能值（包含 ±0、±NaN 等特殊值）。有效范围约为 [0.5, 7.5]（不含非正规数）。

UE8M0 格式：8-bit 无符号整数指数，0-bit 尾数，值为 2^e。4 个 UE8M0 打包为一个 uint32_t。

7.2 Block Scaling#

MegaMoE 的 scaling 是 block-wise 而非 per-tensor 的。对于权重，每个 32×32 的 block 有自己的 scale factor。对于激活，每 32 个元素（在 K 维度上）有一个 scale factor。

这种粒度避免了 per-tensor scaling 的精度问题：不同 expert 的权重分布可能差异很大，per-tensor scaling 会导致某些 expert 的权重在量化后丧失精度。Block-wise scaling 以存储少量额外 scale 为代价，保证了各 expert 的精度。

Block scaling 的 UMMA 指令语义为 D = (A * scale_A) * (B * scale_B) + D。硬件在 MMA 过程中自动应用 scale factor，不增加额外的指令开销。

SF 存储在 TMEM 中，与累加器共享同一块 TMEM 区域（L217）：

1
constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages;
2
constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32;
3
constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32;

7.3 L1 权重 Interleave#

L1 权重（gate + up projection）使用 interleave 布局，而非简单的 [gate | up] 拼接（mega/__init__.py#L75）：

1
# [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...]

这是为了配合 SwiGLU epilogue 的 gate/up pair 消费模式。TMEM load 指令 SM100_TMEM_LOAD_16dp256b1x 一次读出 8 个值，排列为 (0,2) (1,3) (4,6) (5,7) 的 gate/up pair。Interleave 布局保证了 gate 和对应位置的 up 在 TMEM 中相邻，简化了 SwiGLU 的索引计算。