From bd42ee2f0e5a9b4cdd530a49c11fde16fced59a3 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 16 Jun 2026 10:40:17 +0800 Subject: [PATCH 01/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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