DeepGEMM/tests/test_mega_moe_sm90.py

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