Align MegaMoE training with upgraded K-grouped TMA and layouts

csrc/jit_kernels/impls/sm100_bf16_gemm.hpp

@@ -52,6 +52,8 @@ class SM100BF16GemmRuntime final {
         bool combine_reduce = false;
         std::string combine_order_mode = "fixed_topk";
         uint32_t combine_num_extra_threads = 0;
+        bool mask_grouped_k_tail = false;
+        int k_alignment_override = 0;
     };
 
     static void compile_and_launch(const std::string& tag, const Args& args) {
@@ -76,7 +78,7 @@ static void __instantiate_kernel() {{
         {}, {}, {},
         {},
         {},
-        {}, {}, {}, {}
+        {}, {}, {}, {}, {}
     >);
 }};
 )",
@@ -91,15 +93,16 @@ static void __instantiate_kernel() {{
         args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
         args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
         args.gemm_config.launch_config.num_sms,
-        heuristics_runtime->get_mk_alignment_for_contiguous_layout(),
+        args.k_alignment_override != 0 ? args.k_alignment_override :
+            heuristics_runtime->get_mk_alignment_for_contiguous_layout(),
         args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding,
         to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
         to_string(args.gemm_desc.cd_dtype),
         args.epilogue.type,
         args.gemm_desc.tc_util,
         args.combine_num_ranks, args.fuse_combine,
         get_bf16_gemm_combine_order_mode_name(args.combine_order_mode),
-        args.combine_num_extra_threads));
+        args.combine_num_extra_threads, args.mask_grouped_k_tail));
 
         // Launch
         jit->launch(

csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp

@@ -57,23 +57,21 @@ static void sm100_bf16_mega_moe_wgrad_1sm(
     const int kBlockN = deep_jit::get_env<int>(
         "DG_BF16_MEGA_MOE_WGRAD_BLOCK_N",
         n % 256 == 0 ? 256 : 128);
-    // The K-grouped scheduler addresses each expert in the shared physical
-    // pool. Its K tile must divide the forward pool alignment; otherwise the
-    // final tile of one expert reads rows from the next expert. In particular,
-    // BLOCK_M=96 previously contaminated Qwen top-8 wgrads while BLOCK_M=128
-    // happened to pass.
+    // Preserve the physical expert offsets while amortizing TMA/barrier cost
+    // over K=64. For pools such as 96/240, mask whole 16-wide MMA atoms in the
+    // final tile rather than accumulating rows from the following expert.
     const int kBlockK = deep_jit::get_env<int>(
-        "DG_BF16_MEGA_MOE_WGRAD_BLOCK_K",
-        pool_block_m % 64 == 0 ? 64 :
-        pool_block_m % 32 == 0 ? 32 : 16);
+        "DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", 64);
+    DG_HOST_ASSERT(kBlockK == 16 or kBlockK == 32 or kBlockK == 64);
+    DG_HOST_ASSERT(pool_block_m % 16 == 0);
+    const bool mask_grouped_k_tail = pool_block_m % kBlockK != 0;
     const int kNumStages = deep_jit::get_env<int>(
         "DG_BF16_MEGA_MOE_WGRAD_NUM_STAGES", 4);
     const int kSwizzle =
         kBlockK * static_cast<int>(sizeof(cutlass::bfloat16_t));
     const int kStoreBlockN = deep_jit::get_env<int>(
         "DG_BF16_MEGA_MOE_WGRAD_STORE_BLOCK_N", 64);
-    // The output store width is independent of the A/B K tile. In particular,
-    // a 240-row pool requires K=16 but still uses the default 64-column store.
+    // The output store width is independent of the A/B K tile.
     const int kSwizzleCD = kStoreBlockN * sizeof(cutlass::bfloat16_t);
     DG_HOST_ASSERT(kSwizzleCD == 32 or kSwizzleCD == 64 or kSwizzleCD == 128);
     constexpr int kNumNonEpilogueThreads = 128;
@@ -155,9 +153,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,
@@ -184,6 +185,11 @@ static void sm100_bf16_mega_moe_wgrad_1sm(
         .combine_order_mode = combine.order_mode,
         .combine_num_extra_threads =
             static_cast<uint32_t>(num_extra_combine_threads),
+        .mask_grouped_k_tail = mask_grouped_k_tail,
+        // Count-based BF16 groups use explicit physical offsets and no SF.
+        // Do not inherit unrelated global grouped-GEMM alignment (e.g. 224),
+        // which need not be divisible by this kernel's K tile.
+        .k_alignment_override = kBlockK,
     };
     SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_mega_moe_wgrad_1sm", args);
 }

deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh

@@ -37,7 +37,8 @@ template <cute::UMMA::Major kMajorA, cute::UMMA::Major kMajorB,
           uint64_t kTensorCoreUtilControl,
           uint32_t kCombineNumRanks, bool kFuseCombine,
           CombineOrderMode kCombineOrderMode,
-          uint32_t kNumExtraCombineThreads>
+          uint32_t kNumExtraCombineThreads,
+          bool kMaskGroupedKTail = false>
 CUTLASS_GLOBAL void __launch_bounds__(
     kNumNonEpilogueThreads + kNumEpilogueThreads +
         kNumExtraCombineThreads,
@@ -92,6 +93,10 @@ sm100_bf16_gemm_impl(int* grouped_layout,
         "Invalid block K");
     DG_STATIC_ASSERT(BLOCK_K % UMMA_K == 0, "Block K must be divisible by UMMA K");
     DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % BLOCK_K == 0, "K alignment must be divisible by block K");
