chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement sm90 megamoe phase4 l1 epilogue

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 01:09:45 +08:00
parent 2bb1756787
commit f3553f976c
6 changed files with 330 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
csrc/apis/sm90_mega.hpp
View File

@@ -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<int>(sizeof(float)), false);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 128 * static_cast<int>(sizeof(float)), false);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 64 * static_cast<int>(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<int64_t>(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<int>(runtime_workspace.num_max_pool_blocks)},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l2_arrival_mask = torch::from_blob(
reinterpret_cast<int64_t*>(runtime_workspace.get_l2_arrival_mask_ptr()),
{static_cast<int>(runtime_workspace.num_max_pool_blocks)},
torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
auto token_src_metadata = torch::from_blob(
reinterpret_cast<int32_t*>(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<int64_t>(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);

3
csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp
View File

@@ -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,

85
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
View File

@@ -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<uint32_t>(sizeof(float)), false);
constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kL2ActsGranK * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(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<const float*>(l1_weights_sf);
auto l1_accum_debug_ptr = reinterpret_cast<float*>(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<float>();
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<float>() : 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<float>() : 0.0f;
const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) {
if constexpr (kActivationClamp != cute::numeric_limits<float>::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<float>());
const float row_amax_1 = math::warp_reduce<4, false>(local_amax_1, math::ReduceMax<float>());
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;

3
deep_gemm/mega/__init__.py
View File

@@ -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

2
megamoe_dev_test_scripts/phase1/interface_smoke.py
View File

@@ -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

232
megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py Normal file
View File

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