062cb160cf
- 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
215 lines
7.6 KiB
Python
215 lines
7.6 KiB
Python
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()
|