From ef6154a6e5170ede44a57f8dbf1616b1c5462f35 Mon Sep 17 00:00:00 2001 From: morgendave Date: Sun, 20 Sep 2026 20:03:01 +0000 Subject: [PATCH 1/3] Align MegaMoE training with upstream K-grouped TMA and block layouts --- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 9 ++-- tests/test_mega_moe_training.py | 43 +++---------------- tests/test_mega_moe_wgrad_layout.py | 31 +++++++++++++ 3 files changed, 42 insertions(+), 41 deletions(-) create mode 100644 tests/test_mega_moe_wgrad_layout.py diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index a46ead27d..3d5addd6e 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/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(b.stride(0)), 1, kSwizzle); - const auto tensor_map_d = make_tma_cd_desc( - d, m, n, kBlockM, kStoreBlockN, - static_cast(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(d.stride(1)), static_cast(d.stride(0)), + kSwizzleCD); const SM100BF16GemmRuntime::Args args = { .gemm_desc = desc, diff --git a/tests/test_mega_moe_training.py b/tests/test_mega_moe_training.py index 48e8a660c..61c984bff 100644 --- a/tests/test_mega_moe_training.py +++ b/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') diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py new file mode 100644 index 000000000..97faf40df --- /dev/null +++ b/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 From f08dcb0b226d7bb1ef06f4f9a5661de1601833cf Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 21 Sep 2026 02:01:50 +0000 Subject: [PATCH 2/3] perf(mega): use wide wgrad tiles with expert-tail MMA masking --- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 5 +- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 19 ++++---- .../deep_gemm/impls/sm100_bf16_gemm.cuh | 18 ++++++- tests/compile_mega_moe_compat.py | 6 +++ tests/test_mega_moe_wgrad_layout.py | 48 +++++++++++++++++++ 5 files changed, 83 insertions(+), 13 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 47544e74c..2440b3176 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -52,6 +52,7 @@ 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; }; static void compile_and_launch(const std::string& tag, const Args& args) { @@ -76,7 +77,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {}, {}, {}, {} + {}, {}, {}, {}, {} >); }}; )", @@ -99,7 +100,7 @@ static void __instantiate_kernel() {{ 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( diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index 3d5addd6e..99c5f12bc 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/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( "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( - "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( "DG_BF16_MEGA_MOE_WGRAD_NUM_STAGES", 4); const int kSwizzle = kBlockK * static_cast(sizeof(cutlass::bfloat16_t)); const int kStoreBlockN = deep_jit::get_env( "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; @@ -187,6 +185,7 @@ static void sm100_bf16_mega_moe_wgrad_1sm( .combine_order_mode = combine.order_mode, .combine_num_extra_threads = static_cast(num_extra_combine_threads), + .mask_grouped_k_tail = mask_grouped_k_tail, }; SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_mega_moe_wgrad_1sm", args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh index 3d505a262..098771a0b 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh @@ -37,7 +37,8 @@ template + 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( diff --git a/tests/compile_mega_moe_compat.py b/tests/compile_mega_moe_compat.py index c1a5fc566..7cf6963f9 100644 --- a/tests/compile_mega_moe_compat.py +++ b/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", diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py index 97faf40df..1900d3de8 100644 --- a/tests/test_mega_moe_wgrad_layout.py +++ b/tests/test_mega_moe_wgrad_layout.py @@ -29,3 +29,51 @@ def test_wgrad_group_dimension(pool_block_m, output_n): 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) From 0be759300fd5d198ef211e452ccb857217c41a1a Mon Sep 17 00:00:00 2001 From: morgendave Date: Mon, 21 Sep 2026 02:07:12 +0000 Subject: [PATCH 3/3] fix(mega): isolate wgrad tiles from global grouped alignment --- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 4 +++- .../impls/sm100_bf16_mega_moe_wgrad.hpp | 4 ++++ tests/test_mega_moe_wgrad_layout.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 2440b3176..2bb2aa162 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -53,6 +53,7 @@ class SM100BF16GemmRuntime final { 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) { @@ -92,7 +93,8 @@ 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), diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp index 99c5f12bc..67b96d629 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe_wgrad.hpp @@ -186,6 +186,10 @@ static void sm100_bf16_mega_moe_wgrad_1sm( .combine_num_extra_threads = static_cast(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); } diff --git a/tests/test_mega_moe_wgrad_layout.py b/tests/test_mega_moe_wgrad_layout.py index 1900d3de8..848b35f51 100644 --- a/tests/test_mega_moe_wgrad_layout.py +++ b/tests/test_mega_moe_wgrad_layout.py @@ -77,3 +77,18 @@ def test_wgrad_rejects_invalid_k_tile(monkeypatch): 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)