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

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