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

feat: implement sm90 megamoe phase5 l2 scatter

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 15:17:20 +08:00
parent fc8218750c
commit 9bd0519605
5 changed files with 379 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
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, torch::Tensor, torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
@@ -164,11 +164,15 @@ get_symm_buffer_size_for_sm90_mega_moe(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_accum_debug_buffer.base)),
{128, 128},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto combine_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(combine_token_buffer.base)),
{num_topk, num_max_tokens_per_rank, hidden},
torch::TensorOptions().dtype(torch::kBFloat16).device(buffer.device()));
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, l2_arrival_mask, token_src_metadata,
l1_accum_debug);
l1_accum_debug, combine_acts);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
@@ -245,8 +249,9 @@ static void fp8_mega_moe(
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata,
l1_accum_debug] = slice(sym_buffer);
l1_accum_debug, combine_acts] = slice(sym_buffer);
(void)l2_arrival_mask;
(void)combine_acts;
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);

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

@@ -89,6 +89,8 @@ 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 kNumL2WeightSFGroupsN = L2_SHAPE_N / 128;
constexpr uint32_t kNumL2WeightSFGroupsK = L2_SHAPE_K / 128;
constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2;
constexpr uint32_t kL2ActsGranK = 64;
constexpr uint32_t kMathBarrierIdx = 2;
@@ -106,6 +108,9 @@ sm90_fp8_mega_moe_impl(void* y,
cute::prefetch_tma_descriptor(&tensor_map_l1_acts);
cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf);
cute::prefetch_tma_descriptor(&tensor_map_l1_weights);
cute::prefetch_tma_descriptor(&tensor_map_l2_acts);
cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf);
cute::prefetch_tma_descriptor(&tensor_map_l2_weights);
}
__syncwarp();
@@ -121,6 +126,7 @@ sm90_fp8_mega_moe_impl(void* y,
sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);
constexpr auto fp8_token_layout = layout::Data(kHidden);
constexpr auto bf16_token_layout = layout::Data(kHidden * static_cast<uint32_t>(sizeof(nv_bfloat16)));
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 / kL2ActsGranK * static_cast<uint32_t>(sizeof(float)), false);
@@ -156,7 +162,13 @@ sm90_fp8_mega_moe_impl(void* y,
const auto l2_sf_buffer = layout::Buffer(
fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens,
l2_token_buffer.get_end_ptr());
(void)l2_sf_buffer;
constexpr auto l1_accum_debug_layout = layout::Data(128 * static_cast<uint32_t>(sizeof(float)), false);
const auto l1_accum_debug_buffer = layout::Buffer(
l1_accum_debug_layout, 1, 128,
l2_sf_buffer.get_end_ptr());
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, kNumTopk, kNumMaxTokensPerRank,
l1_accum_debug_buffer.get_end_ptr());
constexpr uint32_t kSharedMemoryAlignment = 1024;
extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];
@@ -466,13 +478,16 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
if (block_phase != sched::SM90BlockPhase::Linear1)
return;
const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
if (block_phase == sched::SM90BlockPhase::Linear1) {
const uint32_t valid_m = scheduler.get_valid_m();
const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
while (ptx::ld_acq(arrival_ptr) != valid_m);
} else {
const auto arrival_ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx);
const uint64_t expected = ((1ull << num_k_blocks) << num_k_blocks) - 1;
while (ptx::ld_acq_gpu(arrival_ptr) != expected);
}
#pragma unroll 1
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
@@ -482,6 +497,7 @@ sm90_fp8_mega_moe_impl(void* y,
auto& full_barrier = *full_barriers[stage_idx];
const uint32_t k_idx = k_block_idx * BLOCK_K;
const uint32_t m_idx = pool_block_idx * BLOCK_M;
if (block_phase == sched::SM90BlockPhase::Linear1) {
const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N;
const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
@@ -492,6 +508,15 @@ sm90_fp8_mega_moe_impl(void* y,
&tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
} else {
const uint32_t n_idx = local_expert_idx * L2_SHAPE_N + n_block_idx * BLOCK_N;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
&tensor_map_l2_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
&tensor_map_l2_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
}
}
__syncwarp();
}
@@ -508,6 +533,7 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx;
const uint32_t r_1 = r_0 + 8;
const auto l1_weights_sf_ptr = reinterpret_cast<const float*>(l1_weights_sf);
const auto l2_weights_sf_ptr = reinterpret_cast<const float*>(l2_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>();
@@ -543,11 +569,12 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
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();
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;
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
const auto empty_barrier_arrive = [&](const uint32_t& s) {
@@ -559,6 +586,7 @@ sm90_fp8_mega_moe_impl(void* y,
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
full_barriers[stage_idx]->wait(phase);
if (block_phase == sched::SM90BlockPhase::Linear1) {
const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0);
const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1);
@@ -599,12 +627,80 @@ sm90_fp8_mega_moe_impl(void* y,
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
}
} else {
#pragma unroll
for (uint32_t half = 0; half < 2; ++ half) {
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_arrive();
#pragma unroll
for (uint32_t k = 0; k < 2; ++ k) {
const uint32_t smem_k = (half * 2 + k) * WGMMA::K;
auto desc_a = mma::sm90::make_smem_desc(
smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + smem_k, 1);
auto desc_b = mma::sm90::make_smem_desc(
smem_b[stage_idx] + smem_k, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
}
ptx::warpgroup_commit_batch();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_wait<0>();
const uint32_t sf_group = k_block_idx * 2 + half;
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;
const float scale_a_0 = __ldg(
l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_0);
const float scale_a_1 = __ldg(
l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_1);
const auto sf_base = l2_weights_sf_ptr +
local_expert_idx * kNumL2WeightSFGroupsN * kNumL2WeightSFGroupsK;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col_0 = i * 8 + col_idx * 2;
const uint32_t col_1 = col_0 + 1;
const uint32_t n_group_0 = (n_block_idx * BLOCK_N + col_0) / 128;
const uint32_t n_group_1 = (n_block_idx * BLOCK_N + col_1) / 128;
const float scale_b_0 = __ldg(sf_base + n_group_0 * kNumL2WeightSFGroupsK + k_block_idx);
const float scale_b_1 = __ldg(sf_base + n_group_1 * kNumL2WeightSFGroupsK + k_block_idx);
final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
}
}
empty_barrier_arrive(stage_idx);
}
}
if (block_phase == sched::SM90BlockPhase::Linear2) {
const auto scatter_row = [&](const uint32_t& row, const uint32_t& pool_token_idx, const uint32_t& accum_offset) {
if (row >= valid_m)
return;
const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_token_idx);
const auto dst_token = combine_token_buffer.get_rank_buffer(src_metadata.topk_idx)
.get_data_buffer(src_metadata.token_idx);
auto dst_ptr = dst_token.get_base_ptr<nv_bfloat16>();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t out_col = n_block_idx * BLOCK_N + i * 8 + col_idx * 2;
*sym_buffer.map(dst_ptr + out_col + 0, src_metadata.rank_idx) =
__float2bfloat16(final_accum[i * 4 + accum_offset + 0]);
*sym_buffer.map(dst_ptr + out_col + 1, src_metadata.rank_idx) =
__float2bfloat16(final_accum[i * 4 + accum_offset + 1]);
}
};
scatter_row(row_0, pool_token_idx_0, 0);
scatter_row(row_1, pool_token_idx_1, 2);
return;
}
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 ?

