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()