Phase 0: SM90 MegaMoE design doc, reference baseline, nsys script

- MEGAMOE_SM90_DESIGN.md: complete design document with finalized decisions (fused single kernel, cooperative + single-WG, dynamic BLOCK_M, etc.) - tests/test_mega_moe_sm90.py: PyTorch FP32/BF16 reference implementation for dispatch → L1 GEMM → SwiGLU → L2 GEMM → combine pipeline - scripts/run_nsys_mega_moe_sm90.sh: nsys profiling wrapper script - megamoe-research-reports/: research analysis of PR304/323/347/352/357/360
2026-06-16 18:01:12 +08:00
parent 714dd1a4a9
commit 062cb160cf
11 changed files with 3645 additions and 0 deletions
--- a/tests/test_mega_moe_sm90.py
+++ b/tests/test_mega_moe_sm90.py
@@ -0,0 +1,214 @@
+import argparse
+import os
+import random
+import sys
+import torch
+import torch.distributed as dist
+from typing import Tuple, Optional
+
+from deep_gemm.utils import per_token_cast_to_fp8
+from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather
+from deep_gemm.testing import bench_kineto
+
+
+def reference_swiglu(gate: torch.Tensor, up: torch.Tensor,
+                     activation_clamp: Optional[float] = None) -> torch.Tensor:
+    if activation_clamp is not None:
+        gate = gate.clamp(-activation_clamp, activation_clamp)
+        up = up.clamp(-activation_clamp, activation_clamp)
+    return torch.nn.functional.silu(gate) * up
+
+
+def reference_fp8_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_max = 448.0
+    amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
+    sf = amax / fp8_max
+    x_scaled = x / sf
+    x_fp8 = x_scaled.to(torch.float8_e4m3fn)
+    return x_fp8, sf.squeeze(-1)
+
+
+def reference_fp8_dequantize(x_fp8: torch.Tensor, sf: torch.Tensor,
+                             sf_weights: torch.Tensor,
+                             per_k_gran: int = 128) -> torch.Tensor:
+    x_f32 = x_fp8.to(torch.float32)
+    k = x_f32.shape[-1]
+    num_groups = k // per_k_gran
+    x_f32 = x_f32.reshape(*x_f32.shape[:-1], num_groups, per_k_gran)
+    sf_a = sf.unsqueeze(-1)
+    sf_b = sf_weights.unsqueeze(-2)
+    x_f32 = x_f32 * sf_a * sf_b
+    return x_f32.reshape(*x_f32.shape[:-2], k)
+
+
+def reference_mega_moe(
+    x_bf16: torch.Tensor,
+    topk_idx: torch.Tensor,
+    topk_weights: torch.Tensor,
+    l1_weights_bf16: torch.Tensor,
+    l2_weights_bf16: torch.Tensor,
+    num_experts_per_rank: int,
+    rank_idx: int,
+    num_ranks: int,
+    activation_clamp: Optional[float] = None,
+) -> torch.Tensor:
+    num_tokens, hidden = x_bf16.shape
+    num_topk = topk_idx.shape[1]
+    intermediate_hidden = l2_weights_bf16.shape[2]
+    y = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device=x_bf16.device)
+
+    local_expert_start = rank_idx * num_experts_per_rank
+    local_expert_end = local_expert_start + num_experts_per_rank
+
+    for token_idx in range(num_tokens):
+        for topk_slot in range(num_topk):
+            expert_idx = topk_idx[token_idx, topk_slot].item()
+            if expert_idx < 0:
+                continue
+
+            if expert_idx < local_expert_start or expert_idx >= local_expert_end:
+                continue
+
+            local_expert = expert_idx - local_expert_start
+            weight = topk_weights[token_idx, topk_slot].item()
+
+            token = x_bf16[token_idx].float()
+
+            w1 = l1_weights_bf16[local_expert].float()
+            gate_up = token @ w1.t()
+            gate = gate_up[:intermediate_hidden]
+            up = gate_up[intermediate_hidden:]
+
+            h = reference_swiglu(gate, up, activation_clamp) * weight
+
+            w2 = l2_weights_bf16[local_expert].float()
+            out = h @ w2.t()
+
+            y[token_idx] += out.to(torch.bfloat16)
+
+    return y
+
+
+def reference_mega_moe_batched(
+    x_bf16: torch.Tensor,
+    topk_idx: torch.Tensor,
+    topk_weights: torch.Tensor,
+    l1_weights_bf16: torch.Tensor,
+    l2_weights_bf16: torch.Tensor,
+    num_experts_per_rank: int,
+    rank_idx: int,
+    num_ranks: int,
+    activation_clamp: Optional[float] = None,
+) -> torch.Tensor:
+    num_tokens, hidden = x_bf16.shape
+    num_topk = topk_idx.shape[1]
+    intermediate_hidden = l2_weights_bf16.shape[2]
+    y = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device=x_bf16.device)
+
+    local_expert_start = rank_idx * num_experts_per_rank
+    local_expert_end = local_expert_start + num_experts_per_rank
+
+    for local_expert in range(num_experts_per_rank):
+        expert_idx = local_expert_start + local_expert
+        mask = (topk_idx == expert_idx)
+        token_indices, topk_slots = mask.nonzero(as_tuple=True)
+
+        if token_indices.numel() == 0:
+            continue
+
+        tokens = x_bf16[token_indices].float()
+        weights = topk_weights[token_indices, topk_slots].unsqueeze(-1)
+
+        w1 = l1_weights_bf16[local_expert].float()
+        gate_up = tokens @ w1.t()
+        gate = gate_up[:, :intermediate_hidden]
+        up = gate_up[:, intermediate_hidden:]
+
+        h = reference_swiglu(gate, up, activation_clamp) * weights
+
+        w2 = l2_weights_bf16[local_expert].float()
+        out = h @ w2.t()
+
+        y.index_add_(0, token_indices, out.to(torch.bfloat16))
+
+    return y
+
+
+def test_correctness(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
+    rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks)
+    torch.manual_seed(rank_idx + 42)
+    random.seed(rank_idx + 42)
+
+    num_tokens = args.num_tokens if args.num_tokens > 0 else 32
+    hidden = args.hidden
+    intermediate_hidden = args.intermediate_hidden
+    num_experts = args.num_experts
+    num_topk = args.num_topk
+    num_experts_per_rank = num_experts // num_ranks
+    activation_clamp = args.activation_clamp
+
+    dist_print(f'[SM90 MegaMoE Test] ranks={num_ranks}, tokens={num_tokens}, '
+               f'hidden={hidden}, intermediate={intermediate_hidden}, '
+               f'experts={num_experts}, topk={num_topk}')
+
+    x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
+    l1_weights = torch.randn(
+        (num_experts_per_rank, intermediate_hidden * 2, hidden),
+        dtype=torch.bfloat16, device='cuda') * 0.01
+    l2_weights = torch.randn(
+        (num_experts_per_rank, hidden, intermediate_hidden),
+        dtype=torch.bfloat16, device='cuda') * 0.01
+
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda')
+    topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
+
+    if args.masked_ratio > 0:
+        rand_mask = torch.rand_like(topk_idx, dtype=torch.float)
+        topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1)
+        topk_weights.masked_fill_(topk_idx < 0, 0)
+
+    all_x = uneven_all_gather(x_bf16, group)
+    all_topk_idx = uneven_all_gather(topk_idx, group)
+    all_topk_weights = uneven_all_gather(topk_weights, group)
+
+    ref_y = reference_mega_moe_batched(
+        all_x, all_topk_idx, all_topk_weights,
+        l1_weights, l2_weights,
+        num_experts_per_rank, rank_idx, num_ranks,
+        activation_clamp
+    )
+
+    ref_y_per_rank = ref_y[:num_tokens]
+
+    dist_print(f'[Reference] y norm: {ref_y_per_rank.float().norm().item():.4f}, '
+               f'y abs max: {ref_y_per_rank.float().abs().max().item():.6f}')
+
+    # TODO: Phase 5+ will add SM90 kernel call here and compare with ref_y_per_rank
+    dist_print('[SM90 MegaMoE] Reference baseline computed successfully. '
+               'Kernel comparison will be added in Phase 5.')
+
+    group.barrier()
+    dist_print('[PASSED] Phase 0 reference baseline test')
+
+
+def main():
+    parser = argparse.ArgumentParser(description='SM90 MegaMoE Test')
+    parser.add_argument('--num-tokens', type=int, default=32)
+    parser.add_argument('--num-max-tokens-per-rank', type=int, default=512)
+    parser.add_argument('--hidden', type=int, default=4096)
+    parser.add_argument('--intermediate-hidden', type=int, default=2048)
+    parser.add_argument('--num-experts', type=int, default=16)
+    parser.add_argument('--num-topk', type=int, default=6)
+    parser.add_argument('--masked-ratio', type=float, default=0.0)
+    parser.add_argument('--activation-clamp', type=float, default=None)
+    parser.add_argument('--local-rank', type=int, default=None)
+    args = parser.parse_args()
+
+    num_local_ranks = int(os.environ.get('LOCAL_WORLD_SIZE', '1'))
+    local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0'))
+
+    test_correctness(local_rank, num_local_ranks, args)
+
+
+if __name__ == '__main__':
+    main()