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

feat: implement sm90 megamoe phase4 l1 epilogue

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 01:09:45 +08:00
parent 2bb1756787
commit f3553f976c
6 changed files with 330 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -89,9 +89,13 @@ 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 L1_OUT_BLOCK_N = BLOCK_N / 2;
constexpr uint32_t kL2ActsGranK = 64;
constexpr uint32_t kMathBarrierIdx = 2;
DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");
DG_STATIC_ASSERT(BLOCK_N == WGMMA::N and BLOCK_K % WGMMA::K == 0, "Invalid WGMMA tile shape");
DG_STATIC_ASSERT(L1_OUT_BLOCK_N == kL2ActsGranK, "SM90 Phase 4 expects per-64 L2 activation SF");
const uint32_t thread_idx = threadIdx.x;
const uint32_t sm_idx = blockIdx.x;
@@ -119,7 +123,7 @@ sm90_fp8_mega_moe_impl(void* y,
constexpr auto fp8_token_layout = layout::Data(kHidden);
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 / 128 * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kL2ActsGranK * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(sizeof(float)), false);
@@ -505,6 +509,8 @@ sm90_fp8_mega_moe_impl(void* y,
const uint32_t r_1 = r_0 + 8;
const auto l1_weights_sf_ptr = reinterpret_cast<const float*>(l1_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>();
const auto get_l1_weight_sf_group = [](const uint32_t& interleaved_n) {
constexpr uint32_t kInterleaveGran = 8;
@@ -540,6 +546,7 @@ sm90_fp8_mega_moe_impl(void* y,
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();
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
@@ -594,9 +601,81 @@ sm90_fp8_mega_moe_impl(void* y,
}
}
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;
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 ?
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx_1).get_base_ptr<float>() : 0.0f;
const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) {
if constexpr (kActivationClamp != cute::numeric_limits<float>::infinity()) {
gate = fminf(fmaxf(gate, -kActivationClamp), kActivationClamp);
up = fminf(fmaxf(up, -kActivationClamp), kActivationClamp);
}
const float denom = 1.0f + (kFastMath ? __expf(-gate) : expf(-gate));
const float silu = gate * (kFastMath ? math::fast_rcp(denom) : 1.0f / denom);
return silu * up * topk_weight;
};
float local_amax_0 = 0.0f, local_amax_1 = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) {
const float v_00 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0);
const float v_01 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0);
const float v_10 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1);
const float v_11 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1);
local_amax_0 = fmaxf(local_amax_0, fmaxf(fabsf(v_00), fabsf(v_01)));
local_amax_1 = fmaxf(local_amax_1, fmaxf(fabsf(v_10), fabsf(v_11)));
}
const float row_amax_0 = math::warp_reduce<4, false>(local_amax_0, math::ReduceMax<float>());
const float row_amax_1 = math::warp_reduce<4, false>(local_amax_1, math::ReduceMax<float>());
const float sf_0 = fmaxf(row_amax_0 / 448.0f, 1.0e-12f);
const float sf_1 = fmaxf(row_amax_1 / 448.0f, 1.0e-12f);
const float sf_inv_0 = kFastMath ? math::fast_rcp(sf_0) : 1.0f / sf_0;
const float sf_inv_1 = kFastMath ? math::fast_rcp(sf_1) : 1.0f / sf_1;
if (col_idx == 0) {
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;
if (row_0 < valid_m)
l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_0] = sf_0;
if (row_1 < valid_m)
l2_acts_sf_ptr[n_block_idx * kNumPaddedSFPoolTokens + sf_pool_token_idx_1] = sf_1;
}
const uint32_t out_n_base = n_block_idx * L1_OUT_BLOCK_N;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; i += 2) {
const uint32_t out_col = out_n_base + (i / 2) * 8 + col_idx * 2;
if (row_0 < valid_m) {
const float v_0 = apply_swiglu(final_accum[i * 4 + 0], final_accum[(i + 1) * 4 + 0], topk_weight_0);
const float v_1 = apply_swiglu(final_accum[i * 4 + 1], final_accum[(i + 1) * 4 + 1], topk_weight_0);
l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_0);
l2_acts_ptr[pool_token_idx_0 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_0);
}
if (row_1 < valid_m) {
const float v_0 = apply_swiglu(final_accum[i * 4 + 2], final_accum[(i + 1) * 4 + 2], topk_weight_1);
const float v_1 = apply_swiglu(final_accum[i * 4 + 3], final_accum[(i + 1) * 4 + 3], topk_weight_1);
l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 0] = __nv_fp8_e4m3(v_0 * sf_inv_1);
l2_acts_ptr[pool_token_idx_1 * L2_SHAPE_K + out_col + 1] = __nv_fp8_e4m3(v_1 * sf_inv_1);
}
}
__threadfence();
ptx::sync_aligned(kNumMathThreads, kMathBarrierIdx);
if (math_warp_idx == 0 and cute::elect_one_sync()) {
ptx::red_or_rel_gpu(
workspace.get_l2_arrival_mask_ptr(pool_block_idx),
1ull << n_block_idx);
}
__syncwarp();
if (local_expert_idx == 0 and m_block_idx == 0 and n_block_idx == 0) {
const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0;
const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col = i * 8 + col_idx * 2;

3
deep_gemm/mega/__init__.py
View File

@@ -68,6 +68,7 @@ class SymmBuffer:
self.l2_acts, self.l2_acts_sf,
self.expert_recv_count_sum,
self.l1_arrival_count,
self.l2_arrival_mask,
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
else:
@@ -78,6 +79,7 @@ class SymmBuffer:
self.l1_topk_weights = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
@@ -96,6 +98,7 @@ class SymmBuffer:
self.l2_acts_sf = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None