From 453fc7b04654f149d7c5b5cd543e83c74a58e925 Mon Sep 17 00:00:00 2001 From: Xinyi Liu <94362768+XinyiLiu577086410@users.noreply.github.com> Date: Thu, 18 Jun 2026 17:40:49 +0800 Subject: [PATCH] feat: implement sm90 megamoe phase6 combine --- csrc/apis/sm90_mega.hpp | 2 + .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 41 ++- .../phase6/end_to_end_correctness.py | 274 ++++++++++++++++++ 3 files changed, 316 insertions(+), 1 deletion(-) create mode 100644 megamoe_dev_test_scripts/phase6/end_to_end_correctness.py diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 6e28683..69958ab 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -217,6 +217,8 @@ static void fp8_mega_moe( DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_); DG_HOST_ASSERT(hidden == hidden_); DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(y.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(y.dim() == 2 and y.size(1) == hidden and y.is_contiguous()); DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); // Check weight SF layout: float, natural MN-major, per-128-N and per-128-K. diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh index 8f16fc8..06ded41 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -94,6 +94,7 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; constexpr uint32_t kL2ActsGranK = 64; constexpr uint32_t kMathBarrierIdx = 2; + constexpr uint32_t kDispatchWithMathBarrierIdx = 3; DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts"); DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size"); DG_STATIC_ASSERT(BLOCK_N == WGMMA::N and BLOCK_K % WGMMA::K == 0, "Invalid WGMMA tile shape"); @@ -118,6 +119,7 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kDispatchGridSyncIndex = 0; constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1; constexpr uint32_t kBeforeDispatchPullBarrierTag = 2; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 3; const auto dispatch_sync = []() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }; @@ -454,6 +456,9 @@ sm90_fp8_mega_moe_impl(void* y, } __syncwarp(); } + + if constexpr (BLOCK_M == 128) + ptx::sync_unaligned(kNumDispatchThreads + kNumMathThreads, kDispatchWithMathBarrierIdx); } else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) { if constexpr (BLOCK_M == 128) { if (warp_idx == kNumDispatchWarps) { @@ -525,6 +530,7 @@ sm90_fp8_mega_moe_impl(void* y, } } else { if constexpr (BLOCK_M == 128) { + const uint32_t math_thread_idx = thread_idx - kNumDispatchThreads - kNumTMAThreads; const uint32_t math_warp_idx = warp_idx - kMathWarpStart; const uint32_t math_wg_idx = math_warp_idx / 4; const uint32_t warp_idx_in_wg = math_warp_idx % 4; @@ -707,7 +713,7 @@ sm90_fp8_mega_moe_impl(void* y, *l1_topk_weights_buffer.get_data_buffer(pool_token_idx_1).get_base_ptr() : 0.0f; const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) { - if constexpr (kActivationClamp != cute::numeric_limits::infinity()) { + if constexpr (kActivationClamp < 1.0e30f) { gate = fminf(fmaxf(gate, -kActivationClamp), kActivationClamp); up = fminf(fmaxf(up, -kActivationClamp), kActivationClamp); } @@ -786,6 +792,39 @@ sm90_fp8_mega_moe_impl(void* y, } } }); + + ptx::sync_unaligned(kNumDispatchThreads + kNumMathThreads, kDispatchWithMathBarrierIdx); + __threadfence_system(); + const auto math_sync = []() { + ptx::sync_aligned(kNumMathThreads, kMathBarrierIdx); + }; + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, math_thread_idx, math_sync); + + auto y_ptr = reinterpret_cast(y); + const uint64_t num_output_values = static_cast(num_tokens) * kHidden; + const uint64_t output_stride = static_cast(kNumSMs) * kNumMathThreads; + for (uint64_t elem_idx = static_cast(sm_idx) * kNumMathThreads + math_thread_idx; + elem_idx < num_output_values; + elem_idx += output_stride) { + const uint32_t token_idx = static_cast(elem_idx / kHidden); + const uint32_t hidden_idx = static_cast(elem_idx - static_cast(token_idx) * kHidden); + float reduced = 0.0f; + + #pragma unroll 1 + for (uint32_t topk_slot = 0; topk_slot < kNumTopk; ++ topk_slot) { + const auto expert_idx = static_cast(__ldg( + input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + topk_slot)); + if (expert_idx >= 0 and expert_idx < static_cast(kNumExperts)) { + const auto src_ptr = combine_token_buffer.get_rank_buffer(topk_slot) + .get_data_buffer(token_idx).get_base_ptr(); + reduced += __bfloat162float(src_ptr[hidden_idx]); + } + } + + y_ptr[token_idx * kHidden + hidden_idx] = __float2bfloat16(reduced); + } } } #endif diff --git a/megamoe_dev_test_scripts/phase6/end_to_end_correctness.py b/megamoe_dev_test_scripts/phase6/end_to_end_correctness.py new file mode 100644 index 0000000..a857b2f --- /dev/null +++ b/megamoe_dev_test_scripts/phase6/end_to_end_correctness.py @@ -0,0 +1,274 @@ +import argparse +import inspect +import os +import pathlib +import sys +from typing import Tuple + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import deep_gemm +from deep_gemm.utils.math import ceil_div + + +def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]: + local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0')) + rank = int(os.environ.get('RANK', '0')) + world_size = int(os.environ.get('WORLD_SIZE', '1')) + master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1') + master_port = int(os.environ.get('MASTER_PORT', '8365')) + + torch.cuda.set_device(local_rank) + sig = inspect.signature(dist.init_process_group) + params = { + 'backend': 'nccl', + 'init_method': f'tcp://{master_addr}:{master_port}', + 'world_size': world_size, + 'rank': rank, + } + if 'device_id' in sig.parameters: + params['device_id'] = torch.device(f'cuda:{local_rank}') + dist.init_process_group(**params) + torch.set_default_device('cuda') + return rank, world_size, dist.new_group(list(range(world_size))) + + +def gather_same_shape(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor: + gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))] + dist.all_gather(gathered, tensor.contiguous(), group=group) + return torch.stack(gathered, dim=0) + + +def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int): + torch.manual_seed(4567 + dist.get_rank()) + l1_weights = (torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + l2_weights = (torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + + l1_weights_sf = torch.empty( + (num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128), + dtype=torch.float32, device='cuda') + l2_weights_sf = torch.empty( + (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128), + dtype=torch.float32, device='cuda') + + rank_term = 0.03125 * dist.get_rank() + for expert in range(num_experts_per_rank): + for n_group in range(l1_weights_sf.shape[1]): + for k_group in range(l1_weights_sf.shape[2]): + l1_weights_sf[expert, n_group, k_group] = ( + 0.5 + rank_term + 0.0625 * expert + 0.03125 * n_group + 0.015625 * k_group) + for n_group in range(l2_weights_sf.shape[1]): + for k_group in range(l2_weights_sf.shape[2]): + l2_weights_sf[expert, n_group, k_group] = ( + 0.625 + rank_term + 0.0625 * expert + 0.03125 * n_group + 0.015625 * k_group) + + transformed = deep_gemm.transform_weights_for_mega_moe( + (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf)) + raw = (l1_weights, l1_weights_sf, l2_weights, l2_weights_sf) + return raw, transformed + + +def make_topk(num_tokens: int, + num_experts: int, + num_topk: int, + rank_idx: int, + iteration: int) -> Tuple[torch.Tensor, torch.Tensor]: + token_idx = torch.arange(num_tokens, dtype=torch.long, device='cuda').unsqueeze(1) + topk_slot = torch.arange(num_topk, dtype=torch.long, device='cuda').unsqueeze(0) + slot_stride = max(1, num_experts // max(1, num_topk)) + topk_idx = (token_idx + rank_idx + iteration + topk_slot * slot_stride) % num_experts + + token_term = (token_idx % 7).to(torch.float32) * 0.025 + slot_term = topk_slot.to(torch.float32) * 0.075 + topk_weights = (0.5 + token_term + slot_term).contiguous() + return topk_idx.contiguous(), topk_weights.contiguous() + + +def dequant_input(x_fp8: torch.Tensor, x_sf: torch.Tensor, hidden: int) -> torch.Tensor: + x = x_fp8.to(torch.float32) + out = torch.empty_like(x) + for k_group in range(hidden // 128): + start = k_group * 128 + end = start + 128 + out[:, start:end] = x[:, start:end] * x_sf[:, k_group].to(torch.float32)[:, None] + return out + + +def scaled_fp8_gemm(a: torch.Tensor, + w: torch.Tensor, + w_sf: torch.Tensor, + n_size: int, + k_size: int) -> torch.Tensor: + out = torch.zeros((a.shape[0], n_size), dtype=torch.float32, device='cuda') + for k_group in range(k_size // 128): + k_start = k_group * 128 + k_end = k_start + 128 + partial = a[:, k_start:k_end] @ w[:, k_start:k_end].t() + for n_group in range(n_size // 128): + n_start = n_group * 128 + n_end = n_start + 128 + out[:, n_start:n_end] += partial[:, n_start:n_end] * w_sf[n_group, k_group].to(torch.float32) + return out + + +def quantize_l2_acts(act: torch.Tensor, intermediate_hidden: int) -> torch.Tensor: + out = torch.empty_like(act) + for sf_group in range(intermediate_hidden // 64): + start = sf_group * 64 + end = start + 64 + chunk = act[:, start:end] + sf = (chunk.abs().amax(dim=1).clamp(min=1e-12) / 448.0).to(torch.float32) + out[:, start:end] = (chunk / sf[:, None]).to(torch.float8_e4m3fn).to(torch.float32) * sf[:, None] + return out + + +def reference_output(x_fp8: torch.Tensor, + x_sf: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + all_l1_weights: torch.Tensor, + all_l1_weights_sf: torch.Tensor, + all_l2_weights: torch.Tensor, + all_l2_weights_sf: torch.Tensor, + hidden: int, + intermediate_hidden: int, + num_experts: int, + num_experts_per_rank: int, + activation_clamp: float) -> torch.Tensor: + num_tokens, num_topk = topk_idx.shape + x = dequant_input(x_fp8, x_sf, hidden) + expected = torch.zeros((num_tokens, hidden), dtype=torch.float32, device='cuda') + + for topk_slot in range(num_topk): + for expert_idx in range(num_experts): + mask = topk_idx[:, topk_slot] == expert_idx + if not bool(mask.any()): + continue + + src_rank = expert_idx // num_experts_per_rank + local_expert = expert_idx - src_rank * num_experts_per_rank + l1_accum = scaled_fp8_gemm( + x[mask], + all_l1_weights[src_rank, local_expert], + all_l1_weights_sf[src_rank, local_expert], + intermediate_hidden * 2, + hidden) + gate = l1_accum[:, :intermediate_hidden] + up = l1_accum[:, intermediate_hidden:] + if activation_clamp is not None: + gate = gate.clamp(-activation_clamp, activation_clamp) + up = up.clamp(-activation_clamp, activation_clamp) + l2_input = F.silu(gate) * up * topk_weights[mask, topk_slot].to(torch.float32)[:, None] + l2_input = quantize_l2_acts(l2_input, intermediate_hidden) + contribution = scaled_fp8_gemm( + l2_input, + all_l2_weights[src_rank, local_expert], + all_l2_weights_sf[src_rank, local_expert], + hidden, + intermediate_hidden) + + expected[mask] += contribution.to(torch.bfloat16).to(torch.float32) + + return expected.to(torch.bfloat16) + + +def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None: + hidden = args.hidden + intermediate_hidden = args.intermediate_hidden + num_tokens = args.num_tokens + num_topk = args.num_topk + num_experts = args.num_experts if args.num_experts is not None else max(num_ranks, num_topk) + assert num_experts % num_ranks == 0 + assert num_topk <= num_experts + assert num_tokens * num_ranks * num_topk / num_experts > 64.5, 'Phase 6 test requires BLOCK_M=128' + num_experts_per_rank = num_experts // num_ranks + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, args.num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden) + raw_weights, weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden) + l1_weights, l1_weights_sf, l2_weights, l2_weights_sf = raw_weights + all_l1_weights = gather_same_shape(l1_weights.to(torch.float32), group) + all_l1_weights_sf = gather_same_shape(l1_weights_sf, group) + all_l2_weights = gather_same_shape(l2_weights.to(torch.float32), group) + all_l2_weights_sf = gather_same_shape(l2_weights_sf, group) + + if rank_idx == 0: + print(f'[Phase 6] ranks={num_ranks}, iterations={args.iterations}, tokens={num_tokens}, ' + f'hidden={hidden}, intermediate={intermediate_hidden}, experts={num_experts}, topk={num_topk}', + flush=True) + + for iteration in range(args.iterations): + torch.manual_seed(8901 + rank_idx * 17 + iteration) + x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') * 0.25 + 0.875 + topk_idx, topk_weights = make_topk(num_tokens, num_experts, num_topk, rank_idx, iteration) + + buffer.x[:num_tokens].copy_(x) + buffer.x_sf[:num_tokens].copy_(x_sf) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + torch.cuda.synchronize() + dist.barrier(group=group) + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer, + activation_clamp=args.activation_clamp, + fast_math=False) + torch.cuda.synchronize() + + expected = reference_output( + x, x_sf, topk_idx, topk_weights, + all_l1_weights, all_l1_weights_sf, + all_l2_weights, all_l2_weights_sf, + hidden, intermediate_hidden, + num_experts, num_experts_per_rank, + args.activation_clamp) + torch.testing.assert_close(y.cpu(), expected.cpu(), rtol=args.rtol, atol=args.atol) + + dist.barrier(group=group) + if rank_idx == 0: + max_diff = (y.to(torch.float32) - expected.to(torch.float32)).abs().max().item() + print(f'[PASSED] iteration={iteration}, max_diff={max_diff:.6f}', flush=True) + + buffer.destroy() + + +def main() -> None: + parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 6 end-to-end correctness') + parser.add_argument('--num-tokens', type=int, default=128) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=384) + parser.add_argument('--hidden', type=int, default=256) + parser.add_argument('--intermediate-hidden', type=int, default=128) + parser.add_argument('--num-experts', type=int, default=None) + parser.add_argument('--num-topk', type=int, default=2) + parser.add_argument('--iterations', type=int, default=2) + parser.add_argument('--activation-clamp', type=float, default=None) + parser.add_argument('--local-rank', type=int, default=None) + parser.add_argument('--atol', type=float, default=5e-2) + parser.add_argument('--rtol', type=float, default=1e-1) + args = parser.parse_args() + + rank_idx, num_ranks, group = init_test_dist(args.local_rank) + assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9 + assert args.num_tokens == 128 + assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0 + assert args.intermediate_hidden % 64 == 0 + + run_case(args, group, rank_idx, num_ranks) + dist.destroy_process_group() + + +if __name__ == '__main__': + main()