From 680874d06790c70e529549a6165e3123ff090a83 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 12 May 2026 07:08:08 +0000 Subject: [PATCH] NVFP4 L1 epilogue: group_size=16 SF layout - 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 --- .../impls/sm100_fp8_nvfp4_mega_moe.cuh | 56 +++++++++---------- 1 file changed, 25 insertions(+), 31 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh index 36a1b1b..be0410c 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh @@ -281,10 +281,12 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y, }); // Epilogue amax reduction shared memory - auto smem_amax_reduction = reinterpret_cast(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(smem_amax_reduction + STORE_BLOCK_M * kNumEpilogueWarps / 2); + auto barrier_start_ptr = reinterpret_cast(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()); - 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()); - 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()); } // 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(); 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(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(); @@ -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(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(); }