+    DG_STATIC_ASSERT(not kMaskGroupedKTail or
+                     (kGemmType == GemmType::KGroupedContiguous and
+                      kMajorA == cute::UMMA::Major::MN and kMajorB == cute::UMMA::Major::MN),
+                     "Masked K tails require physical count-based MN-major groups");
     DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast");
     DG_STATIC_ASSERT((kSwapAB and BLOCK_N == LAYOUT_AD_M) or
                      (not kSwapAB and (BLOCK_M == 32 or BLOCK_M == 64 or BLOCK_M == LAYOUT_AD_M)), "Invalid block size");
@@ -321,6 +326,17 @@ sm100_bf16_gemm_impl(int* grouped_layout,
                 if (cute::elect_one_sync()) {
                     #pragma unroll
                     for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / UMMA_K; ++ umma_k_idx) {
+                        // MegaMoE groups are padded to at least UMMA_K, not
+                        // necessarily BLOCK_K (e.g. 240-row pools). TMA may
+                        // read the next group's rows in the final tile, but
+                        // those entire 16-wide MMA atoms must not contribute.
+                        // Keep the existing empty-group initialization path;
+                        // the epilogue explicitly writes zeros for that group.
+                        if constexpr (kMaskGroupedKTail) {
+                            if (scheduler.current_shape_k != 0 and
+                                k_block_idx * BLOCK_K + umma_k_idx * UMMA_K >= scheduler.current_shape_k)
+                                continue;
+                        }
                         const uint32_t atom_k_idx = umma_k_idx * UMMA_K / BLOCK_ATOM_K;
                         const uint32_t inner_k_idx = umma_k_idx * UMMA_K % BLOCK_ATOM_K;
                         a_desc.lo = mma::sm100::advance_umma_desc_lo<kMajorA, LOAD_BLOCK_M, kSwizzleAMode, cutlass::bfloat16_t>(

tests/compile_mega_moe_compat.py

@@ -55,6 +55,12 @@ def cases():
                 "cutlass::bfloat16_t", "epilogue::transform::EpilogueIdentity", 100,
                 2, combine, "CombineOrderMode::DeepEPV1", 64 if combine else 0]
         yield f"wgrad_combine{int(combine)}", "sm100_bf16_gemm", "sm100_bf16_gemm_impl", args
+        masked_args = args.copy()
+        masked_args[7] = 64
+        masked_args[9:11] = [128, 128]
+        masked_args[18] = 128
+        masked_args.append(True)
+        yield f"wgrad_masked_tail_combine{int(combine)}", "sm100_bf16_gemm", "sm100_bf16_gemm_impl", masked_args
 
     for layout in (1, 2):
         yield (f"psum_layout{layout}", "smxx_layout", "transpose_and_pack_strided_fp32_into_ue8m0",

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,94 @@
+"""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
+
+
+@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_wide_k_masks_expert_tail(monkeypatch, pool_block_m, output_n):
+    """Wide loads may read the next expert; its MMA atoms must never execute."""
+    if torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("requires SM100/SM103")
+    torch.manual_seed(941)
+    # Prefix offsets and tails cover all 16-row residues modulo 64. Empty
+    # first/middle/last groups exercise accumulator initialization and OOB TMA.
+    sizes = [0, pool_block_m, 2 * pool_block_m, 0, 3 * pool_block_m, pool_block_m, 0]
+    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
+    # Deliberately make adjacent experts very different, so a cross-expert
+    # contribution cannot hide under the approximate reference tolerance.
+    a[pool_block_m:3 * pool_block_m].mul_(32)
+    b[pool_block_m:3 * pool_block_m].mul_(32)
+    shape = (len(sizes), 512, output_n)
+    baseline = torch.full(shape, float("nan"), dtype=torch.bfloat16, device="cuda")
+    monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", "16")
+    deep_gemm.bf16_mega_moe_backward_w13(baseline, a, b, counts, pool_block_m)
+    torch.cuda.synchronize()
+    for block_k in (32, 64):
+        monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", str(block_k))
+        output = torch.full_like(baseline, float("nan"))
+        for _ in range(2):
+            deep_gemm.bf16_mega_moe_backward_w13(output, a, b, counts, pool_block_m)
+            torch.cuda.synchronize()
+            assert torch.isfinite(output).all()
+            assert torch.equal(output, baseline)
+            for expert, rows in enumerate(sizes):
+                if rows == 0:
+                    assert torch.count_nonzero(output[expert]) == 0
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU")
+def test_wgrad_rejects_invalid_k_tile(monkeypatch):
+    if torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("requires SM100/SM103")
+    monkeypatch.setenv("DG_BF16_MEGA_MOE_WGRAD_BLOCK_K", "48")
+    a = torch.ones(240, 128, dtype=torch.bfloat16, device="cuda")
+    counts = torch.tensor([240], dtype=torch.int32, device="cuda")
+    output = torch.empty(1, 128, 128, dtype=torch.bfloat16, device="cuda")
+    with pytest.raises(RuntimeError, match="kBlockK"):
+        deep_gemm.bf16_mega_moe_backward_w13(output, a, a, counts, 240)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a Blackwell GPU")
+def test_wgrad_independent_of_global_grouped_alignment(monkeypatch):
+    if torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("requires SM100/SM103")
+    previous = deep_gemm.get_mk_alignment_for_contiguous_layout()
+    try:
+        # Dense/grouped BF16 users may select 224; MegaMoE's physical pool
+        # must remain independent, and its wgrad must not mutate that setting.
+        deep_gemm.set_mk_alignment_for_contiguous_layout(224)
+        test_wgrad_wide_k_masks_expert_tail(monkeypatch, 240, 128)
+        assert deep_gemm.get_mk_alignment_for_contiguous_layout() == 224
+    finally:
+        deep_gemm.set_mk_alignment_for_contiguous_layout(previous)