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')