From 79fcfd6abf69ec502de26d2124577503768baa9e Mon Sep 17 00:00:00 2001 From: LuminolT Date: Wed, 8 Jul 2026 18:29:09 +0800 Subject: [PATCH] feat(megamoe): add nvfp4 group16 capability gate Allow SM100 FP4 scale layout transforms to accept group16 and thread weight granularity through the MegaMoE Python wrapper, API checks, and synthetic benchmark entrypoint. Keep fused SM100 MegaMoE compute behind an explicit group16 capability gate until the SFB/TMEM/MMA scale path is updated and validated. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/mega/__init__.py tests/test_mega_moe.py tests/generators.py Tested: git diff --check Not-tested: CUDA build and SM100/B300 runtime validation are not available locally. --- csrc/apis/layout.hpp | 10 +++-- csrc/apis/mega.hpp | 15 +++++-- .../impls/sm100_fp8_fp4_mega_moe.hpp | 15 ++++--- csrc/jit_kernels/impls/smxx_layout.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe.cuh | 8 ++-- deep_gemm/mega/__init__.py | 12 ++--- .../glm52_nvfp4_group16_notes.md | 44 +++++++++++++++++++ tests/generators.py | 4 +- tests/test_mega_moe.py | 20 ++++++--- 9 files changed, 99 insertions(+), 31 deletions(-) create mode 100644 megamoe-research-reports/glm52_nvfp4_group16_notes.md diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index b404241..4bc95c4 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -44,8 +44,10 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, if (sf.scalar_type() == torch::kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, false, true, torch::kFloat); - // (FP32, x, gran_k) on SM100: transform to (INT, 1, gran_k), TMA-aligned and MN-major - if (sf.scalar_type() == torch::kFloat and (gran_k == 32 or gran_k == 128) and arch_major == 10) { + // (FP32, x, gran_k) on SM100: transform to (INT, 1, gran_k), TMA-aligned and MN-major. + // GLM-5.2 NVFP4 checkpoints use weight granularity 16, while the original + // SM100 MegaMoE path only exercised 32. + if (sf.scalar_type() == torch::kFloat and (gran_k == 16 or gran_k == 32 or gran_k == 128) and arch_major == 10) { DG_HOST_ASSERT(not disable_ue8m0_cast); const auto broadcasted = gran_mn == 1 ? sf : sf.index_select(-2, torch::arange(mn, at::TensorOptions().device(sf.device())).floor_divide_(gran_mn)); @@ -53,7 +55,7 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, } // (INT, 1, gran_k) on SM100: transform to TMA-aligned and MN-major - if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 10) + if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 16 or gran_k == 32 or gran_k == 128) and arch_major == 10) return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, true, false, torch::kInt); DG_HOST_UNREACHABLE("Unknown SF transformation"); @@ -94,7 +96,7 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te DG_HOST_ASSERT(std::get<0>(recipe) == 1 and std::get<1>(recipe) == 1); const int gran_k = std::get<2>(recipe); - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); + DG_HOST_ASSERT(gran_k == 16 or gran_k == 32 or gran_k == 128); const auto arch_major = device_runtime->get_arch_major(); diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index efc3a78..f97e151 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -150,7 +150,7 @@ static void fp8_fp4_mega_moe( // Config checks const auto num_tokens = static_cast(y.size(0)); const auto [rm, rn, rk] = recipe; - DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32); + DG_HOST_ASSERT(rm == 1 and rn == 1 and (rk == 16 or rk == 32)); DG_HOST_ASSERT(activation == "swiglu"); // Activation checks @@ -173,11 +173,17 @@ static void fp8_fp4_mega_moe( DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); // Check weight SF layout for UE8M0 packing, MN-major, and TMA alignment - constexpr int kGranMN = 1, kGranK = 32; - check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, + constexpr int kGranMN = 1; + const int weight_gran_k = rk; + check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, weight_gran_k, num_experts_per_rank, true, false, torch::kInt); - check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, + check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, weight_gran_k, num_experts_per_rank, true, false, torch::kInt); + if (weight_gran_k == 16) { + DG_HOST_UNREACHABLE( + "SM100 FP8xFP4 MegaMoE weight granularity 16 requires kernel support for " + "NVFP4 group16; the current fused compute path still uses mxf4.block_scale.block32"); + } // Check stats counter if (cumulative_local_expert_recv_stats.has_value()) { @@ -213,6 +219,7 @@ static void fp8_fp4_mega_moe( num_experts_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, + weight_gran_k, activation_clamp, fast_math); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 4d91256..a569e1e 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -22,6 +22,7 @@ public: int num_max_tokens_per_rank; int hidden, intermediate_hidden; int num_experts, num_topk; + int weight_gran_k; int num_ranks; float activation_clamp; bool fast_math; @@ -60,6 +61,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, + {}, {}, {}, {}, {}, {}, {}, @@ -75,6 +77,7 @@ static void __instantiate_kernel() {{ )", args.num_max_tokens_per_rank, args.hidden, args.intermediate_hidden, args.num_experts, args.num_topk, + args.weight_gran_k, args.config.num_experts_per_wave, args.config.block_m, args.config.block_n, args.config.block_k, args.config.store_block_m, @@ -120,6 +123,7 @@ static void sm100_fp8_fp4_mega_moe( const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, + const int& weight_gran_k, const float& activation_clamp, const bool& fast_math ) { @@ -133,7 +137,7 @@ static void sm100_fp8_fp4_mega_moe( num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens); // Make tensormap - constexpr int kGranK = 32; + constexpr int kActivationGranK = 32; const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, hidden, config.num_max_pool_tokens, config.block_k, config.load_block_m, @@ -141,7 +145,7 @@ static void sm100_fp8_fp4_mega_moe( config.swizzle_acts_mode); const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, config.num_padded_sf_pool_tokens, hidden, - config.sf_block_m, kGranK, + config.sf_block_m, kActivationGranK, 1, 0); const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights, hidden, num_experts_per_rank * intermediate_hidden * 2, @@ -150,7 +154,7 @@ static void sm100_fp8_fp4_mega_moe( config.swizzle_weights_mode); const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf, intermediate_hidden * 2, hidden, - config.block_n, kGranK, + config.block_n, weight_gran_k, num_experts_per_rank, 0); // NOTES: L1 output and L2 activations are essentially the same tensor. // Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile), @@ -167,7 +171,7 @@ static void sm100_fp8_fp4_mega_moe( config.swizzle_acts_mode); const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, config.num_padded_sf_pool_tokens, intermediate_hidden, - config.sf_block_m, kGranK, + config.sf_block_m, kActivationGranK, 1, 0); const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, intermediate_hidden, num_experts_per_rank * hidden, @@ -176,7 +180,7 @@ static void sm100_fp8_fp4_mega_moe( config.swizzle_weights_mode); const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf, hidden, intermediate_hidden, - config.block_n, kGranK, + config.block_n, weight_gran_k, num_experts_per_rank, 0); // Stats can be optional @@ -190,6 +194,7 @@ static void sm100_fp8_fp4_mega_moe( .num_max_tokens_per_rank = num_max_tokens_per_rank, .hidden = hidden, .intermediate_hidden = intermediate_hidden, .num_experts = num_experts, .num_topk = num_topk, + .weight_gran_k = weight_gran_k, .num_ranks = num_ranks, .activation_clamp = activation_clamp, .fast_math = fast_math, diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 5d1f17b..6b54eab 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -228,7 +228,7 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons const torch::Tensor& ks_tensor, const std::vector& ks, const int gran_k) { - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); + DG_HOST_ASSERT(gran_k == 16 or gran_k == 32 or gran_k == 128); const auto [sf_k, mn] = get_shape<2>(sf); const auto num_groups = static_cast(ks.size()); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh index b2adc6c..7a7c09d 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh @@ -22,6 +22,7 @@ template < uint32_t kNumMaxTokensPerRank, uint32_t kHidden, uint32_t kIntermediateHidden, uint32_t kNumExperts, uint32_t kNumTopk, + uint32_t kWeightGranK, uint32_t kNumExpertsPerWave, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, uint32_t STORE_BLOCK_M, @@ -119,7 +120,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, input_topk_idx_buffer.get_end_ptr()); // SF and its buffer configs - constexpr uint32_t kGranK = 32; + constexpr uint32_t kActivationGranK = 32; + DG_STATIC_ASSERT(kWeightGranK == 16 or kWeightGranK == 32, "Invalid FP4 weight scale granularity"); constexpr uint32_t kNumUTCCPAlignedElems = 128; DG_STATIC_ASSERT(SF_BLOCK_M == math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems), "Invalid SF_BLOCK_M"); DG_STATIC_ASSERT(SF_BLOCK_N == BLOCK_N, "No padding is needed for SFB"); @@ -673,7 +675,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; - const auto shape_sfa_k = math::ceil_div(shape_k, kGranK * 4u); + const auto shape_sfa_k = math::ceil_div(shape_k, kActivationGranK * 4u); // Compute pool block offset for this expert const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; @@ -743,7 +745,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; const auto shape_n = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; - const auto shape_sfb_k = math::ceil_div(shape_k, kGranK * 4u); + const auto shape_sfb_k = math::ceil_div(shape_k, kWeightGranK * 4u); for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { // Wait consumer release diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index b82f25a..9fb5c7f 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -142,10 +142,11 @@ def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tup return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1]) -def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: +def _transpose_sf_for_utccp(sf: torch.Tensor, gran_k: int = 32) -> torch.Tensor: num_groups, mn, packed_sf_k = sf.shape assert sf.dtype == torch.int and mn % 128 == 0 - result = (sf.reshape(num_groups, -1, 4, 32, packed_sf_k) + assert 128 % gran_k == 0 + result = (sf.reshape(num_groups, -1, 128 // gran_k, gran_k, packed_sf_k) .transpose(2, 3) .reshape(num_groups, mn, packed_sf_k)) return torch.empty_like(sf).copy_(result) @@ -163,14 +164,15 @@ def transform_weights_for_mega_moe_sm90( def transform_weights_for_mega_moe( l1_weights: Tuple[torch.Tensor, torch.Tensor], - l2_weights: Tuple[torch.Tensor, torch.Tensor] + l2_weights: Tuple[torch.Tensor, torch.Tensor], + weight_gran_k: int = 32, ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: if _is_sm90(): return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights) # SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose l1_interleaved = _interleave_l1_weights(l1_weights) - l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1])) - l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1])) + l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1], weight_gran_k)) + l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1], weight_gran_k)) return l1_weights, l2_weights diff --git a/megamoe-research-reports/glm52_nvfp4_group16_notes.md b/megamoe-research-reports/glm52_nvfp4_group16_notes.md new file mode 100644 index 0000000..3bfa775 --- /dev/null +++ b/megamoe-research-reports/glm52_nvfp4_group16_notes.md @@ -0,0 +1,44 @@ +# GLM-5.2 NVFP4 group16 MegaMoE notes + +## Context + +GLM-5.2 NVFP4 checkpoints use FP4 weight scales with `group_size=16`. +The existing SM100 FP8xFP4 MegaMoE path was developed and tested with +`recipe=(1, 1, 32)`, so both the weight scale layout and fused compute path +carried an implicit group32 assumption. + +## Current patch scope + +- `transform_sf_into_required_layout` now accepts SM100 packed UE8M0 scale + transforms with `gran_k=16`. +- `transform_weights_for_mega_moe(..., weight_gran_k=...)` can apply the UTCCP + scale transpose with a group16-aware 128-element tiling. +- `tests/test_mega_moe.py` exposes `--weight-gran-k 16|32` so synthetic runs can + reproduce GLM-style group16 inputs without loading model weights. +- The fused SM100 MegaMoE compute API now performs an explicit capability check + for `recipe=(1, 1, 16)` instead of failing earlier with + `Unknown SF transformation`. + +## Remaining kernel work + +The fused compute kernel still uses the SM100 MXF4 block-scale path and its +current shared-memory/TMEM scale layout is group32-equivalent. Supporting +group16 correctly requires auditing at least: + +- weight scale TMA width per K block; +- SFB shared-memory and tensor-memory column allocation; +- scale id selection passed to the MMA instruction; +- the UTCCP scale transpose layout consumed by `SM100_UTCCP_4x32dp128bit_2cta`. + +Until that kernel work is complete and validated on B300/SM100, group16 should +be treated as layout-supported but fused-compute unsupported. + +## Validation target + +After kernel support is added, validate with: + +- existing group32 MegaMoE tests unchanged; +- `tests/test_layout.py` on SM100 for `gran_k=16`; +- `tests/test_mega_moe.py --weight-gran-k 16 --ncu-profile-only` for synthetic + fused execution; +- SGLang GLM-5.2 NVFP4 real-weight layout build and 8-card e2e smoke. diff --git a/tests/generators.py b/tests/generators.py index 989e984..3ab5761 100644 --- a/tests/generators.py +++ b/tests/generators.py @@ -202,7 +202,7 @@ def enumerate_k_grouped_contiguous(dtype: torch.dtype): def enumerate_sf_layout(): - gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128) + gran_k_list = (128, ) if get_arch_major() == 9 else (16, 32, 128) for use_ue8m0 in (False, True): for with_transpose in (True, False): for mn in (4096, 4097, 8192): @@ -214,7 +214,7 @@ def enumerate_sf_layout(): def enumerate_k_grouped_sf_layout(): - gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128) + gran_k_list = (128, ) if get_arch_major() == 9 else (16, 32, 128) for mn in (4096, 7168): for num_groups, avg_k in ((16, 2048), (8, 4096), (72, 384), (128, 256)): for gran_k in gran_k_list: diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 83e8d62..b25a236 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -82,7 +82,8 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): assert intermediate_hidden % 128 == 0 assert l1_weights.shape[2] % 128 == 0 and l2_weights.shape[2] % 128 == 0 - # Cast inputs to FP8 with per-32 UE8M0 SF + # Cast inputs to FP8 with per-32 UE8M0 SF. GLM NVFP4 group16 applies to + # FP4 weights only; activation dispatch remains per-32 in this test. x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) # Cast grouped BF16 weights to FP4 with MN-major SF @@ -90,15 +91,17 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): def cast_grouped_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: num_groups, n, k = bf16_weights.shape w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) - w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + w_sf = torch.empty((num_groups, n, k // args.weight_gran_k), device='cuda', dtype=torch.float) for i in range(num_groups): - w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + w[i], w_sf[i] = per_token_cast_to_fp4( + bf16_weights[i], use_ue8m0=True, gran_k=args.weight_gran_k) + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, args.weight_gran_k), num_groups) return w, w_sf l1_weights = cast_grouped_weights_to_fp4(l1_weights) l2_weights = cast_grouped_weights_to_fp4(l2_weights) - transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights) + transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe( + l1_weights, l2_weights, weight_gran_k=args.weight_gran_k) # Run fused mega MoE # NOTES: copy x into buffer before each call because debug mode zeros the entire buffer @@ -115,6 +118,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): transformed_l1_weights, transformed_l2_weights, buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, + recipe=(1, 1, args.weight_gran_k), activation_clamp=args.activation_clamp, fast_math=bool(args.fast_math) ) @@ -166,7 +170,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l1_y = torch.empty((n, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda') deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, 32)) + use_psum_layout=True, recipe=(1, 1, args.weight_gran_k)) # noinspection PyCallingNonCallable l1_y = tilelang_ops.swiglu_apply_weight_to_fp8( x=l1_y, @@ -183,7 +187,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, 32)) + use_psum_layout=True, recipe=(1, 1, args.weight_gran_k)) return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline # Check correctness (must be bitwise identical) @@ -275,6 +279,8 @@ if __name__ == '__main__': parser.add_argument('--num-topk', type=int, default=6, help='Number of expert selections') parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections') parser.add_argument('--fast-math', type=int, default=1, help='Enable fast math (0 or 1, default: 1)') + parser.add_argument('--weight-gran-k', type=int, default=32, choices=(16, 32), + help='FP4 weight scale granularity along K') # Test settings parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test')