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

NVFP4 L1 epilogue: group_size=16 SF layout

Browse Source 
- Single amax per warp (16 N-elements = 1 SF group, no warp-pair reduction)
- Single sf_val instead of sf.x/sf.y split
- All 4 warps write SF (k_idx = n_block_idx*4 + warp_idx_in_wg)
- Remove dead SMEM amax storage, reclaim barrier offset space
- Remove dead __syncwarp after register-local amax
This commit is contained in:
biondizzle
2026-05-12 07:08:08 +00:00
parent c0850a6859
commit 680874d067
1 changed files with 25 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
deep_gemm/include/deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh
View File

@@ -281,10 +281,12 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
});
// Epilogue amax reduction shared memory
auto smem_amax_reduction = reinterpret_cast<float2*>(smem_sfb[kNumStages]);
// NVFP4: scalar amax per warp (group_size=16, no warp-pair reduction)
// No SMEM amax storage needed — each warp computes its own amax independently
auto smem_after_sfb = smem_sfb[kNumStages];
// Barriers and tensor memory pointer
auto barrier_start_ptr = reinterpret_cast<Barrier*>(smem_amax_reduction + STORE_BLOCK_M * kNumEpilogueWarps / 2);
auto barrier_start_ptr = reinterpret_cast<Barrier*>(smem_after_sfb);
auto dispatch_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); });
auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumDispatchWarps + i); });
auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumDispatchWarps + kNumStages + i); });
@@ -988,7 +990,7 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
// Iterate all atoms in the store block
float2 swiglu_values[kNumAtomsPerStore * 2];
float2 amax_values[kNumAtomsPerStore];
float amax_values[kNumAtomsPerStore];
#pragma unroll
for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) {
const uint32_t j = s * kNumAtomsPerStore + i;
@@ -1052,20 +1054,17 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
swiglu_values[i * 2 + k] = __fmul2_rn(__fmul2_rn(gate, up), weights);
}
// Amax reduction
amax_values[i].x = math::warp_reduce<4, true>(
cute::max(cute::abs(swiglu_values[i * 2 + 0].x), cute::abs(swiglu_values[i * 2 + 1].x)),
math::ReduceMax<float>());
amax_values[i].y = math::warp_reduce<4, true>(
cute::max(cute::abs(swiglu_values[i * 2 + 0].y), cute::abs(swiglu_values[i * 2 + 1].y)),
math::ReduceMax<float>());
if (lane_idx < 4)
smem_amax_reduction[epilogue_warp_idx * (STORE_BLOCK_M / 2) + i * (ATOM_M / 2) + lane_idx] = amax_values[i];
__syncwarp();
// Amax reduction — NVFP4 group_size=16: single amax per warp
// Each warp covers 16 N-elements = exactly one SF group
const float lane_amax = cute::max(
cute::max(cute::abs(swiglu_values[i * 2 + 0].x),
cute::abs(swiglu_values[i * 2 + 0].y)),
cute::max(cute::abs(swiglu_values[i * 2 + 1].x),
cute::abs(swiglu_values[i * 2 + 1].y)));
amax_values[i] = math::warp_reduce<4, true>(lane_amax, math::ReduceMax<float>());
}
// Wait shared memory release from previous TMA store
// And fence `smem_amax_reduction`
const uint32_t tma_stage_idx = s % kNumTMAStoreStages;
ptx::tma_store_wait<kNumTMAStoreStages - 1>();
ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx);
@@ -1097,22 +1096,16 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
#pragma unroll
for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) {
// Reduce amax across warp pair
const float2 wp_amax =
smem_amax_reduction[(epilogue_warp_idx ^ 1) * (STORE_BLOCK_M / 2) + i * (ATOM_M / 2) + lane_idx % 4];
amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x);
amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y);
// NVFP4 group_size=16: no warp-pair amax reduction needed
// Each warp owns 16 N-elements = exactly one SF group
// E2M1: scale = amax / 6.0 (E2M1 max magnitude)
float2 sf, sf_inv;
sf.x = fmaxf(amax_values[i].x / 6.0f, 1e-8f);
sf.y = fmaxf(amax_values[i].y / 6.0f, 1e-8f);
sf_inv.x = 1.0f / sf.x;
sf_inv.y = 1.0f / sf.y;
float sf_val = fmaxf(amax_values[i] / 6.0f, 1e-8f);
float sf_inv = 1.0f / sf_val;
// 4 SwiGLU floats → 4 E2M1 nibbles → 2 bytes (uint16) for this lane
const float2 upper = __fmul2_rn(swiglu_values[i * 2 + 0], sf_inv);
const float2 lower = __fmul2_rn(swiglu_values[i * 2 + 1], sf_inv);
const float2 upper = __fmul2_rn(swiglu_values[i * 2 + 0], {sf_inv, sf_inv});
const float2 lower = __fmul2_rn(swiglu_values[i * 2 + 1], {sf_inv, sf_inv});
const uint8_t e0 = quant_e2m1_byte(upper.x);
const uint8_t e1 = quant_e2m1_byte(upper.y);
const uint8_t e2 = quant_e2m1_byte(lower.x);
@@ -1129,9 +1122,10 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
+ col_pair * 2;
*reinterpret_cast<uint16_t*>(smem_byte_ptr) = packed16;
// SF store to l2_sf_buffer as UE4M3 (MN-major layout)
if (warp_idx_in_wg % 2 == 0 and lane_idx < 4) {
const uint32_t k_idx = n_block_idx * 2 + warp_idx_in_wg / 2;
// SF store — NVFP4 group_size=16: all 4 warps warps write, one K position each
// k_idx = n_block_idx * 4 + warp_idx_in_wg → 4 K positions per atom
if (lane_idx < 4) {
const uint32_t k_idx = n_block_idx * 4 + warp_idx_in_wg;
const uint32_t k_uint_idx = k_idx / 4, byte_idx = k_idx % 4;
const uint32_t mn_stride = kNumPaddedSFPoolTokens * sizeof(uint32_t);
const auto sf_base_ptr = l2_sf_buffer.get_base_ptr<uint8_t>();
@@ -1145,8 +1139,8 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
cutlass::float_e4m3_t e = cutlass::float_e4m3_t(v);
return reinterpret_cast<uint8_t&>(e) & 0x7F;
};
sf_base_ptr[sf_addr] = to_ue4m3(sf.x);
sf_base_ptr[sf_addr + 4 * uint32_t(sizeof(uint32_t))] = to_ue4m3(sf.y);
sf_base_ptr[sf_addr] = to_ue4m3(sf_val);
sf_base_ptr[sf_addr + 4 * uint32_t(sizeof(uint32_t))] = to_ue4m3(sf_val);
}
__syncwarp();
}