Align MegaMoE training with upgraded K-grouped TMA and layouts

csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp

@@ -155,9 +155,12 @@ static void sm100_bf16_mega_moe_wgrad_1sm(
     const auto tensor_map_b = make_tma_b_desc(
         cute::UMMA::Major::MN, b, n, pool_rows_b,
         kBlockN, kBlockK, static_cast<int>(b.stride(0)), 1, kSwizzle);
-    const auto tensor_map_d = make_tma_cd_desc(
-        d, m, n, kBlockM, kStoreBlockN,
-        static_cast<int>(d.stride(1)), num_groups, kSwizzleCD);
+    // Upstream K-grouped epilogues address [N, M, group] with a 3D TMA
+    // store. A flattened 2D descriptor compiles but traps on the first store.
+    const auto tensor_map_d = make_tma_3d_desc(
+        d, n, m, num_groups, kStoreBlockN, kBlockM, 1,
+        static_cast<int>(d.stride(1)), static_cast<int>(d.stride(0)),
+        kSwizzleCD);
 
     const SM100BF16GemmRuntime::Args args = {
         .gemm_desc = desc,

tests/test_mega_moe_training.py

@@ -434,21 +434,9 @@ def create_inputs():
                 op=dist.ReduceOp.MAX,
                 group=group)
         num_config_tokens = int(num_config_tokens_tensor.item())
-        expected_tokens_per_expert = (
-            num_config_tokens * num_ranks * num_topk /
-            num_experts)
-        if expected_tokens_per_expert <= 8.5:
-            pool_block_m = 16
-        elif expected_tokens_per_expert <= 16.5:
-            pool_block_m = 32
-        elif expected_tokens_per_expert <= 32.5:
-            pool_block_m = 64
-        elif expected_tokens_per_expert <= 64.5:
-            pool_block_m = 96
-        elif expected_tokens_per_expert <= 96.5:
-            pool_block_m = 128
-        else:
-            pool_block_m = 192
+        pool_block_m = deep_gemm._C.get_block_m_for_mega_moe(
+            num_ranks, num_experts, buffer.num_max_tokens_per_rank,
+            num_config_tokens, num_topk, args.mma_type)
         if args.expect_block_m:
             assert pool_block_m == args.expect_block_m, (
                 f'expected BLOCK_M={args.expect_block_m}, '
@@ -1550,28 +1538,7 @@ def check_native_forward_repeatability(
     def run_bf16_backward_test():
         backward_base_allocated = torch.cuda.memory_allocated()
         torch.cuda.reset_peak_memory_stats()
-        config_num_tokens = torch.tensor(
-            num_tokens, dtype=torch.int32, device='cuda')
-        if num_ranks > 1:
-            dist.all_reduce(
-                config_num_tokens,
-                op=dist.ReduceOp.MAX,
-                group=group)
-        expected_tokens_per_expert = (
-            int(config_num_tokens.item()) *
-            num_ranks * num_topk / num_experts)
-        if expected_tokens_per_expert <= 8.5:
-            block_m = 16
-        elif expected_tokens_per_expert <= 16.5:
-            block_m = 32
-        elif expected_tokens_per_expert <= 32.5:
-            block_m = 64
-        elif expected_tokens_per_expert <= 64.5:
-            block_m = 96
-        elif expected_tokens_per_expert <= 96.5:
-            block_m = 128
-        else:
-            block_m = 192
+        block_m = pool_block_m
 
         all_topk_idx = gather_rank_padded(topk_idx, -1)
         all_topk_weights = gather_rank_padded(topk_weights, 0)
@@ -3194,7 +3161,7 @@ def run_baseline():
         '--expect-block-m',
         type=int,
         default=0,
-        choices=[0, 16, 32, 64, 96, 128, 192],
+        choices=[0, 16, 32, 64, 96, 128, 192, 240],
         help='Assert the rank-uniform forward/backward BLOCK_M selection')
     parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections')
     parser.add_argument('--routing', choices=['random', 'balanced', 'skew', 'extreme'], default='random')

tests/test_mega_moe_wgrad_layout.py

@@ -0,0 +1,31 @@
+"""MegaMoE wgrad must follow the upstream K-grouped 3D output ABI."""
+
+import pytest
+import torch
+
+import deep_gemm
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU")
+@pytest.mark.parametrize("pool_block_m", [16, 32, 64, 96, 128, 192, 240])
+@pytest.mark.parametrize("output_n", [128, 256])
+def test_wgrad_group_dimension(pool_block_m, output_n):
+    if torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("requires SM100/SM103")
+    torch.manual_seed(123)
+    # Include an empty expert and unequal group extents to catch a descriptor
+    # that silently flattens (or aliases) the output group coordinate.
+    sizes = [pool_block_m, 0, 2 * pool_block_m, 3 * pool_block_m]
+    counts = torch.tensor(sizes, dtype=torch.int32, device="cuda")
+    a = torch.randn(sum(sizes), 512, dtype=torch.bfloat16, device="cuda") * 0.1
+    b = torch.randn(sum(sizes), output_n, dtype=torch.bfloat16, device="cuda") * 0.1
+    output = torch.full((4, 512, output_n), float("nan"), dtype=torch.bfloat16, device="cuda")
+    deep_gemm.bf16_mega_moe_backward_w13(output, a, b, counts, pool_block_m)
+    torch.cuda.synchronize()
+    offset = 0
+    for expert, rows in enumerate(sizes):
+        reference = a[offset:offset + rows].float().T @ b[offset:offset + rows].float()
+        assert torch.isfinite(output[expert]).all()
+        torch.testing.assert_close(output[expert].float(), reference, rtol=0.004, atol=0.002)
+        offset += rows
+    assert torch.count_nonzero(output[1]) == 0