chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(megamoe): add nvfp4 group16 capability gate

Browse Source 
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.
This commit is contained in:
LuminolT
2026-07-08 18:29:09 +08:00
parent 0ff91f1285
commit 79fcfd6abf
9 changed files with 99 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh
+5 -3
View File
@@ -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
+7 -5
View File
@@ -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