5
deep_gemm/mega/__init__.py
View File

@@ -70,7 +70,8 @@ class SymmBuffer:
self.l1_arrival_count,
self.l2_arrival_mask,
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
self.l1_accum_debug,
self.combine_acts) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
@@ -82,6 +83,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def destroy(self):
self.handle = None
@@ -101,6 +103,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,

2
megamoe_dev_test_scripts/phase1/interface_smoke.py
View File

@@ -85,6 +85,8 @@ def main() -> None:
assert buffer.l2_acts.shape[1] == args.intermediate_hidden
assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 64
assert buffer.l2_acts_sf.dtype == torch.float32
assert buffer.combine_acts.shape == (args.num_topk, buffer.num_max_tokens_per_rank, args.hidden)
assert buffer.combine_acts.dtype == torch.bfloat16
num_tokens = args.num_tokens
buffer.x[:num_tokens].copy_(torch.randn((num_tokens, args.hidden), device='cuda').to(torch.float8_e4m3fn))

201
megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py Normal file
View File

@@ -0,0 +1,201 @@
import argparse
import inspect
import os
import pathlib
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', '8364'))
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 make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
torch.manual_seed(2345 + 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')
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 + 0.125 * expert + 0.0625 * n_group + 0.03125 * 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.75 + 0.125 * expert + 0.0625 * n_group + 0.03125 * k_group
return deep_gemm.transform_weights_for_mega_moe(
(l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
def dequant_l2_acts(buffer: deep_gemm.SymmBuffer,
start: int,
count: int,
intermediate_hidden: int) -> torch.Tensor:
x = buffer.l2_acts[start:start + count, :intermediate_hidden].to(torch.float32)
out = torch.empty_like(x)
for sf_group in range(intermediate_hidden // 64):
col_start = sf_group * 64
col_end = col_start + 64
sf = buffer.l2_acts_sf[start:start + count, sf_group].to(torch.float32)
out[:, col_start:col_end] = x[:, col_start:col_end] * sf[:, None]
return out
def verify_l2_scatter(buffer: deep_gemm.SymmBuffer,
l2_weights: torch.Tensor,
l2_weights_sf: torch.Tensor,
hidden: int,
intermediate_hidden: int,
num_tokens: int,
atol: float,
rtol: float) -> None:
block_m = 128
ref = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
counts = [int(v) & 0xffffffff for v in buffer.expert_recv_count_sum.to(torch.int64).cpu().tolist()]
pool_block_offset = 0
expected_mask = (1 << (intermediate_hidden // 64)) - 1
for expert, count in enumerate(counts):
if count == 0:
continue
pool_start = pool_block_offset * block_m
x = dequant_l2_acts(buffer, pool_start, count, intermediate_hidden)
expected = torch.zeros((count, hidden), dtype=torch.float32, device='cuda')
for k_group in range(intermediate_hidden // 128):
k_start = k_group * 128
k_end = k_start + 128
w = l2_weights[expert, :, k_start:k_end].to(torch.float32)
partial = x[:, k_start:k_end] @ w.t()
for n_group in range(hidden // 128):
n_start = n_group * 128
n_end = n_start + 128
sfb = l2_weights_sf[expert, n_group, k_group].to(torch.float32)
expected[:, n_start:n_end] += partial[:, n_start:n_end] * sfb
metadata = buffer.token_src_metadata[pool_start:pool_start + count].to(torch.int64)
for row in range(count):
rank_idx, token_idx, topk_idx = [int(v) for v in metadata[row].tolist()]
assert rank_idx == dist.get_rank()
assert topk_idx == 0
ref[token_idx] = expected[row].to(torch.bfloat16)
for block in range(ceil_div(count, block_m)):
mask = int(buffer.l2_arrival_mask[pool_block_offset + block].item())
assert mask == expected_mask, (expert, block, mask, expected_mask)
pool_block_offset += ceil_div(count, block_m)
actual = buffer.combine_acts[0, :num_tokens, :hidden]
torch.testing.assert_close(actual.cpu(), ref.cpu(), rtol=rtol, atol=atol)
def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None:
assert num_ranks == 1, 'Phase 5 milestone verifies multi-expert single-rank scatter first'
hidden = args.hidden
intermediate_hidden = args.intermediate_hidden
num_tokens = args.num_tokens
num_topk = 1
num_experts = args.num_experts
num_experts_per_rank = num_experts // num_ranks
assert num_experts % num_ranks == 0
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(6789 + 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
topk_idx = (torch.arange(num_tokens, dtype=torch.long, device='cuda') // 128).reshape(num_tokens, 1)
topk_idx = torch.clamp(topk_idx, max=num_experts - 1)
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_l2_scatter(buffer, weights[1][0], weights[1][1], hidden, intermediate_hidden,
num_tokens, args.atol, args.rtol)
dist.barrier(group=group)
if rank_idx == 0:
print('[PASSED] Phase 5 L2 GEMM scatter correctness', flush=True)
buffer.destroy()
def main() -> None:
parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 5 L2 GEMM scatter correctness')
parser.add_argument('--num-tokens', type=int, default=256)
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=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=3e-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 == 256
assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
run_case(args, group, rank_idx, num_ranks)
dist.destroy_process_group()
if __name__ == '__main__':
main()