From ed1a32318e901ee0f058ebb9ec995cedb269d1ba Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Tue, 29 Sep 2026 16:06:33 +0300 Subject: [PATCH 1/4] Add native block32 FP8 GEMM support on SM90 Support recipe=(1,32,32) in dense NT FP8 GEMM using four independently scaled K32 partial sums per K128 TMA stage. Adapt scale layouts, accumulator scheduling, tile heuristics, and N/K specialization while retaining the default block128 path. --- csrc/apis/gemm.hpp | 6 +- csrc/apis/layout.hpp | 12 +- csrc/jit_kernels/heuristics/sm90.hpp | 22 ++- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 26 ++-- .../deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh | 135 +++++++++++++++--- 5 files changed, 165 insertions(+), 36 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index e058a4922..ee9f0948f 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -195,16 +195,18 @@ static void fp8_fp4_gemm_nt(const std::pair& a, if (arch_major == 9 or arch_major == 10) { // SM90/SM100 share the "transform scaling factors, then dispatch" flow. const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( - a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, disable_ue8m0_cast); + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, disable_ue8m0_cast, std::nullopt, true); if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) { DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100"); const int gran_n = recipe.has_value() ? std::get<1>(recipe.value()) : std::get<0>(recipe_b.value()); + DG_HOST_ASSERT(gran_k_a == gran_k_b); + DG_HOST_ASSERT((gran_n == 32 and gran_k_a == 32) or (gran_k_a == 128 and (gran_n == 1 or gran_n == 128))); if (gran_n == 1) { sm90_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims); } else { const auto major_sfb = get_major_type_ab(sfb); - sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims); + sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims, std::nullopt, gran_k_a); } } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) { sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, gran_k_a, gran_k_b, diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index caed591d9..bda553393 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -35,11 +35,11 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups); // (FP32, 1, 128) on SM90: transform to TMA-aligned and MN-major - if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 128 or gran_k == 32) and (arch_major == 9 or disable_ue8m0_cast)) return get_mn_major_tma_aligned_tensor(sf); // (FP32, 128, 128) on SM90: no need to transform, check SFB requirements - if (sf.scalar_type() == torch::kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + if (sf.scalar_type() == torch::kFloat and ((gran_mn == 128 and gran_k == 128) or (gran_mn == 32 and gran_k == 32)) and (arch_major == 9 or disable_ue8m0_cast)) return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, false, true, torch::kFloat); // (FP32, x, gran_k) on SM100/SM120: transform to (INT, 1, gran_k), TMA-aligned and MN-major @@ -67,7 +67,8 @@ static std::tuple transform_sf_pair_into const std::optional& num_groups_b, const bool& disable_ue8m0_cast = false, // PSUM layout only applies to SFA (B is weights with no gaps) - const std::optional& psum_layout = std::nullopt) { + const std::optional& psum_layout = std::nullopt, + const bool allow_sm90_block32 = false) { // Use default recipe, if none is specified if (not recipe_a.has_value() and not recipe.has_value()) recipe = get_default_recipe(sfa.scalar_type(), sfb.scalar_type()); @@ -76,6 +77,11 @@ static std::tuple transform_sf_pair_into DG_HOST_ASSERT(recipe_a.has_value() == recipe_b.has_value()); DG_HOST_ASSERT(recipe_a.has_value() != recipe.has_value()); + if (jit->device.get_arch_major() == 9 and not allow_sm90_block32) { + DG_HOST_ASSERT((recipe.has_value() ? std::get<2>(*recipe) : std::get<1>(*recipe_a)) == 128); + DG_HOST_ASSERT((recipe.has_value() ? std::get<2>(*recipe) : std::get<1>(*recipe_b)) == 128); + } + // Transform SFA and SFB layout (PSUM only for SFA) const auto transformed_sfa = recipe.has_value() ? transform_sf_into_required_layout(sfa, m, k, recipe.value(), num_groups_a, true, disable_ue8m0_cast, psum_layout) : transform_sf_into_required_layout(sfa, m, k, recipe_a.value(), num_groups_a, std::nullopt, disable_ue8m0_cast, psum_layout); diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index 5a08cacc8..aa7b4acfb 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -14,8 +14,11 @@ namespace deep_gemm { struct SM90ArchSpec { static constexpr int smem_capacity = 232448; + static constexpr int block32_short_k = 8 * 128; + static constexpr int block32_small_m_k = 2 * block32_short_k; static std::vector get_layout_candidates(const GemmDesc& desc) { + const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32; // Block M candidates std::vector block_m_candidates; if (desc.gemm_type == GemmType::Normal or @@ -24,8 +27,8 @@ struct SM90ArchSpec { // TODO: check 256's performance block_m_candidates = {64, 128}; // NOTES: smaller block M can avoid TMA L2 OOB bound - if (desc.m <= 16) block_m_candidates.push_back(16); - if (desc.m <= 32) block_m_candidates.push_back(32); + if (desc.m <= 16 and (not block32 or desc.k <= block32_small_m_k)) block_m_candidates.push_back(16); + if (desc.m <= 32 and (not block32 or desc.k <= block32_small_m_k)) block_m_candidates.push_back(32); // BF16 output GEMM supports 256 if (desc.cd_dtype != torch::kFloat) @@ -39,7 +42,7 @@ struct SM90ArchSpec { // Block N candidates std::vector block_n_candidates; - int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of()); + int step = std::lcm(block32 ? 32 : 16, heuristics_runtime->get_block_n_multiple_of()); int start = step; // Avoid bank conflicts for 1D1D kernel FP32 output if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { @@ -51,7 +54,7 @@ struct SM90ArchSpec { // Register spills int end = 256; if (desc.kernel_type == KernelType::Kernel1D2D) - end = 192; + end = block32 ? 128 : 192; if (desc.kernel_type == KernelType::Kernel1D1D) end = 160; // Enumerate @@ -66,7 +69,7 @@ struct SM90ArchSpec { // The number of k-groups is large (a heuristic) (desc.gemm_type == GemmType::KGroupedContiguous and desc.num_groups > 4) or // Not supported - (desc.gemm_type == GemmType::Batched); + (desc.gemm_type == GemmType::Batched) or block32; // Enumerate all candidates std::vector candidates; @@ -82,6 +85,10 @@ struct SM90ArchSpec { for (int block_m: block_m_candidates) { for (int block_n: block_n_candidates) { + // Long K favors overlapping wide single-wave tiles; short K can + // amortize B loads across two M waves without spilling. + if (block32 and block_m > (desc.k <= block32_short_k ? 256 : 128)) + continue; // 1D2D kernel unroll requirement if (desc.kernel_type == KernelType::Kernel1D2D and block_n > block_k and (block_n % (block_n - block_k) != 0 and block_k % (block_n - block_k) != 0)) continue; @@ -162,15 +169,16 @@ struct SM90ArchSpec { const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype); // Calculate SF A/B per stages + const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32; const int smem_sfa_per_stage = desc.kernel_type == KernelType::KernelNoSF ? - 0 : align(layout.block_m * static_cast(sizeof(float)), 128); + 0 : align(layout.block_m * (block32 ? 4 : 1) * static_cast(sizeof(float)), 128); const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ? 0 : align(layout.block_n * static_cast(sizeof(float)), 128); // Extra SFB sizes for 1D2D kernels const int use_uniform_sfb = layout.block_k % layout.block_n == 0 ? 1 : 2; const int smem_extra_sfb = desc.kernel_type != KernelType::Kernel1D2D ? - 0 : align(ceil_div(desc.k, layout.block_k) * static_cast(sizeof(float)) * use_uniform_sfb, 8); + 0 : align(ceil_div(desc.k, block32 ? 32 : layout.block_k) * static_cast(sizeof(float)) * (block32 ? layout.block_n / 32 : use_uniform_sfb), 8); // Extra tensormap for 1D1D kernels const int smem_tensormap = diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 215941a6b..10d868da0 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -27,6 +27,7 @@ class SM90FP8Gemm1D2DRuntime final { CUtensorMap tensor_map_b; CUtensorMap tensor_map_d; CUtensorMap tensor_map_sfa; + int sf_granularity = 128; }; static void compile_and_launch(const std::string& tag, const Args& args) { @@ -34,6 +35,10 @@ class SM90FP8Gemm1D2DRuntime final { const bool specialize_n = args.gemm_desc.gemm_type == GemmType::Normal && args.gemm_desc.compiled_dims.empty() && args.gemm_config.layout.block_m == 128 && args.gemm_config.layout.block_n > 128 && args.gemm_config.layout.get_cluster_size() == 1; + // Block32 scale addressing benefits from fixed weight dimensions. + // M remains dynamic, so one specialization serves decode and prefill. + const auto compiled_dims = args.sf_granularity == 32 and args.gemm_desc.compiled_dims.empty() ? + std::string("nk") : args.gemm_desc.compiled_dims; const auto kernel = jit->compile(tag, std::format(R"( #include @@ -51,14 +56,14 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, - {} + {}, {} >); }}; )", to_string(args.major_sfb), - get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), - get_compiled_dim(args.gemm_desc.n, 'n', specialize_n ? "n" : args.gemm_desc.compiled_dims), - get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.m, 'm', compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', specialize_n ? "n" : compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', compiled_dims), args.gemm_desc.num_groups, args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k, args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode, @@ -67,7 +72,7 @@ static void __instantiate_kernel() {{ args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), to_string(args.gemm_desc.cd_dtype), - get_default_epilogue_type(args.epilogue_type))); + get_default_epilogue_type(args.epilogue_type), args.sf_granularity)); // Launch jit->launch( @@ -87,8 +92,11 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, const int& m, const int& n, const int& k, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb, const std::string& compiled_dims, - const std::optional& epilogue_type = std::nullopt) { + const std::optional& epilogue_type = std::nullopt, + const int sf_granularity = 128) { DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(sf_granularity == 128 or sf_granularity == 32); + DG_HOST_ASSERT(sf_granularity != 32 or (n % 8 == 0 and k % 32 == 0)); DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); const auto desc = GemmDesc { @@ -101,7 +109,8 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .with_accumulation = c.has_value(), .num_sms = runtime->get_num_sms(), .tc_util = runtime->get_tc_util(), - .compiled_dims = compiled_dims + .compiled_dims = compiled_dims, + .max_gran_k = sf_granularity }; const auto config = get_best_config(desc); @@ -124,7 +133,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, - config.layout.block_m, config.layout.block_k, 1, 0); + config.layout.block_m, sf_granularity, 1, 0, 0, false, config.layout.block_k / sf_granularity); // Compile and launch SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_fp8_gemm_1d2d", { @@ -144,6 +153,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, .tensor_map_sfa = tensor_map_sfa, + .sf_granularity = sf_granularity, }); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh index ffd1c2dd5..8d0f42199 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh @@ -45,7 +45,7 @@ template + typename epilogue_op_t, uint32_t SF_GRAN = 128> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, @@ -55,7 +55,10 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) // Scaling checks - DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT(BLOCK_K == 128, "TMA stages use K=128"); + DG_STATIC_ASSERT(SF_GRAN == 32 or SF_GRAN == 128, "Unsupported scale granularity"); + DG_STATIC_ASSERT(SF_GRAN != 32 or (BLOCK_N % 32 == 0 and kGemmType == GemmType::Normal), + "Block32 currently supports dense GEMM with aligned N tiles"); DG_STATIC_ASSERT( math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); @@ -78,11 +81,12 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, static constexpr uint32_t SMEM_D_SIZE = math::constexpr_align(BLOCK_M * BLOCK_N * static_cast(sizeof(__nv_bfloat16)), 1024u); static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3); - static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * (BLOCK_K / SF_GRAN) * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); - const uint32_t shape_k_scales = math::ceil_div(shape_k, BLOCK_K); - const uint32_t shape_n_sfb = math::ceil_div(shape_n, BLOCK_K); - const uint32_t smem_sfb_size = math::align(shape_k_scales * (kMustUseUniformedScaleB ? 1 : 2) * sizeof(float), sizeof(Barrier)); + const uint32_t shape_k_scales = math::ceil_div(shape_k, SF_GRAN); + const uint32_t shape_n_sfb = math::ceil_div(shape_n, SF_GRAN); + constexpr uint32_t kScaleBRows = SF_GRAN == 32 ? BLOCK_N / 32 : (kMustUseUniformedScaleB ? 1 : 2); + const uint32_t smem_sfb_size = math::align(shape_k_scales * kScaleBRows * sizeof(float), sizeof(Barrier)); // NOTES: Make sure we have enough shared memory for WGMMA padding static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); @@ -194,8 +198,8 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, tma::copy(&tensor_map_a, &full_barrier, smem_a[stage_idx], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx), num_tma_multicast_a, batch_idx); - tma::copy(&tensor_map_sfa, &full_barrier, - smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx), + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx * (BLOCK_K / SF_GRAN)), num_tma_multicast_a); // Issue TMA B @@ -234,7 +238,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8; } - uint32_t num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2); + uint32_t num_sfb = shape_k_scales * (SF_GRAN == 32 ? kScaleBRows : (num_former_iters >= num_full_iters ? 1 : 2)); // Load B scales with math warp-groups // NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks @@ -242,18 +246,32 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, auto previous_group_offset = scheduler.template get_global_idx(shape_n_sfb * shape_k_scales, 0, 0, m_block_idx); const uint32_t stride_n_sfb = kMajorSFB == cute::UMMA::Major::MN ? 1 : shape_k_scales; const uint32_t stride_k_sfb = kMajorSFB == cute::UMMA::Major::MN ? shape_n_sfb : 1; - auto local_sfb = sfb + previous_group_offset + ((n_block_idx * BLOCK_N) / BLOCK_K) * stride_n_sfb; - - #pragma unroll - for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32) - ptx::st_shared(smem_sfb + i, i < shape_k_scales ? local_sfb[i * stride_k_sfb] : local_sfb[(i - shape_k_scales) * stride_k_sfb + stride_n_sfb]); + auto local_sfb = sfb + previous_group_offset + ((n_block_idx * BLOCK_N) / SF_GRAN) * stride_n_sfb; + + if constexpr (SF_GRAN == 32) { + #pragma unroll + for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32) { + const auto row = i / shape_k_scales, col = i % shape_k_scales; + const bool valid = n_block_idx * BLOCK_N / SF_GRAN + row < shape_n_sfb; + ptx::st_shared(smem_sfb + i, valid ? local_sfb[row * stride_n_sfb + col * stride_k_sfb] : 0.0f); + } + } else { + #pragma unroll + for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32) + ptx::st_shared(smem_sfb + i, i < shape_k_scales ? local_sfb[i * stride_k_sfb] : local_sfb[(i - shape_k_scales) * stride_k_sfb + stride_n_sfb]); + } } cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); // Accumulation for WGMMA or CUDA promotion constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); - float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; + // Full-M tiles reuse B with one bank; wide single-wave tiles overlap + // promotion of the previous partial with WGMMA on the other bank. + constexpr uint32_t kAccumBanks = SF_GRAN == 32 ? + (BLOCK_M > 128 ? 1 : (BLOCK_N <= 64 ? 4 : 2)) : 1; + constexpr bool kOverlapPromotion = SF_GRAN == 32 and kAccumBanks == 2; + float accum[WGMMA::kNumAccum * kAccumBanks], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; // Pick threads whose WGMMA results are to be stored in shared memory DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for `BLOCK_M < 64`"); @@ -273,7 +291,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, // Skip useless computations if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { // The compiler must know the dynamic variable `num_former_iters`'s real value - constexpr bool kShouldOptimize = BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr bool kShouldOptimize = SF_GRAN == 128 and BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; @@ -284,6 +302,90 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); const auto b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16); + if constexpr (SF_GRAN == 32) { + full_barriers[stage_idx]->wait(phase); + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + const auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + float sa0[4], sa1[4], sb[4][kScaleBRows]; + // Capture all scales before arrive: non-storing warps can preload + // the next tile immediately after the final WGMMA retires. + #pragma unroll + for (uint32_t sub_k = 0; sub_k < 4; ++ sub_k) { + const auto sf_a = smem_sfa[stage_idx] + sub_k * BLOCK_M + m_offset; + const uint32_t sf_k = k_block_idx * 4 + sub_k; + sa0[sub_k] = do_wgmma_store ? ptx::ld_shared(sf_a + r_0) : 0; + sa1[sub_k] = do_wgmma_store ? ptx::ld_shared(sf_a + r_1) : 0; + #pragma unroll + for (uint32_t i = 0; i < kScaleBRows; ++ i) + sb[sub_k][i] = do_wgmma_store and sf_k < shape_k_scales ? + ptx::ld_shared(smem_sfb + i * shape_k_scales + sf_k) : 0; + } + auto promote = [&](const uint32_t sub_k, float* partial) { + if (not do_wgmma_store) return; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float scale0 = sa0[sub_k] * sb[sub_k][i / 4]; + const float scale1 = sa1[sub_k] * sb[sub_k][i / 4]; + shifted_accum[i * 4 + 0] += scale0 * partial[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale0 * partial[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale1 * partial[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale1 * partial[i * 4 + 3]; + } + }; + if constexpr (kOverlapPromotion) { + #pragma unroll + for (uint32_t sub_k = 0; sub_k < 4; ++ sub_k) { + auto partial = accum + (sub_k % 2) * WGMMA::kNumAccum; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(partial[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * 32) / 16; + b_desc.reg32_[0] = b_desc_base_lo + sub_k * 32 / 16; + WGMMA::wgmma(a_desc, b_desc, partial, 0); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(partial[i]); + if (sub_k > 0) { + ptx::warpgroup_wait<1>(); + promote(sub_k - 1, accum + ((sub_k - 1) % 2) * WGMMA::kNumAccum); + } + } + ptx::warpgroup_wait<0>(); + if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1) + empty_barrier_arrive(); + promote(3, accum + WGMMA::kNumAccum); + } else { + #pragma unroll + for (uint32_t batch = 0; batch < 4 / kAccumBanks; ++ batch) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum * kAccumBanks; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t bank = 0; bank < kAccumBanks; ++ bank) { + const auto sub_k = batch * kAccumBanks + bank; + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * 32) / 16; + b_desc.reg32_[0] = b_desc_base_lo + sub_k * 32 / 16; + WGMMA::wgmma(a_desc, b_desc, accum + bank * WGMMA::kNumAccum, 0); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum * kAccumBanks; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + if (batch == 4 / kAccumBanks - 1 and local_idx == BLOCK_M / WAVE_BLOCK_M - 1) + empty_barrier_arrive(); + #pragma unroll + for (uint32_t bank = 0; bank < kAccumBanks; ++ bank) + promote(batch * kAccumBanks + bank, accum + bank * WGMMA::kNumAccum); + } + } + } + } else { // Read B scales float scale_b_0 = ptx::ld_shared(smem_sfb + k_block_idx), scale_b_1; // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks @@ -346,6 +448,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3]; } } + } } }); } else { From e462f4daaf18f261377847d600c2cf4107d66dd3 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Tue, 29 Sep 2026 17:38:59 +0300 Subject: [PATCH 2/4] Clean up SM90 block32 configuration and validation Move block32 capability checks from the shared scale-layout helper to GEMM dispatch. Derive scale counts and offsets from block parameters, centralize accumulator-bank selection, clarify N/K specialization, and fix block128 branch indentation. Validated on H200: 65 existing gates and 8 API checks passed; all 360 dense block32 points match the previous commit bit for bit. Paired full-path timings across 40 points give a 1.003x geometric mean baseline/cleanup ratio. --- csrc/apis/gemm.hpp | 4 +- csrc/apis/layout.hpp | 12 +- csrc/jit_kernels/heuristics/sm90.hpp | 36 ++-- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 9 +- .../deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh | 182 +++++++++--------- 5 files changed, 128 insertions(+), 115 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index ee9f0948f..42bed3e18 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -195,7 +195,7 @@ static void fp8_fp4_gemm_nt(const std::pair& a, if (arch_major == 9 or arch_major == 10) { // SM90/SM100 share the "transform scaling factors, then dispatch" flow. const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( - a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, disable_ue8m0_cast, std::nullopt, true); + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, disable_ue8m0_cast); if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) { DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100"); @@ -321,6 +321,7 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair transform_sf_pair_into const std::optional& num_groups_b, const bool& disable_ue8m0_cast = false, // PSUM layout only applies to SFA (B is weights with no gaps) - const std::optional& psum_layout = std::nullopt, - const bool allow_sm90_block32 = false) { + const std::optional& psum_layout = std::nullopt) { // Use default recipe, if none is specified if (not recipe_a.has_value() and not recipe.has_value()) recipe = get_default_recipe(sfa.scalar_type(), sfb.scalar_type()); @@ -77,11 +76,6 @@ static std::tuple transform_sf_pair_into DG_HOST_ASSERT(recipe_a.has_value() == recipe_b.has_value()); DG_HOST_ASSERT(recipe_a.has_value() != recipe.has_value()); - if (jit->device.get_arch_major() == 9 and not allow_sm90_block32) { - DG_HOST_ASSERT((recipe.has_value() ? std::get<2>(*recipe) : std::get<1>(*recipe_a)) == 128); - DG_HOST_ASSERT((recipe.has_value() ? std::get<2>(*recipe) : std::get<1>(*recipe_b)) == 128); - } - // Transform SFA and SFB layout (PSUM only for SFA) const auto transformed_sfa = recipe.has_value() ? transform_sf_into_required_layout(sfa, m, k, recipe.value(), num_groups_a, true, disable_ue8m0_cast, psum_layout) : transform_sf_into_required_layout(sfa, m, k, recipe_a.value(), num_groups_a, std::nullopt, disable_ue8m0_cast, psum_layout); diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index aa7b4acfb..a2575c395 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -14,11 +14,17 @@ namespace deep_gemm { struct SM90ArchSpec { static constexpr int smem_capacity = 232448; - static constexpr int block32_short_k = 8 * 128; - static constexpr int block32_small_m_k = 2 * block32_short_k; - static std::vector get_layout_candidates(const GemmDesc& desc) { + constexpr int short_k_stages = 8; + // Block K is fixed by the MMA operand width. + const int block_k = 128 * 8 / get_num_element_bits(desc.get_mma_kind()); const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32; + const int n_tile_multiple = block32 ? desc.max_gran_k : 16; + const int max_1d2d_block_n = block32 ? block_k : 192; + const int short_k = short_k_stages * block_k; + const int max_block_m = desc.k <= short_k ? 2 * block_k : block_k; + const bool allow_small_m = not block32 or desc.k <= 2 * short_k; + // Block M candidates std::vector block_m_candidates; if (desc.gemm_type == GemmType::Normal or @@ -27,8 +33,8 @@ struct SM90ArchSpec { // TODO: check 256's performance block_m_candidates = {64, 128}; // NOTES: smaller block M can avoid TMA L2 OOB bound - if (desc.m <= 16 and (not block32 or desc.k <= block32_small_m_k)) block_m_candidates.push_back(16); - if (desc.m <= 32 and (not block32 or desc.k <= block32_small_m_k)) block_m_candidates.push_back(32); + if (desc.m <= 16 and allow_small_m) block_m_candidates.push_back(16); + if (desc.m <= 32 and allow_small_m) block_m_candidates.push_back(32); // BF16 output GEMM supports 256 if (desc.cd_dtype != torch::kFloat) @@ -42,7 +48,7 @@ struct SM90ArchSpec { // Block N candidates std::vector block_n_candidates; - int step = std::lcm(block32 ? 32 : 16, heuristics_runtime->get_block_n_multiple_of()); + int step = std::lcm(n_tile_multiple, heuristics_runtime->get_block_n_multiple_of()); int start = step; // Avoid bank conflicts for 1D1D kernel FP32 output if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { @@ -54,16 +60,13 @@ struct SM90ArchSpec { // Register spills int end = 256; if (desc.kernel_type == KernelType::Kernel1D2D) - end = block32 ? 128 : 192; + end = max_1d2d_block_n; if (desc.kernel_type == KernelType::Kernel1D1D) end = 160; // Enumerate for (int i = start; i <= end; i += step) block_n_candidates.push_back(i); - // Block K is always in a fixed manner - const int block_k = 128 * 8 / get_num_element_bits(desc.get_mma_kind()); - // Disable multicast for performance const bool disable_multicast = // The number of k-groups is large (a heuristic) @@ -87,7 +90,7 @@ struct SM90ArchSpec { for (int block_n: block_n_candidates) { // Long K favors overlapping wide single-wave tiles; short K can // amortize B loads across two M waves without spilling. - if (block32 and block_m > (desc.k <= block32_short_k ? 256 : 128)) + if (block32 and block_m > max_block_m) continue; // 1D2D kernel unroll requirement if (desc.kernel_type == KernelType::Kernel1D2D and block_n > block_k and (block_n % (block_n - block_k) != 0 and block_k % (block_n - block_k) != 0)) @@ -156,6 +159,11 @@ struct SM90ArchSpec { static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) { constexpr int kNumMaxStages = 16; + const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32; + const int scale_gran_k = block32 ? desc.max_gran_k : layout.block_k; + const int scales_per_stage = layout.block_k / scale_gran_k; + const int scale_b_rows = block32 ? layout.block_n / scale_gran_k : + (layout.block_k % layout.block_n == 0 ? 1 : 2); // TODO: consider swap AB // C/D for TMA stores @@ -169,16 +177,14 @@ struct SM90ArchSpec { const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype); // Calculate SF A/B per stages - const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32; const int smem_sfa_per_stage = desc.kernel_type == KernelType::KernelNoSF ? - 0 : align(layout.block_m * (block32 ? 4 : 1) * static_cast(sizeof(float)), 128); + 0 : align(layout.block_m * scales_per_stage * static_cast(sizeof(float)), 128); const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ? 0 : align(layout.block_n * static_cast(sizeof(float)), 128); // Extra SFB sizes for 1D2D kernels - const int use_uniform_sfb = layout.block_k % layout.block_n == 0 ? 1 : 2; const int smem_extra_sfb = desc.kernel_type != KernelType::Kernel1D2D ? - 0 : align(ceil_div(desc.k, block32 ? 32 : layout.block_k) * static_cast(sizeof(float)) * (block32 ? layout.block_n / 32 : use_uniform_sfb), 8); + 0 : align(ceil_div(desc.k, scale_gran_k) * scale_b_rows * static_cast(sizeof(float)), 8); // Extra tensormap for 1D1D kernels const int smem_tensormap = diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 10d868da0..a527ab422 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -35,10 +35,11 @@ class SM90FP8Gemm1D2DRuntime final { const bool specialize_n = args.gemm_desc.gemm_type == GemmType::Normal && args.gemm_desc.compiled_dims.empty() && args.gemm_config.layout.block_m == 128 && args.gemm_config.layout.block_n > 128 && args.gemm_config.layout.get_cluster_size() == 1; - // Block32 scale addressing benefits from fixed weight dimensions. - // M remains dynamic, so one specialization serves decode and prefill. - const auto compiled_dims = args.sf_granularity == 32 and args.gemm_desc.compiled_dims.empty() ? - std::string("nk") : args.gemm_desc.compiled_dims; + // Specialize fixed N/K to simplify scale indexing. + // Keep M dynamic to avoid specialization for each token count. + auto compiled_dims = args.gemm_desc.compiled_dims; + if (args.sf_granularity == 32 and compiled_dims.empty()) + compiled_dims = "nk"; const auto kernel = jit->compile(tag, std::format(R"( #include diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh index 8d0f42199..dc25fbf71 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh @@ -70,6 +70,21 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, using WGMMA = typename mma::sm90::FP8MMASelector::type; using Barrier = cutlass::arch::ClusterTransactionBarrier; DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M, "Invalid block size"); + constexpr bool kIsBlock32 = SF_GRAN == 32; + constexpr uint32_t kScaleGroupsPerStage = BLOCK_K / SF_GRAN; + constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); + constexpr uint32_t kFourBankMaxBlockN = 64; + // Multiple M waves reuse B with one bank; narrower single-wave tiles can + // hold every partial, while wider tiles overlap promotion with two banks. + constexpr uint32_t kAccumBanks = [] { + if constexpr (not kIsBlock32 or BLOCK_M > WAVE_BLOCK_M) + return 1u; + if constexpr (BLOCK_N <= kFourBankMaxBlockN) + return kScaleGroupsPerStage; + return 2u; + }(); + constexpr bool kOverlapPromotion = kIsBlock32 and kAccumBanks == 2; // Overwrite shape constants if the compiler gives shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; @@ -81,11 +96,11 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, static constexpr uint32_t SMEM_D_SIZE = math::constexpr_align(BLOCK_M * BLOCK_N * static_cast(sizeof(__nv_bfloat16)), 1024u); static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3); - static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * (BLOCK_K / SF_GRAN) * sizeof(float); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * kScaleGroupsPerStage * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); const uint32_t shape_k_scales = math::ceil_div(shape_k, SF_GRAN); const uint32_t shape_n_sfb = math::ceil_div(shape_n, SF_GRAN); - constexpr uint32_t kScaleBRows = SF_GRAN == 32 ? BLOCK_N / 32 : (kMustUseUniformedScaleB ? 1 : 2); + constexpr uint32_t kScaleBRows = kIsBlock32 ? BLOCK_N / SF_GRAN : (kMustUseUniformedScaleB ? 1 : 2); const uint32_t smem_sfb_size = math::align(shape_k_scales * kScaleBRows * sizeof(float), sizeof(Barrier)); // NOTES: Make sure we have enough shared memory for WGMMA padding @@ -198,8 +213,8 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, tma::copy(&tensor_map_a, &full_barrier, smem_a[stage_idx], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx), num_tma_multicast_a, batch_idx); - tma::copy(&tensor_map_sfa, &full_barrier, - smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx * (BLOCK_K / SF_GRAN)), + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx * kScaleGroupsPerStage), num_tma_multicast_a); // Issue TMA B @@ -238,7 +253,9 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8; } - uint32_t num_sfb = shape_k_scales * (SF_GRAN == 32 ? kScaleBRows : (num_former_iters >= num_full_iters ? 1 : 2)); + uint32_t num_sfb = shape_k_scales * kScaleBRows; + if constexpr (not kIsBlock32) + num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2); // Load B scales with math warp-groups // NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks @@ -248,7 +265,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, const uint32_t stride_k_sfb = kMajorSFB == cute::UMMA::Major::MN ? shape_n_sfb : 1; auto local_sfb = sfb + previous_group_offset + ((n_block_idx * BLOCK_N) / SF_GRAN) * stride_n_sfb; - if constexpr (SF_GRAN == 32) { + if constexpr (kIsBlock32) { #pragma unroll for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32) { const auto row = i / shape_k_scales, col = i % shape_k_scales; @@ -264,13 +281,6 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); // Accumulation for WGMMA or CUDA promotion - constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; - DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); - // Full-M tiles reuse B with one bank; wide single-wave tiles overlap - // promotion of the previous partial with WGMMA on the other bank. - constexpr uint32_t kAccumBanks = SF_GRAN == 32 ? - (BLOCK_M > 128 ? 1 : (BLOCK_N <= 64 ? 4 : 2)) : 1; - constexpr bool kOverlapPromotion = SF_GRAN == 32 and kAccumBanks == 2; float accum[WGMMA::kNumAccum * kAccumBanks], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; // Pick threads whose WGMMA results are to be stored in shared memory @@ -291,7 +301,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, // Skip useless computations if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { // The compiler must know the dynamic variable `num_former_iters`'s real value - constexpr bool kShouldOptimize = SF_GRAN == 128 and BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr bool kShouldOptimize = not kIsBlock32 and BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; @@ -302,19 +312,19 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); const auto b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16); - if constexpr (SF_GRAN == 32) { + if constexpr (kIsBlock32) { full_barriers[stage_idx]->wait(phase); #pragma unroll for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { const auto m_offset = local_idx * WAVE_BLOCK_M; auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; - float sa0[4], sa1[4], sb[4][kScaleBRows]; + float sa0[kScaleGroupsPerStage], sa1[kScaleGroupsPerStage], sb[kScaleGroupsPerStage][kScaleBRows]; // Capture all scales before arrive: non-storing warps can preload // the next tile immediately after the final WGMMA retires. #pragma unroll - for (uint32_t sub_k = 0; sub_k < 4; ++ sub_k) { + for (uint32_t sub_k = 0; sub_k < kScaleGroupsPerStage; ++ sub_k) { const auto sf_a = smem_sfa[stage_idx] + sub_k * BLOCK_M + m_offset; - const uint32_t sf_k = k_block_idx * 4 + sub_k; + const uint32_t sf_k = k_block_idx * kScaleGroupsPerStage + sub_k; sa0[sub_k] = do_wgmma_store ? ptx::ld_shared(sf_a + r_0) : 0; sa1[sub_k] = do_wgmma_store ? ptx::ld_shared(sf_a + r_1) : 0; #pragma unroll @@ -336,14 +346,14 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, }; if constexpr (kOverlapPromotion) { #pragma unroll - for (uint32_t sub_k = 0; sub_k < 4; ++ sub_k) { - auto partial = accum + (sub_k % 2) * WGMMA::kNumAccum; + for (uint32_t sub_k = 0; sub_k < kScaleGroupsPerStage; ++ sub_k) { + auto partial = accum + (sub_k % kAccumBanks) * WGMMA::kNumAccum; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(partial[i]); ptx::warpgroup_arrive(); - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * 32) / 16; - b_desc.reg32_[0] = b_desc_base_lo + sub_k * 32 / 16; + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * SF_GRAN) / 16; + b_desc.reg32_[0] = b_desc_base_lo + sub_k * SF_GRAN / 16; WGMMA::wgmma(a_desc, b_desc, partial, 0); ptx::warpgroup_commit_batch(); #pragma unroll @@ -351,16 +361,16 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, ptx::warpgroup_fence_operand(partial[i]); if (sub_k > 0) { ptx::warpgroup_wait<1>(); - promote(sub_k - 1, accum + ((sub_k - 1) % 2) * WGMMA::kNumAccum); + promote(sub_k - 1, accum + ((sub_k - 1) % kAccumBanks) * WGMMA::kNumAccum); } } ptx::warpgroup_wait<0>(); if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1) empty_barrier_arrive(); - promote(3, accum + WGMMA::kNumAccum); + promote(kScaleGroupsPerStage - 1, accum + (kAccumBanks - 1) * WGMMA::kNumAccum); } else { #pragma unroll - for (uint32_t batch = 0; batch < 4 / kAccumBanks; ++ batch) { + for (uint32_t batch = 0; batch < kScaleGroupsPerStage / kAccumBanks; ++ batch) { #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum * kAccumBanks; ++ i) ptx::warpgroup_fence_operand(accum[i]); @@ -368,8 +378,8 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, #pragma unroll for (uint32_t bank = 0; bank < kAccumBanks; ++ bank) { const auto sub_k = batch * kAccumBanks + bank; - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * 32) / 16; - b_desc.reg32_[0] = b_desc_base_lo + sub_k * 32 / 16; + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + sub_k * SF_GRAN) / 16; + b_desc.reg32_[0] = b_desc_base_lo + sub_k * SF_GRAN / 16; WGMMA::wgmma(a_desc, b_desc, accum + bank * WGMMA::kNumAccum, 0); } ptx::warpgroup_commit_batch(); @@ -377,7 +387,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, for (uint32_t i = 0; i < WGMMA::kNumAccum * kAccumBanks; ++ i) ptx::warpgroup_fence_operand(accum[i]); ptx::warpgroup_wait<0>(); - if (batch == 4 / kAccumBanks - 1 and local_idx == BLOCK_M / WAVE_BLOCK_M - 1) + if (batch == kScaleGroupsPerStage / kAccumBanks - 1 and local_idx == BLOCK_M / WAVE_BLOCK_M - 1) empty_barrier_arrive(); #pragma unroll for (uint32_t bank = 0; bank < kAccumBanks; ++ bank) @@ -386,69 +396,69 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, } } } else { - // Read B scales - float scale_b_0 = ptx::ld_shared(smem_sfb + k_block_idx), scale_b_1; - // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks - if constexpr (not kMustUseUniformedScaleB) - scale_b_1 = ptx::ld_shared(smem_sfb + k_block_idx + shape_k_scales); - - // Wait TMA arrivals - full_barriers[stage_idx]->wait(phase); - - // TODO: remove some useless computation for unaligned Ms - #pragma unroll - for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { - auto m_offset = local_idx * WAVE_BLOCK_M; - - // Read A scales - // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results - auto scale_a_0 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_0 + m_offset) : 0; - auto scale_a_1 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_1 + m_offset) : 0; - - // Commit WGMMA instructions - #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) { - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16; - b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, 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>(); - - // Notify barrier arrival at the last warpgroup wave - if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1) - empty_barrier_arrive(); - - // Skip promotion for the unfilled parts - if (not do_wgmma_store) - continue; - - // Promote with scales - // NOTES: making it as predicates is very important for performance, comparing to two loops - float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0; - float scale_0_1, scale_1_1; + // Read B scales + float scale_b_0 = ptx::ld_shared(smem_sfb + k_block_idx), scale_b_1; + // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks if constexpr (not kMustUseUniformedScaleB) - scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1; + scale_b_1 = ptx::ld_shared(smem_sfb + k_block_idx + shape_k_scales); + + // Wait TMA arrivals + full_barriers[stage_idx]->wait(phase); - auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + // TODO: remove some useless computation for unaligned Ms #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant - const bool predicate = kMustUseUniformedScaleB or i < num_former_iters; - shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0]; - shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1]; - shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2]; - shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3]; + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto m_offset = local_idx * WAVE_BLOCK_M; + + // Read A scales + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + auto scale_a_0 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_0 + m_offset) : 0; + auto scale_a_1 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_1 + m_offset) : 0; + + // Commit WGMMA instructions + #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) { + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, 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>(); + + // Notify barrier arrival at the last warpgroup wave + if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1) + empty_barrier_arrive(); + + // Skip promotion for the unfilled parts + if (not do_wgmma_store) + continue; + + // Promote with scales + // NOTES: making it as predicates is very important for performance, comparing to two loops + float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0; + float scale_0_1, scale_1_1; + if constexpr (not kMustUseUniformedScaleB) + scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1; + + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant + const bool predicate = kMustUseUniformedScaleB or i < num_former_iters; + shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0]; + shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1]; + shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2]; + shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3]; + } } } - } } }); } else { From a14a476399baae462415be587f8ff94214f60e74 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Wed, 30 Sep 2026 15:43:48 +0300 Subject: [PATCH 3/4] Validate SM90 activation scales and cover block32 GEMM Require one activation scale per row in the SM90 dense FP8 dispatch for both block32 and block128. Add Hopper regression coverage with an independent dequantized reference, K-stage fragments, M/N tails, both B-scale layouts, all three accumulator-bank schedules, and invalid recipe rejection. --- csrc/apis/gemm.hpp | 3 + sgl_deep_gemm/run_tests.sh | 6 + sgl_deep_gemm/tests/test_sm90_block32.py | 150 +++++++++++++++++++++++ 3 files changed, 159 insertions(+) create mode 100644 sgl_deep_gemm/tests/test_sm90_block32.py diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42bed3e18..b566de0a4 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -199,7 +199,10 @@ static void fp8_fp4_gemm_nt(const std::pair& a, if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) { DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100"); + // SM90 kernels read one activation scale per row. + const int gran_m = recipe.has_value() ? std::get<0>(recipe.value()) : std::get<0>(recipe_a.value()); const int gran_n = recipe.has_value() ? std::get<1>(recipe.value()) : std::get<0>(recipe_b.value()); + DG_HOST_ASSERT(gran_m == 1); DG_HOST_ASSERT(gran_k_a == gran_k_b); DG_HOST_ASSERT((gran_n == 32 and gran_k_a == 32) or (gran_k_a == 128 and (gran_n == 1 or gran_n == 128))); if (gran_n == 1) { diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index e6e4d9f1f..ab1c96037 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -150,6 +150,12 @@ for t in "${SINGLE_GPU_TESTS[@]}"; do fi done +if [ "${ARCH_MAJOR}" -eq 9 ]; then + run_test test_sm90_block32.py +else + skip_test test_sm90_block32.py "SM90-only, arch major ${ARCH_MAJOR}" +fi + for t in test_mega_gate.py test_mega_mhc.py test_clean_logits_bounds.py; do if [ "${ARCH_MAJOR}" -eq 10 ]; then run_test "${t}" diff --git a/sgl_deep_gemm/tests/test_sm90_block32.py b/sgl_deep_gemm/tests/test_sm90_block32.py new file mode 100644 index 000000000..ae7d82dbb --- /dev/null +++ b/sgl_deep_gemm/tests/test_sm90_block32.py @@ -0,0 +1,150 @@ +"""SM90 FP8 block32 correctness and activation-scale contract tests.""" + +import os +import re +import sys +import tempfile + +import torch + +import deep_gemm + + +def _ceil_div(x, y): + return (x + y - 1) // y + + +def _run(a, sa, b, sb, *, capture_config=False, **recipe): + out = torch.empty((a.shape[0], b.shape[0]), device="cuda", dtype=torch.bfloat16) + config_log = "" + if capture_config: + sys.stdout.flush() + with tempfile.TemporaryFile(mode="w+") as trace: + stdout_fd = os.dup(1) + try: + os.dup2(trace.fileno(), 1) + deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) + finally: + os.dup2(stdout_fd, 1) + os.close(stdout_fd) + trace.seek(0) + config_log = trace.read() + else: + deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) + torch.cuda.synchronize() + return out, config_log + + +def _reference(a, sa, b, sb): + # Independent FP32 dequantized reference: each K32 product has its own scales. + n, k = b.shape + out = torch.zeros((a.shape[0], n), device="cuda", dtype=torch.float32) + for group in range(_ceil_div(k, 32)): + start, end = group * 32, min((group + 1) * 32, k) + b_scale = sb[:, group].repeat_interleave(32)[:n] + product = a[:, start:end].float() @ b[:, start:end].float().T + out += product * sa[:, group, None] * b_scale[None, :] + return out.to(torch.bfloat16) + + +def _selected_banks(config_log): + match = re.search(r"block_m=(\d+), block_n=(\d+)", config_log) + assert match is not None, f"SM90 did not report a selected tile: {config_log}" + block_m, block_n = map(int, match.groups()) + # WGMMA M=64; the kernel uses a 128-row wave for larger M tiles. + wave_m = min(block_m, 128) + banks = 1 if block_m > wave_m else 4 if block_n <= 64 else 2 + return banks, block_m, block_n + + +def test_invalid_activation_granularity(): + m = n = 32 + k = 128 + a = torch.ones((m, k), device="cuda").to(torch.float8_e4m3fn) + b = torch.ones((n, k), device="cuda").to(torch.float8_e4m3fn) + invalid = ( + (32, 32, 32, False), + (32, 32, 32, True), + (128, 128, 128, False), + (128, 128, 128, True), + (128, 1, 128, False), + ) + for gran_m, gran_n, gran_k, separate in invalid: + # Allocate enough backing for every row if the guard regresses. + backing = torch.ones((m, _ceil_div(k, gran_k)), device="cuda") + sa = backing[:_ceil_div(m, gran_m)] + sb = torch.ones((_ceil_div(n, gran_n), _ceil_div(k, gran_k)), device="cuda") + recipe = (dict(recipe_a=(gran_m, gran_k), recipe_b=(gran_n, gran_k)) + if separate else dict(recipe=(gran_m, gran_n, gran_k))) + try: + _run(a, sa, b, sb, **recipe) + except RuntimeError as error: + assert "gran_m == 1" in str(error), (recipe, str(error)) + else: + raise AssertionError(f"SM90 accepted unsupported A-scale granularity: {recipe}") + + +def test_block32_correctness(): + # Cover K128-stage fragments, M/N tails, and both physical B-scale layouts. + # The remaining shapes select all three bank schedules on H200; fallbacks + # keep this check useful when another Hopper selects different tiles. + cases = ( + (1, 8, 32, False, False, False), + (17, 40, 96, True, True, False), + (65, 72, 160, False, True, False), + (129, 96, 224, True, False, False), + (128, 32, 2048, False, False, True), + (2048, 1024, 512, True, True, True), + (2048, 2048, 512, False, False, True), + (2048, 1024, 2048, True, True, True), + (4096, 2048, 512, False, False, True), + ) + observed_banks = set() + for m, n, k, transpose_sb, separate_recipe, check_banks in cases: + a = torch.randint(0, 4, (m, k), device="cuda").float().to(torch.float8_e4m3fn) + b = torch.randint(0, 4, (n, k), device="cuda").float().to(torch.float8_e4m3fn) + sa = torch.pow(2.0, torch.randint(-3, 1, (m, _ceil_div(k, 32)), device="cuda").float()) + sb = torch.pow(2.0, torch.randint(-3, 1, (_ceil_div(n, 32), _ceil_div(k, 32)), device="cuda").float()) + if transpose_sb: + sb = sb.T.contiguous().T + recipe = (dict(recipe_a=(1, 32), recipe_b=(32, 32)) + if separate_recipe else dict(recipe=(1, 32, 32))) + actual, config_log = _run(a, sa, b, sb, capture_config=check_banks, **recipe) + expected = _reference(a, sa, b, sb) + assert torch.equal(actual, expected), ( + f"block32 mismatch at {m=}, {n=}, {k=}, {transpose_sb=}, {separate_recipe=}; " + f"max error={(actual.float() - expected.float()).abs().max().item()}") + bank_note = "" + if check_banks: + banks, block_m, block_n = _selected_banks(config_log) + observed_banks.add(banks) + bank_note = f" banks={banks} tile=({block_m},{block_n})" + print(f"PASS block32 M={m} N={n} K={k} transposed_B_scales={transpose_sb}{bank_note}", flush=True) + if observed_banks == {1, 2, 4}: + break + assert observed_banks == {1, 2, 4}, observed_banks + + +def test_block128_control(): + m = n = 32 + k = 128 + a = torch.ones((m, k), device="cuda").to(torch.float8_e4m3fn) + b = torch.ones((n, k), device="cuda").to(torch.float8_e4m3fn) + sa = torch.ones((m, 1), device="cuda") + sb = torch.ones((1, 1), device="cuda") + for recipe in (dict(recipe=(1, 128, 128)), + dict(recipe_a=(1, 128), recipe_b=(128, 128))): + actual, _ = _run(a, sa, b, sb, **recipe) + assert torch.all(actual == 128), recipe + + +if __name__ == "__main__": + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + print("SKIP: SM90 GPU required") + else: + os.environ["DG_PRINT_CONFIGS"] = "1" + torch.manual_seed(32) + test_invalid_activation_granularity() + test_block32_correctness() + test_block128_control() + print("PASS: SM90 block32/activation-scale cases") From cbac37fe007ae36cad127222b6a245a4c8f916a7 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Wed, 30 Sep 2026 16:11:05 +0300 Subject: [PATCH 4/4] Fold SM90 block32 regressions into FP8 GEMM tests Run native block32 correctness, tail, scale-layout, bank-schedule, and invalid activation-scale cases from the existing FP8 GEMM test entrypoint. Remove the dedicated SM90 test file and runner dispatch; keep the existing block128 positive coverage in the general suite. --- sgl_deep_gemm/run_tests.sh | 6 - sgl_deep_gemm/tests/test_fp8_fp4.py | 133 ++++++++++++++++++++ sgl_deep_gemm/tests/test_sm90_block32.py | 150 ----------------------- 3 files changed, 133 insertions(+), 156 deletions(-) delete mode 100644 sgl_deep_gemm/tests/test_sm90_block32.py diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index ab1c96037..e6e4d9f1f 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -150,12 +150,6 @@ for t in "${SINGLE_GPU_TESTS[@]}"; do fi done -if [ "${ARCH_MAJOR}" -eq 9 ]; then - run_test test_sm90_block32.py -else - skip_test test_sm90_block32.py "SM90-only, arch major ${ARCH_MAJOR}" -fi - for t in test_mega_gate.py test_mega_mhc.py test_clean_logits_bounds.py; do if [ "${ARCH_MAJOR}" -eq 10 ]; then run_test "${t}" diff --git a/sgl_deep_gemm/tests/test_fp8_fp4.py b/sgl_deep_gemm/tests/test_fp8_fp4.py index 1b079d99c..482b8cd2d 100644 --- a/sgl_deep_gemm/tests/test_fp8_fp4.py +++ b/sgl_deep_gemm/tests/test_fp8_fp4.py @@ -1,4 +1,9 @@ +import os import random +import re +import sys +import tempfile + import torch import deep_gemm @@ -372,6 +377,133 @@ def test_func(): print() +def _sm90_run(a, sa, b, sb, *, capture_config=False, **recipe): + out = torch.empty((a.shape[0], b.shape[0]), device="cuda", dtype=torch.bfloat16) + config_log = "" + if capture_config: + sys.stdout.flush() + with tempfile.TemporaryFile(mode="w+") as trace: + stdout_fd = os.dup(1) + try: + os.dup2(trace.fileno(), 1) + deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) + finally: + os.dup2(stdout_fd, 1) + os.close(stdout_fd) + trace.seek(0) + config_log = trace.read() + else: + deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) + torch.cuda.synchronize() + return out, config_log + + +def _sm90_reference(a, sa, b, sb): + # Independent FP32 dequantized reference: each K32 product has its own scales. + n, k = b.shape + out = torch.zeros((a.shape[0], n), device="cuda", dtype=torch.float32) + for group in range(ceil_div(k, 32)): + start, end = group * 32, min((group + 1) * 32, k) + b_scale = sb[:, group].repeat_interleave(32)[:n] + product = a[:, start:end].float() @ b[:, start:end].float().T + out += product * sa[:, group, None] * b_scale[None, :] + return out.to(torch.bfloat16) + + +def _sm90_selected_banks(config_log): + match = re.search(r"block_m=(\d+), block_n=(\d+)", config_log) + assert match is not None, f"SM90 did not report a selected tile: {config_log}" + block_m, block_n = map(int, match.groups()) + # WGMMA M=64; the kernel uses a 128-row wave for larger M tiles. + wave_m = min(block_m, 128) + banks = 1 if block_m > wave_m else 4 if block_n <= 64 else 2 + return banks, block_m, block_n + + +def _check_sm90_invalid_activation_granularity(): + m = n = 32 + k = 128 + a = torch.ones((m, k), device="cuda").to(torch.float8_e4m3fn) + b = torch.ones((n, k), device="cuda").to(torch.float8_e4m3fn) + invalid = ( + (32, 32, 32, False), + (32, 32, 32, True), + (128, 128, 128, False), + (128, 128, 128, True), + (128, 1, 128, False), + ) + for gran_m, gran_n, gran_k, separate in invalid: + # Allocate enough backing for every row if the guard regresses. + backing = torch.ones((m, ceil_div(k, gran_k)), device="cuda") + sa = backing[:ceil_div(m, gran_m)] + sb = torch.ones((ceil_div(n, gran_n), ceil_div(k, gran_k)), device="cuda") + recipe = (dict(recipe_a=(gran_m, gran_k), recipe_b=(gran_n, gran_k)) + if separate else dict(recipe=(gran_m, gran_n, gran_k))) + try: + _sm90_run(a, sa, b, sb, **recipe) + except RuntimeError as error: + assert "gran_m == 1" in str(error), (recipe, str(error)) + else: + raise AssertionError(f"SM90 accepted unsupported A-scale granularity: {recipe}") + + +def _check_sm90_block32_correctness(): + # Cover K128-stage fragments, M/N tails, and both physical B-scale layouts. + # The remaining shapes select all three bank schedules on H200; fallbacks + # keep this check useful when another Hopper selects different tiles. + cases = ( + (1, 8, 32, False, False, False), + (17, 40, 96, True, True, False), + (65, 72, 160, False, True, False), + (129, 96, 224, True, False, False), + (128, 32, 2048, False, False, True), + (2048, 1024, 512, True, True, True), + (2048, 2048, 512, False, False, True), + (2048, 1024, 2048, True, True, True), + (4096, 2048, 512, False, False, True), + ) + observed_banks = set() + for m, n, k, transpose_sb, separate_recipe, check_banks in cases: + a = torch.randint(0, 4, (m, k), device="cuda").float().to(torch.float8_e4m3fn) + b = torch.randint(0, 4, (n, k), device="cuda").float().to(torch.float8_e4m3fn) + sa = torch.pow(2.0, torch.randint(-3, 1, (m, ceil_div(k, 32)), device="cuda").float()) + sb = torch.pow(2.0, torch.randint(-3, 1, (ceil_div(n, 32), ceil_div(k, 32)), device="cuda").float()) + if transpose_sb: + sb = sb.T.contiguous().T + recipe = (dict(recipe_a=(1, 32), recipe_b=(32, 32)) + if separate_recipe else dict(recipe=(1, 32, 32))) + actual, config_log = _sm90_run(a, sa, b, sb, capture_config=check_banks, **recipe) + expected = _sm90_reference(a, sa, b, sb) + assert torch.equal(actual, expected), ( + f"block32 mismatch at {m=}, {n=}, {k=}, {transpose_sb=}, {separate_recipe=}; " + f"max error={(actual.float() - expected.float()).abs().max().item()}") + bank_note = "" + if check_banks: + banks, block_m, block_n = _sm90_selected_banks(config_log) + observed_banks.add(banks) + bank_note = f" banks={banks} tile=({block_m},{block_n})" + print(f"PASS block32 M={m} N={n} K={k} transposed_B_scales={transpose_sb}{bank_note}", flush=True) + if observed_banks == {1, 2, 4}: + break + assert observed_banks == {1, 2, 4}, observed_banks + + +def test_sm90_block32() -> None: + if get_arch_major() != 9: + return + previous_print_configs = os.environ.get("DG_PRINT_CONFIGS") + os.environ["DG_PRINT_CONFIGS"] = "1" + try: + torch.manual_seed(32) + _check_sm90_invalid_activation_granularity() + _check_sm90_block32_correctness() + finally: + if previous_print_configs is None: + os.environ.pop("DG_PRINT_CONFIGS", None) + else: + os.environ["DG_PRINT_CONFIGS"] = previous_print_configs + + if __name__ == '__main__': torch.manual_seed(0) random.seed(0) @@ -383,3 +515,4 @@ def test_func(): test_m_grouped_gemm_contiguous() test_m_grouped_gemm_masked() test_k_grouped_gemm_contiguous() + test_sm90_block32() diff --git a/sgl_deep_gemm/tests/test_sm90_block32.py b/sgl_deep_gemm/tests/test_sm90_block32.py deleted file mode 100644 index ae7d82dbb..000000000 --- a/sgl_deep_gemm/tests/test_sm90_block32.py +++ /dev/null @@ -1,150 +0,0 @@ -"""SM90 FP8 block32 correctness and activation-scale contract tests.""" - -import os -import re -import sys -import tempfile - -import torch - -import deep_gemm - - -def _ceil_div(x, y): - return (x + y - 1) // y - - -def _run(a, sa, b, sb, *, capture_config=False, **recipe): - out = torch.empty((a.shape[0], b.shape[0]), device="cuda", dtype=torch.bfloat16) - config_log = "" - if capture_config: - sys.stdout.flush() - with tempfile.TemporaryFile(mode="w+") as trace: - stdout_fd = os.dup(1) - try: - os.dup2(trace.fileno(), 1) - deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) - finally: - os.dup2(stdout_fd, 1) - os.close(stdout_fd) - trace.seek(0) - config_log = trace.read() - else: - deep_gemm.fp8_gemm_nt((a, sa), (b, sb), out, **recipe) - torch.cuda.synchronize() - return out, config_log - - -def _reference(a, sa, b, sb): - # Independent FP32 dequantized reference: each K32 product has its own scales. - n, k = b.shape - out = torch.zeros((a.shape[0], n), device="cuda", dtype=torch.float32) - for group in range(_ceil_div(k, 32)): - start, end = group * 32, min((group + 1) * 32, k) - b_scale = sb[:, group].repeat_interleave(32)[:n] - product = a[:, start:end].float() @ b[:, start:end].float().T - out += product * sa[:, group, None] * b_scale[None, :] - return out.to(torch.bfloat16) - - -def _selected_banks(config_log): - match = re.search(r"block_m=(\d+), block_n=(\d+)", config_log) - assert match is not None, f"SM90 did not report a selected tile: {config_log}" - block_m, block_n = map(int, match.groups()) - # WGMMA M=64; the kernel uses a 128-row wave for larger M tiles. - wave_m = min(block_m, 128) - banks = 1 if block_m > wave_m else 4 if block_n <= 64 else 2 - return banks, block_m, block_n - - -def test_invalid_activation_granularity(): - m = n = 32 - k = 128 - a = torch.ones((m, k), device="cuda").to(torch.float8_e4m3fn) - b = torch.ones((n, k), device="cuda").to(torch.float8_e4m3fn) - invalid = ( - (32, 32, 32, False), - (32, 32, 32, True), - (128, 128, 128, False), - (128, 128, 128, True), - (128, 1, 128, False), - ) - for gran_m, gran_n, gran_k, separate in invalid: - # Allocate enough backing for every row if the guard regresses. - backing = torch.ones((m, _ceil_div(k, gran_k)), device="cuda") - sa = backing[:_ceil_div(m, gran_m)] - sb = torch.ones((_ceil_div(n, gran_n), _ceil_div(k, gran_k)), device="cuda") - recipe = (dict(recipe_a=(gran_m, gran_k), recipe_b=(gran_n, gran_k)) - if separate else dict(recipe=(gran_m, gran_n, gran_k))) - try: - _run(a, sa, b, sb, **recipe) - except RuntimeError as error: - assert "gran_m == 1" in str(error), (recipe, str(error)) - else: - raise AssertionError(f"SM90 accepted unsupported A-scale granularity: {recipe}") - - -def test_block32_correctness(): - # Cover K128-stage fragments, M/N tails, and both physical B-scale layouts. - # The remaining shapes select all three bank schedules on H200; fallbacks - # keep this check useful when another Hopper selects different tiles. - cases = ( - (1, 8, 32, False, False, False), - (17, 40, 96, True, True, False), - (65, 72, 160, False, True, False), - (129, 96, 224, True, False, False), - (128, 32, 2048, False, False, True), - (2048, 1024, 512, True, True, True), - (2048, 2048, 512, False, False, True), - (2048, 1024, 2048, True, True, True), - (4096, 2048, 512, False, False, True), - ) - observed_banks = set() - for m, n, k, transpose_sb, separate_recipe, check_banks in cases: - a = torch.randint(0, 4, (m, k), device="cuda").float().to(torch.float8_e4m3fn) - b = torch.randint(0, 4, (n, k), device="cuda").float().to(torch.float8_e4m3fn) - sa = torch.pow(2.0, torch.randint(-3, 1, (m, _ceil_div(k, 32)), device="cuda").float()) - sb = torch.pow(2.0, torch.randint(-3, 1, (_ceil_div(n, 32), _ceil_div(k, 32)), device="cuda").float()) - if transpose_sb: - sb = sb.T.contiguous().T - recipe = (dict(recipe_a=(1, 32), recipe_b=(32, 32)) - if separate_recipe else dict(recipe=(1, 32, 32))) - actual, config_log = _run(a, sa, b, sb, capture_config=check_banks, **recipe) - expected = _reference(a, sa, b, sb) - assert torch.equal(actual, expected), ( - f"block32 mismatch at {m=}, {n=}, {k=}, {transpose_sb=}, {separate_recipe=}; " - f"max error={(actual.float() - expected.float()).abs().max().item()}") - bank_note = "" - if check_banks: - banks, block_m, block_n = _selected_banks(config_log) - observed_banks.add(banks) - bank_note = f" banks={banks} tile=({block_m},{block_n})" - print(f"PASS block32 M={m} N={n} K={k} transposed_B_scales={transpose_sb}{bank_note}", flush=True) - if observed_banks == {1, 2, 4}: - break - assert observed_banks == {1, 2, 4}, observed_banks - - -def test_block128_control(): - m = n = 32 - k = 128 - a = torch.ones((m, k), device="cuda").to(torch.float8_e4m3fn) - b = torch.ones((n, k), device="cuda").to(torch.float8_e4m3fn) - sa = torch.ones((m, 1), device="cuda") - sb = torch.ones((1, 1), device="cuda") - for recipe in (dict(recipe=(1, 128, 128)), - dict(recipe_a=(1, 128), recipe_b=(128, 128))): - actual, _ = _run(a, sa, b, sb, **recipe) - assert torch.all(actual == 128), recipe - - -if __name__ == "__main__": - if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: - print("SKIP: SM90 GPU required") - else: - os.environ["DG_PRINT_CONFIGS"] = "1" - torch.manual_seed(32) - test_invalid_activation_granularity() - test_block32_correctness() - test_block128_control() - print("PASS: SM90 block32/activation-scale cases")