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

@@ -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
+    // (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) {
+    // 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();
 

csrc/apis/mega.hpp

@@ -150,7 +150,7 @@ static void fp8_fp4_mega_moe(
     // Config checks
     const auto num_tokens = static_cast<int>(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;
+    constexpr int kGranMN = 1;
-    check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK,
+    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");

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,

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<int>& 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<int>(ks.size());
 

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

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]))
+    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]))
+    l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1], weight_gran_k))
     return l1_weights, l2_weights
 
 

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.

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:

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[i], w_sf[i] = per_token_cast_to_fp4(
-            w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups)
+                    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')