feat: implement sm90 megamoe phase6 combine

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.

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<float>() : 0.0f;
 
                 const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) {
-                    if constexpr (kActivationClamp != cute::numeric_limits<float>::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<kNumRanks, kNumSMs, kNumMathThreads,
+                                 kDispatchGridSyncIndex, kBeforeCombineReduceBarrierTag>(
+                workspace, sym_buffer, sm_idx, math_thread_idx, math_sync);
+
+            auto y_ptr = reinterpret_cast<nv_bfloat16*>(y);
+            const uint64_t num_output_values = static_cast<uint64_t>(num_tokens) * kHidden;
+            const uint64_t output_stride = static_cast<uint64_t>(kNumSMs) * kNumMathThreads;
+            for (uint64_t elem_idx = static_cast<uint64_t>(sm_idx) * kNumMathThreads + math_thread_idx;
+                 elem_idx < num_output_values;
+                 elem_idx += output_stride) {
+                const uint32_t token_idx = static_cast<uint32_t>(elem_idx / kHidden);
+                const uint32_t hidden_idx = static_cast<uint32_t>(elem_idx - static_cast<uint64_t>(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<int64_t>(__ldg(
+                        input_topk_idx_buffer.get_base_ptr<int64_t>() + token_idx * kNumTopk + topk_slot));
+                    if (expert_idx >= 0 and expert_idx < static_cast<int64_t>(kNumExperts)) {
+                        const auto src_ptr = combine_token_buffer.get_rank_buffer(topk_slot)
+                                             .get_data_buffer(token_idx).get_base_ptr<nv_bfloat16>();
+                        reduced += __bfloat162float(src_ptr[hidden_idx]);
+                    }
+                }
+
+                y_ptr[token_idx * kHidden + hidden_idx] = __float2bfloat16(reduced);
+            }
         }
     }
 #endif

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