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

feat: implement sm90 megamoe phase5 l2 scatter

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 15:17:20 +08:00
parent fc8218750c
commit 9bd0519605
5 changed files with 379 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
View File

@@ -49,14 +49,14 @@ sm90_fp8_mega_moe_impl(void* y,
const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
const void* l1_weights_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
const void* l2_weights_sf,
void* l1_accum_debug) {
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
const void* l1_weights_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
const void* l2_weights_sf,
void* l1_accum_debug) {
DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
@@ -89,6 +89,8 @@ sm90_fp8_mega_moe_impl(void* y,
constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t);
constexpr uint32_t kNumL1WeightSFGroupsN = L1_SHAPE_N / 128;
constexpr uint32_t kNumL1WeightSFGroupsK = L1_SHAPE_K / 128;
constexpr uint32_t kNumL2WeightSFGroupsN = L2_SHAPE_N / 128;
constexpr uint32_t kNumL2WeightSFGroupsK = L2_SHAPE_K / 128;
constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2;
constexpr uint32_t kL2ActsGranK = 64;
constexpr uint32_t kMathBarrierIdx = 2;
@@ -106,6 +108,9 @@ sm90_fp8_mega_moe_impl(void* y,
cute::prefetch_tma_descriptor(&tensor_map_l1_acts);
cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf);
cute::prefetch_tma_descriptor(&tensor_map_l1_weights);
cute::prefetch_tma_descriptor(&tensor_map_l2_acts);
cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf);
cute::prefetch_tma_descriptor(&tensor_map_l2_weights);
}
__syncwarp();
@@ -121,6 +126,7 @@ sm90_fp8_mega_moe_impl(void* y,
sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);
constexpr auto fp8_token_layout = layout::Data(kHidden);
constexpr auto bf16_token_layout = layout::Data(kHidden * static_cast<uint32_t>(sizeof(nv_bfloat16)));
constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kL2ActsGranK * static_cast<uint32_t>(sizeof(float)), false);
@@ -156,7 +162,13 @@ sm90_fp8_mega_moe_impl(void* y,
const auto l2_sf_buffer = layout::Buffer(
fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens,
l2_token_buffer.get_end_ptr());
(void)l2_sf_buffer;
constexpr auto l1_accum_debug_layout = layout::Data(128 * static_cast<uint32_t>(sizeof(float)), false);
const auto l1_accum_debug_buffer = layout::Buffer(
l1_accum_debug_layout, 1, 128,
l2_sf_buffer.get_end_ptr());
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, kNumTopk, kNumMaxTokensPerRank,
l1_accum_debug_buffer.get_end_ptr());
constexpr uint32_t kSharedMemoryAlignment = 1024;
extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];
@@ -466,13 +478,16 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
if (block_phase != sched::SM90BlockPhase::Linear1)
return;
const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
const uint32_t valid_m = scheduler.get_valid_m();
const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
while (ptx::ld_acq(arrival_ptr) != valid_m);
if (block_phase == sched::SM90BlockPhase::Linear1) {
const uint32_t valid_m = scheduler.get_valid_m();
const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
while (ptx::ld_acq(arrival_ptr) != valid_m);
} else {
const auto arrival_ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx);
const uint64_t expected = ((1ull << num_k_blocks) << num_k_blocks) - 1;
while (ptx::ld_acq_gpu(arrival_ptr) != expected);
}
#pragma unroll 1
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
@@ -482,16 +497,26 @@ sm90_fp8_mega_moe_impl(void* y,
auto& full_barrier = *full_barriers[stage_idx];
const uint32_t k_idx = k_block_idx * BLOCK_K;
const uint32_t m_idx = pool_block_idx * BLOCK_M;
const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N;
const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
&tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
tma::copy<SF_BLOCK_M, 1, 0>(
&tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
&tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
if (block_phase == sched::SM90BlockPhase::Linear1) {
const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N;
const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
&tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
tma::copy<SF_BLOCK_M, 1, 0>(
&tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
&tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
} else {
const uint32_t n_idx = local_expert_idx * L2_SHAPE_N + n_block_idx * BLOCK_N;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
&tensor_map_l2_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
&tensor_map_l2_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
}
}
__syncwarp();
}
@@ -508,6 +533,7 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx;
const uint32_t r_1 = r_0 + 8;
const auto l1_weights_sf_ptr = reinterpret_cast<const float*>(l1_weights_sf);
const auto l2_weights_sf_ptr = reinterpret_cast<const float*>(l2_weights_sf);
auto l1_accum_debug_ptr = reinterpret_cast<float*>(l1_accum_debug);
auto l2_acts_ptr = l2_token_buffer.get_base_ptr<__nv_fp8_e4m3>();
auto l2_acts_sf_ptr = l2_sf_buffer.get_base_ptr<float>();
@@ -543,11 +569,12 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
if (block_phase != sched::SM90BlockPhase::Linear1)
return;
const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
const uint32_t valid_m = scheduler.get_valid_m();
const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0;
const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1;
const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0;
const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1;
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
const auto empty_barrier_arrive = [&](const uint32_t& s) {
@@ -559,52 +586,121 @@ sm90_fp8_mega_moe_impl(void* y,
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
full_barriers[stage_idx]->wait(phase);
const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0);
const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1);
if (block_phase == sched::SM90BlockPhase::Linear1) {
const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0);
const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1);
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_arrive();
#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
auto desc_a = mma::sm90::make_smem_desc(
smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
auto desc_b = mma::sm90::make_smem_desc(
smem_b[stage_idx] + k * WGMMA::K, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
}
ptx::warpgroup_commit_batch();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_wait<0>();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_arrive();
#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
auto desc_a = mma::sm90::make_smem_desc(
smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
auto desc_b = mma::sm90::make_smem_desc(
smem_b[stage_idx] + k * WGMMA::K, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
}
ptx::warpgroup_commit_batch();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_wait<0>();
empty_barrier_arrive(stage_idx);
empty_barrier_arrive(stage_idx);
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col_0 = i * 8 + col_idx * 2;
const uint32_t col_1 = col_0 + 1;
const uint32_t n_0 = n_block_idx * BLOCK_N + col_0;
const uint32_t n_1 = n_block_idx * BLOCK_N + col_1;
const auto sf_base = l1_weights_sf_ptr +
local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK;
const float scale_b_0 = __ldg(
sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx);
const float scale_b_1 = __ldg(
sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx);
final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col_0 = i * 8 + col_idx * 2;
const uint32_t col_1 = col_0 + 1;
const uint32_t n_0 = n_block_idx * BLOCK_N + col_0;
const uint32_t n_1 = n_block_idx * BLOCK_N + col_1;
const auto sf_base = l1_weights_sf_ptr +
local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK;
const float scale_b_0 = __ldg(
sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx);
const float scale_b_1 = __ldg(
sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx);
final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
}
} else {
#pragma unroll
for (uint32_t half = 0; half < 2; ++ half) {
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_arrive();
#pragma unroll
for (uint32_t k = 0; k < 2; ++ k) {
const uint32_t smem_k = (half * 2 + k) * WGMMA::K;
auto desc_a = mma::sm90::make_smem_desc(
smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + smem_k, 1);
auto desc_b = mma::sm90::make_smem_desc(
smem_b[stage_idx] + smem_k, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
}
ptx::warpgroup_commit_batch();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_wait<0>();
const uint32_t sf_group = k_block_idx * 2 + half;
const uint32_t sf_pool_token_idx_0 = pool_block_idx * SF_BLOCK_M + row_0;
const uint32_t sf_pool_token_idx_1 = pool_block_idx * SF_BLOCK_M + row_1;
const float scale_a_0 = __ldg(
l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_0);
const float scale_a_1 = __ldg(
l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_1);
const auto sf_base = l2_weights_sf_ptr +
local_expert_idx * kNumL2WeightSFGroupsN * kNumL2WeightSFGroupsK;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col_0 = i * 8 + col_idx * 2;
const uint32_t col_1 = col_0 + 1;
const uint32_t n_group_0 = (n_block_idx * BLOCK_N + col_0) / 128;
const uint32_t n_group_1 = (n_block_idx * BLOCK_N + col_1) / 128;
const float scale_b_0 = __ldg(sf_base + n_group_0 * kNumL2WeightSFGroupsK + k_block_idx);
const float scale_b_1 = __ldg(sf_base + n_group_1 * kNumL2WeightSFGroupsK + k_block_idx);
final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
}
}
empty_barrier_arrive(stage_idx);
}
}
const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0;
const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1;
const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0;
const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1;
if (block_phase == sched::SM90BlockPhase::Linear2) {
const auto scatter_row = [&](const uint32_t& row, const uint32_t& pool_token_idx, const uint32_t& accum_offset) {
if (row >= valid_m)
return;
const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_token_idx);
const auto dst_token = combine_token_buffer.get_rank_buffer(src_metadata.topk_idx)
.get_data_buffer(src_metadata.token_idx);
auto dst_ptr = dst_token.get_base_ptr<nv_bfloat16>();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t out_col = n_block_idx * BLOCK_N + i * 8 + col_idx * 2;
*sym_buffer.map(dst_ptr + out_col + 0, src_metadata.rank_idx) =
__float2bfloat16(final_accum[i * 4 + accum_offset + 0]);
*sym_buffer.map(dst_ptr + out_col + 1, src_metadata.rank_idx) =
__float2bfloat16(final_accum[i * 4 + accum_offset + 1]);
}
};
scatter_row(row_0, pool_token_idx_0, 0);
scatter_row(row_1, pool_token_idx_1, 2);
return;
}
const float topk_weight_0 = row_0 < valid_m ?
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx_0).get_base_ptr<float>() : 0.0f;
const float topk_weight_1 = row_1 < valid_m ?

5
deep_gemm/mega/__init__.py
View File

@@ -70,7 +70,8 @@ class SymmBuffer:
self.l1_arrival_count,
self.l2_arrival_mask,
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
self.l1_accum_debug,
self.combine_acts) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
@@ -82,6 +83,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def destroy(self):
self.handle = None
@@ -101,6 +103,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,