From bd42ee2f0e5a9b4cdd530a49c11fde16fced59a3 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 10:40:17 +0800 Subject: [PATCH 01/30] Add SM90 MXFP8 FP8 grouped kernels --- csrc/apis/gemm.hpp | 78 +++++ .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 213 ++++++++++++ deep_gemm/__init__.py | 2 + .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 315 ++++++++++++++++++ tests/test_sm90_mxfp8_fp8.py | 100 ++++++ 5 files changed, 708 insertions(+) create mode 100644 csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh create mode 100644 tests/test_sm90_mxfp8_fp8.py diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42622df7d8..753853deb7 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -5,6 +5,8 @@ #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" +#include "../jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp" +#include "../jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" @@ -268,6 +270,76 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& grouped_layout, + const std::string& compiled_dims) { + (void) compiled_dims; + const auto major_a = get_major_type_ab(a.first); + const auto major_b = get_major_type_ab(b.first); + DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); + DG_HOST_ASSERT(grouped_layout.is_contiguous()); + + const auto arch_major = device_runtime->get_arch_major(); + const auto [m, k] = check_ab_fp8_fp4(a.first, major_a, arch_major); + const auto [num_groups, n, k_] = check_grouped_ab_fp8_fp4(b.first, major_b, arch_major); + const auto [m_, n_] = get_shape<2>(d); + const auto [m__] = get_shape<1>(grouped_layout); + DG_HOST_ASSERT(arch_major == 9); + DG_HOST_ASSERT(m == m_ and m == m__ and n == n_ and k == k_); + DG_HOST_ASSERT(k % 32 == 0 and num_groups > 0); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); + DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); + + if (m == 0) + return; + + const std::variant, std::tuple> sfa_recipe = std::make_tuple(1, 128); + const auto sfa = layout::transform_sf_into_required_layout( + a.second, m, k, sfa_recipe, std::nullopt, std::nullopt, false); + sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( + a.first, sfa, b.first, b.second, d, grouped_layout, num_groups, m, n, k, compiled_dims); +} + +static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& expected_m, + const std::string& compiled_dims) { + (void) expected_m; + (void) compiled_dims; + const auto major_a = get_major_type_ab(a.first); + const auto major_b = get_major_type_ab(b.first); + DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); + DG_HOST_ASSERT(masked_m.is_contiguous()); + + const auto arch_major = device_runtime->get_arch_major(); + const auto [num_groups, m, k] = check_grouped_ab_fp8_fp4(a.first, major_a, arch_major); + const auto [num_groups_, n, k_] = check_grouped_ab_fp8_fp4(b.first, major_b, arch_major); + const auto [num_groups__, m_, n_] = get_shape<3>(d); + DG_HOST_ASSERT(arch_major == 9); + DG_HOST_ASSERT(num_groups == num_groups_ and num_groups == num_groups__); + DG_HOST_ASSERT(masked_m.numel() == num_groups); + DG_HOST_ASSERT(m == m_ and n == n_ and k == k_); + DG_HOST_ASSERT(k % 32 == 0 and m > 0 and n > 0); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); + DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); + + const std::variant, std::tuple> sfa_recipe = std::make_tuple(1, 128); + const auto sfa = layout::transform_sf_into_required_layout( + a.second, m, k, sfa_recipe, num_groups, std::nullopt, false); + sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( + a.first, sfa, b.first, b.second, d, masked_m, num_groups, m, n, k, compiled_dims); +} + static void k_grouped_fp8_gemm_tn_contiguous(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -644,6 +716,12 @@ static void register_apis(pybind11::module_& m) { py::arg("expected_m"), py::arg("recipe") = std::nullopt, py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "nk", py::arg("disable_ue8m0_cast") = false); + m.def("m_grouped_mxfp8_fp8_gemm_nt_contiguous", &m_grouped_mxfp8_fp8_gemm_nt_contiguous, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), + py::arg("compiled_dims") = "nk"); + m.def("m_grouped_mxfp8_fp8_gemm_nt_masked", &m_grouped_mxfp8_fp8_gemm_nt_masked, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"), + py::arg("expected_m"), py::arg("compiled_dims") = "nk"); m.def("k_grouped_fp8_gemm_tn_contiguous", &k_grouped_fp8_gemm_tn_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"), py::arg("ks_tensor"), py::arg("c") = std::nullopt, diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp new file mode 100644 index 0000000000..26873144cb --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -0,0 +1,213 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "../heuristics/sm90.hpp" + +#include "runtime_utils.hpp" + +namespace deep_gemm { + +template +class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime> { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + void *sfb, *grouped_layout; + uint32_t sfb_stride_group, sfb_stride_n, sfb_stride_k; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + CUtensorMap tensor_map_sfa; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, {}, {}, + {}, {}, + {}, {} + >); +}} +)", + kMasked ? "true" : "false", + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.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, + args.gemm_config.pipeline_config.num_stages, + args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, + 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)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.sfb, args.grouped_layout, + args.sfb_stride_group, args.sfb_stride_n, args.sfb_stride_k, + args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.tensor_map_sfa)); + } +}; + +static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const torch::Tensor& grouped_layout, + const int& num_groups, const int& m, const int& n, const int& k, + const std::string& compiled_dims) { + DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and grouped_layout.is_contiguous()); + DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); + + const auto desc = GemmDesc { + .gemm_type = GemmType::MGroupedContiguous, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(), + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = 1 + }; + const auto config = get_best_config(desc); + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(0)), 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(1)), num_groups, + config.storage_config.swizzle_b_mode); + const auto tensor_map_d = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + 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); + + const typename SM90MXFP8FP8Gemm1D2DRuntime::Args& args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .sfb = sfb.data_ptr(), + .grouped_layout = grouped_layout.data_ptr(), + .sfb_stride_group = static_cast(sfb.stride(0)), + .sfb_stride_n = static_cast(sfb.stride(1)), + .sfb_stride_k = static_cast(sfb.stride(2)), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .tensor_map_sfa = tensor_map_sfa, + }; + const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d", code); + SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); +} + +static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const torch::Tensor& masked_m, + const int& num_groups, const int& m, const int& n, const int& k, + const std::string& compiled_dims) { + DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt and masked_m.is_contiguous()); + DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); + + const auto desc = GemmDesc { + .gemm_type = GemmType::MGroupedMasked, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(), + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups + }; + const auto config = get_best_config(desc); + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(1)), num_groups, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(1)), num_groups, + config.storage_config.swizzle_b_mode); + const auto tensor_map_d = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), num_groups, + 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, num_groups, 0); + + const typename SM90MXFP8FP8Gemm1D2DRuntime::Args& args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .sfb = sfb.data_ptr(), + .grouped_layout = masked_m.data_ptr(), + .sfb_stride_group = static_cast(sfb.stride(0)), + .sfb_stride_n = static_cast(sfb.stride(1)), + .sfb_stride_k = static_cast(sfb.stride(2)), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .tensor_map_sfa = tensor_map_sfa, + }; + const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d", code); + SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a9542e2f44..ddea7c59a7 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -47,6 +47,8 @@ m_grouped_fp8_gemm_nt_contiguous, m_grouped_fp8_gemm_nn_contiguous, m_grouped_fp8_gemm_nt_masked, + m_grouped_mxfp8_fp8_gemm_nt_contiguous, + m_grouped_mxfp8_fp8_gemm_nt_masked, k_grouped_fp8_gemm_nt_contiguous, k_grouped_fp8_gemm_tn_contiguous, # BF16 GEMMs diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh new file mode 100644 index 0000000000..04e945ef7b --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -0,0 +1,315 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +namespace mxfp8_fp8_detail { + +CUTLASS_DEVICE float e8m0_to_float(uint8_t scale) { + return __uint_as_float(static_cast(scale) << 23); +} + +} // namespace mxfp8_fp8_detail + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void +sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, + uint32_t sfb_stride_group, uint32_t sfb_stride_n, uint32_t sfb_stride_k, + uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, + const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_d, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) + DG_STATIC_ASSERT(BLOCK_K == 128, "A scale is fixed to per-128 K"); + DG_STATIC_ASSERT(kNumStages >= 1, "Invalid pipeline stages"); + + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + + 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 ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + + const uint32_t num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K); + const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane_idx = ptx::get_lane_idx(); + + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_d); + } + __syncwarp(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + }); + + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); + auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); + + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + cutlass::arch::fence_barrier_init(); + } + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + constexpr uint32_t kNumTMARegisters = 40; + constexpr uint32_t kNumMathRegisters = kNumMathThreads == 128 ? 248 : 232; + + cudaGridDependencySynchronize(); + + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + if (warp_idx >= kNumMathThreads / 32) { + cutlass::arch::warpgroup_reg_dealloc(); + if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); + const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); + + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; + auto& full_barrier = *full_barriers[stage_idx]; + const uint32_t k_idx = k_block_idx * BLOCK_K; + 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); + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(num_total_k_blocks, 1, k_block_idx), + num_tma_multicast_a); + + tma::copy(&tensor_map_b, &full_barrier, + smem_b[stage_idx], k_idx, scheduler.get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b); + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + } + } + + if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; advance_pipeline(i)) + empty_barriers[stage_idx]->wait(phase ^ 1); + } + } + } else { + cutlass::arch::warpgroup_reg_alloc(); + + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto r_0 = warp_idx * 16 + lane_idx / 4, r_1 = r_0 + 8; + + auto a_desc = mma::sm90::make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1); + auto b_desc = mma::sm90::make_smem_desc(smem_b[0], 1); + const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0); + const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0); + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + 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}; + + DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for BLOCK_M < 64"); + constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + const bool do_wgmma_store = BLOCK_M >= WGMMA::M or warp_idx < kNumWGMMAStoreThreads / 32; + + auto empty_barrier_arrive = [&]() { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void(); + } + }; + + auto load_sfb = [&](uint32_t n_idx, uint32_t k_scale_idx) { + if (n_idx >= shape_n) + return 1.0f; + const uint32_t offset = scheduler.current_group_idx * sfb_stride_group + + n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; + return mxfp8_fp8_detail::e8m0_to_float(sfb[offset]); + }; + + if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + 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); + full_barriers[stage_idx]->wait(phase); + + #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; + 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; + + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, accum, false); + 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>(); + + if (not do_wgmma_store) + continue; + + const uint32_t global_k_scale_idx = k_block_idx * (BLOCK_K / WGMMA::K) + kk; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t n_scale_idx = n_block_idx * BLOCK_N + i * 8; + const float scale_b = load_sfb(n_scale_idx, global_k_scale_idx); + shifted_accum[i * 4 + 0] += scale_a_0 * scale_b * accum[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_a_0 * scale_b * accum[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_a_1 * scale_b * accum[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_a_1 * scale_b * accum[i * 4 + 3]; + } + } + + if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1) + empty_barrier_arrive(); + } + } + } else { + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + empty_barrier_arrive(); + } + } + + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes); + constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4; + DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom"); + DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32, + "Unaligned TMA store or too many TMA store instructions"); + DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N"); + + if (not do_wgmma_store) + continue; + + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + #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; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumBankGroupBytes = 16; + auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8); + auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes); + constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8; + auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8); + auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8); + col ^= row % (kSwizzleDMode / 16); + smem_ptr = reinterpret_cast(smem_d) + + warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + + m_offset * kSwizzleDMode + + atom_offset * BLOCK_M * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; + } else { + smem_ptr = reinterpret_cast(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8); + } + + ptx::SM90_U32x2_STSM_N::copy( + __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}), + __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}), + smem_ptr + ); + } + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked; + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) { + auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N; + auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M; + auto n_idx = n_block_idx * BLOCK_N + in_block_n_offset; + auto m_idx = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); + cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr, n_idx, m_idx); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_90a"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py new file mode 100644 index 0000000000..a79928020b --- /dev/null +++ b/tests/test_sm90_mxfp8_fp8.py @@ -0,0 +1,100 @@ +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.math import per_token_cast_to_fp8 + + +def _require_sm90() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for SM90 MXFP8FP8 tests") + major, _ = torch.cuda.get_device_capability() + if major != 9: + pytest.skip(f"SM90 MXFP8FP8 tests require sm_90, got sm_{major}x") + + +def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int) -> torch.Tensor: + group_idx = torch.arange(x.size(-1), device=x.device) // gran_k + return x.float() * sf[..., group_idx] + + +def _e8m0_from_fp32_pow2(sf: torch.Tensor) -> torch.Tensor: + assert sf.dtype == torch.float32 or sf.dtype == torch.float + return ((sf.view(torch.int32) >> 23) & 0xFF).to(torch.uint8) + + +def _make_contiguous_case(groups: int, m_per_group: int, n: int, k: int): + m = groups * m_per_group + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) + b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + for group_id in range(groups): + b_data[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=128) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + start = group_id * m_per_group + end = start + m_per_group + b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_sf_fp32[group_id], gran_k=32) + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + return a, (b_data, _e8m0_from_fp32_pow2(b_sf_fp32)), grouped_layout, ref + + +def test_m_grouped_mxfp8_fp8_contiguous_e8m0_scale_accuracy(): + _require_sm90() + groups, m_per_group, n, k = 2, 17, 48, 64 + a, b, grouped_layout, ref = _make_contiguous_case(groups, m_per_group, n, k) + d = torch.empty_like(ref) + + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous(a, b, d, grouped_layout) + diff = calc_diff(d, ref) + assert diff < 0.03 + + +def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): + _require_sm90() + groups, max_m, n, k = 2, 32, 48, 64 + masked_m = torch.tensor([7, 19], device="cuda", dtype=torch.int32) + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, 1), device="cuda", dtype=torch.float32) + b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref[group_id], use_ue8m0=False, gran_k=128 + ) + b_data[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + + a = (a_data, a_sf) + b = (b_data, _e8m0_from_fp32_pow2(b_sf_fp32)) + a_dequant = _cast_back_from_fp8_1d(a_data, a_sf, gran_k=128) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for group_id, valid_m in enumerate(masked_m.tolist()): + b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_sf_fp32[group_id], gran_k=32) + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + + d = torch.empty_like(ref) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_masked(a, b, d, masked_m, expected_m=max_m) + diff = max( + calc_diff(d[group_id, :valid_m], ref[group_id, :valid_m]) + for group_id, valid_m in enumerate(masked_m.tolist()) + ) + assert diff < 0.03 From 5e8cd3d55708c9d3900728383a69fa7f06e28fa2 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 10:49:27 +0800 Subject: [PATCH 02/30] Fix MXFP8 FP8 per-column B scaling --- .../deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 04e945ef7b..12a6b42a1a 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -223,12 +223,13 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint32_t n_scale_idx = n_block_idx * BLOCK_N + i * 8; - const float scale_b = load_sfb(n_scale_idx, global_k_scale_idx); - shifted_accum[i * 4 + 0] += scale_a_0 * scale_b * accum[i * 4 + 0]; - shifted_accum[i * 4 + 1] += scale_a_0 * scale_b * accum[i * 4 + 1]; - shifted_accum[i * 4 + 2] += scale_a_1 * scale_b * accum[i * 4 + 2]; - shifted_accum[i * 4 + 3] += scale_a_1 * scale_b * accum[i * 4 + 3]; + const uint32_t n_scale_idx = n_block_idx * BLOCK_N + i * 8 + (lane_idx % 4) * 2; + const float scale_b_0 = load_sfb(n_scale_idx, global_k_scale_idx); + const float scale_b_1 = load_sfb(n_scale_idx + 1, global_k_scale_idx); + shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; } } From 4365d1ddf49fd6c8a6292899ec7356e244455aac Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 10:52:48 +0800 Subject: [PATCH 03/30] Fix MXFP8 contiguous accuracy test constraints --- .../include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 2 +- tests/test_sm90_mxfp8_fp8.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 12a6b42a1a..1248a2282e 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -182,7 +182,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, }; auto load_sfb = [&](uint32_t n_idx, uint32_t k_scale_idx) { - if (n_idx >= shape_n) + if (n_idx >= shape_n or k_scale_idx >= math::ceil_div(shape_k, 32u)) return 1.0f; const uint32_t offset = scheduler.current_group_idx * sfb_stride_group + n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index a79928020b..16c4c9b085 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -55,7 +55,8 @@ def _make_contiguous_case(groups: int, m_per_group: int, n: int, k: int): def test_m_grouped_mxfp8_fp8_contiguous_e8m0_scale_accuracy(): _require_sm90() - groups, m_per_group, n, k = 2, 17, 48, 64 + # SM90 grouped-contiguous WGMMA/TMA maps one B group per M tile. + groups, m_per_group, n, k = 2, 128, 48, 128 a, b, grouped_layout, ref = _make_contiguous_case(groups, m_per_group, n, k) d = torch.empty_like(ref) @@ -66,7 +67,7 @@ def test_m_grouped_mxfp8_fp8_contiguous_e8m0_scale_accuracy(): def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): _require_sm90() - groups, max_m, n, k = 2, 32, 48, 64 + groups, max_m, n, k = 2, 32, 48, 128 masked_m = torch.tensor([7, 19], device="cuda", dtype=torch.int32) a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) From 5608f23cc6a9d85d023b79aa22e93c6e1f4eadfd Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 11:00:57 +0800 Subject: [PATCH 04/30] Add MXFP8 FP8 performance comparison test --- tests/test_sm90_mxfp8_fp8.py | 139 ++++++++++++++++++++++++++++++++++- 1 file changed, 138 insertions(+), 1 deletion(-) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 16c4c9b085..7c24cb7285 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -1,5 +1,6 @@ import sys from pathlib import Path +from typing import Callable import pytest import torch @@ -7,7 +8,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import deep_gemm -from deep_gemm.testing import calc_diff +from deep_gemm.testing import bench, calc_diff from deep_gemm.utils.math import per_token_cast_to_fp8 @@ -29,6 +30,15 @@ def _e8m0_from_fp32_pow2(sf: torch.Tensor) -> torch.Tensor: return ((sf.view(torch.int32) >> 23) & 0xFF).to(torch.uint8) +def _tflops(m: int, n: int, k: int, elapsed: float) -> float: + return 2.0 * m * n * k / elapsed / 1e12 + + +def _time_kernel(fn: Callable[[], None]) -> float: + fn() + return bench(fn, num_warmups=5, num_tests=10) + + def _make_contiguous_case(groups: int, m_per_group: int, n: int, k: int): m = groups * m_per_group a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) @@ -99,3 +109,130 @@ def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): for group_id, valid_m in enumerate(masked_m.tolist()) ) assert diff < 0.03 + + +def test_m_grouped_mxfp8_vs_fp8_perf_contiguous_and_masked(): + _require_sm90() + groups, n, k = 4, 1024, 1024 + + # Contiguous: one B group per M tile. + m_per_group = 128 + m = groups * m_per_group + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + a = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) + + b_mx_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_mx_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // 128), device="cuda", dtype=torch.float32) + for group_id in range(groups): + b_mx_data[group_id], b_mx_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=False, gran_k=128 + ) + b_mx = (b_mx_data, _e8m0_from_fp32_pow2(b_mx_sf_fp32)) + b_fp8 = (b_fp8_data, b_fp8_sf) + d_mx = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + d_fp8 = torch.empty_like(d_mx) + + def run_mx_contiguous(): + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous(a, b_mx, d_mx, grouped_layout) + + def run_fp8_contiguous(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, b_fp8, d_fp8, grouped_layout, recipe_a=(1, 128), recipe_b=(1, 128) + ) + + mx_contiguous_elapsed = _time_kernel(run_mx_contiguous) + fp8_contiguous_elapsed = _time_kernel(run_fp8_contiguous) + contiguous_diff = float(calc_diff(d_mx, d_fp8)) + + # Masked: same shape class, but allow uneven active rows per group. + max_m = 128 + masked_m = torch.tensor([128, 96, 64, 32], device="cuda", dtype=torch.int32) + a_ref_masked = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_masked = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + a_masked_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_masked_sf = torch.empty((groups, max_m, k // 128), device="cuda", dtype=torch.float32) + b_mx_masked_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_mx_masked_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + b_fp8_masked_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_masked_sf = torch.empty((groups, n, k // 128), device="cuda", dtype=torch.float32) + for group_id in range(groups): + a_masked_data[group_id], a_masked_sf[group_id] = per_token_cast_to_fp8( + a_ref_masked[group_id], use_ue8m0=False, gran_k=128 + ) + b_mx_masked_data[group_id], b_mx_masked_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref_masked[group_id], use_ue8m0=True, gran_k=32 + ) + b_fp8_masked_data[group_id], b_fp8_masked_sf[group_id] = per_token_cast_to_fp8( + b_ref_masked[group_id], use_ue8m0=False, gran_k=128 + ) + a_masked = (a_masked_data, a_masked_sf) + b_mx_masked = (b_mx_masked_data, _e8m0_from_fp32_pow2(b_mx_masked_sf_fp32)) + b_fp8_masked = (b_fp8_masked_data, b_fp8_masked_sf) + d_mx_masked = torch.empty((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + d_fp8_masked = torch.empty_like(d_mx_masked) + + def run_mx_masked(): + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_masked( + a_masked, b_mx_masked, d_mx_masked, masked_m, expected_m=max_m + ) + + def run_fp8_masked(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a_masked, + b_fp8_masked, + d_fp8_masked, + masked_m, + expected_m=max_m, + recipe_a=(1, 128), + recipe_b=(1, 128), + ) + + mx_masked_elapsed = _time_kernel(run_mx_masked) + fp8_masked_elapsed = _time_kernel(run_fp8_masked) + masked_diff = max( + float(calc_diff(d_mx_masked[group_id, :valid_m], d_fp8_masked[group_id, :valid_m])) + for group_id, valid_m in enumerate(masked_m.tolist()) + ) + + masked_active_m = int(masked_m.sum().item()) + rows = [ + ( + "contiguous", + m, + mx_contiguous_elapsed, + fp8_contiguous_elapsed, + _tflops(m, n, k, mx_contiguous_elapsed), + _tflops(m, n, k, fp8_contiguous_elapsed), + contiguous_diff, + ), + ( + "masked", + masked_active_m, + mx_masked_elapsed, + fp8_masked_elapsed, + _tflops(masked_active_m, n, k, mx_masked_elapsed), + _tflops(masked_active_m, n, k, fp8_masked_elapsed), + masked_diff, + ), + ] + print("kernel | active M | MXFP8 us | FP8 us | MXFP8 TFLOPS | FP8 TFLOPS | speedup | diff") + print("-- | -- | -- | -- | -- | -- | -- | --") + for name, active_m, mx_elapsed, fp8_elapsed, mx_tflops, fp8_tflops, diff in rows: + print( + f"{name} | {active_m} | {mx_elapsed * 1e6:.0f} | {fp8_elapsed * 1e6:.0f} | " + f"{mx_tflops:.1f} | {fp8_tflops:.1f} | {fp8_elapsed / mx_elapsed:.2f}x | {diff:.4f}" + ) + + assert mx_contiguous_elapsed > 0 + assert fp8_contiguous_elapsed > 0 + assert mx_masked_elapsed > 0 + assert fp8_masked_elapsed > 0 + assert contiguous_diff == contiguous_diff + assert masked_diff == masked_diff From aaecde60cec5bc7b66a13e591e5400bdabe7741d Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 11:09:41 +0800 Subject: [PATCH 05/30] Stage MXFP8 B scales in shared memory --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 23 +++++- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 71 ++++++++++++------- 2 files changed, 68 insertions(+), 26 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 26873144cb..899fc0d2bb 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include #include "../../jit/compiler.hpp" @@ -72,6 +74,21 @@ static void __instantiate_kernel() {{ } }; +static void tune_mxfp8_fp8_smem_config(GemmConfig& config, const GemmDesc& desc) { + const int orig_num_stages = config.pipeline_config.num_stages; + const int original_per_stage = + config.storage_config.load_block_m * config.layout.block_k * c10::elementSize(desc.a_dtype) + + config.storage_config.load_block_n * config.layout.block_k * c10::elementSize(desc.b_dtype) + + align(config.layout.block_m * static_cast(sizeof(float)), 128); + const int sfb_per_stage = align(config.layout.block_n * (config.layout.block_k / 32) * static_cast(sizeof(uint8_t)), 128); + const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage; + const int merged_per_stage = original_per_stage + sfb_per_stage; + int chosen_stages = std::min(orig_num_stages, (SM90ArchSpec::smem_capacity - smem_extra) / merged_per_stage); + DG_HOST_ASSERT(chosen_stages >= 1); + config.pipeline_config.num_stages = chosen_stages; + config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; +} + static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, @@ -98,7 +115,8 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = 1 }; - const auto config = get_best_config(desc); + auto config = get_best_config(desc); + tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); @@ -167,7 +185,8 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; - const auto config = get_best_config(desc); + auto config = get_best_config(desc); + tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 1248a2282e..0dafb6dd4b 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -63,6 +63,9 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, 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 ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + static constexpr uint32_t SHAPE_K_SFB_PER_STAGE = BLOCK_K / 32; + static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * SHAPE_K_SFB_PER_STAGE * sizeof(uint8_t); + static constexpr uint32_t ALIGNED_SMEM_SFB_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFB_SIZE_PER_STAGE, 128u); const uint32_t num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K); const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); @@ -88,8 +91,13 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); }); + auto smem_sfb = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE + + i * ALIGNED_SMEM_SFB_SIZE_PER_STAGE); + }); - auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + auto barrier_start_ptr = reinterpret_cast( + smem_buffer + SMEM_SF_OFFSET + kNumStages * (ALIGNED_SMEM_SFA_SIZE_PER_STAGE + ALIGNED_SMEM_SFB_SIZE_PER_STAGE)); auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); @@ -120,7 +128,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, if (warp_idx >= kNumMathThreads / 32) { cutlass::arch::warpgroup_reg_dealloc(); - if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + if (warp_idx == kNumMathThreads / 32 + 2) { while (scheduler.get_next_block(m_block_idx, n_block_idx)) { const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; @@ -128,22 +136,41 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { - empty_barriers[stage_idx]->wait(phase ^ 1); + const bool is_producer_leader = cute::elect_one_sync(); + if (is_producer_leader) + empty_barriers[stage_idx]->wait(phase ^ 1); + __syncwarp(); constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; auto& full_barrier = *full_barriers[stage_idx]; const uint32_t k_idx = k_block_idx * BLOCK_K; - 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); - tma::copy(&tensor_map_sfa, &full_barrier, - smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(num_total_k_blocks, 1, k_block_idx), - num_tma_multicast_a); - - tma::copy(&tensor_map_b, &full_barrier, - smem_b[stage_idx], k_idx, scheduler.get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), - num_tma_multicast_b); - full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + if (is_producer_leader) { + 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); + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(num_total_k_blocks, 1, k_block_idx), + num_tma_multicast_a); + + tma::copy(&tensor_map_b, &full_barrier, + smem_b[stage_idx], k_idx, scheduler.get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b); + } + + for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { + const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; + const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; + const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; + const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFB_PER_STAGE + k_scale_offset; + const bool is_valid = n_idx < shape_n and k_scale_idx < math::ceil_div(shape_k, 32u); + const uint32_t gmem_offset = scheduler.current_group_idx * sfb_stride_group + + n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; + smem_sfb[stage_idx][i] = is_valid ? sfb[gmem_offset] : static_cast(127); + } + __syncwarp(); + + if (is_producer_leader) + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); } } @@ -181,12 +208,9 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, } }; - auto load_sfb = [&](uint32_t n_idx, uint32_t k_scale_idx) { - if (n_idx >= shape_n or k_scale_idx >= math::ceil_div(shape_k, 32u)) - return 1.0f; - const uint32_t offset = scheduler.current_group_idx * sfb_stride_group + - n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; - return mxfp8_fp8_detail::e8m0_to_float(sfb[offset]); + auto load_sfb = [&](uint32_t n_offset, uint32_t k_scale_offset) { + const uint32_t offset = n_offset * SHAPE_K_SFB_PER_STAGE + k_scale_offset; + return mxfp8_fp8_detail::e8m0_to_float(smem_sfb[stage_idx][offset]); }; if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { @@ -219,13 +243,12 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, if (not do_wgmma_store) continue; - const uint32_t global_k_scale_idx = k_block_idx * (BLOCK_K / WGMMA::K) + kk; auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint32_t n_scale_idx = n_block_idx * BLOCK_N + i * 8 + (lane_idx % 4) * 2; - const float scale_b_0 = load_sfb(n_scale_idx, global_k_scale_idx); - const float scale_b_1 = load_sfb(n_scale_idx + 1, global_k_scale_idx); + const uint32_t n_scale_offset = i * 8 + (lane_idx % 4) * 2; + const float scale_b_0 = load_sfb(n_scale_offset, kk); + const float scale_b_1 = load_sfb(n_scale_offset + 1, kk); shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; From cb8e933e8c04ae6e17ad7aee6417a9ec1c7e96aa Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 11:21:20 +0800 Subject: [PATCH 06/30] Fix main-based MXFP8 include list --- csrc/apis/gemm.hpp | 1 - 1 file changed, 1 deletion(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 753853deb7..4d4dc8fb0f 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -6,7 +6,6 @@ #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp" -#include "../jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" From dcb101703fe2abb9e52aacfe3d974afcdf62c731 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 11:25:47 +0800 Subject: [PATCH 07/30] Fence staged MXFP8 B scales before consume --- deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 0dafb6dd4b..65d8f6d861 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -167,6 +167,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; smem_sfb[stage_idx][i] = is_valid ? sfb[gmem_offset] : static_cast(127); } + __threadfence_block(); __syncwarp(); if (is_producer_leader) From c1125fca70f700ad6d5e33464fd71962eed7fdd7 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 15:19:02 +0800 Subject: [PATCH 08/30] Support MXFP8 A scales on SM90 grouped kernels --- csrc/apis/gemm.hpp | 16 ++++---- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 33 ++++++++------- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 40 ++++++++++++------- tests/test_sm90_mxfp8_fp8.py | 35 ++++++++++------ 4 files changed, 74 insertions(+), 50 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 4d4dc8fb0f..d8e57697f9 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -290,18 +290,18 @@ static void m_grouped_mxfp8_fp8_gemm_nt_contiguous(const std::pair 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 and a.second.is_contiguous()); DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + const auto [m_sfa, k_sfa] = get_shape<2>(a.second); + DG_HOST_ASSERT(m == m_sfa and k_sfa == ceil_div(k, 32)); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); if (m == 0) return; - const std::variant, std::tuple> sfa_recipe = std::make_tuple(1, 128); - const auto sfa = layout::transform_sf_into_required_layout( - a.second, m, k, sfa_recipe, std::nullopt, std::nullopt, false); sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( - a.first, sfa, b.first, b.second, d, grouped_layout, num_groups, m, n, k, compiled_dims); + a.first, a.second, b.first, b.second, d, grouped_layout, num_groups, m, n, k, compiled_dims); } static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair& a, @@ -328,15 +328,15 @@ static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair 0 and n > 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 and a.second.is_contiguous()); DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + const auto [num_groups_sfa, m_sfa, k_sfa] = get_shape<3>(a.second); + DG_HOST_ASSERT(num_groups == num_groups_sfa and m == m_sfa and k_sfa == ceil_div(k, 32)); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); - const std::variant, std::tuple> sfa_recipe = std::make_tuple(1, 128); - const auto sfa = layout::transform_sf_into_required_layout( - a.second, m, k, sfa_recipe, num_groups, std::nullopt, false); sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( - a.first, sfa, b.first, b.second, d, masked_m, num_groups, m, n, k, compiled_dims); + a.first, a.second, b.first, b.second, d, masked_m, num_groups, m, n, k, compiled_dims); } static void k_grouped_fp8_gemm_tn_contiguous(const std::pair& a, diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 899fc0d2bb..0135abb345 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -23,12 +23,12 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime(sizeof(float)), 128); + const int sfa_per_stage = align(config.layout.block_m * (config.layout.block_k / 32) * static_cast(sizeof(uint8_t)), 128); const int sfb_per_stage = align(config.layout.block_n * (config.layout.block_k / 32) * static_cast(sizeof(uint8_t)), 128); const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage; - const int merged_per_stage = original_per_stage + sfb_per_stage; + const int merged_per_stage = + config.storage_config.load_block_m * config.layout.block_k * c10::elementSize(desc.a_dtype) + + config.storage_config.load_block_n * config.layout.block_k * c10::elementSize(desc.b_dtype) + + sfa_per_stage + sfb_per_stage; int chosen_stages = std::min(orig_num_stages, (SM90ArchSpec::smem_capacity - smem_extra) / merged_per_stage); DG_HOST_ASSERT(chosen_stages >= 1); config.pipeline_config.num_stages = chosen_stages; @@ -97,7 +102,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( const std::string& compiled_dims) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8); DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and grouped_layout.is_contiguous()); @@ -135,24 +140,23 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( config.storage_config.store_block_n, 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); - const typename SM90MXFP8FP8Gemm1D2DRuntime::Args& args = { .gemm_desc = desc, .gemm_config = config, .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, config.pipeline_config.smem_size, config.layout.get_cluster_size()), + .sfa = sfa.data_ptr(), .sfb = sfb.data_ptr(), .grouped_layout = grouped_layout.data_ptr(), + .sfa_stride_m = static_cast(sfa.stride(0)), + .sfa_stride_k = static_cast(sfa.stride(1)), .sfb_stride_group = static_cast(sfb.stride(0)), .sfb_stride_n = static_cast(sfb.stride(1)), .sfb_stride_k = static_cast(sfb.stride(2)), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, - .tensor_map_sfa = tensor_map_sfa, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d", code); @@ -167,7 +171,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( const std::string& compiled_dims) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8); DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt and masked_m.is_contiguous()); @@ -205,24 +209,23 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( config.storage_config.store_block_n, static_cast(d.stride(-2)), num_groups, 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, num_groups, 0); - const typename SM90MXFP8FP8Gemm1D2DRuntime::Args& args = { .gemm_desc = desc, .gemm_config = config, .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, config.pipeline_config.smem_size, config.layout.get_cluster_size()), + .sfa = sfa.data_ptr(), .sfb = sfb.data_ptr(), .grouped_layout = masked_m.data_ptr(), + .sfa_stride_m = static_cast(sfa.stride(-2)), + .sfa_stride_k = static_cast(sfa.stride(-1)), .sfb_stride_group = static_cast(sfb.stride(0)), .sfb_stride_n = static_cast(sfb.stride(1)), .sfb_stride_k = static_cast(sfb.stride(2)), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, - .tensor_map_sfa = tensor_map_sfa, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d", code); diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 65d8f6d861..b751b621b2 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -40,15 +40,15 @@ template CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void -sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, +sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, + uint32_t sfa_stride_m, uint32_t sfa_stride_k, uint32_t sfb_stride_group, uint32_t sfb_stride_n, uint32_t sfb_stride_k, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, const __grid_constant__ cute::TmaDescriptor tensor_map_a, const __grid_constant__ cute::TmaDescriptor tensor_map_b, - const __grid_constant__ cute::TmaDescriptor tensor_map_d, - const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) { + const __grid_constant__ cute::TmaDescriptor tensor_map_d) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) - DG_STATIC_ASSERT(BLOCK_K == 128, "A scale is fixed to per-128 K"); + DG_STATIC_ASSERT(BLOCK_K == 128, "MXFP8 scale stage assumes 4 K/32 scale groups"); DG_STATIC_ASSERT(kNumStages >= 1, "Invalid pipeline stages"); using WGMMA = typename mma::sm90::FP8MMASelector::type; @@ -61,7 +61,8 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* 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 SHAPE_K_SFA_PER_STAGE = BLOCK_K / 32; + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * SHAPE_K_SFA_PER_STAGE * sizeof(uint8_t); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); static constexpr uint32_t SHAPE_K_SFB_PER_STAGE = BLOCK_K / 32; static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * SHAPE_K_SFB_PER_STAGE * sizeof(uint8_t); @@ -74,7 +75,6 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { cute::prefetch_tma_descriptor(&tensor_map_a); cute::prefetch_tma_descriptor(&tensor_map_b); - cute::prefetch_tma_descriptor(&tensor_map_sfa); cute::prefetch_tma_descriptor(&tensor_map_d); } __syncwarp(); @@ -89,7 +89,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, }); constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); }); auto smem_sfb = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE + @@ -148,15 +148,23 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* 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); - tma::copy(&tensor_map_sfa, &full_barrier, - smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(num_total_k_blocks, 1, k_block_idx), - num_tma_multicast_a); - tma::copy(&tensor_map_b, &full_barrier, smem_b[stage_idx], k_idx, scheduler.get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), num_tma_multicast_b); } + const uint32_t sfa_base_m = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); + for (uint32_t i = lane_idx; i < BLOCK_M * SHAPE_K_SFA_PER_STAGE; i += 32) { + const uint32_t m_offset = i / SHAPE_K_SFA_PER_STAGE; + const uint32_t k_scale_offset = i % SHAPE_K_SFA_PER_STAGE; + const uint32_t m_idx = sfa_base_m + m_offset; + const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFA_PER_STAGE + k_scale_offset; + const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and + k_scale_idx < math::ceil_div(shape_k, 32u); + smem_sfa[stage_idx][i] = is_valid ? sfa[m_idx * sfa_stride_m + k_scale_idx * sfa_stride_k] : + static_cast(127); + } + for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; @@ -171,7 +179,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, __syncwarp(); if (is_producer_leader) - full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); } } @@ -213,6 +221,10 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, const uint32_t offset = n_offset * SHAPE_K_SFB_PER_STAGE + k_scale_offset; return mxfp8_fp8_detail::e8m0_to_float(smem_sfb[stage_idx][offset]); }; + auto load_sfa = [&](uint32_t m_offset, uint32_t k_scale_offset) { + const uint32_t offset = m_offset * SHAPE_K_SFA_PER_STAGE + k_scale_offset; + return mxfp8_fp8_detail::e8m0_to_float(smem_sfa[stage_idx][offset]); + }; if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { @@ -223,8 +235,6 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, #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; - 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; #pragma unroll for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { @@ -244,6 +254,8 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfb, int* grouped_layout, if (not do_wgmma_store) continue; + const float scale_a_0 = load_sfa(r_0 + m_offset, kk); + const float scale_a_1 = load_sfa(r_1 + m_offset, kk); auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 7c24cb7285..94c5ae5717 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -44,7 +44,7 @@ def _make_contiguous_case(groups: int, m_per_group: int, n: int, k: int): a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) - a = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) + a_data, a_sf_fp32 = per_token_cast_to_fp8(a_ref, use_ue8m0=True, gran_k=32) b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) for group_id in range(groups): @@ -53,7 +53,8 @@ def _make_contiguous_case(groups: int, m_per_group: int, n: int, k: int): ) grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) - a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=128) + a = (a_data, _e8m0_from_fp32_pow2(a_sf_fp32)) + a_dequant = _cast_back_from_fp8_1d(a_data, a_sf_fp32, gran_k=32) ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) for group_id in range(groups): start = group_id * m_per_group @@ -83,20 +84,20 @@ def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) - a_sf = torch.empty((groups, max_m, 1), device="cuda", dtype=torch.float32) + a_sf_fp32 = torch.empty((groups, max_m, k // 32), device="cuda", dtype=torch.float32) b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) for group_id in range(groups): - a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( - a_ref[group_id], use_ue8m0=False, gran_k=128 + a_data[group_id], a_sf_fp32[group_id] = per_token_cast_to_fp8( + a_ref[group_id], use_ue8m0=True, gran_k=32 ) b_data[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp8( b_ref[group_id], use_ue8m0=True, gran_k=32 ) - a = (a_data, a_sf) + a = (a_data, _e8m0_from_fp32_pow2(a_sf_fp32)) b = (b_data, _e8m0_from_fp32_pow2(b_sf_fp32)) - a_dequant = _cast_back_from_fp8_1d(a_data, a_sf, gran_k=128) + a_dequant = _cast_back_from_fp8_1d(a_data, a_sf_fp32, gran_k=32) ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) for group_id, valid_m in enumerate(masked_m.tolist()): b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_sf_fp32[group_id], gran_k=32) @@ -120,7 +121,9 @@ def test_m_grouped_mxfp8_vs_fp8_perf_contiguous_and_masked(): m = groups * m_per_group a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) - a = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) + a_data, a_mx_sf_fp32 = per_token_cast_to_fp8(a_ref, use_ue8m0=True, gran_k=32) + a = (a_data, _e8m0_from_fp32_pow2(a_mx_sf_fp32)) + a_fp8 = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) b_mx_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) @@ -144,7 +147,7 @@ def run_mx_contiguous(): def run_fp8_contiguous(): deep_gemm.m_grouped_fp8_gemm_nt_contiguous( - a, b_fp8, d_fp8, grouped_layout, recipe_a=(1, 128), recipe_b=(1, 128) + a_fp8, b_fp8, d_fp8, grouped_layout, recipe_a=(1, 128), recipe_b=(1, 128) ) mx_contiguous_elapsed = _time_kernel(run_mx_contiguous) @@ -157,13 +160,18 @@ def run_fp8_contiguous(): a_ref_masked = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) b_ref_masked = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) a_masked_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) - a_masked_sf = torch.empty((groups, max_m, k // 128), device="cuda", dtype=torch.float32) + a_masked_sf_fp32 = torch.empty((groups, max_m, k // 32), device="cuda", dtype=torch.float32) + a_fp8_masked_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_fp8_masked_sf = torch.empty((groups, max_m, k // 128), device="cuda", dtype=torch.float32) b_mx_masked_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) b_mx_masked_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) b_fp8_masked_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) b_fp8_masked_sf = torch.empty((groups, n, k // 128), device="cuda", dtype=torch.float32) for group_id in range(groups): - a_masked_data[group_id], a_masked_sf[group_id] = per_token_cast_to_fp8( + a_masked_data[group_id], a_masked_sf_fp32[group_id] = per_token_cast_to_fp8( + a_ref_masked[group_id], use_ue8m0=True, gran_k=32 + ) + a_fp8_masked_data[group_id], a_fp8_masked_sf[group_id] = per_token_cast_to_fp8( a_ref_masked[group_id], use_ue8m0=False, gran_k=128 ) b_mx_masked_data[group_id], b_mx_masked_sf_fp32[group_id] = per_token_cast_to_fp8( @@ -172,7 +180,8 @@ def run_fp8_contiguous(): b_fp8_masked_data[group_id], b_fp8_masked_sf[group_id] = per_token_cast_to_fp8( b_ref_masked[group_id], use_ue8m0=False, gran_k=128 ) - a_masked = (a_masked_data, a_masked_sf) + a_masked = (a_masked_data, _e8m0_from_fp32_pow2(a_masked_sf_fp32)) + a_fp8_masked = (a_fp8_masked_data, a_fp8_masked_sf) b_mx_masked = (b_mx_masked_data, _e8m0_from_fp32_pow2(b_mx_masked_sf_fp32)) b_fp8_masked = (b_fp8_masked_data, b_fp8_masked_sf) d_mx_masked = torch.empty((groups, max_m, n), device="cuda", dtype=torch.bfloat16) @@ -185,7 +194,7 @@ def run_mx_masked(): def run_fp8_masked(): deep_gemm.m_grouped_fp8_gemm_nt_masked( - a_masked, + a_fp8_masked, b_fp8_masked, d_fp8_masked, masked_m, From 405be71da6d373731d3b3df24114a58e780ef079 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 15:24:06 +0800 Subject: [PATCH 09/30] Load SM90 MXFP8 A scales from global memory --- .../deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index b751b621b2..b610152409 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -221,16 +221,22 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, const uint32_t offset = n_offset * SHAPE_K_SFB_PER_STAGE + k_scale_offset; return mxfp8_fp8_detail::e8m0_to_float(smem_sfb[stage_idx][offset]); }; - auto load_sfa = [&](uint32_t m_offset, uint32_t k_scale_offset) { - const uint32_t offset = m_offset * SHAPE_K_SFA_PER_STAGE + k_scale_offset; - return mxfp8_fp8_detail::e8m0_to_float(smem_sfa[stage_idx][offset]); - }; - if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { 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); full_barriers[stage_idx]->wait(phase); + constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; + const uint32_t sfa_base_m = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); + auto load_sfa = [&](uint32_t m_offset, uint32_t k_scale_offset) { + const uint32_t m_idx = sfa_base_m + m_offset; + const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFA_PER_STAGE + k_scale_offset; + const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and + k_scale_idx < math::ceil_div(shape_k, 32u); + return mxfp8_fp8_detail::e8m0_to_float( + is_valid ? sfa[m_idx * sfa_stride_m + k_scale_idx * sfa_stride_k] : + static_cast(127)); + }; #pragma unroll for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { From c2b8716c0511837ca3922c933013ad9df7242c7a Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 15:28:20 +0800 Subject: [PATCH 10/30] Load SM90 MXFP8 B scales from global memory --- .../deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index b610152409..ac83a603a5 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -217,10 +217,6 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, } }; - auto load_sfb = [&](uint32_t n_offset, uint32_t k_scale_offset) { - const uint32_t offset = n_offset * SHAPE_K_SFB_PER_STAGE + k_scale_offset; - return mxfp8_fp8_detail::e8m0_to_float(smem_sfb[stage_idx][offset]); - }; if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) { for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); @@ -237,6 +233,15 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, is_valid ? sfa[m_idx * sfa_stride_m + k_scale_idx * sfa_stride_k] : static_cast(127)); }; + auto load_sfb = [&](uint32_t n_offset, uint32_t k_scale_offset) { + const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; + const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFB_PER_STAGE + k_scale_offset; + const bool is_valid = n_idx < shape_n and k_scale_idx < math::ceil_div(shape_k, 32u); + const uint32_t gmem_offset = scheduler.current_group_idx * sfb_stride_group + + n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; + return mxfp8_fp8_detail::e8m0_to_float( + is_valid ? sfb[gmem_offset] : static_cast(127)); + }; #pragma unroll for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { From ffe17ba69e078dedb3a912c786a1729a05179649 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 15:39:12 +0800 Subject: [PATCH 11/30] Avoid SM90 masked MXFP8 cross-group stores --- csrc/jit_kernels/heuristics/sm90.hpp | 11 ++++++++++- csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 5 +++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index c411fb7e01..7b53c8ba24 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -32,7 +32,16 @@ struct SM90ArchSpec { desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) { block_m_candidates = std::vector{heuristics_runtime->get_mk_alignment_for_contiguous_layout()}; } else if (desc.gemm_type == GemmType::MGroupedMasked) { - block_m_candidates = {64, 128}; + // Masked grouped outputs are laid out as [num_groups, expected_m, n]. + // Keep BLOCK_M aligned with expected_m when possible, otherwise a + // full-block TMA store can cross into the next group's rows. + const int expected_m = desc.get_expected_m(); + for (int candidate: {16, 32, 64, 128}) { + if (expected_m % candidate == 0) + block_m_candidates.push_back(candidate); + } + if (block_m_candidates.empty()) + block_m_candidates = {64, 128}; } // Block N candidates diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 0135abb345..baf5ef3537 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -36,6 +36,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; +static constexpr int kSm90MXFP8FP8K32E8M0JitVersion = 4; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -159,7 +160,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_k32e8m0_v4", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -228,7 +229,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_k32e8m0_v4", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } From 44017de1376bee09e27813fcfce2c62b9a1dbd1a Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 16:00:46 +0800 Subject: [PATCH 12/30] Pack-load SM90 MXFP8 scales from SMEM in consumer The MXFP8 consumer was issuing a fresh global load for every (kk, accum) pair, even though the producer warp had already staged SFA/SFB into shared memory. Replace the per-iteration GMEM lookups with two LDS loads (one 32-bit pack per SFA row, one 64-bit pack per pair of SFB rows) outside the wgmma kk loop, and decode the per-kk byte with a register-side shift. Keeps the result bit-identical while removing the GMEM dependency from the inner wgmma pipeline. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 56 ++++++++++--------- 1 file changed, 31 insertions(+), 25 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index ac83a603a5..f5fb0368ac 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -222,31 +222,34 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* 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); full_barriers[stage_idx]->wait(phase); - constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; - const uint32_t sfa_base_m = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); - auto load_sfa = [&](uint32_t m_offset, uint32_t k_scale_offset) { - const uint32_t m_idx = sfa_base_m + m_offset; - const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFA_PER_STAGE + k_scale_offset; - const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and - k_scale_idx < math::ceil_div(shape_k, 32u); - return mxfp8_fp8_detail::e8m0_to_float( - is_valid ? sfa[m_idx * sfa_stride_m + k_scale_idx * sfa_stride_k] : - static_cast(127)); - }; - auto load_sfb = [&](uint32_t n_offset, uint32_t k_scale_offset) { - const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; - const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFB_PER_STAGE + k_scale_offset; - const bool is_valid = n_idx < shape_n and k_scale_idx < math::ceil_div(shape_k, 32u); - const uint32_t gmem_offset = scheduler.current_group_idx * sfb_stride_group + - n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; - return mxfp8_fp8_detail::e8m0_to_float( - is_valid ? sfb[gmem_offset] : static_cast(127)); - }; #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; + // Pack-load all SFA/SFB bytes for this wave_block from SMEM: + // SFA row stride is SHAPE_K_SFA_PER_STAGE (== BLOCK_K/32 == 4) bytes, + // so 4 SFA bytes (one per kk) are loaded with a single 32-bit LDS. + // SFB has the same layout; two adjacent N rows (n_base, n_base+1) form + // 8 contiguous bytes that we fetch with a single ld.shared.v2.u32. + uint32_t sfa_pack_0 = 0, sfa_pack_1 = 0; + uint32_t sfb_pack[WGMMA::kNumAccum / 4][2]; + if (do_wgmma_store) { + sfa_pack_0 = ptx::ld_shared(reinterpret_cast( + smem_sfa[stage_idx] + (r_0 + m_offset) * SHAPE_K_SFA_PER_STAGE)); + sfa_pack_1 = ptx::ld_shared(reinterpret_cast( + smem_sfa[stage_idx] + (r_1 + m_offset) * SHAPE_K_SFA_PER_STAGE)); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t n_scale_offset = i * 8 + (lane_idx % 4) * 2; + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];\n" + : "=r"(sfb_pack[i][0]), "=r"(sfb_pack[i][1]) + : "l"(__cvta_generic_to_shared( + smem_sfb[stage_idx] + + n_scale_offset * SHAPE_K_SFB_PER_STAGE))); + } + } + #pragma unroll for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { #pragma unroll @@ -265,14 +268,17 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, if (not do_wgmma_store) continue; - const float scale_a_0 = load_sfa(r_0 + m_offset, kk); - const float scale_a_1 = load_sfa(r_1 + m_offset, kk); + const float scale_a_0 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfa_pack_0 >> (kk * 8)) & 0xff)); + const float scale_a_1 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfa_pack_1 >> (kk * 8)) & 0xff)); auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint32_t n_scale_offset = i * 8 + (lane_idx % 4) * 2; - const float scale_b_0 = load_sfb(n_scale_offset, kk); - const float scale_b_1 = load_sfb(n_scale_offset + 1, kk); + const float scale_b_0 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff)); + const float scale_b_1 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff)); shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; From 6de214d13121687a4ef8d78b2bab57aed03ec4fa Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 18:00:50 +0800 Subject: [PATCH 13/30] Support packed UE8M0 scales in SM90 MXFP8 GEMM --- csrc/apis/gemm.hpp | 25 ++++++---- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 44 ++++++++++++++--- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 40 ++++++++++++---- tests/test_sm90_mxfp8_fp8.py | 47 +++++++++++++++++++ 4 files changed, 131 insertions(+), 25 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index d8e57697f9..1fbfb6687c 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -290,12 +290,16 @@ static void m_grouped_mxfp8_fp8_gemm_nt_contiguous(const std::pair 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); - DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 and a.second.is_contiguous()); - DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT((a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt) and + a.second.is_contiguous()); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8 or b.second.scalar_type() == torch::kInt); const auto [m_sfa, k_sfa] = get_shape<2>(a.second); - DG_HOST_ASSERT(m == m_sfa and k_sfa == ceil_div(k, 32)); + const auto num_sfa_scales = k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(m == m_sfa and k % num_sfa_scales == 0 and (k / num_sfa_scales == 32 or k / num_sfa_scales == 128)); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); - DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); + const auto num_sfb_scales = k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k % num_sfb_scales == 0 and + (k / num_sfb_scales == 32 or k / num_sfb_scales == 128)); if (m == 0) return; @@ -328,12 +332,17 @@ static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair 0 and n > 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); - DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 and a.second.is_contiguous()); - DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT((a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt) and + a.second.is_contiguous()); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8 or b.second.scalar_type() == torch::kInt); const auto [num_groups_sfa, m_sfa, k_sfa] = get_shape<3>(a.second); - DG_HOST_ASSERT(num_groups == num_groups_sfa and m == m_sfa and k_sfa == ceil_div(k, 32)); + const auto num_sfa_scales = k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(num_groups == num_groups_sfa and m == m_sfa and k % num_sfa_scales == 0 and + (k / num_sfa_scales == 32 or k / num_sfa_scales == 128)); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); - DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k_sfb == ceil_div(k, 32)); + const auto num_sfb_scales = k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k % num_sfb_scales == 0 and + (k / num_sfb_scales == 32 or k / num_sfb_scales == 128)); sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( a.first, a.second, b.first, b.second, d, masked_m, num_groups, m, n, k, compiled_dims); diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index baf5ef3537..49735f728e 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -25,7 +25,11 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8K32E8M0JitVersion = 4; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 6; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -70,7 +74,9 @@ static void __instantiate_kernel() {{ DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.sfa, args.sfb, args.grouped_layout, args.sfa_stride_m, args.sfa_stride_k, + args.sfa_gran_k, args.sfa_packed_int32, args.sfb_stride_group, args.sfb_stride_n, args.sfb_stride_k, + args.sfb_gran_k, args.sfb_packed_int32, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d)); } @@ -103,8 +109,8 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( const std::string& compiled_dims) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8); - DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8 or sfa.scalar_type() == torch::kInt); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8 or sfb.scalar_type() == torch::kInt); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and grouped_layout.is_contiguous()); DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); @@ -125,6 +131,14 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + const auto sfa_num_scales = static_cast(sfa.size(1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(sfa_num_scales > 0 and k % sfa_num_scales == 0); + const auto sfa_gran_k = k / sfa_num_scales; + DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); + const auto sfb_num_scales = static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(sfb_num_scales > 0 and k % sfb_num_scales == 0); + const auto sfb_gran_k = k / sfb_num_scales; + DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, config.storage_config.load_block_m, @@ -152,15 +166,19 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .grouped_layout = grouped_layout.data_ptr(), .sfa_stride_m = static_cast(sfa.stride(0)), .sfa_stride_k = static_cast(sfa.stride(1)), + .sfa_gran_k = static_cast(sfa_gran_k), + .sfa_packed_int32 = sfa.scalar_type() == torch::kInt, .sfb_stride_group = static_cast(sfb.stride(0)), .sfb_stride_n = static_cast(sfb.stride(1)), .sfb_stride_k = static_cast(sfb.stride(2)), + .sfb_gran_k = static_cast(sfb_gran_k), + .sfb_packed_int32 = sfb.scalar_type() == torch::kInt, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_k32e8m0_v4", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v6", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -172,8 +190,8 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( const std::string& compiled_dims) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8); - DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8 or sfa.scalar_type() == torch::kInt); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kUInt8 or sfb.scalar_type() == torch::kInt); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt and masked_m.is_contiguous()); DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); @@ -194,6 +212,14 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + const auto sfa_num_scales = static_cast(sfa.size(-1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(sfa_num_scales > 0 and k % sfa_num_scales == 0); + const auto sfa_gran_k = k / sfa_num_scales; + DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); + const auto sfb_num_scales = static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1); + DG_HOST_ASSERT(sfb_num_scales > 0 and k % sfb_num_scales == 0); + const auto sfb_gran_k = k / sfb_num_scales; + DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, config.storage_config.load_block_m, @@ -221,15 +247,19 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .grouped_layout = masked_m.data_ptr(), .sfa_stride_m = static_cast(sfa.stride(-2)), .sfa_stride_k = static_cast(sfa.stride(-1)), + .sfa_gran_k = static_cast(sfa_gran_k), + .sfa_packed_int32 = sfa.scalar_type() == torch::kInt, .sfb_stride_group = static_cast(sfb.stride(0)), .sfb_stride_n = static_cast(sfb.stride(1)), .sfb_stride_k = static_cast(sfb.stride(2)), + .sfb_gran_k = static_cast(sfb_gran_k), + .sfb_packed_int32 = sfb.scalar_type() == torch::kInt, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_k32e8m0_v4", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v6", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index f5fb0368ac..0d2af3b6b2 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -28,6 +28,16 @@ CUTLASS_DEVICE float e8m0_to_float(uint8_t scale) { return __uint_as_float(static_cast(scale) << 23); } +CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, + uint32_t k_scale_idx, uint32_t stride_k, + bool packed_int32) { + if (packed_int32) { + const auto packed = reinterpret_cast(ptr)[base_offset + (k_scale_idx / 4) * stride_k]; + return static_cast((packed >> ((k_scale_idx % 4) * 8)) & 0xff); + } + return reinterpret_cast(ptr)[base_offset + k_scale_idx * stride_k]; +} + } // namespace mxfp8_fp8_detail template CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void -sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, +sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, uint32_t sfa_stride_m, uint32_t sfa_stride_k, + uint32_t sfa_gran_k, bool sfa_packed_int32, uint32_t sfb_stride_group, uint32_t sfb_stride_n, uint32_t sfb_stride_k, + uint32_t sfb_gran_k, bool sfb_packed_int32, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, const __grid_constant__ cute::TmaDescriptor tensor_map_a, const __grid_constant__ cute::TmaDescriptor tensor_map_b, @@ -158,22 +170,30 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(uint8_t* sfa, uint8_t* sfb, int* grouped_layout, const uint32_t m_offset = i / SHAPE_K_SFA_PER_STAGE; const uint32_t k_scale_offset = i % SHAPE_K_SFA_PER_STAGE; const uint32_t m_idx = sfa_base_m + m_offset; - const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFA_PER_STAGE + k_scale_offset; + const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; + const uint32_t k_scale_idx = k_idx_for_scale / sfa_gran_k; const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and - k_scale_idx < math::ceil_div(shape_k, 32u); - smem_sfa[stage_idx][i] = is_valid ? sfa[m_idx * sfa_stride_m + k_scale_idx * sfa_stride_k] : - static_cast(127); + k_idx_for_scale < shape_k; + const uint32_t sfa_base_offset = m_idx * sfa_stride_m; + smem_sfa[stage_idx][i] = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale( + sfa, sfa_base_offset, k_scale_idx, sfa_stride_k, sfa_packed_int32) : + static_cast(127); } for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; - const uint32_t k_scale_idx = k_block_idx * SHAPE_K_SFB_PER_STAGE + k_scale_offset; - const bool is_valid = n_idx < shape_n and k_scale_idx < math::ceil_div(shape_k, 32u); - const uint32_t gmem_offset = scheduler.current_group_idx * sfb_stride_group + - n_idx * sfb_stride_n + k_scale_idx * sfb_stride_k; - smem_sfb[stage_idx][i] = is_valid ? sfb[gmem_offset] : static_cast(127); + const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; + const uint32_t k_scale_idx = k_idx_for_scale / sfb_gran_k; + const bool is_valid = n_idx < shape_n and k_idx_for_scale < shape_k; + const uint32_t sfb_base_offset = scheduler.current_group_idx * sfb_stride_group + + n_idx * sfb_stride_n; + smem_sfb[stage_idx][i] = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale( + sfb, sfb_base_offset, k_scale_idx, sfb_stride_k, sfb_packed_int32) : + static_cast(127); } __threadfence_block(); __syncwarp(); diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 94c5ae5717..7eaa0948ae 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -30,6 +30,18 @@ def _e8m0_from_fp32_pow2(sf: torch.Tensor) -> torch.Tensor: return ((sf.view(torch.int32) >> 23) & 0xFF).to(torch.uint8) +def _pack_ue8m0_u8_to_i32(sf: torch.Tensor) -> torch.Tensor: + assert sf.dtype == torch.uint8 + assert sf.shape[-1] % 4 == 0 + sf_i32 = sf.contiguous().view(*sf.shape[:-1], sf.shape[-1] // 4, 4).to(torch.int32) + return ( + sf_i32[..., 0] + | torch.bitwise_left_shift(sf_i32[..., 1], 8) + | torch.bitwise_left_shift(sf_i32[..., 2], 16) + | torch.bitwise_left_shift(sf_i32[..., 3], 24) + ).contiguous() + + def _tflops(m: int, n: int, k: int, elapsed: float) -> float: return 2.0 * m * n * k / elapsed / 1e12 @@ -112,6 +124,41 @@ def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): assert diff < 0.03 +def test_m_grouped_mxfp8_fp8_contiguous_packed_int32_scale_accuracy(): + _require_sm90() + groups, m_per_group, n, k = 2, 128, 48, 512 + m = groups * m_per_group + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data, a_sf_fp32 = per_token_cast_to_fp8(a_ref, use_ue8m0=True, gran_k=128) + b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + for group_id in range(groups): + b_data[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + + a = (a_data, _pack_ue8m0_u8_to_i32(_e8m0_from_fp32_pow2(a_sf_fp32))) + b = (b_data, _pack_ue8m0_u8_to_i32(_e8m0_from_fp32_pow2(b_sf_fp32))) + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave( + m_per_group + ) + + a_dequant = _cast_back_from_fp8_1d(a_data, a_sf_fp32, gran_k=128) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + start = group_id * m_per_group + end = start + m_per_group + b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_sf_fp32[group_id], gran_k=32) + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + + d = torch.empty_like(ref) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous(a, b, d, grouped_layout) + diff = calc_diff(d, ref) + assert diff < 0.03 + + def test_m_grouped_mxfp8_vs_fp8_perf_contiguous_and_masked(): _require_sm90() groups, n, k = 4, 1024, 1024 From 06b2ac57fd2ecded8a3d846fd776b3235d3bf49f Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 20:18:41 +0800 Subject: [PATCH 14/30] Pass explicit recipes to SM90 MXFP8 GEMM --- csrc/apis/gemm.hpp | 66 ++++++++++++++----- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 40 ++++++----- tests/test_sm90_mxfp8_fp8.py | 15 ++++- 3 files changed, 85 insertions(+), 36 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 1fbfb6687c..30167934af 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -273,7 +273,9 @@ static void m_grouped_mxfp8_fp8_gemm_nt_contiguous(const std::pair& b, const torch::Tensor& d, const torch::Tensor& grouped_layout, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional>& recipe_a, + const std::optional>& recipe_b) { (void) compiled_dims; const auto major_a = get_major_type_ab(a.first); const auto major_b = get_major_type_ab(b.first); @@ -293,19 +295,32 @@ static void m_grouped_mxfp8_fp8_gemm_nt_contiguous(const std::pair(recipe_a.value()) == 1 and + (std::get<1>(recipe_a.value()) == 32 or std::get<1>(recipe_a.value()) == 128)); + if (recipe_b.has_value()) + DG_HOST_ASSERT(std::get<0>(recipe_b.value()) == 1 and + (std::get<1>(recipe_b.value()) == 32 or std::get<1>(recipe_b.value()) == 128)); const auto [m_sfa, k_sfa] = get_shape<2>(a.second); - const auto num_sfa_scales = k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(m == m_sfa and k % num_sfa_scales == 0 and (k / num_sfa_scales == 32 or k / num_sfa_scales == 128)); + const auto gran_k_a = recipe_a.has_value() + ? std::get<1>(recipe_a.value()) + : k / (k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(m == m_sfa and (gran_k_a == 32 or gran_k_a == 128) and + k_sfa == ceil_div(k, gran_k_a * (a.second.scalar_type() == torch::kInt ? 4 : 1))); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); - const auto num_sfb_scales = k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k % num_sfb_scales == 0 and - (k / num_sfb_scales == 32 or k / num_sfb_scales == 128)); + const auto gran_k_b = recipe_b.has_value() + ? std::get<1>(recipe_b.value()) + : k / (k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and + (gran_k_b == 32 or gran_k_b == 128) and + k_sfb == ceil_div(k, gran_k_b * (b.second.scalar_type() == torch::kInt ? 4 : 1))); if (m == 0) return; sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( - a.first, a.second, b.first, b.second, d, grouped_layout, num_groups, m, n, k, compiled_dims); + a.first, a.second, b.first, b.second, d, grouped_layout, num_groups, m, n, k, + compiled_dims, recipe_a, recipe_b); } static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair& a, @@ -313,7 +328,9 @@ static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair>& recipe_a, + const std::optional>& recipe_b) { (void) expected_m; (void) compiled_dims; const auto major_a = get_major_type_ab(a.first); @@ -335,17 +352,30 @@ static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair(recipe_a.value()) == 1 and + (std::get<1>(recipe_a.value()) == 32 or std::get<1>(recipe_a.value()) == 128)); + if (recipe_b.has_value()) + DG_HOST_ASSERT(std::get<0>(recipe_b.value()) == 1 and + (std::get<1>(recipe_b.value()) == 32 or std::get<1>(recipe_b.value()) == 128)); const auto [num_groups_sfa, m_sfa, k_sfa] = get_shape<3>(a.second); - const auto num_sfa_scales = k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(num_groups == num_groups_sfa and m == m_sfa and k % num_sfa_scales == 0 and - (k / num_sfa_scales == 32 or k / num_sfa_scales == 128)); + const auto gran_k_a = recipe_a.has_value() + ? std::get<1>(recipe_a.value()) + : k / (k_sfa * (a.second.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(num_groups == num_groups_sfa and m == m_sfa and + (gran_k_a == 32 or gran_k_a == 128) and + k_sfa == ceil_div(k, gran_k_a * (a.second.scalar_type() == torch::kInt ? 4 : 1))); const auto [num_groups_sfb, n_sfb, k_sfb] = get_shape<3>(b.second); - const auto num_sfb_scales = k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and k % num_sfb_scales == 0 and - (k / num_sfb_scales == 32 or k / num_sfb_scales == 128)); + const auto gran_k_b = recipe_b.has_value() + ? std::get<1>(recipe_b.value()) + : k / (k_sfb * (b.second.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(num_groups == num_groups_sfb and n == n_sfb and + (gran_k_b == 32 or gran_k_b == 128) and + k_sfb == ceil_div(k, gran_k_b * (b.second.scalar_type() == torch::kInt ? 4 : 1))); sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( - a.first, a.second, b.first, b.second, d, masked_m, num_groups, m, n, k, compiled_dims); + a.first, a.second, b.first, b.second, d, masked_m, num_groups, m, n, k, + compiled_dims, recipe_a, recipe_b); } static void k_grouped_fp8_gemm_tn_contiguous(const std::pair& a, @@ -726,10 +756,12 @@ static void register_apis(pybind11::module_& m) { py::arg("compiled_dims") = "nk", py::arg("disable_ue8m0_cast") = false); m.def("m_grouped_mxfp8_fp8_gemm_nt_contiguous", &m_grouped_mxfp8_fp8_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), - py::arg("compiled_dims") = "nk"); + py::arg("compiled_dims") = "nk", py::arg("recipe_a") = std::nullopt, + py::arg("recipe_b") = std::nullopt); m.def("m_grouped_mxfp8_fp8_gemm_nt_masked", &m_grouped_mxfp8_fp8_gemm_nt_masked, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"), - py::arg("expected_m"), py::arg("compiled_dims") = "nk"); + py::arg("expected_m"), py::arg("compiled_dims") = "nk", + py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt); m.def("k_grouped_fp8_gemm_tn_contiguous", &k_grouped_fp8_gemm_tn_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"), py::arg("ks_tensor"), py::arg("c") = std::nullopt, diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 49735f728e..ef9a50ca95 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -106,7 +106,9 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( const torch::Tensor& b, const torch::Tensor& sfb, const torch::Tensor& d, const torch::Tensor& grouped_layout, const int& num_groups, const int& m, const int& n, const int& k, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional>& recipe_a, + const std::optional>& recipe_b) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8 or sfa.scalar_type() == torch::kInt); @@ -131,14 +133,16 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); - const auto sfa_num_scales = static_cast(sfa.size(1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(sfa_num_scales > 0 and k % sfa_num_scales == 0); - const auto sfa_gran_k = k / sfa_num_scales; + const auto sfa_gran_k = recipe_a.has_value() + ? std::get<1>(recipe_a.value()) + : k / (static_cast(sfa.size(1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1)); DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); - const auto sfb_num_scales = static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(sfb_num_scales > 0 and k % sfb_num_scales == 0); - const auto sfb_gran_k = k / sfb_num_scales; + DG_HOST_ASSERT(sfa.size(1) == ceil_div(k, sfa_gran_k * (sfa.scalar_type() == torch::kInt ? 4 : 1))); + const auto sfb_gran_k = recipe_b.has_value() + ? std::get<1>(recipe_b.value()) + : k / (static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1)); DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); + DG_HOST_ASSERT(sfb.size(-1) == ceil_div(k, sfb_gran_k * (sfb.scalar_type() == torch::kInt ? 4 : 1))); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, config.storage_config.load_block_m, @@ -178,7 +182,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v6", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v7", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -187,7 +191,9 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( const torch::Tensor& b, const torch::Tensor& sfb, const torch::Tensor& d, const torch::Tensor& masked_m, const int& num_groups, const int& m, const int& n, const int& k, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional>& recipe_a, + const std::optional>& recipe_b) { DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(sfa.scalar_type() == torch::kUInt8 or sfa.scalar_type() == torch::kInt); @@ -212,14 +218,16 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); - const auto sfa_num_scales = static_cast(sfa.size(-1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(sfa_num_scales > 0 and k % sfa_num_scales == 0); - const auto sfa_gran_k = k / sfa_num_scales; + const auto sfa_gran_k = recipe_a.has_value() + ? std::get<1>(recipe_a.value()) + : k / (static_cast(sfa.size(-1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1)); DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); - const auto sfb_num_scales = static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1); - DG_HOST_ASSERT(sfb_num_scales > 0 and k % sfb_num_scales == 0); - const auto sfb_gran_k = k / sfb_num_scales; + DG_HOST_ASSERT(sfa.size(-1) == ceil_div(k, sfa_gran_k * (sfa.scalar_type() == torch::kInt ? 4 : 1))); + const auto sfb_gran_k = recipe_b.has_value() + ? std::get<1>(recipe_b.value()) + : k / (static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1)); DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); + DG_HOST_ASSERT(sfb.size(-1) == ceil_div(k, sfb_gran_k * (sfb.scalar_type() == torch::kInt ? 4 : 1))); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, config.storage_config.load_block_m, @@ -259,7 +267,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v6", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v7", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 7eaa0948ae..03cdd5e590 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -32,7 +32,14 @@ def _e8m0_from_fp32_pow2(sf: torch.Tensor) -> torch.Tensor: def _pack_ue8m0_u8_to_i32(sf: torch.Tensor) -> torch.Tensor: assert sf.dtype == torch.uint8 - assert sf.shape[-1] % 4 == 0 + if sf.shape[-1] % 4 != 0: + padded = torch.zeros( + (*sf.shape[:-1], ((sf.shape[-1] + 3) // 4) * 4), + device=sf.device, + dtype=sf.dtype, + ) + padded[..., : sf.shape[-1]] = sf + sf = padded sf_i32 = sf.contiguous().view(*sf.shape[:-1], sf.shape[-1] // 4, 4).to(torch.int32) return ( sf_i32[..., 0] @@ -126,7 +133,7 @@ def test_m_grouped_mxfp8_fp8_masked_e8m0_scale_accuracy(): def test_m_grouped_mxfp8_fp8_contiguous_packed_int32_scale_accuracy(): _require_sm90() - groups, m_per_group, n, k = 2, 128, 48, 512 + groups, m_per_group, n, k = 2, 128, 48, 640 m = groups * m_per_group a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) @@ -154,7 +161,9 @@ def test_m_grouped_mxfp8_fp8_contiguous_packed_int32_scale_accuracy(): ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) d = torch.empty_like(ref) - deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous(a, b, d, grouped_layout) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous( + a, b, d, grouped_layout, recipe_a=(1, 128), recipe_b=(1, 32) + ) diff = calc_diff(d, ref) assert diff < 0.03 From a5ef61348a679f91be18286a7b33ad328c0b15f2 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 22:04:28 +0800 Subject: [PATCH 15/30] Fix SM90 masked packed scale group stride --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 16 ++++--- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 6 ++- tests/test_sm90_mxfp8_fp8.py | 43 +++++++++++++++++++ 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index ef9a50ca95..c50def739b 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -24,7 +24,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 6; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 8; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -73,7 +73,7 @@ static void __instantiate_kernel() {{ static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.sfa, args.sfb, args.grouped_layout, - args.sfa_stride_m, args.sfa_stride_k, + args.sfa_stride_group, args.sfa_stride_m, args.sfa_stride_k, args.sfa_gran_k, args.sfa_packed_int32, args.sfb_stride_group, args.sfb_stride_n, args.sfb_stride_k, args.sfb_gran_k, args.sfb_packed_int32, @@ -168,6 +168,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .sfa = sfa.data_ptr(), .sfb = sfb.data_ptr(), .grouped_layout = grouped_layout.data_ptr(), + .sfa_stride_group = 0, .sfa_stride_m = static_cast(sfa.stride(0)), .sfa_stride_k = static_cast(sfa.stride(1)), .sfa_gran_k = static_cast(sfa_gran_k), @@ -182,7 +183,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v7", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v8", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -253,8 +254,9 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .sfa = sfa.data_ptr(), .sfb = sfb.data_ptr(), .grouped_layout = masked_m.data_ptr(), - .sfa_stride_m = static_cast(sfa.stride(-2)), - .sfa_stride_k = static_cast(sfa.stride(-1)), + .sfa_stride_group = static_cast(sfa.stride(0)), + .sfa_stride_m = static_cast(sfa.stride(1)), + .sfa_stride_k = static_cast(sfa.stride(2)), .sfa_gran_k = static_cast(sfa_gran_k), .sfa_packed_int32 = sfa.scalar_type() == torch::kInt, .sfb_stride_group = static_cast(sfb.stride(0)), @@ -267,7 +269,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v7", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v8", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 0d2af3b6b2..db360f3ec4 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -51,7 +51,7 @@ template CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, - uint32_t sfa_stride_m, uint32_t sfa_stride_k, + uint32_t sfa_stride_group, uint32_t sfa_stride_m, uint32_t sfa_stride_k, uint32_t sfa_gran_k, bool sfa_packed_int32, uint32_t sfb_stride_group, uint32_t sfb_stride_n, uint32_t sfb_stride_k, uint32_t sfb_gran_k, bool sfb_packed_int32, @@ -174,7 +174,9 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, const uint32_t k_scale_idx = k_idx_for_scale / sfa_gran_k; const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and k_idx_for_scale < shape_k; - const uint32_t sfa_base_offset = m_idx * sfa_stride_m; + const uint32_t sfa_local_m = kMasked ? (m_idx - scheduler.current_group_idx * shape_m) : m_idx; + const uint32_t sfa_base_offset = (kMasked ? scheduler.current_group_idx * sfa_stride_group : 0) + + sfa_local_m * sfa_stride_m; smem_sfa[stage_idx][i] = is_valid ? mxfp8_fp8_detail::load_e8m0_scale( sfa, sfa_base_offset, k_scale_idx, sfa_stride_k, sfa_packed_int32) : diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 03cdd5e590..a5dbd7a14e 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -49,6 +49,11 @@ def _pack_ue8m0_u8_to_i32(sf: torch.Tensor) -> torch.Tensor: ).contiguous() +def _pack_ue8m0_u8_to_i32_mn_major(sf: torch.Tensor) -> torch.Tensor: + packed = _pack_ue8m0_u8_to_i32(sf) + return packed.transpose(-1, -2).contiguous().transpose(-1, -2) + + def _tflops(m: int, n: int, k: int, elapsed: float) -> float: return 2.0 * m * n * k / elapsed / 1e12 @@ -168,6 +173,44 @@ def test_m_grouped_mxfp8_fp8_contiguous_packed_int32_scale_accuracy(): assert diff < 0.03 +def test_m_grouped_mxfp8_fp8_masked_packed_int32_mn_major_scale_accuracy(): + _require_sm90() + groups, max_m, n, k = 3, 128, 64, 640 + masked_m = torch.tensor([7, 65, 113], device="cuda", dtype=torch.int32) + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_sf_fp32 = torch.empty((groups, max_m, k // 128), device="cuda", dtype=torch.float32) + b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + for group_id in range(groups): + a_data[group_id], a_sf_fp32[group_id] = per_token_cast_to_fp8( + a_ref[group_id], use_ue8m0=True, gran_k=128 + ) + b_data[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + + a = (a_data, _pack_ue8m0_u8_to_i32_mn_major(_e8m0_from_fp32_pow2(a_sf_fp32))) + b = (b_data, _e8m0_from_fp32_pow2(b_sf_fp32)) + a_dequant = _cast_back_from_fp8_1d(a_data, a_sf_fp32, gran_k=128) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for group_id, valid_m in enumerate(masked_m.tolist()): + b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_sf_fp32[group_id], gran_k=32) + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + + d = torch.empty_like(ref) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_masked( + a, b, d, masked_m, expected_m=max_m, recipe_a=(1, 128), recipe_b=(1, 32) + ) + diff = max( + calc_diff(d[group_id, :valid_m], ref[group_id, :valid_m]) + for group_id, valid_m in enumerate(masked_m.tolist()) + ) + assert diff < 0.03 + + def test_m_grouped_mxfp8_vs_fp8_perf_contiguous_and_masked(): _require_sm90() groups, n, k = 4, 1024, 1024 From 2b5b958c4ee987d898c68545e794f5ef7d19f71e Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 22:13:14 +0800 Subject: [PATCH 16/30] Allow non-contiguous SM90 MXFP8 scales --- csrc/apis/gemm.hpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 30167934af..9675cf6e8c 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -292,8 +292,7 @@ static void m_grouped_mxfp8_fp8_gemm_nt_contiguous(const std::pair 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); - DG_HOST_ASSERT((a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt) and - a.second.is_contiguous()); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt); DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8 or b.second.scalar_type() == torch::kInt); if (recipe_a.has_value()) DG_HOST_ASSERT(std::get<0>(recipe_a.value()) == 1 and @@ -349,8 +348,7 @@ static void m_grouped_mxfp8_fp8_gemm_nt_masked(const std::pair 0 and n > 0); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 and d.is_contiguous()); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); - DG_HOST_ASSERT((a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt) and - a.second.is_contiguous()); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kUInt8 or a.second.scalar_type() == torch::kInt); DG_HOST_ASSERT(b.second.scalar_type() == torch::kUInt8 or b.second.scalar_type() == torch::kInt); if (recipe_a.has_value()) DG_HOST_ASSERT(std::get<0>(recipe_a.value()) == 1 and From 9862891134057e2fd916df26055bedd6718389e5 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 17 Jun 2026 12:02:29 +0800 Subject: [PATCH 17/30] Test UE8M0 int32 packing byte order --- tests/test_sm90_mxfp8_fp8.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index a5dbd7a14e..852807d747 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -54,6 +54,21 @@ def _pack_ue8m0_u8_to_i32_mn_major(sf: torch.Tensor) -> torch.Tensor: return packed.transpose(-1, -2).contiguous().transpose(-1, -2) +def test_packed_ue8m0_i32_byte_order_matches_sm100_layout(): + _require_sm90() + import deep_gemm.utils.layout + + sf = torch.tensor( + [[2.0, 4.0, 8.0, 16.0], [32.0, 64.0, 128.0, 256.0]], + device="cuda", + dtype=torch.float32, + ) + expected = _pack_ue8m0_u8_to_i32(_e8m0_from_fp32_pow2(sf)) + packed = deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) + + assert torch.equal(packed.cpu(), expected.cpu()) + + def _tflops(m: int, n: int, k: int, elapsed: float) -> float: return 2.0 * m * n * k / elapsed / 1e12 From 65871a518b45abbea0219fa163de1941c396b16a Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 24 Jun 2026 16:00:04 +0800 Subject: [PATCH 18/30] Fix SM90 MXFP8 contiguous RHS scale group --- .../include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index db360f3ec4..bd1ea3e56a 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -190,7 +190,10 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; const uint32_t k_scale_idx = k_idx_for_scale / sfb_gran_k; const bool is_valid = n_idx < shape_n and k_idx_for_scale < shape_k; - const uint32_t sfb_base_offset = scheduler.current_group_idx * sfb_stride_group + + const uint32_t sfb_group_idx = kMasked ? + scheduler.current_group_idx : + static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + const uint32_t sfb_base_offset = sfb_group_idx * sfb_stride_group + n_idx * sfb_stride_n; smem_sfb[stage_idx][i] = is_valid ? mxfp8_fp8_detail::load_e8m0_scale( From 203e3b3e8873246078dc23663b9fa08004417bb7 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Thu, 25 Jun 2026 17:30:38 +0800 Subject: [PATCH 19/30] Add SM90 MXFP8 DeepEP scale layout test --- tests/test_sm90_mxfp8_fp8.py | 65 ++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 852807d747..bd7f4bb303 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -54,6 +54,10 @@ def _pack_ue8m0_u8_to_i32_mn_major(sf: torch.Tensor) -> torch.Tensor: return packed.transpose(-1, -2).contiguous().transpose(-1, -2) +def _fp32_from_e8m0_u8(sf: torch.Tensor) -> torch.Tensor: + return torch.bitwise_left_shift(sf.to(torch.int32), 23).contiguous().view(torch.float32) + + def test_packed_ue8m0_i32_byte_order_matches_sm100_layout(): _require_sm90() import deep_gemm.utils.layout @@ -188,6 +192,67 @@ def test_m_grouped_mxfp8_fp8_contiguous_packed_int32_scale_accuracy(): assert diff < 0.03 +def test_m_grouped_mxfp8_fp8_contiguous_deepep_normal_scale_layout_accuracy(): + _require_sm90() + torch.manual_seed(0) + # Matches SGLang DeepEP normal layout: + # A scale: packed int32 MN-major non-contiguous view, gran_k=128 + # B scale: raw uint8 [expert, n, k/32], gran_k=32 + groups, m_per_group, n, k = 3, 128, 80, 640 + m = groups * m_per_group + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + a_data, _ = per_token_cast_to_fp8(a_ref, use_ue8m0=True, gran_k=128) + b_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + for group_id in range(groups): + b_data[group_id], _ = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + + a_exp = ( + 124 + + (torch.arange(m, device="cuda", dtype=torch.uint8).view(m, 1) % 5) + + (torch.arange(k // 128, device="cuda", dtype=torch.uint8).view(1, -1) % 3) + ) + b_exp = ( + 123 + + (torch.arange(groups, device="cuda", dtype=torch.uint8).view(groups, 1, 1) % 3) + + (torch.arange(n, device="cuda", dtype=torch.uint8).view(1, n, 1) % 5) + + (torch.arange(k // 32, device="cuda", dtype=torch.uint8).view(1, 1, -1) % 2) + ) + + a_scale_i32 = _pack_ue8m0_u8_to_i32_mn_major(a_exp) + b_scale_u8 = b_exp.contiguous() + a = (a_data, a_scale_i32) + b = (b_data, b_scale_u8) + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave( + m_per_group + ) + + a_dequant = _cast_back_from_fp8_1d(a_data, _fp32_from_e8m0_u8(a_exp), gran_k=128) + b_scale_fp32 = _fp32_from_e8m0_u8(b_exp) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + start = group_id * m_per_group + end = start + m_per_group + b_dequant = _cast_back_from_fp8_1d(b_data[group_id], b_scale_fp32[group_id], gran_k=32) + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + + d = torch.empty_like(ref) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous( + a, b, d, grouped_layout, recipe_a=(1, 128), recipe_b=(1, 32) + ) + diff = calc_diff(d, ref) + max_abs_diff = (d.float() - ref.float()).abs().max().item() + print( + "DeepEP-normal scale layout diff: " + f"calc_diff={diff:.6f}, max_abs_diff={max_abs_diff:.6f}, " + f"a_scale_shape={tuple(a_scale_i32.shape)}, a_scale_stride={tuple(a_scale_i32.stride())}, " + f"b_scale_shape={tuple(b_scale_u8.shape)}, b_scale_stride={tuple(b_scale_u8.stride())}" + ) + assert diff < 0.03 + + def test_m_grouped_mxfp8_fp8_masked_packed_int32_mn_major_scale_accuracy(): _require_sm90() groups, max_m, n, k = 3, 128, 64, 640 From d2032ea95868e48411c9590c5a56d22985c2e515 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Thu, 25 Jun 2026 17:37:57 +0800 Subject: [PATCH 20/30] Tighten SM90 MXFP8 DeepEP scale layout test --- tests/test_sm90_mxfp8_fp8.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index bd7f4bb303..ac1bddd5e6 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -209,15 +209,18 @@ def test_m_grouped_mxfp8_fp8_contiguous_deepep_normal_scale_layout_accuracy(): b_ref[group_id], use_ue8m0=True, gran_k=32 ) + # Keep exponents near 127 (scale 1.0). Wider synthetic ranges produce very + # large BF16 outputs, where a normal one-ULP BF16 difference has a misleading + # absolute error while relative/cosine error is still essentially zero. a_exp = ( - 124 - + (torch.arange(m, device="cuda", dtype=torch.uint8).view(m, 1) % 5) - + (torch.arange(k // 128, device="cuda", dtype=torch.uint8).view(1, -1) % 3) + 126 + + (torch.arange(m, device="cuda", dtype=torch.uint8).view(m, 1) % 2) + + (torch.arange(k // 128, device="cuda", dtype=torch.uint8).view(1, -1) % 2) ) b_exp = ( - 123 - + (torch.arange(groups, device="cuda", dtype=torch.uint8).view(groups, 1, 1) % 3) - + (torch.arange(n, device="cuda", dtype=torch.uint8).view(1, n, 1) % 5) + 126 + + (torch.arange(groups, device="cuda", dtype=torch.uint8).view(groups, 1, 1) % 2) + + (torch.arange(n, device="cuda", dtype=torch.uint8).view(1, n, 1) % 2) + (torch.arange(k // 32, device="cuda", dtype=torch.uint8).view(1, 1, -1) % 2) ) @@ -244,9 +247,12 @@ def test_m_grouped_mxfp8_fp8_contiguous_deepep_normal_scale_layout_accuracy(): ) diff = calc_diff(d, ref) max_abs_diff = (d.float() - ref.float()).abs().max().item() + ref_absmax = ref.float().abs().max().item() + max_rel_diff = max_abs_diff / max(ref_absmax, 1.0) print( "DeepEP-normal scale layout diff: " f"calc_diff={diff:.6f}, max_abs_diff={max_abs_diff:.6f}, " + f"ref_absmax={ref_absmax:.6f}, max_rel_diff={max_rel_diff:.6f}, " f"a_scale_shape={tuple(a_scale_i32.shape)}, a_scale_stride={tuple(a_scale_i32.stride())}, " f"b_scale_shape={tuple(b_scale_u8.shape)}, b_scale_stride={tuple(b_scale_u8.stride())}" ) From 664fa78c10d2716f224aa1bac320e69aca90d69c Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Thu, 25 Jun 2026 17:49:03 +0800 Subject: [PATCH 21/30] Add SM90 MXFP8 dense raw scale test --- tests/test_sm90_mxfp8_fp8.py | 69 ++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index ac1bddd5e6..6c2ba76f4c 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -259,6 +259,75 @@ def test_m_grouped_mxfp8_fp8_contiguous_deepep_normal_scale_layout_accuracy(): assert diff < 0.03 +def test_m_grouped_mxfp8_fp8_contiguous_dense_linear_raw_u8_scale_accuracy(): + _require_sm90() + torch.manual_seed(1) + # Matches SGLang dense linear through the SM90 grouped-contiguous wrapper: + # one logical RHS group, raw uint8 UE8M0 scales on both A and B, and padded M. + m, padded_m, n, k = 137, 256, 96, 640 + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + a_data, _ = per_token_cast_to_fp8(a_ref, use_ue8m0=True, gran_k=32) + b_data, _ = per_token_cast_to_fp8(b_ref, use_ue8m0=True, gran_k=32) + + a_exp = ( + 124 + + (torch.arange(m, device="cuda", dtype=torch.uint8).view(m, 1) % 5) + + (torch.arange(k // 32, device="cuda", dtype=torch.uint8).view(1, -1) % 3) + ) + b_exp = ( + 124 + + (torch.arange(n, device="cuda", dtype=torch.uint8).view(n, 1) % 5) + + (torch.arange(k // 32, device="cuda", dtype=torch.uint8).view(1, -1) % 3) + ) + + kernel_a = torch.zeros((padded_m, k), device="cuda", dtype=torch.float8_e4m3fn) + kernel_a[:m] = a_data + kernel_a_scale = torch.zeros((padded_m, k // 32), device="cuda", dtype=torch.uint8) + kernel_a_scale[:m] = a_exp + kernel_b = b_data.unsqueeze(0).contiguous() + kernel_b_scale = b_exp.unsqueeze(0).contiguous() + m_indices = torch.full((padded_m,), -1, device="cuda", dtype=torch.int32) + m_indices[:m] = 0 + + a_dequant = _cast_back_from_fp8_1d(a_data, _fp32_from_e8m0_u8(a_exp), gran_k=32) + b_dequant = _cast_back_from_fp8_1d(b_data, _fp32_from_e8m0_u8(b_exp), gran_k=32) + ref = (a_dequant @ b_dequant.t()).to(torch.bfloat16) + d_padded = torch.empty((padded_m, n), device="cuda", dtype=torch.bfloat16) + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous( + (kernel_a, kernel_a_scale), + (kernel_b, kernel_b_scale), + d_padded, + m_indices, + recipe_a=(1, 32), + recipe_b=(1, 32), + ) + d = d_padded[:m] + + inv_a_scale = _fp32_from_e8m0_u8((254 - a_exp).to(torch.uint8)) + inv_b_scale = _fp32_from_e8m0_u8((254 - b_exp).to(torch.uint8)) + inv_ref = ( + _cast_back_from_fp8_1d(a_data, inv_a_scale, gran_k=32) + @ _cast_back_from_fp8_1d(b_data, inv_b_scale, gran_k=32).t() + ).to(torch.bfloat16) + + diff = calc_diff(d, ref) + inverse_diff = calc_diff(d, inv_ref) + max_abs_diff = (d.float() - ref.float()).abs().max().item() + ref_absmax = ref.float().abs().max().item() + max_rel_diff = max_abs_diff / max(ref_absmax, 1.0) + print( + "Dense raw-u8 scale layout diff: " + f"calc_diff={diff:.6f}, inverse_scale_calc_diff={inverse_diff:.6f}, " + f"max_abs_diff={max_abs_diff:.6f}, ref_absmax={ref_absmax:.6f}, " + f"max_rel_diff={max_rel_diff:.6f}, " + f"a_scale_shape={tuple(kernel_a_scale.shape)}, " + f"b_scale_shape={tuple(kernel_b_scale.shape)}" + ) + assert diff < 0.03 + assert inverse_diff > diff + 0.1 + + def test_m_grouped_mxfp8_fp8_masked_packed_int32_mn_major_scale_accuracy(): _require_sm90() groups, max_m, n, k = 3, 128, 64, 640 From b3897f906f28cc8cbb55b22918a75f3c2201938f Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 14 Jul 2026 22:44:44 +0800 Subject: [PATCH 22/30] test: add mxfp8 contiguous large-m perf check --- tests/test_sm90_mxfp8_fp8.py | 102 +++++++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/tests/test_sm90_mxfp8_fp8.py b/tests/test_sm90_mxfp8_fp8.py index 6c2ba76f4c..fe5775c7df 100644 --- a/tests/test_sm90_mxfp8_fp8.py +++ b/tests/test_sm90_mxfp8_fp8.py @@ -499,3 +499,105 @@ def run_fp8_masked(): assert fp8_masked_elapsed > 0 assert contiguous_diff == contiguous_diff assert masked_diff == masked_diff + + +def test_m_grouped_mxfp8_fp8_contiguous_perf_large_m_scaling(): + _require_sm90() + torch.manual_seed(2) + + groups, n, k = 4, 1024, 1024 + cases = [ + ("small_m", 128), + ("large_m", 2048), + ] + rows = [] + + for case_name, m_per_group in cases: + m = groups * m_per_group + a_ref = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_mx_data, a_mx_sf_fp32 = per_token_cast_to_fp8( + a_ref, use_ue8m0=True, gran_k=32 + ) + a_mx = (a_mx_data, _e8m0_from_fp32_pow2(a_mx_sf_fp32)) + a_fp8 = per_token_cast_to_fp8(a_ref, use_ue8m0=False, gran_k=128) + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave( + m_per_group + ) + + b_mx_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_mx_sf_fp32 = torch.empty((groups, n, k // 32), device="cuda", dtype=torch.float32) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // 128), device="cuda", dtype=torch.float32) + for group_id in range(groups): + b_mx_data[group_id], b_mx_sf_fp32[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=True, gran_k=32 + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_ref[group_id], use_ue8m0=False, gran_k=128 + ) + b_mx = (b_mx_data, _e8m0_from_fp32_pow2(b_mx_sf_fp32)) + b_fp8 = (b_fp8_data, b_fp8_sf) + d_mx = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + d_fp8 = torch.empty_like(d_mx) + + def run_mx_contiguous(): + deep_gemm.m_grouped_mxfp8_fp8_gemm_nt_contiguous( + a_mx, b_mx, d_mx, grouped_layout + ) + + def run_fp8_contiguous(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a_fp8, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, 128), + recipe_b=(1, 128), + ) + + mx_elapsed = _time_kernel(run_mx_contiguous) + fp8_elapsed = _time_kernel(run_fp8_contiguous) + speedup = fp8_elapsed / mx_elapsed + rows.append( + { + "case_name": case_name, + "m": m, + "mx_elapsed": mx_elapsed, + "fp8_elapsed": fp8_elapsed, + "mx_tflops": _tflops(m, n, k, mx_elapsed), + "fp8_tflops": _tflops(m, n, k, fp8_elapsed), + "speedup": speedup, + } + ) + + assert mx_elapsed > 0 + assert fp8_elapsed > 0 + assert speedup == speedup + + print("case | M | MXFP8 us | FP8 us | MXFP8 TFLOPS | FP8 TFLOPS | speedup") + print("-- | -- | -- | -- | -- | -- | --") + for row in rows: + print( + f"{row['case_name']} | {row['m']} | " + f"{row['mx_elapsed'] * 1e6:.0f} | {row['fp8_elapsed'] * 1e6:.0f} | " + f"{row['mx_tflops']:.1f} | {row['fp8_tflops']:.1f} | " + f"{row['speedup']:.2f}x" + ) + + large_m_row = rows[-1] + print( + "When m increases, contiguous kernel performance is poor; " + f"at m = {large_m_row['m']}, MXFP8 contiguous is " + f"{large_m_row['mx_elapsed'] * 1e6:.0f} us versus FP8 contiguous " + f"{large_m_row['fp8_elapsed'] * 1e6:.0f} us -- only " + f"{large_m_row['speedup']:.2f}x the speed." + ) + assert large_m_row["m"] == 8192 + assert large_m_row["speedup"] >= 0.25, ( + "MXFP8 contiguous kernel regressed badly at m=8192: " + f"{large_m_row['mx_elapsed'] * 1e6:.0f} us vs " + f"{large_m_row['fp8_elapsed'] * 1e6:.0f} us, " + f"speedup={large_m_row['speedup']:.2f}x" + ) From e5f029ca248456c9b53407260aaa3fda1f1504f4 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 14 Jul 2026 22:54:33 +0800 Subject: [PATCH 23/30] perf: optimize sm90 mxfp8 scale handling --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 6 +- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 106 +++++++++++++----- 2 files changed, 81 insertions(+), 31 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index c50def739b..e4e796418b 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 8; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 9; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -183,7 +183,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v8", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v9", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -269,7 +269,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v8", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v9", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index bd1ea3e56a..3a63b66552 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -38,6 +38,38 @@ CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, return reinterpret_cast(ptr)[base_offset + k_scale_idx * stride_k]; } +CUTLASS_DEVICE uint32_t load_e8m0_scale_quad(const void* ptr, uint32_t base_offset, + uint32_t k_block_idx, uint32_t stride_k, + uint32_t gran_k, bool packed_int32, + uint32_t shape_k) { + constexpr uint32_t kIdentityScalePack = 0x7f7f7f7f; + const uint32_t first_k_scale_idx = (k_block_idx * 128) / gran_k; + const uint32_t first_k = k_block_idx * 128; + + if (first_k >= shape_k) + return kIdentityScalePack; + + if (gran_k == 32 and first_k + 128 <= shape_k) { + if (packed_int32 and stride_k == 1) + return reinterpret_cast(ptr)[base_offset + first_k_scale_idx / 4]; + if (not packed_int32 and stride_k == 1 and base_offset % 4 == 0) + return *reinterpret_cast( + reinterpret_cast(ptr) + base_offset + first_k_scale_idx); + } + + uint32_t packed = 0; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + const uint32_t k = first_k + kk * 32; + const uint32_t k_scale_idx = k / gran_k; + const uint8_t scale = k < shape_k ? + load_e8m0_scale(ptr, base_offset, k_scale_idx, stride_k, packed_int32) : + static_cast(127); + packed |= static_cast(scale) << (kk * 8); + } + return packed; +} + } // namespace mxfp8_fp8_detail template (shape_m, BLOCK_M, m_block_idx); - for (uint32_t i = lane_idx; i < BLOCK_M * SHAPE_K_SFA_PER_STAGE; i += 32) { - const uint32_t m_offset = i / SHAPE_K_SFA_PER_STAGE; - const uint32_t k_scale_offset = i % SHAPE_K_SFA_PER_STAGE; + for (uint32_t m_offset = lane_idx; m_offset < BLOCK_M; m_offset += 32) { const uint32_t m_idx = sfa_base_m + m_offset; - const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; - const uint32_t k_scale_idx = k_idx_for_scale / sfa_gran_k; - const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and - k_idx_for_scale < shape_k; + const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1); const uint32_t sfa_local_m = kMasked ? (m_idx - scheduler.current_group_idx * shape_m) : m_idx; const uint32_t sfa_base_offset = (kMasked ? scheduler.current_group_idx * sfa_stride_group : 0) + sfa_local_m * sfa_stride_m; - smem_sfa[stage_idx][i] = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale( - sfa, sfa_base_offset, k_scale_idx, sfa_stride_k, sfa_packed_int32) : - static_cast(127); + const uint32_t scales = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale_quad( + sfa, sfa_base_offset, k_block_idx, sfa_stride_k, + sfa_gran_k, sfa_packed_int32, shape_k) : + 0x7f7f7f7f; + *reinterpret_cast( + smem_sfa[stage_idx] + m_offset * SHAPE_K_SFA_PER_STAGE) = scales; } - for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { - const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; - const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; + const uint32_t sfb_group_idx = kMasked ? + scheduler.current_group_idx : + static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + for (uint32_t n_offset = lane_idx; n_offset < BLOCK_N; n_offset += 32) { const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; - const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; - const uint32_t k_scale_idx = k_idx_for_scale / sfb_gran_k; - const bool is_valid = n_idx < shape_n and k_idx_for_scale < shape_k; - const uint32_t sfb_group_idx = kMasked ? - scheduler.current_group_idx : - static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + const bool is_valid = n_idx < shape_n; const uint32_t sfb_base_offset = sfb_group_idx * sfb_stride_group + n_idx * sfb_stride_n; - smem_sfb[stage_idx][i] = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale( - sfb, sfb_base_offset, k_scale_idx, sfb_stride_k, sfb_packed_int32) : - static_cast(127); + const uint32_t scales = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale_quad( + sfb, sfb_base_offset, k_block_idx, sfb_stride_k, + sfb_gran_k, sfb_packed_int32, shape_k) : + 0x7f7f7f7f; + *reinterpret_cast( + smem_sfb[stage_idx] + n_offset * SHAPE_K_SFB_PER_STAGE) = scales; } __threadfence_block(); __syncwarp(); @@ -227,7 +256,9 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, while (scheduler.get_next_block(m_block_idx, n_block_idx)) { 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}; + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, "MXFP8 pairwise WGMMA promotion assumes four K/32 chunks"); + float accum[WGMMA::kNumAccum], accum_1[WGMMA::kNumAccum], + final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for BLOCK_M < 64"); constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); @@ -276,18 +307,25 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, } #pragma unroll - for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; kk += 2) { #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_1[i]); + } ptx::warpgroup_arrive(); a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; WGMMA::wgmma(a_desc, b_desc, accum, false); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + (kk + 1) * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + (kk + 1) * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, accum_1, false); ptx::warpgroup_commit_batch(); #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_1[i]); + } ptx::warpgroup_wait<0>(); if (not do_wgmma_store) @@ -297,6 +335,10 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, static_cast((sfa_pack_0 >> (kk * 8)) & 0xff)); const float scale_a_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfa_pack_1 >> (kk * 8)) & 0xff)); + const float scale_a_next_0 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfa_pack_0 >> ((kk + 1) * 8)) & 0xff)); + const float scale_a_next_1 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfa_pack_1 >> ((kk + 1) * 8)) & 0xff)); auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { @@ -304,10 +346,18 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff)); const float scale_b_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff)); + const float scale_b_next_0 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfb_pack[i][0] >> ((kk + 1) * 8)) & 0xff)); + const float scale_b_next_1 = mxfp8_fp8_detail::e8m0_to_float( + static_cast((sfb_pack[i][1] >> ((kk + 1) * 8)) & 0xff)); shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + shifted_accum[i * 4 + 0] += scale_a_next_0 * scale_b_next_0 * accum_1[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_a_next_0 * scale_b_next_1 * accum_1[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_a_next_1 * scale_b_next_0 * accum_1[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_a_next_1 * scale_b_next_1 * accum_1[i * 4 + 3]; } } From 39e756ca6463b55b0fda15279293236a8b1d80b4 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 14 Jul 2026 23:03:24 +0800 Subject: [PATCH 24/30] Revert "perf: optimize sm90 mxfp8 scale handling" This reverts commit e5f029ca248456c9b53407260aaa3fda1f1504f4. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 6 +- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 106 +++++------------- 2 files changed, 31 insertions(+), 81 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index e4e796418b..c50def739b 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 9; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 8; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -183,7 +183,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v9", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v8", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -269,7 +269,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v9", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v8", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 3a63b66552..bd1ea3e56a 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -38,38 +38,6 @@ CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, return reinterpret_cast(ptr)[base_offset + k_scale_idx * stride_k]; } -CUTLASS_DEVICE uint32_t load_e8m0_scale_quad(const void* ptr, uint32_t base_offset, - uint32_t k_block_idx, uint32_t stride_k, - uint32_t gran_k, bool packed_int32, - uint32_t shape_k) { - constexpr uint32_t kIdentityScalePack = 0x7f7f7f7f; - const uint32_t first_k_scale_idx = (k_block_idx * 128) / gran_k; - const uint32_t first_k = k_block_idx * 128; - - if (first_k >= shape_k) - return kIdentityScalePack; - - if (gran_k == 32 and first_k + 128 <= shape_k) { - if (packed_int32 and stride_k == 1) - return reinterpret_cast(ptr)[base_offset + first_k_scale_idx / 4]; - if (not packed_int32 and stride_k == 1 and base_offset % 4 == 0) - return *reinterpret_cast( - reinterpret_cast(ptr) + base_offset + first_k_scale_idx); - } - - uint32_t packed = 0; - #pragma unroll - for (uint32_t kk = 0; kk < 4; ++ kk) { - const uint32_t k = first_k + kk * 32; - const uint32_t k_scale_idx = k / gran_k; - const uint8_t scale = k < shape_k ? - load_e8m0_scale(ptr, base_offset, k_scale_idx, stride_k, packed_int32) : - static_cast(127); - packed |= static_cast(scale) << (kk * 8); - } - return packed; -} - } // namespace mxfp8_fp8_detail template (shape_m, BLOCK_M, m_block_idx); - for (uint32_t m_offset = lane_idx; m_offset < BLOCK_M; m_offset += 32) { + for (uint32_t i = lane_idx; i < BLOCK_M * SHAPE_K_SFA_PER_STAGE; i += 32) { + const uint32_t m_offset = i / SHAPE_K_SFA_PER_STAGE; + const uint32_t k_scale_offset = i % SHAPE_K_SFA_PER_STAGE; const uint32_t m_idx = sfa_base_m + m_offset; - const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1); + const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; + const uint32_t k_scale_idx = k_idx_for_scale / sfa_gran_k; + const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and + k_idx_for_scale < shape_k; const uint32_t sfa_local_m = kMasked ? (m_idx - scheduler.current_group_idx * shape_m) : m_idx; const uint32_t sfa_base_offset = (kMasked ? scheduler.current_group_idx * sfa_stride_group : 0) + sfa_local_m * sfa_stride_m; - const uint32_t scales = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale_quad( - sfa, sfa_base_offset, k_block_idx, sfa_stride_k, - sfa_gran_k, sfa_packed_int32, shape_k) : - 0x7f7f7f7f; - *reinterpret_cast( - smem_sfa[stage_idx] + m_offset * SHAPE_K_SFA_PER_STAGE) = scales; + smem_sfa[stage_idx][i] = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale( + sfa, sfa_base_offset, k_scale_idx, sfa_stride_k, sfa_packed_int32) : + static_cast(127); } - const uint32_t sfb_group_idx = kMasked ? - scheduler.current_group_idx : - static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); - for (uint32_t n_offset = lane_idx; n_offset < BLOCK_N; n_offset += 32) { + for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { + const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; + const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; - const bool is_valid = n_idx < shape_n; + const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; + const uint32_t k_scale_idx = k_idx_for_scale / sfb_gran_k; + const bool is_valid = n_idx < shape_n and k_idx_for_scale < shape_k; + const uint32_t sfb_group_idx = kMasked ? + scheduler.current_group_idx : + static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); const uint32_t sfb_base_offset = sfb_group_idx * sfb_stride_group + n_idx * sfb_stride_n; - const uint32_t scales = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale_quad( - sfb, sfb_base_offset, k_block_idx, sfb_stride_k, - sfb_gran_k, sfb_packed_int32, shape_k) : - 0x7f7f7f7f; - *reinterpret_cast( - smem_sfb[stage_idx] + n_offset * SHAPE_K_SFB_PER_STAGE) = scales; + smem_sfb[stage_idx][i] = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale( + sfb, sfb_base_offset, k_scale_idx, sfb_stride_k, sfb_packed_int32) : + static_cast(127); } __threadfence_block(); __syncwarp(); @@ -256,9 +227,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, while (scheduler.get_next_block(m_block_idx, n_block_idx)) { 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"); - DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, "MXFP8 pairwise WGMMA promotion assumes four K/32 chunks"); - float accum[WGMMA::kNumAccum], accum_1[WGMMA::kNumAccum], - final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for BLOCK_M < 64"); constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); @@ -307,25 +276,18 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, } #pragma unroll - for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; kk += 2) { + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_fence_operand(accum_1[i]); - } ptx::warpgroup_arrive(); a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; WGMMA::wgmma(a_desc, b_desc, accum, false); - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + (kk + 1) * WGMMA::K) / 16; - b_desc.reg32_[0] = b_desc_base_lo + (kk + 1) * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, accum_1, false); ptx::warpgroup_commit_batch(); #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_fence_operand(accum_1[i]); - } ptx::warpgroup_wait<0>(); if (not do_wgmma_store) @@ -335,10 +297,6 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, static_cast((sfa_pack_0 >> (kk * 8)) & 0xff)); const float scale_a_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfa_pack_1 >> (kk * 8)) & 0xff)); - const float scale_a_next_0 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfa_pack_0 >> ((kk + 1) * 8)) & 0xff)); - const float scale_a_next_1 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfa_pack_1 >> ((kk + 1) * 8)) & 0xff)); auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { @@ -346,18 +304,10 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff)); const float scale_b_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff)); - const float scale_b_next_0 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfb_pack[i][0] >> ((kk + 1) * 8)) & 0xff)); - const float scale_b_next_1 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfb_pack[i][1] >> ((kk + 1) * 8)) & 0xff)); shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; - shifted_accum[i * 4 + 0] += scale_a_next_0 * scale_b_next_0 * accum_1[i * 4 + 0]; - shifted_accum[i * 4 + 1] += scale_a_next_0 * scale_b_next_1 * accum_1[i * 4 + 1]; - shifted_accum[i * 4 + 2] += scale_a_next_1 * scale_b_next_0 * accum_1[i * 4 + 2]; - shifted_accum[i * 4 + 3] += scale_a_next_1 * scale_b_next_1 * accum_1[i * 4 + 3]; } } From e1ccec7dc49159ae66e874f890a2c3b95c6585c5 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 14 Jul 2026 23:51:42 +0800 Subject: [PATCH 25/30] perf: pack sm90 mxfp8 producer scale loads without touching wgmma loop Re-apply only the safe half of the reverted e5f029c: the producer warp now loads the four per-32-K UE8M0 scales covering one 128-K block with a single 32-bit load (gran_k==32 unit-stride fast path) and writes them to SMEM with one 32-bit store per row, instead of four dependent per-byte global loads/stores. The math-warpgroup WGMMA loop is left on the v8 baseline structure, so this carries no extra accumulator registers and cannot spill (the pairwise accum_1 change from e5f029c was the cause of the 0.03x collapse). Bump JIT cache to v10 to avoid reusing the bad v9 kernel. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 6 +- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 82 +++++++++++++------ 2 files changed, 61 insertions(+), 27 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index c50def739b..981f07e23c 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 8; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 10; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -183,7 +183,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v8", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v10", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -269,7 +269,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v8", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v10", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index bd1ea3e56a..3b76b9b5b0 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -38,6 +38,43 @@ CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, return reinterpret_cast(ptr)[base_offset + k_scale_idx * stride_k]; } +// Load the four UE8M0 scales that cover one 128-wide K block (chunks kk = 0..3) +// and pack them into a single 32-bit word (byte kk = scale for chunk kk). This +// replaces four dependent per-byte global loads with a single 32-bit load on the +// common `gran_k == 32`, unit-stride path. Bounds/granularity are handled here so +// callers can drop their per-element K guards. +CUTLASS_DEVICE uint32_t load_e8m0_scale_quad(const void* ptr, uint32_t base_offset, + uint32_t k_block_idx, uint32_t stride_k, + uint32_t gran_k, bool packed_int32, + uint32_t shape_k) { + constexpr uint32_t kIdentityScalePack = 0x7f7f7f7f; + const uint32_t first_k_scale_idx = (k_block_idx * 128) / gran_k; + const uint32_t first_k = k_block_idx * 128; + + if (first_k >= shape_k) + return kIdentityScalePack; + + if (gran_k == 32 and first_k + 128 <= shape_k) { + if (packed_int32 and stride_k == 1) + return reinterpret_cast(ptr)[base_offset + first_k_scale_idx / 4]; + if (not packed_int32 and stride_k == 1 and base_offset % 4 == 0) + return *reinterpret_cast( + reinterpret_cast(ptr) + base_offset + first_k_scale_idx); + } + + uint32_t packed = 0; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + const uint32_t k = first_k + kk * 32; + const uint32_t k_scale_idx = k / gran_k; + const uint8_t scale = k < shape_k ? + load_e8m0_scale(ptr, base_offset, k_scale_idx, stride_k, packed_int32) : + static_cast(127); + packed |= static_cast(scale) << (kk * 8); + } + return packed; +} + } // namespace mxfp8_fp8_detail template (shape_m, BLOCK_M, m_block_idx); - for (uint32_t i = lane_idx; i < BLOCK_M * SHAPE_K_SFA_PER_STAGE; i += 32) { - const uint32_t m_offset = i / SHAPE_K_SFA_PER_STAGE; - const uint32_t k_scale_offset = i % SHAPE_K_SFA_PER_STAGE; + for (uint32_t m_offset = lane_idx; m_offset < BLOCK_M; m_offset += 32) { const uint32_t m_idx = sfa_base_m + m_offset; - const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; - const uint32_t k_scale_idx = k_idx_for_scale / sfa_gran_k; - const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1) and - k_idx_for_scale < shape_k; + const bool is_valid = m_idx < shape_m * (kMasked ? kNumGroups : 1); const uint32_t sfa_local_m = kMasked ? (m_idx - scheduler.current_group_idx * shape_m) : m_idx; const uint32_t sfa_base_offset = (kMasked ? scheduler.current_group_idx * sfa_stride_group : 0) + sfa_local_m * sfa_stride_m; - smem_sfa[stage_idx][i] = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale( - sfa, sfa_base_offset, k_scale_idx, sfa_stride_k, sfa_packed_int32) : - static_cast(127); + const uint32_t scales = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale_quad( + sfa, sfa_base_offset, k_block_idx, sfa_stride_k, + sfa_gran_k, sfa_packed_int32, shape_k) : + 0x7f7f7f7f; + *reinterpret_cast( + smem_sfa[stage_idx] + m_offset * SHAPE_K_SFA_PER_STAGE) = scales; } - for (uint32_t i = lane_idx; i < BLOCK_N * SHAPE_K_SFB_PER_STAGE; i += 32) { - const uint32_t n_offset = i / SHAPE_K_SFB_PER_STAGE; - const uint32_t k_scale_offset = i % SHAPE_K_SFB_PER_STAGE; + const uint32_t sfb_group_idx = kMasked ? + scheduler.current_group_idx : + static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + for (uint32_t n_offset = lane_idx; n_offset < BLOCK_N; n_offset += 32) { const uint32_t n_idx = n_block_idx * BLOCK_N + n_offset; - const uint32_t k_idx_for_scale = k_block_idx * BLOCK_K + k_scale_offset * 32; - const uint32_t k_scale_idx = k_idx_for_scale / sfb_gran_k; - const bool is_valid = n_idx < shape_n and k_idx_for_scale < shape_k; - const uint32_t sfb_group_idx = kMasked ? - scheduler.current_group_idx : - static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + const bool is_valid = n_idx < shape_n; const uint32_t sfb_base_offset = sfb_group_idx * sfb_stride_group + n_idx * sfb_stride_n; - smem_sfb[stage_idx][i] = is_valid ? - mxfp8_fp8_detail::load_e8m0_scale( - sfb, sfb_base_offset, k_scale_idx, sfb_stride_k, sfb_packed_int32) : - static_cast(127); + const uint32_t scales = is_valid ? + mxfp8_fp8_detail::load_e8m0_scale_quad( + sfb, sfb_base_offset, k_block_idx, sfb_stride_k, + sfb_gran_k, sfb_packed_int32, shape_k) : + 0x7f7f7f7f; + *reinterpret_cast( + smem_sfb[stage_idx] + n_offset * SHAPE_K_SFB_PER_STAGE) = scales; } __threadfence_block(); __syncwarp(); From 7816268928b5bd6996599a1303a675b47f45d712 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 15 Jul 2026 11:01:19 +0800 Subject: [PATCH 26/30] perf: double-buffer sm90 mxfp8 wgmma/promotion for large M MXFP8 promotes each 32-wide K chunk with its own UE8M0 scale, so unlike FP8 it cannot hardware-accumulate across the 4 chunks of a 128-K block; each WGMMA is drained (warpgroup_wait<0>) before its FP32 promotion, serializing the main loop and leaving MXFP8 ~4x slower than FP8 at large M. Add a 2-deep software pipeline that overlaps chunk (kk+1)'s WGMMA with chunk kk's promotion via a ping-pong accumulator and warpgroup_wait<1>. It is guarded by `if constexpr (kNumAccum <= 48 && BLOCK_M >= WGMMA::M)` so it only compiles when both accumulators fit under the 256-thread / 232-register budget; otherwise the proven serial single-buffer path runs (avoids the local-memory spill that crashed the earlier pairwise attempt). To actually take the pipeline, clamp the MXFP8 layout to BLOCK_N <= 96 (kNumAccum <= 48) by re-selecting the fastest narrow candidate from the shared SM90 heuristic, without polluting the FP8 heuristic. Bump JIT version 10 -> 11 to avoid stale kernel caches. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 54 +++++++++- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 102 ++++++++++++++---- 2 files changed, 128 insertions(+), 28 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 981f07e23c..ff558d1350 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 10; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 11; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -101,6 +101,50 @@ static void tune_mxfp8_fp8_smem_config(GemmConfig& config, const GemmDesc& desc) config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; } +// The MXFP8 kernel promotes every 32-wide K chunk with a distinct UE8M0 scale, so it cannot +// hardware-accumulate across the 4 chunks of a 128-K block the way FP8 does. To hide the +// per-chunk WGMMA drain it runs a 2-deep software pipeline with a ping-pong FP32 accumulator, +// which needs `kNumAccum = block_m * block_n / 128 <= 48` (i.e. block_n <= 96) to fit both +// buffers under the 256-thread / 232-register budget. The shared SM90 heuristic maximizes MMA +// throughput and typically picks block_n up to 192, which would force the serial fallback. +// Re-select the best layout among candidates constrained to block_n <= 96 so MXFP8 gets the +// pipeline; other dtypes keep using the unconstrained heuristic. +static constexpr int kMXFP8PipelineMaxBlockN = 96; + +static GemmConfig get_mxfp8_fp8_best_config(const GemmDesc& desc) { + auto config = get_best_config(desc); + if (config.layout.block_n <= kMXFP8PipelineMaxBlockN) + return config; + + // Filter the candidate layouts down to those that keep the ping-pong accumulator small + // enough, then pick the fastest among them using the same cost model. + const auto layout_candidates = SM90ArchSpec::get_layout_candidates(desc); + bool found = false; + Layout best_layout{}; + LayoutInfo best_info{}; + for (const auto& candidate: layout_candidates) { + if (candidate.block_n > kMXFP8PipelineMaxBlockN) + continue; + const auto info = SM90ArchSpec::get_layout_info(desc, candidate); + if (not found or SM90ArchSpec::compare(info, best_info)) { + best_layout = candidate; + best_info = info; + found = true; + } + } + + // No narrow candidate survived the heuristic's stage/swizzle filters: keep the original + // config so the kernel falls back to the (correct) serial path. + if (not found) + return config; + + config.layout = best_layout; + config.storage_config = SM90ArchSpec::get_storage_config(desc, best_layout); + config.pipeline_config = SM90ArchSpec::get_pipeline_config(desc, best_layout, config.storage_config); + config.launch_config = SM90ArchSpec::get_launch_config(desc, best_layout); + return config; +} + static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, @@ -129,7 +173,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = 1 }; - auto config = get_best_config(desc); + auto config = get_mxfp8_fp8_best_config(desc); tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); @@ -183,7 +227,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v10", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v11", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -215,7 +259,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; - auto config = get_best_config(desc); + auto config = get_mxfp8_fp8_best_config(desc); tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); @@ -269,7 +313,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v10", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v11", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 3b76b9b5b0..c1c2c39dbc 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -261,7 +261,16 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, while (scheduler.get_next_block(m_block_idx, n_block_idx)) { 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}; + // MXFP8 promotes each 32-wide K chunk with its own scale, so unlike FP8 we cannot + // hardware-accumulate across chunks: each WGMMA must be drained before its FP32 + // promotion. When the accumulator is small enough to double-buffer without spilling + // (BLOCK_N <= 96 => kNumAccum <= 48) and every math warp stores (BLOCK_M >= WGMMA::M, + // so `do_wgmma_store` is warpgroup-uniform and the WGMMA stays collective), run a + // 2-deep software pipeline that overlaps chunk kk's WGMMA with chunk kk-1's promotion. + // Otherwise fall back to the serial single-buffer path to stay under the 232-register + // budget (a second full-width accumulator at BLOCK_N=128 spills to local memory). + constexpr bool kPipelinePromotion = WGMMA::kNumAccum <= 48 and BLOCK_M >= WGMMA::M; + float final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for BLOCK_M < 64"); constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); @@ -309,39 +318,86 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, } } - #pragma unroll - for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; - b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, accum, false); - 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>(); - - if (not do_wgmma_store) - continue; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + // FP32 promotion of one drained 32-K chunk `kk` from accumulator `src`. + auto promote_chunk = [&](const float* src, const uint32_t& kk) { const float scale_a_0 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfa_pack_0 >> (kk * 8)) & 0xff)); const float scale_a_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfa_pack_1 >> (kk * 8)) & 0xff)); - auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { const float scale_b_0 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff)); const float scale_b_1 = mxfp8_fp8_detail::e8m0_to_float( static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff)); - shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; - shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; - shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; - shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * src[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * src[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * src[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * src[i * 4 + 3]; + } + }; + + if constexpr (kPipelinePromotion) { + // 2-deep software pipeline: overlap chunk (kk+1)'s WGMMA with chunk kk's + // FP32 promotion using a ping-pong accumulator. Enabled only when the + // accumulator is small (kNumAccum <= 48, guaranteed by the MXFP8 BLOCK_N + // clamp) so both buffers fit under the 232-register budget. + constexpr uint32_t kNumChunks = BLOCK_K / WGMMA::K; + float accum_buf[2][WGMMA::kNumAccum]; + + auto issue_chunk = [&](float* dst, const uint32_t& kk) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(dst[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, dst, false); + ptx::warpgroup_commit_batch(); + }; + + // Prologue: launch the first chunk, then in each step launch the next + // chunk before draining the current one. + issue_chunk(accum_buf[0], 0); + #pragma unroll + for (uint32_t kk = 0; kk < kNumChunks; ++ kk) { + const uint32_t cur = kk & 1; + if (kk + 1 < kNumChunks) + issue_chunk(accum_buf[(kk + 1) & 1], kk + 1); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum_buf[cur][i]); + // Non-last chunks keep the freshly issued (kk+1) group in flight and + // only drain kk; the last chunk drains everything. + if (kk + 1 < kNumChunks) + ptx::warpgroup_wait<1>(); + else + ptx::warpgroup_wait<0>(); + promote_chunk(accum_buf[cur], kk); + } + } else { + float accum[WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, accum, false); + 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>(); + + if (not do_wgmma_store) + continue; + + promote_chunk(accum, kk); } } From 8b27d66cb24f14566a8378666547830b119f280a Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 15 Jul 2026 11:21:20 +0800 Subject: [PATCH 27/30] perf: full-depth sm90 mxfp8 wgmma pipeline at BLOCK_N=64 The 2-deep pipeline lifted large-M MXFP8 from 0.31x to 0.47x of FP8, but it still keeps only two WGMMAs in flight because BLOCK_N=96 (kNumAccum=48) leaves no register room for more accumulators. Clamp the MXFP8 layout to BLOCK_N=64 (kNumAccum=32), which for n=1024 also tiles evenly (16 blocks, no wasted columns), and pick the deepest pipeline that fits the FP32 register budget at compile time: * kFullPipeline - one accumulator per 32-K chunk; issue all 4 WGMMAs as a single warpgroup batch and drain once, mirroring FP8's fully-loaded tensor pipeline (BLOCK_N<=64); * kTwoDeepPipeline - ping-pong overlap (BLOCK_N==96); * serial fallback - single buffer, no spill (BLOCK_N>=128). Depth is gated by `(depth + promote_waves) * kNumAccum <= 168` so it never spills. Bump JIT version 11 -> 12. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 21 ++-- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 106 ++++++++++++------ 2 files changed, 82 insertions(+), 45 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index ff558d1350..b330d9ebee 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 11; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 12; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -103,13 +103,14 @@ static void tune_mxfp8_fp8_smem_config(GemmConfig& config, const GemmDesc& desc) // The MXFP8 kernel promotes every 32-wide K chunk with a distinct UE8M0 scale, so it cannot // hardware-accumulate across the 4 chunks of a 128-K block the way FP8 does. To hide the -// per-chunk WGMMA drain it runs a 2-deep software pipeline with a ping-pong FP32 accumulator, -// which needs `kNumAccum = block_m * block_n / 128 <= 48` (i.e. block_n <= 96) to fit both -// buffers under the 256-thread / 232-register budget. The shared SM90 heuristic maximizes MMA -// throughput and typically picks block_n up to 192, which would force the serial fallback. -// Re-select the best layout among candidates constrained to block_n <= 96 so MXFP8 gets the -// pipeline; other dtypes keep using the unconstrained heuristic. -static constexpr int kMXFP8PipelineMaxBlockN = 96; +// per-chunk WGMMA drain it runs a software pipeline over a per-chunk FP32 accumulator whose +// depth is bounded by the register budget. Clamping to BLOCK_N == 64 gives kNumAccum == 32, +// which is small enough to hold one accumulator per chunk (full, FP8-style pipeline) under the +// 256-thread / 232-register budget, and (for n == 1024) tiles evenly with no wasted columns. +// The shared SM90 heuristic maximizes MMA throughput and typically picks block_n up to 192, +// which would force the serial fallback; re-select the fastest candidate with block_n <= 64 so +// MXFP8 gets the deep pipeline. Other dtypes keep using the unconstrained heuristic. +static constexpr int kMXFP8PipelineMaxBlockN = 64; static GemmConfig get_mxfp8_fp8_best_config(const GemmDesc& desc) { auto config = get_best_config(desc); @@ -227,7 +228,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v11", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v12", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -313,7 +314,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v11", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v12", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index c1c2c39dbc..a7946d906b 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -261,15 +261,26 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, while (scheduler.get_next_block(m_block_idx, n_block_idx)) { 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"); - // MXFP8 promotes each 32-wide K chunk with its own scale, so unlike FP8 we cannot - // hardware-accumulate across chunks: each WGMMA must be drained before its FP32 - // promotion. When the accumulator is small enough to double-buffer without spilling - // (BLOCK_N <= 96 => kNumAccum <= 48) and every math warp stores (BLOCK_M >= WGMMA::M, - // so `do_wgmma_store` is warpgroup-uniform and the WGMMA stays collective), run a - // 2-deep software pipeline that overlaps chunk kk's WGMMA with chunk kk-1's promotion. - // Otherwise fall back to the serial single-buffer path to stay under the 232-register - // budget (a second full-width accumulator at BLOCK_N=128 spills to local memory). - constexpr bool kPipelinePromotion = WGMMA::kNumAccum <= 48 and BLOCK_M >= WGMMA::M; + // MXFP8 promotes each 32-wide K chunk with its own UE8M0 scale, so unlike FP8 it + // cannot hardware-accumulate the 4 chunks of a 128-K block into one accumulator: + // every WGMMA must be drained before its FP32 promotion. To hide that drain we run a + // software pipeline whose depth is bounded by the FP32 register budget. We reserve + // ~168 of the 232 math registers for accumulators (the rest cover scales, descriptors + // and indices) and pick the deepest form that fits at compile time: + // * kFullPipeline - one independent accumulator per chunk; all WGMMAs are issued + // back-to-back and drained once (mirrors FP8's full pipeline); + // * kTwoDeepPipeline - 2-deep ping-pong overlapping chunk kk's WGMMA with kk-1's + // promotion (used when only two accumulators fit); + // * otherwise - serial single-buffer fallback (no spill). + // All pipelined paths require BLOCK_M >= WGMMA::M so `do_wgmma_store` is warpgroup- + // uniform and the collective WGMMA is always issued. + constexpr uint32_t kNumKChunks = BLOCK_K / WGMMA::K; + constexpr uint32_t kPromoteWaves = BLOCK_M / WAVE_BLOCK_M; + constexpr uint32_t kAccumRegBudget = 168; + constexpr bool kFullPipeline = BLOCK_M >= WGMMA::M and + (kNumKChunks + kPromoteWaves) * WGMMA::kNumAccum <= kAccumRegBudget; + constexpr bool kTwoDeepPipeline = BLOCK_M >= WGMMA::M and not kFullPipeline and + (2 + kPromoteWaves) * WGMMA::kNumAccum <= kAccumRegBudget; float final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for BLOCK_M < 64"); @@ -339,56 +350,81 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, } }; - if constexpr (kPipelinePromotion) { - // 2-deep software pipeline: overlap chunk (kk+1)'s WGMMA with chunk kk's - // FP32 promotion using a ping-pong accumulator. Enabled only when the - // accumulator is small (kNumAccum <= 48, guaranteed by the MXFP8 BLOCK_N - // clamp) so both buffers fit under the 232-register budget. - constexpr uint32_t kNumChunks = BLOCK_K / WGMMA::K; - float accum_buf[2][WGMMA::kNumAccum]; + // Issue one 32-K chunk's collective WGMMA into accumulator `dst`. + auto issue_chunk = [&](float* dst, const uint32_t& kk) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(dst[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, dst, false); + ptx::warpgroup_commit_batch(); + }; - auto issue_chunk = [&](float* dst, const uint32_t& kk) { + if constexpr (kFullPipeline) { + // Full pipeline: one dedicated accumulator per chunk. Issue every + // WGMMA back-to-back as a single warpgroup batch and drain once + // (`warpgroup_wait<0>`), matching FP8's fully-loaded tensor pipeline, + // then promote each chunk with its own scale. Fits only when + // kNumKChunks accumulators (+ the final_accum wave) stay under the + // register budget (BLOCK_N <= 64 => kNumAccum 32). + float accum_buf[kNumKChunks][WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(dst[i]); - ptx::warpgroup_arrive(); + ptx::warpgroup_fence_operand(accum_buf[kk][i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) { a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, dst, false); - ptx::warpgroup_commit_batch(); - }; + WGMMA::wgmma(a_desc, b_desc, accum_buf[kk], false); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum_buf[kk][i]); + ptx::warpgroup_wait<0>(); + if (do_wgmma_store) { + #pragma unroll + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) + promote_chunk(accum_buf[kk], kk); + } + } else if constexpr (kTwoDeepPipeline) { + // 2-deep software pipeline: overlap chunk (kk+1)'s WGMMA with chunk kk's + // FP32 promotion using a ping-pong accumulator. Used when only two + // accumulators fit under the register budget (BLOCK_N == 96). + float accum_buf[2][WGMMA::kNumAccum]; // Prologue: launch the first chunk, then in each step launch the next // chunk before draining the current one. issue_chunk(accum_buf[0], 0); #pragma unroll - for (uint32_t kk = 0; kk < kNumChunks; ++ kk) { + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) { const uint32_t cur = kk & 1; - if (kk + 1 < kNumChunks) + if (kk + 1 < kNumKChunks) issue_chunk(accum_buf[(kk + 1) & 1], kk + 1); #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum_buf[cur][i]); // Non-last chunks keep the freshly issued (kk+1) group in flight and // only drain kk; the last chunk drains everything. - if (kk + 1 < kNumChunks) + if (kk + 1 < kNumKChunks) ptx::warpgroup_wait<1>(); else ptx::warpgroup_wait<0>(); - promote_chunk(accum_buf[cur], kk); + if (do_wgmma_store) + promote_chunk(accum_buf[cur], kk); } } else { float accum[WGMMA::kNumAccum]; #pragma unroll - for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + kk * WGMMA::K) / 16; - b_desc.reg32_[0] = b_desc_base_lo + kk * WGMMA::K / 16; - WGMMA::wgmma(a_desc, b_desc, accum, false); - ptx::warpgroup_commit_batch(); + for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) { + issue_chunk(accum, kk); #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); From 5a57e93d8c3ac9190e363e4fd6ea33243097a3a7 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 15 Jul 2026 11:40:29 +0800 Subject: [PATCH 28/30] perf: revert sm90 mxfp8 clamp to BLOCK_N=96 (2-deep pipeline) BLOCK_N=64 full-depth pipeline measured slower at large M (0.41x vs 0.47x, 85.3 vs 96.7 TFLOPS): the 64x64x32 WGMMA loses more tensor-core efficiency than the deeper pipeline recovers. Revert the launcher clamp to BLOCK_N=96, which auto-selects the 2-deep ping-pong path ((2+1)*48=144 <= 168 reg budget) and restores the known-best 0.47x. BLOCK_M=256 was considered to amortize the drain but rejected: SM90 grouped-contiguous selects a B group per M-block at 128-alignment only (get_mk_alignment_for_contiguous_layout=128), so BLOCK_M=256 would let a block straddle two groups and silently multiply half its rows by the wrong B -- and the large_m perf test only asserts speedup, not accuracy, so it would pass while producing wrong results. It also pushes kPromoteWaves=2, blowing the accumulator register budget. Bump JIT recipe version 12 -> 13. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index b330d9ebee..b8ac2da118 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 12; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 13; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -104,13 +104,14 @@ static void tune_mxfp8_fp8_smem_config(GemmConfig& config, const GemmDesc& desc) // The MXFP8 kernel promotes every 32-wide K chunk with a distinct UE8M0 scale, so it cannot // hardware-accumulate across the 4 chunks of a 128-K block the way FP8 does. To hide the // per-chunk WGMMA drain it runs a software pipeline over a per-chunk FP32 accumulator whose -// depth is bounded by the register budget. Clamping to BLOCK_N == 64 gives kNumAccum == 32, -// which is small enough to hold one accumulator per chunk (full, FP8-style pipeline) under the -// 256-thread / 232-register budget, and (for n == 1024) tiles evenly with no wasted columns. -// The shared SM90 heuristic maximizes MMA throughput and typically picks block_n up to 192, -// which would force the serial fallback; re-select the fastest candidate with block_n <= 64 so -// MXFP8 gets the deep pipeline. Other dtypes keep using the unconstrained heuristic. -static constexpr int kMXFP8PipelineMaxBlockN = 64; +// depth is bounded by the register budget. Clamping to BLOCK_N == 96 keeps kNumAccum == 48, +// small enough for a 2-deep ping-pong pipeline under the 256-thread / 232-register budget. +// (BLOCK_N == 64 was measured slower at large M: the 64x64x32 WGMMA loses more tensor-core +// efficiency than the deeper pipeline recovers.) The shared SM90 heuristic maximizes MMA +// throughput and typically picks block_n up to 192, which would force the serial fallback; +// re-select the fastest candidate with block_n <= 96 so MXFP8 gets the pipeline. Other dtypes +// keep using the unconstrained heuristic. +static constexpr int kMXFP8PipelineMaxBlockN = 96; static GemmConfig get_mxfp8_fp8_best_config(const GemmDesc& desc) { auto config = get_best_config(desc); @@ -228,7 +229,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v12", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v13", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -314,7 +315,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v12", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v13", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } From b3436277720dd690d1202d9d0bbe92f680e31644 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 15 Jul 2026 12:02:19 +0800 Subject: [PATCH 29/30] perf: combine sm90 mxfp8 ue8m0 scales in the integer exponent domain MXFP8 promotion is FP32-pipe bound: each 4-element group did two FP32 scale multiplies (scale_a * scale_b * src). Since UE8M0 scales are pure powers of two, 2^(Ea-127) * 2^(Eb-127) == the float whose biased exponent byte is Ea+Eb-127 (exact, mantissa stays zero). Combine the scales via one integer add + clamp + shift on the integer ALU, leaving only the `* src` FMA on the FP32 pipe -- directly attacking the promotion bottleneck without touching the WGMMA/pipeline structure or the accumulator register budget. Clamp the combined exponent to [0, 255] so overflow saturates to +inf and underflow flushes to +0 (matching IEEE multiply of the normal, finite scales amax-quantization emits) and so `<< 23` can never corrupt sign/mantissa bits with an out-of-range exponent -- no silent miscompute. Bit-exact with the old two-multiply path for the normal [1, 254] exponent range. Bump JIT recipe version 13 -> 14. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 6 +-- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 43 +++++++++++++------ 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index b8ac2da118..8b9569ca59 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -40,7 +40,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 13; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 14; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -229,7 +229,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v13", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v14", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -315,7 +315,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v13", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v14", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index a7946d906b..2476595863 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -28,6 +28,21 @@ CUTLASS_DEVICE float e8m0_to_float(uint8_t scale) { return __uint_as_float(static_cast(scale) << 23); } +// Multiply two UE8M0 scales in the integer exponent domain instead of via two FP32 +// multiplies. Both are pure powers of two, so e8m0_to_float(Ea) * e8m0_to_float(Eb) +// == 2^(Ea-127) * 2^(Eb-127) == 2^((Ea+Eb-127)-127); the product is exactly the float +// whose biased exponent byte is Ea+Eb-127 (no rounding, mantissa stays zero). This moves +// the scale combination off the FP32 pipe (which is the promotion bottleneck) onto the +// integer ALU, leaving only the `* src` FMA on the FP32 pipe. Clamp the combined exponent +// to [0, 255] so overflow saturates to +inf and underflow flushes to +0 (matching IEEE +// multiply of the finite, normal scales MXFP8 amax-quantization produces) and, critically, +// so the `<< 23` never corrupts the sign/mantissa bits with an out-of-range exponent. +CUTLASS_DEVICE float e8m0_mul_to_float(uint8_t scale_a, uint8_t scale_b) { + const int biased = static_cast(scale_a) + static_cast(scale_b) - 127; + const uint32_t clamped = static_cast(max(0, min(255, biased))); + return __uint_as_float(clamped << 23); +} + CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, uint32_t k_scale_idx, uint32_t stride_k, bool packed_int32) { @@ -333,20 +348,24 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, // FP32 promotion of one drained 32-K chunk `kk` from accumulator `src`. auto promote_chunk = [&](const float* src, const uint32_t& kk) { - const float scale_a_0 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfa_pack_0 >> (kk * 8)) & 0xff)); - const float scale_a_1 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfa_pack_1 >> (kk * 8)) & 0xff)); + const uint8_t scale_a_0 = static_cast((sfa_pack_0 >> (kk * 8)) & 0xff); + const uint8_t scale_a_1 = static_cast((sfa_pack_1 >> (kk * 8)) & 0xff); #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const float scale_b_0 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff)); - const float scale_b_1 = mxfp8_fp8_detail::e8m0_to_float( - static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff)); - shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * src[i * 4 + 0]; - shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * src[i * 4 + 1]; - shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * src[i * 4 + 2]; - shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * src[i * 4 + 3]; + const uint8_t scale_b_0 = static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff); + const uint8_t scale_b_1 = static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff); + // Combine the A/B UE8M0 scales in the integer exponent domain (one + // int add + clamp + shift each) so the FP32 pipe only sees the `* src` + // FMA below. Bit-exact with two IEEE FP32 scale multiplies for the + // normal exponent range [1, 254] that MXFP8 amax-quantization emits. + const float scale_00 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_0, scale_b_0); + const float scale_01 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_0, scale_b_1); + const float scale_10 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_1, scale_b_0); + const float scale_11 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_1, scale_b_1); + shifted_accum[i * 4 + 0] += scale_00 * src[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_01 * src[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_10 * src[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_11 * src[i * 4 + 3]; } }; From 632467b5267eaf8dbe78019fe4f287fff09c2181 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Sat, 18 Jul 2026 11:37:35 +0800 Subject: [PATCH 30/30] perf: pipeline sm90 mxfp8 fp32 scale staging and launch state Move UE8M0-to-FP32 expansion into the producer warp and stage A/B scales as FP32. Math warpgroups now hoist their two A-scale rows with LDS.128 and load B scales only at promotion, avoiding the register spill caused by preloading the full B tile. This mirrors the exact gran-32 H20 pipeline: promotion is LDS + FMUL/FFMA rather than packed-byte unpacking and integer exponent reconstruction in the WGMMA path. Split scale readiness from the TMA full barrier with one scale mbarrier per math warpgroup and stage. A/B TMA completion can now launch chunk-0 WGMMA while producer-side scale expansion is still in flight; promotion waits for its own scale barrier before reading staged scales. Account for FP32 scale storage and scale barriers in the host SMEM stage calculation. Add a thread-local, single-entry contiguous launch-state cache keyed by device, tensor storage addresses, shapes, strides, scale recipes, and compiled dimensions. Stable serving buffers reuse config selection, TMA maps, generated code, and the JIT runtime without changing output-buffer ownership. Any key change rebuilds the launch state. Bump the MXFP8 JIT recipe version to v15. --- .../impls/sm90_mxfp8_fp8_gemm_1d2d.hpp | 103 ++++++++++--- .../impls/sm90_mxfp8_fp8_gemm_1d2d.cuh | 135 ++++++++++-------- deep_gemm/include/deep_gemm/ptx/ld_st.cuh | 6 + 3 files changed, 168 insertions(+), 76 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp index 8b9569ca59..ecf82f61e4 100644 --- a/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_mxfp8_fp8_gemm_1d2d.hpp @@ -1,6 +1,10 @@ #pragma once #include +#include +#include +#include +#include #include @@ -40,7 +44,7 @@ class SM90MXFP8FP8Gemm1D2DRuntime final: public LaunchRuntime using namespace deep_gemm; -static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 14; +static constexpr int kSm90MXFP8FP8ScaleRecipeJitVersion = 15; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm90_mxfp8_fp8_gemm_1d2d_impl< @@ -82,19 +86,62 @@ static void __instantiate_kernel() {{ } }; +struct SM90MXFP8FP8ContiguousLaunchKey { + int device, num_groups, m, n, k; + const void *a, *sfa, *b, *sfb, *d, *grouped_layout; + int64_t a_stride_m, b_stride_group, b_stride_n, d_stride_m; + int64_t sfa_stride_m, sfa_stride_k, sfb_stride_group, sfb_stride_n, sfb_stride_k; + uint32_t sfa_gran_k, sfb_gran_k; + bool sfa_packed_int32, sfb_packed_int32; + std::string compiled_dims; + + bool operator==(const SM90MXFP8FP8ContiguousLaunchKey& other) const { + return device == other.device and num_groups == other.num_groups and + m == other.m and n == other.n and k == other.k and + a == other.a and sfa == other.sfa and b == other.b and sfb == other.sfb and + d == other.d and grouped_layout == other.grouped_layout and + a_stride_m == other.a_stride_m and b_stride_group == other.b_stride_group and + b_stride_n == other.b_stride_n and d_stride_m == other.d_stride_m and + sfa_stride_m == other.sfa_stride_m and sfa_stride_k == other.sfa_stride_k and + sfb_stride_group == other.sfb_stride_group and sfb_stride_n == other.sfb_stride_n and + sfb_stride_k == other.sfb_stride_k and + sfa_gran_k == other.sfa_gran_k and sfb_gran_k == other.sfb_gran_k and + sfa_packed_int32 == other.sfa_packed_int32 and + sfb_packed_int32 == other.sfb_packed_int32 and compiled_dims == other.compiled_dims; + } +}; + +struct SM90MXFP8FP8ContiguousLaunchState { + SM90MXFP8FP8ContiguousLaunchKey key; + SM90MXFP8FP8Gemm1D2DRuntime::Args args; + std::shared_ptr runtime; +}; + +// The normal serving path repeatedly launches the same GEMM over stable buffers. Cache the +// complete host launch state for that one hot call site: config selection, TMA map encoding, +// generated code and JIT runtime are all invariant while this key matches. The cache is +// thread-local so independent Python execution threads never race on tensor lifetimes. +static thread_local std::optional sm90_mxfp8_fp8_contiguous_launch_state; + static void tune_mxfp8_fp8_smem_config(GemmConfig& config, const GemmDesc& desc) { const int orig_num_stages = config.pipeline_config.num_stages; const int original_per_stage = config.storage_config.load_block_m * config.layout.block_k * c10::elementSize(desc.a_dtype) + config.storage_config.load_block_n * config.layout.block_k * c10::elementSize(desc.b_dtype) + align(config.layout.block_m * static_cast(sizeof(float)), 128); - const int sfa_per_stage = align(config.layout.block_m * (config.layout.block_k / 32) * static_cast(sizeof(uint8_t)), 128); - const int sfb_per_stage = align(config.layout.block_n * (config.layout.block_k / 32) * static_cast(sizeof(uint8_t)), 128); + // The producer expands the four UE8M0 scales for every A/B row into FP32 + // before publishing a stage, so the math warpgroup only loads scales and + // performs the exact scale product during promotion. + const int sfa_per_stage = align(config.layout.block_m * (config.layout.block_k / 32) * static_cast(sizeof(float)), 128); + const int sfb_per_stage = align(config.layout.block_n * (config.layout.block_k / 32) * static_cast(sizeof(float)), 128); const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage; + // One 64-bit scale-ready mbarrier per math warpgroup and stage. The existing + // heuristic reserves the full/empty TMA barriers separately in smem_extra. + const int scale_barriers_per_stage = (config.launch_config.num_math_threads / 128) * static_cast(sizeof(uint64_t)); const int merged_per_stage = config.storage_config.load_block_m * config.layout.block_k * c10::elementSize(desc.a_dtype) + config.storage_config.load_block_n * config.layout.block_k * c10::elementSize(desc.b_dtype) + - sfa_per_stage + sfb_per_stage; + sfa_per_stage + sfb_per_stage + scale_barriers_per_stage; int chosen_stages = std::min(orig_num_stages, (SM90ArchSpec::smem_capacity - smem_extra) / merged_per_stage); DG_HOST_ASSERT(chosen_stages >= 1); config.pipeline_config.num_stages = chosen_stages; @@ -163,6 +210,35 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and grouped_layout.is_contiguous()); DG_HOST_ASSERT(a.is_contiguous() and b.is_contiguous() and d.is_contiguous()); + const auto sfa_gran_k = recipe_a.has_value() + ? std::get<1>(recipe_a.value()) + : k / (static_cast(sfa.size(1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); + DG_HOST_ASSERT(sfa.size(1) == ceil_div(k, sfa_gran_k * (sfa.scalar_type() == torch::kInt ? 4 : 1))); + const auto sfb_gran_k = recipe_b.has_value() + ? std::get<1>(recipe_b.value()) + : k / (static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); + DG_HOST_ASSERT(sfb.size(-1) == ceil_div(k, sfb_gran_k * (sfb.scalar_type() == torch::kInt ? 4 : 1))); + + const SM90MXFP8FP8ContiguousLaunchKey launch_key = { + .device = a.get_device(), .num_groups = num_groups, .m = m, .n = n, .k = k, + .a = a.data_ptr(), .sfa = sfa.data_ptr(), .b = b.data_ptr(), .sfb = sfb.data_ptr(), + .d = d.data_ptr(), .grouped_layout = grouped_layout.data_ptr(), + .a_stride_m = a.stride(0), .b_stride_group = b.stride(0), .b_stride_n = b.stride(1), + .d_stride_m = d.stride(-2), .sfa_stride_m = sfa.stride(0), .sfa_stride_k = sfa.stride(1), + .sfb_stride_group = sfb.stride(0), .sfb_stride_n = sfb.stride(1), .sfb_stride_k = sfb.stride(2), + .sfa_gran_k = static_cast(sfa_gran_k), .sfb_gran_k = static_cast(sfb_gran_k), + .sfa_packed_int32 = sfa.scalar_type() == torch::kInt, .sfb_packed_int32 = sfb.scalar_type() == torch::kInt, + .compiled_dims = compiled_dims, + }; + if (sm90_mxfp8_fp8_contiguous_launch_state.has_value() and + sm90_mxfp8_fp8_contiguous_launch_state->key == launch_key) { + const auto& cached = sm90_mxfp8_fp8_contiguous_launch_state.value(); + SM90MXFP8FP8Gemm1D2DRuntime::launch(cached.runtime, cached.args); + return; + } + const auto desc = GemmDesc { .gemm_type = GemmType::MGroupedContiguous, .kernel_type = KernelType::Kernel1D2D, @@ -179,16 +255,6 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( tune_mxfp8_fp8_smem_config(config, desc); DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); - const auto sfa_gran_k = recipe_a.has_value() - ? std::get<1>(recipe_a.value()) - : k / (static_cast(sfa.size(1)) * (sfa.scalar_type() == torch::kInt ? 4 : 1)); - DG_HOST_ASSERT(sfa_gran_k == 32 or sfa_gran_k == 128); - DG_HOST_ASSERT(sfa.size(1) == ceil_div(k, sfa_gran_k * (sfa.scalar_type() == torch::kInt ? 4 : 1))); - const auto sfb_gran_k = recipe_b.has_value() - ? std::get<1>(recipe_b.value()) - : k / (static_cast(sfb.size(-1)) * (sfb.scalar_type() == torch::kInt ? 4 : 1)); - DG_HOST_ASSERT(sfb_gran_k == 32 or sfb_gran_k == 128); - DG_HOST_ASSERT(sfb.size(-1) == ceil_div(k, sfb_gran_k * (sfb.scalar_type() == torch::kInt ? 4 : 1))); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, config.storage_config.load_block_m, @@ -229,7 +295,12 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v14", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_contiguous_1d2d_scale_recipe_v15", code); + sm90_mxfp8_fp8_contiguous_launch_state = SM90MXFP8FP8ContiguousLaunchState { + .key = launch_key, + .args = args, + .runtime = runtime, + }; SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } @@ -315,7 +386,7 @@ static void sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d( .tensor_map_d = tensor_map_d, }; const auto code = SM90MXFP8FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v14", code); + const auto runtime = compiler->build("sm90_m_grouped_mxfp8_fp8_gemm_masked_1d2d_scale_recipe_v15", code); SM90MXFP8FP8Gemm1D2DRuntime::launch(runtime, args); } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh index 2476595863..0e381a2876 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mxfp8_fp8_gemm_1d2d.cuh @@ -28,21 +28,6 @@ CUTLASS_DEVICE float e8m0_to_float(uint8_t scale) { return __uint_as_float(static_cast(scale) << 23); } -// Multiply two UE8M0 scales in the integer exponent domain instead of via two FP32 -// multiplies. Both are pure powers of two, so e8m0_to_float(Ea) * e8m0_to_float(Eb) -// == 2^(Ea-127) * 2^(Eb-127) == 2^((Ea+Eb-127)-127); the product is exactly the float -// whose biased exponent byte is Ea+Eb-127 (no rounding, mantissa stays zero). This moves -// the scale combination off the FP32 pipe (which is the promotion bottleneck) onto the -// integer ALU, leaving only the `* src` FMA on the FP32 pipe. Clamp the combined exponent -// to [0, 255] so overflow saturates to +inf and underflow flushes to +0 (matching IEEE -// multiply of the finite, normal scales MXFP8 amax-quantization produces) and, critically, -// so the `<< 23` never corrupts the sign/mantissa bits with an out-of-range exponent. -CUTLASS_DEVICE float e8m0_mul_to_float(uint8_t scale_a, uint8_t scale_b) { - const int biased = static_cast(scale_a) + static_cast(scale_b) - 127; - const uint32_t clamped = static_cast(max(0, min(255, biased))); - return __uint_as_float(clamped << 23); -} - CUTLASS_DEVICE uint8_t load_e8m0_scale(const void* ptr, uint32_t base_offset, uint32_t k_scale_idx, uint32_t stride_k, bool packed_int32) { @@ -117,6 +102,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, using WGMMA = typename mma::sm90::FP8MMASelector::type; using Barrier = cutlass::arch::ClusterTransactionBarrier; + static constexpr uint32_t kNumMathWarpgroups = kNumMathThreads / 128; shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; @@ -126,10 +112,12 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, 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 SHAPE_K_SFA_PER_STAGE = BLOCK_K / 32; - static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * SHAPE_K_SFA_PER_STAGE * sizeof(uint8_t); + // Materialize UE8M0 as FP32 in the producer warp. This shifts exponent unpacking out + // of the math warpgroups' promotion loop, where instruction issue is the bottleneck. + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * SHAPE_K_SFA_PER_STAGE * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); static constexpr uint32_t SHAPE_K_SFB_PER_STAGE = BLOCK_K / 32; - static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * SHAPE_K_SFB_PER_STAGE * sizeof(uint8_t); + static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * SHAPE_K_SFB_PER_STAGE * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFB_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFB_SIZE_PER_STAGE, 128u); const uint32_t num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K); @@ -153,23 +141,29 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, }); constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); }); auto smem_sfb = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE + - i * ALIGNED_SMEM_SFB_SIZE_PER_STAGE); + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE + + i * ALIGNED_SMEM_SFB_SIZE_PER_STAGE); }); auto barrier_start_ptr = reinterpret_cast( smem_buffer + SMEM_SF_OFFSET + kNumStages * (ALIGNED_SMEM_SFA_SIZE_PER_STAGE + ALIGNED_SMEM_SFB_SIZE_PER_STAGE)); auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); + auto scale_full_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return barrier_start_ptr + kNumStages * 2 + i; + }); if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { full_barriers[i]->init(1); empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + #pragma unroll + for (uint32_t wg = 0; wg < kNumMathWarpgroups; ++ wg) + scale_full_barriers[i * kNumMathWarpgroups + wg]->init(1); } cutlass::arch::fence_barrier_init(); } @@ -215,6 +209,9 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, tma::copy(&tensor_map_b, &full_barrier, smem_b[stage_idx], k_idx, scheduler.get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), num_tma_multicast_b); + // Publish A/B readiness immediately. Scale staging is guarded by + // scale_full_barriers below, allowing chunk-0 WGMMA to overlap it. + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); } const uint32_t sfa_base_m = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); @@ -229,8 +226,12 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, sfa, sfa_base_offset, k_block_idx, sfa_stride_k, sfa_gran_k, sfa_packed_int32, shape_k) : 0x7f7f7f7f; - *reinterpret_cast( - smem_sfa[stage_idx] + m_offset * SHAPE_K_SFA_PER_STAGE) = scales; + ptx::st_shared(reinterpret_cast( + smem_sfa[stage_idx] + m_offset * SHAPE_K_SFA_PER_STAGE), make_float4( + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 8)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 16)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 24)))); } const uint32_t sfb_group_idx = kMasked ? @@ -246,14 +247,21 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, sfb, sfb_base_offset, k_block_idx, sfb_stride_k, sfb_gran_k, sfb_packed_int32, shape_k) : 0x7f7f7f7f; - *reinterpret_cast( - smem_sfb[stage_idx] + n_offset * SHAPE_K_SFB_PER_STAGE) = scales; + ptx::st_shared(reinterpret_cast( + smem_sfb[stage_idx] + n_offset * SHAPE_K_SFB_PER_STAGE), make_float4( + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 8)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 16)), + mxfp8_fp8_detail::e8m0_to_float(static_cast(scales >> 24)))); } __threadfence_block(); __syncwarp(); - if (is_producer_leader) - full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + if (is_producer_leader) { + #pragma unroll + for (uint32_t wg = 0; wg < kNumMathWarpgroups; ++ wg) + scale_full_barriers[stage_idx * kNumMathWarpgroups + wg]->arrive(); + } } } @@ -321,47 +329,37 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { auto m_offset = local_idx * WAVE_BLOCK_M; - // Pack-load all SFA/SFB bytes for this wave_block from SMEM: - // SFA row stride is SHAPE_K_SFA_PER_STAGE (== BLOCK_K/32 == 4) bytes, - // so 4 SFA bytes (one per kk) are loaded with a single 32-bit LDS. - // SFB has the same layout; two adjacent N rows (n_base, n_base+1) form - // 8 contiguous bytes that we fetch with a single ld.shared.v2.u32. - uint32_t sfa_pack_0 = 0, sfa_pack_1 = 0; - uint32_t sfb_pack[WGMMA::kNumAccum / 4][2]; - if (do_wgmma_store) { - sfa_pack_0 = ptx::ld_shared(reinterpret_cast( - smem_sfa[stage_idx] + (r_0 + m_offset) * SHAPE_K_SFA_PER_STAGE)); - sfa_pack_1 = ptx::ld_shared(reinterpret_cast( - smem_sfa[stage_idx] + (r_1 + m_offset) * SHAPE_K_SFA_PER_STAGE)); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint32_t n_scale_offset = i * 8 + (lane_idx % 4) * 2; - asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];\n" - : "=r"(sfb_pack[i][0]), "=r"(sfb_pack[i][1]) - : "l"(__cvta_generic_to_shared( - smem_sfb[stage_idx] + - n_scale_offset * SHAPE_K_SFB_PER_STAGE))); + // The producer materializes the four per-32-K scales as FP32. Hoist the + // two A rows with LDS.128, but leave B scales in SMEM until promotion: + // preloading every B scale would consume ~96 FP32 registers at BLOCK_N=96 + // and spill the 2-deep accumulator pipeline. + float4 sfa_pack_0 = {}, sfa_pack_1 = {}; + auto load_scales = [&]() { + if (do_wgmma_store) { + sfa_pack_0 = ptx::ld_shared(reinterpret_cast( + smem_sfa[stage_idx] + (r_0 + m_offset) * SHAPE_K_SFA_PER_STAGE)); + sfa_pack_1 = ptx::ld_shared(reinterpret_cast( + smem_sfa[stage_idx] + (r_1 + m_offset) * SHAPE_K_SFA_PER_STAGE)); } - } + }; auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; // FP32 promotion of one drained 32-K chunk `kk` from accumulator `src`. auto promote_chunk = [&](const float* src, const uint32_t& kk) { - const uint8_t scale_a_0 = static_cast((sfa_pack_0 >> (kk * 8)) & 0xff); - const uint8_t scale_a_1 = static_cast((sfa_pack_1 >> (kk * 8)) & 0xff); + const float* scale_a_0 = reinterpret_cast(&sfa_pack_0); + const float* scale_a_1 = reinterpret_cast(&sfa_pack_1); #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint8_t scale_b_0 = static_cast((sfb_pack[i][0] >> (kk * 8)) & 0xff); - const uint8_t scale_b_1 = static_cast((sfb_pack[i][1] >> (kk * 8)) & 0xff); - // Combine the A/B UE8M0 scales in the integer exponent domain (one - // int add + clamp + shift each) so the FP32 pipe only sees the `* src` - // FMA below. Bit-exact with two IEEE FP32 scale multiplies for the - // normal exponent range [1, 254] that MXFP8 amax-quantization emits. - const float scale_00 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_0, scale_b_0); - const float scale_01 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_0, scale_b_1); - const float scale_10 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_1, scale_b_0); - const float scale_11 = mxfp8_fp8_detail::e8m0_mul_to_float(scale_a_1, scale_b_1); + const uint32_t n_scale_offset = i * 8 + (lane_idx % 4) * 2; + const float scale_b_0 = ptx::ld_shared( + smem_sfb[stage_idx] + n_scale_offset * SHAPE_K_SFB_PER_STAGE + kk); + const float scale_b_1 = ptx::ld_shared( + smem_sfb[stage_idx] + (n_scale_offset + 1) * SHAPE_K_SFB_PER_STAGE + kk); + const float scale_00 = scale_a_0[kk] * scale_b_0; + const float scale_01 = scale_a_0[kk] * scale_b_1; + const float scale_10 = scale_a_1[kk] * scale_b_0; + const float scale_11 = scale_a_1[kk] * scale_b_1; shifted_accum[i * 4 + 0] += scale_00 * src[i * 4 + 0]; shifted_accum[i * 4 + 1] += scale_01 * src[i * 4 + 1]; shifted_accum[i * 4 + 2] += scale_10 * src[i * 4 + 2]; @@ -381,6 +379,18 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, ptx::warpgroup_commit_batch(); }; + // TMA data is ready after full_barrier, but FP32 scale staging is + // independent. Defer this wait until after chunk-0 has been issued so + // WGMMA can cover producer-side scale expansion. + bool scales_loaded = false; + auto ensure_scales_loaded = [&]() { + if (do_wgmma_store and not scales_loaded) { + scale_full_barriers[stage_idx * kNumMathWarpgroups + math_wg_idx]->wait(phase); + load_scales(); + scales_loaded = true; + } + }; + if constexpr (kFullPipeline) { // Full pipeline: one dedicated accumulator per chunk. Issue every // WGMMA back-to-back as a single warpgroup batch and drain once @@ -402,6 +412,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, WGMMA::wgmma(a_desc, b_desc, accum_buf[kk], false); } ptx::warpgroup_commit_batch(); + ensure_scales_loaded(); #pragma unroll for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) #pragma unroll @@ -420,8 +431,11 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, float accum_buf[2][WGMMA::kNumAccum]; // Prologue: launch the first chunk, then in each step launch the next - // chunk before draining the current one. + // chunk before draining the current one. Defer scale loading until + // chunk-0 is in flight; issuing chunks 0+1 before the scale wait was + // measured to increase scoreboard pressure. issue_chunk(accum_buf[0], 0); + ensure_scales_loaded(); #pragma unroll for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) { const uint32_t cur = kk & 1; @@ -444,6 +458,7 @@ sm90_mxfp8_fp8_gemm_1d2d_impl(void* sfa, void* sfb, int* grouped_layout, #pragma unroll for (uint32_t kk = 0; kk < kNumKChunks; ++ kk) { issue_chunk(accum, kk); + ensure_scales_loaded(); #pragma unroll for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) ptx::warpgroup_fence_operand(accum[i]); diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index c3e03bec73..e11429c661 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -122,6 +122,12 @@ CUTLASS_DEVICE void st_shared(const float2* ptr, float2 val) { asm volatile("st.shared.v2.f32 [%0], {%1, %2};" :: "l"(__cvta_generic_to_shared(ptr)), "f"(val.x), "f"(val.y)); } +CUTLASS_DEVICE void st_shared(const float4* ptr, float4 val) { + asm volatile("st.shared.v4.f32 [%0], {%1, %2, %3, %4};" :: + "l"(__cvta_generic_to_shared(ptr)), + "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)); +} + CUTLASS_DEVICE void st_shared(const uint32_t* ptr, uint32_t val) { asm volatile("st.shared.u32 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "r"(val)); }