diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 203e39e..0bd32f5 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -16,7 +16,7 @@ using SM90MegaMoEBufferViews = std::tuple< torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, - torch::Tensor>; + torch::Tensor, torch::Tensor>; static int get_token_alignment_for_sm90_mega_moe() { return layout::kLCMCandidateBlockM; @@ -38,7 +38,7 @@ get_symm_buffer_size_for_sm90_mega_moe( const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); const auto bf16_token_layout = layout::Data(hidden * 2); const auto fp8_sf_layout = layout::Data(hidden / 128 * static_cast(sizeof(float)), false); - const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 128 * static_cast(sizeof(float)), false); + const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 64 * static_cast(sizeof(float)), false); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); @@ -141,7 +141,7 @@ get_symm_buffer_size_for_sm90_mega_moe( torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); auto l2_acts_sf = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), - {num_max_padded_sf_pool_tokens, intermediate_hidden / 128}, + {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, {1, num_max_padded_sf_pool_tokens}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); auto expert_recv_count_sum = torch::from_blob( @@ -152,6 +152,10 @@ get_symm_buffer_size_for_sm90_mega_moe( runtime_workspace.get_l1_arrival_count_ptr(), {static_cast(runtime_workspace.num_max_pool_blocks)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); + auto l2_arrival_mask = torch::from_blob( + reinterpret_cast(runtime_workspace.get_l2_arrival_mask_ptr()), + {static_cast(runtime_workspace.num_max_pool_blocks)}, + torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); auto token_src_metadata = torch::from_blob( reinterpret_cast(runtime_workspace.get_token_src_metadata_ptr()), {num_max_pool_tokens, 3}, @@ -163,7 +167,7 @@ get_symm_buffer_size_for_sm90_mega_moe( return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l1_topk_weights, l2_acts, l2_acts_sf, - expert_recv_count_sum, l1_arrival_count, token_src_metadata, + expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata, l1_accum_debug); }; return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; @@ -240,8 +244,9 @@ static void fp8_mega_moe( const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l1_topk_weights, l2_acts, l2_acts_sf, - expert_recv_count_sum, l1_arrival_count, token_src_metadata, + expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata, l1_accum_debug] = slice(sym_buffer); + (void)l2_arrival_mask; // Dispatch into SM90 path DG_HOST_ASSERT(arch_major == 9); diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index 6043cca..611b585 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -135,6 +135,7 @@ static void sm90_fp8_mega_moe( // Make tensormap constexpr int kGranK = 128; + constexpr int kL2ActsGranK = 64; const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, hidden, config.num_max_pool_tokens, config.block_k, config.block_m, @@ -162,7 +163,7 @@ static void sm90_fp8_mega_moe( 128); const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, config.num_padded_sf_pool_tokens, intermediate_hidden, - config.block_m, kGranK, + config.block_m, kL2ActsGranK, 1, 0); const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, intermediate_hidden, num_experts_per_rank * hidden, 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 6bd9d21..e112d7c 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 @@ -89,9 +89,13 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); constexpr uint32_t kNumL1WeightSFGroupsN = L1_SHAPE_N / 128; constexpr uint32_t kNumL1WeightSFGroupsK = L1_SHAPE_K / 128; + constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; + constexpr uint32_t kL2ActsGranK = 64; + constexpr uint32_t kMathBarrierIdx = 2; 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"); + DG_STATIC_ASSERT(L1_OUT_BLOCK_N == kL2ActsGranK, "SM90 Phase 4 expects per-64 L2 activation SF"); const uint32_t thread_idx = threadIdx.x; const uint32_t sm_idx = blockIdx.x; @@ -119,7 +123,7 @@ sm90_fp8_mega_moe_impl(void* y, constexpr auto fp8_token_layout = layout::Data(kHidden); constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast(sizeof(float)), false); constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); - constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast(sizeof(float)), false); + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kL2ActsGranK * static_cast(sizeof(float)), false); constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast(sizeof(int64_t)), false); constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast(sizeof(float)), false); constexpr auto l1_topk_weights_layout = layout::Data(static_cast(sizeof(float)), false); @@ -505,6 +509,8 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t r_1 = r_0 + 8; const auto l1_weights_sf_ptr = reinterpret_cast(l1_weights_sf); auto l1_accum_debug_ptr = reinterpret_cast(l1_accum_debug); + auto l2_acts_ptr = l2_token_buffer.get_base_ptr<__nv_fp8_e4m3>(); + auto l2_acts_sf_ptr = l2_sf_buffer.get_base_ptr(); const auto get_l1_weight_sf_group = [](const uint32_t& interleaved_n) { constexpr uint32_t kInterleaveGran = 8; @@ -540,6 +546,7 @@ sm90_fp8_mega_moe_impl(void* y, if (block_phase != sched::SM90BlockPhase::Linear1) return; + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; const uint32_t valid_m = scheduler.get_valid_m(); float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0}; @@ -594,9 +601,81 @@ sm90_fp8_mega_moe_impl(void* y, } } + const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0; + const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1; + const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0; + const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1; + const float topk_weight_0 = row_0 < valid_m ? + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx_0).get_base_ptr() : 0.0f; + const float topk_weight_1 = row_1 < valid_m ? + *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()) { + gate = fminf(fmaxf(gate, -kActivationClamp), kActivationClamp); + up = fminf(fmaxf(up, -kActivationClamp), kActivationClamp); + } + + const float denom = 1.0f + (kFastMath ? __expf(-gate) : expf(-gate)); + const float silu = gate * (kFastMath ? math::fast_rcp(denom) : 1.0f / denom); + return silu * up * topk_weight; + }; + + float local_amax_0 = 0.0f, local_amax_1 = 0.0f; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) { + const float v_00 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0); + const float v_01 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0); + const float v_10 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1); + const float v_11 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1); + local_amax_0 = fmaxf(local_amax_0, fmaxf(fabsf(v_00), fabsf(v_01))); + local_amax_1 = fmaxf(local_amax_1, fmaxf(fabsf(v_10), fabsf(v_11))); + } + + const float row_amax_0 = math::warp_reduce<4, false>(local_amax_0, math::ReduceMax()); + const float row_amax_1 = math::warp_reduce<4, false>(local_amax_1, math::ReduceMax()); + const float sf_0 = fmaxf(row_amax_0 / 448.0f, 1.0e-12f); + const float sf_1 = fmaxf(row_amax_1 / 448.0f, 1.0e-12f); + const float sf_inv_0 = kFastMath ? math::fast_rcp(sf_0) : 1.0f / sf_0; + const float sf_inv_1 = kFastMath ? math::fast_rcp(sf_1) : 1.0f / sf_1; + + if (col_idx == 0) { + const uint32_t sf_pool_token_idx_0 = pool_block_idx * SF_BLOCK_M + row_0; + const uint32_t sf_pool_token_idx_1 = pool_block_idx * SF_BLOCK_M + row_1; + if (row_0 < valid_m) + l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_0] = sf_0; + if (row_1 < valid_m) + l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_1] = sf_1; + } + + const uint32_t out_n_base = n_block_idx * L1_OUT_BLOCK_N; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) { + const uint32_t out_col = out_n_base + (i / 2) * 8 + col_idx * 2; + if (row_0 < valid_m) { + const float v_0 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0); + const float v_1 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0); + l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_0); + l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_0); + } + if (row_1 < valid_m) { + const float v_0 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1); + const float v_1 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1); + l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_1); + l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_1); + } + } + + __threadfence(); + ptx::sync_aligned(kNumMathThreads, kMathBarrierIdx); + if (math_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu( + workspace.get_l2_arrival_mask_ptr(pool_block_idx), + 1ull << n_block_idx); + } + __syncwarp(); + if (local_expert_idx == 0 and m_block_idx == 0 and n_block_idx == 0) { - const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0; - const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { const uint32_t col = i * 8 + col_idx * 2; diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 428bb7c..c048205 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -68,6 +68,7 @@ class SymmBuffer: self.l2_acts, self.l2_acts_sf, self.expert_recv_count_sum, self.l1_arrival_count, + self.l2_arrival_mask, self.token_src_metadata, self.l1_accum_debug) = buffer_views else: @@ -78,6 +79,7 @@ class SymmBuffer: self.l1_topk_weights = None self.expert_recv_count_sum = None self.l1_arrival_count = None + self.l2_arrival_mask = None self.token_src_metadata = None self.l1_accum_debug = None @@ -96,6 +98,7 @@ class SymmBuffer: self.l2_acts_sf = None self.expert_recv_count_sum = None self.l1_arrival_count = None + self.l2_arrival_mask = None self.token_src_metadata = None self.l1_accum_debug = None diff --git a/megamoe_dev_test_scripts/phase1/interface_smoke.py b/megamoe_dev_test_scripts/phase1/interface_smoke.py index b118456..3d71db1 100644 --- a/megamoe_dev_test_scripts/phase1/interface_smoke.py +++ b/megamoe_dev_test_scripts/phase1/interface_smoke.py @@ -83,7 +83,7 @@ def main() -> None: assert buffer.l1_acts_sf.shape[1] == args.hidden // 128 assert buffer.l1_acts_sf.dtype == torch.float32 assert buffer.l2_acts.shape[1] == args.intermediate_hidden - assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 128 + assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 64 assert buffer.l2_acts_sf.dtype == torch.float32 num_tokens = args.num_tokens diff --git a/megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py b/megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py new file mode 100644 index 0000000..aeae976 --- /dev/null +++ b/megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py @@ -0,0 +1,232 @@ +import argparse +import inspect +import os +import pathlib +import random +import sys +from typing import Tuple + +import torch +import torch.distributed as dist + + +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', '8363')) + + 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 interleaved_to_natural_n(n: torch.Tensor, half_n: int, gran: int = 8) -> torch.Tensor: + pair_group = n // (2 * gran) + offset = n - pair_group * (2 * gran) + gate_n = pair_group * gran + offset + up_n = half_n + pair_group * gran + offset - gran + return torch.where(offset < gran, gate_n, up_n) + + +def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int): + torch.manual_seed(9017 + 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') + 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.625 + 0.125 * n_group + 0.03125 * k_group + + l2_weights_sf = torch.ones( + (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128), + dtype=torch.float32, device='cuda') + return deep_gemm.transform_weights_for_mega_moe( + (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf)) + + +def reference_accum_tile(buffer: deep_gemm.SymmBuffer, + l1_weights: torch.Tensor, + l1_weights_sf: torch.Tensor, + hidden: int, + intermediate_hidden: int, + n_block_idx: int) -> torch.Tensor: + block_m = 128 + block_n = 128 + n_start = n_block_idx * block_n + x_fp8 = buffer.l1_acts[:block_m] + x_sf = buffer.l1_acts_sf[:block_m, :hidden // 128] + w_fp8 = l1_weights[0, n_start:n_start + block_n] + n = torch.arange(n_start, n_start + block_n, device='cuda') + natural_n = interleaved_to_natural_n(n, intermediate_hidden) + n_groups = natural_n // 128 + + expected = torch.zeros((block_m, block_n), dtype=torch.float32, device='cuda') + for k_group in range(hidden // 128): + k_start = k_group * 128 + k_end = k_start + 128 + x = x_fp8[:, k_start:k_end].to(torch.float32) + w = w_fp8[:, k_start:k_end].to(torch.float32) + sfa = x_sf[:, k_group].to(torch.float32) + sfb = l1_weights_sf[0, n_groups, k_group].to(torch.float32) + expected += (x @ w.t()) * sfa[:, None] * sfb[None, :] + return expected + + +def reference_swiglu_block(accum: torch.Tensor, topk_weights: torch.Tensor, + activation_clamp: float = None) -> torch.Tensor: + pieces = [] + for group_idx in range(accum.shape[1] // 16): + base = group_idx * 16 + gate = accum[:, base:base + 8] + up = accum[:, base + 8:base + 16] + if activation_clamp is not None: + gate = gate.clamp(-activation_clamp, activation_clamp) + up = up.clamp(-activation_clamp, activation_clamp) + pieces.append(torch.nn.functional.silu(gate) * up) + return torch.cat(pieces, dim=1) * topk_weights[:, None] + + +def verify_l1_epilogue(buffer: deep_gemm.SymmBuffer, + l1_weights: torch.Tensor, + l1_weights_sf: torch.Tensor, + hidden: int, + intermediate_hidden: int, + activation_clamp: float, + atol: float, + rtol: float) -> None: + block_m = 128 + num_l1_n_blocks = intermediate_hidden // 64 + topk_weights = buffer.l1_topk_weights[:block_m].to(torch.float32) + + for n_block_idx in range(num_l1_n_blocks): + accum = reference_accum_tile(buffer, l1_weights, l1_weights_sf, + hidden, intermediate_hidden, n_block_idx) + ref = reference_swiglu_block(accum, topk_weights, activation_clamp) + ref_sf = (ref.abs().amax(dim=1).clamp(min=1e-12) / 448.0).to(torch.float32) + ref_fp8 = (ref / ref_sf[:, None]).to(torch.float8_e4m3fn) + ref_dequant = ref_fp8.to(torch.float32) * ref_sf[:, None] + + col_start = n_block_idx * 64 + col_end = col_start + 64 + actual_sf = buffer.l2_acts_sf[:block_m, n_block_idx].to(torch.float32) + actual_dequant = buffer.l2_acts[:block_m, col_start:col_end].to(torch.float32) * actual_sf[:, None] + + torch.testing.assert_close(actual_sf.cpu(), ref_sf.cpu(), rtol=1e-3, atol=5e-6) + + diff = (actual_dequant - ref_dequant).abs() + base_tol = torch.maximum(torch.full_like(diff, atol), ref_dequant.abs() * rtol) + fp8_step_tol = torch.maximum(actual_sf, ref_sf)[:, None] * 32.0 + tol = torch.maximum(base_tol, fp8_step_tol + 1e-6) + if not torch.all(diff <= tol): + idx = torch.nonzero(diff > tol, as_tuple=False)[0] + row = int(idx[0].item()) + col = int(idx[1].item()) + raise AssertionError( + f'n_block={n_block_idx}, row={row}, col={col}, ' + f'actual={float(actual_dequant[row, col].item())}, ' + f'ref={float(ref_dequant[row, col].item())}, ' + f'diff={float(diff[row, col].item())}, ' + f'tol={float(tol[row, col].item())}') + + mask = int(buffer.l2_arrival_mask[0].item()) + expected_mask = (1 << num_l1_n_blocks) - 1 + assert mask == expected_mask, (mask, expected_mask) + + +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 = 1 + num_experts = args.num_experts if args.num_experts is not None else num_ranks + num_experts_per_rank = num_experts // num_ranks + assert num_experts % num_ranks == 0 + assert num_experts_per_rank >= 1 + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, args.num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden) + weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden) + + torch.manual_seed(3456 + rank_idx) + 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 + local_global_expert = rank_idx * num_experts_per_rank + topk_idx = torch.full((num_tokens, num_topk), local_global_expert, dtype=torch.long, device='cuda') + topk_weights = torch.linspace(0.75, 1.25, num_tokens, dtype=torch.float32, device='cuda').reshape(num_tokens, 1) + + 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() + + verify_l1_epilogue(buffer, weights[0][0], weights[0][1], hidden, intermediate_hidden, + args.activation_clamp, args.atol, args.rtol) + + dist.barrier(group=group) + if rank_idx == 0: + print('[PASSED] Phase 4 L1 epilogue correctness', flush=True) + buffer.destroy() + + +def main() -> None: + parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 4 L1 epilogue 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('--activation-clamp', type=float, default=None) + parser.add_argument('--local-rank', type=int, default=None) + parser.add_argument('--atol', type=float, default=2e-2) + parser.add_argument('--rtol', type=float, default=5e-2) + 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 + + random.seed(7890 + rank_idx) + run_case(args, group, rank_idx, num_ranks) + dist.destroy_process_group() + + +if __name__ == '__main__': + main()