diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42622df7d8..d1b18da5f6 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -5,6 +5,7 @@ #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_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" @@ -140,6 +141,27 @@ static void fp8_fp4_gemm_tt(const std::pair& a, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); } +static void fp8_fp4_gemm_nt_sm90_fused_wgmma(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional& c, + const int& gran_k, + const std::string& compiled_dims) { + const auto [m, k] = get_shape<2>(a.first); + const auto [n, half_k] = get_shape<2>(b.first); + DG_HOST_ASSERT(half_k * 2 == k); + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k); + std::optional> recipe_b = std::make_tuple(1, gran_k); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, false); + DG_HOST_ASSERT(gran_k_a == gran_k and gran_k_b == gran_k); + + sm90_fp8_fp4_gemm_1d1d_fused( + {a.first, sfa}, {b.first, sfb}, d, c, gran_k, compiled_dims); +} + static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -624,6 +646,37 @@ static void register_apis(pybind11::module_& m) { py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "mn", py::arg("disable_ue8m0_cast") = false); + m.def("fp8_fp4_gemm_nt_sm90_fused_wgmma", &fp8_fp4_gemm_nt_sm90_fused_wgmma, + py::arg("a"), py::arg("b"), py::arg("d"), + py::arg("c") = std::nullopt, + py::arg("gran_k") = 128, + py::arg("compiled_dims") = "nk"); + m.def("m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma", + &sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), + py::arg("gran_k") = 128, + py::arg("compiled_dims") = "nk", + py::arg("use_psum_layout") = false, + py::arg("expected_m_for_psum_layout") = std::nullopt, + py::arg("block_m_override") = std::nullopt, + py::arg("block_n_override") = std::nullopt, + py::arg("decode_stub") = false); + m.def("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", + &sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"), + py::arg("expected_m"), + py::arg("gran_k") = 128, + py::arg("gran_k_a") = std::nullopt, + py::arg("gran_k_b") = std::nullopt, + py::arg("compiled_dims") = "nk", + py::arg("block_m_override") = std::nullopt, + py::arg("block_n_override") = std::nullopt, + py::arg("decode_stub") = false, + py::arg("b_is_int4_sym") = false, + py::arg("masked_m_max_hint") = std::nullopt, + py::arg("active_groups_hint") = std::nullopt); + m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") = + m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma"); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index b404241a66..a28468cf92 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -36,8 +36,23 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, // Pre-transform checks check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups); - // (FP32, 1, 128) on SM90: transform to TMA-aligned and MN-major - if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + // (BF16, 1, 32/128) on SM90 SFB fast path: path-A (k128) 与 path-B fast-path (k32) + // 都支持 bf16 sfb，体积砍半。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,8) + // 构造，这里跳过 fp32 only 的 transpose，直接复用 align 路径。 + if (sf.scalar_type() == torch::kBFloat16 and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9) + return get_mn_major_tma_aligned_tensor(sf); + // (UInt8/E8M0, 1, 32) on SM90 SFB fast path：path-B 专用，每元素 1B 即 fp32 的 + // 8 位指数。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,16) 构造。 + if (sf.scalar_type() == torch::kUInt8 and gran_mn == 1 and gran_k == 32 and arch_major == 9) + return get_mn_major_tma_aligned_tensor(sf); + + // (FP32, 1, 32/128) on SM90: transform to TMA-aligned and MN-major + if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and + (arch_major == 9 or disable_ue8m0_cast)) + return get_mn_major_tma_aligned_tensor(sf); + + // (INT packed UE8M0, 1, 32/128) on SM90: transform to TMA-aligned and MN-major. + if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9) return get_mn_major_tma_aligned_tensor(sf); // (FP32, 128, 128) on SM90: no need to transform, check SFB requirements diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp new file mode 100644 index 0000000000..febe4bbb26 --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -0,0 +1,1410 @@ +#pragma once + +#include + +#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/layout.hpp" +#include "epilogue.hpp" +#include "runtime_utils.hpp" +#include "sm90_fp8_fp4_gemm_1d2d_rs.hpp" + +namespace deep_gemm { + +class SM90FP8FP4Gemm1D1DRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + void *gmem_b_ptr; + void *gmem_d_ptr; + void *grouped_layout; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_sfa; + CUtensorMap tensor_map_sfb; + CUtensorMap tensor_map_cd; + }; + + 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_fp8_fp4_gemm_1d1d_impl< + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, + {}, {} + >); +}}; +)", + 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), + to_string(args.gemm_desc.cd_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + nullptr, args.gmem_b_ptr, + args.gmem_d_ptr, args.grouped_layout, + nullptr, + args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.tensor_map_a, + args.tensor_map_sfa, args.tensor_map_sfb, + args.tensor_map_cd)); + } +}; + +class SM90FP8FP4Gemm1D2DRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + cute::UMMA::Major major_sfb; + bool decode_stub; + void *gmem_b_ptr; + void *gmem_d_ptr; + void *sfb; + void *grouped_layout; + 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_fp8_fp4_gemm_1d2d_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {} + >); +}}; +)", + to_string(args.major_sfb), + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', 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), + get_default_epilogue_type(std::nullopt), + args.decode_stub ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.gmem_b_ptr, args.sfb, args.grouped_layout, args.gmem_d_ptr, + 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_fp8_fp4_gemm_1d1d_fused(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional& c, + const int& gran_k, + const std::string& compiled_dims) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(gran_k == 128); + DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat or b.second.scalar_type() == torch::kInt); + DG_HOST_ASSERT(a.first.is_contiguous()); + DG_HOST_ASSERT(b.first.is_contiguous()); + DG_HOST_ASSERT(d.is_contiguous() and c->is_contiguous()); + + const auto [m, k] = get_shape<2>(a.first); + const auto [n, half_k] = get_shape<2>(b.first); + const auto [m_d, n_d] = get_shape<2>(d); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(a.second.size(0) == m and a.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(b.second.size(0) == n and b.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(c->sizes() == d.sizes()); + + if (m == 0 or n == 0) { + return; + } + if (c->data_ptr() != d.data_ptr()) { + d.copy_(*c); + } + + auto desc = GemmDesc { + .gemm_type = GemmType::Normal, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = k, .num_groups = 1, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = c.has_value(), + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims + }; + auto config = get_best_config(desc); + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, a.second, m, k, + config.layout.block_m, config.layout.block_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, b.second, n, k, + config.layout.block_n, config.layout.block_k, 1, 0); + const auto tensor_map_cd = 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 SM90FP8FP4Gemm1D1DRuntime::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()), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .grouped_layout = nullptr, + .tensor_map_a = tensor_map_a, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd, + }; + const auto code = SM90FP8FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm90_fp8_fp4_gemm_1d1d_fused", code); + SM90FP8FP4Gemm1D1DRuntime::launch(runtime, args); +} + +static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( + const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& grouped_layout, + const int& gran_k, + const std::string& compiled_dims, + const bool& use_psum_layout, + const std::optional& expected_m_for_psum_layout, + const std::optional& block_m_override, + const std::optional& block_n_override, + const bool& decode_stub) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(gran_k == 128); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); + 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.first.is_contiguous() and b.first.is_contiguous() and d.is_contiguous()); + + const auto [m, k] = get_shape<2>(a.first); + const auto [num_groups, n, half_k] = get_shape<3>(b.first); + const auto [m_d, n_d] = get_shape<2>(d); + const auto [layout_size] = get_shape<1>(grouped_layout); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(use_psum_layout ? (layout_size == num_groups) : (layout_size == m)); + if (expected_m_for_psum_layout) { + DG_HOST_ASSERT(use_psum_layout); + } + DG_HOST_ASSERT(a.second.size(0) == m and a.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(b.second.size(0) == num_groups and b.second.size(1) == n and b.second.size(2) == ceil_div(k, gran_k)); + + if (m == 0 or n == 0) { + return; + } + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k); + std::optional> recipe_b = std::make_tuple(1, gran_k); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, + std::nullopt, num_groups, false); + DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128); + + const auto gemm_type = use_psum_layout ? + GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous; + // NOTE: psum layout previously always took the 1d1d fallback path. With the + // 1d2d psum scheduler / SFB indexing aligned, we let psum also flow into the + // common 1d2d code below. This relies on: + // - All groups in `grouped_layout` having the same K (current_shape_k == shape_k); + // true for current call sites where K is shared across groups. + // - SFB physical layout `[num_groups, n, shape_k_scales]` already matches + // 1d2d cooperative ld.global indexing. + // - 0-size groups handled by the scheduler's while-loop fallthrough. + if (false /* use_psum_layout disabled: psum now goes through the 1d2d common path */) { + auto desc = GemmDesc { + .gemm_type = gemm_type, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .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 = expected_m_for_psum_layout.value_or(m), + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups + }; + 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.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_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 auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, + config.layout.block_n, config.layout.block_k, num_groups, 0); + const auto tensor_map_cd = 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 SM90FP8FP4Gemm1D1DRuntime::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()), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .grouped_layout = grouped_layout.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd, + }; + const auto code = SM90FP8FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused", code); + SM90FP8FP4Gemm1D1DRuntime::launch(runtime, args); + return; + } + + auto desc = GemmDesc { + .gemm_type = gemm_type, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .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 = expected_m_for_psum_layout.value_or(m), + .expected_n = n, + .expected_k = k, + .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 + }; + auto rebuild_config = [&](Layout layout) { + const auto storage_config = SM90ArchSpec::get_storage_config(desc, layout); + const auto pipeline_config = SM90ArchSpec::get_pipeline_config(desc, layout, storage_config); + DG_HOST_ASSERT(pipeline_config.num_stages >= 3); + const auto launch_config = SM90ArchSpec::get_launch_config(desc, layout); + return GemmConfig{layout, storage_config, pipeline_config, launch_config}; + }; + + const auto layout_candidates = SM90ArchSpec::get_layout_candidates(desc); + DG_HOST_ASSERT(not layout_candidates.empty()); + auto layout = layout_candidates[0]; + auto layout_info = SM90ArchSpec::get_layout_info(desc, layout); + bool found_layout = false; + // FP4 B is decoded by each CTA, so only A multicast (cluster_n) is enabled. + for (const auto& candidate: layout_candidates) { + if (candidate.cluster_m != 1) { + continue; + } + const auto candidate_info = SM90ArchSpec::get_layout_info(desc, candidate); + if (not found_layout or SM90ArchSpec::compare(candidate_info, layout_info)) { + layout = candidate; + layout_info = candidate_info; + found_layout = true; + } + } + DG_HOST_ASSERT(found_layout); + // Psum grouped_layout is encoded with 128-row alignment. Keep BLOCK_M fixed + // at 128 so scheduler psum boundaries match the physical layout. + if (use_psum_layout) { + layout.block_m = 128; + layout.cluster_m = 1; + } + if (not use_psum_layout and n >= 1024 and num_groups > 0 and m % num_groups == 0) { + const int expected_m_per_group = m / num_groups; + if (expected_m_per_group >= 256) { + layout.block_m = 256; + layout.block_n = 64; + } else if (expected_m_per_group == 128) { + // Tuned on the SM90 FP8xFP4 contiguous fallback benchmark at + // N=4096,K=7168. Earlier code carved out BN=128 for 16/24 groups, + // but this consistently underperformed BN=64 (16g:0.74x, 24g:0.78x + // vs the BN=64 baseline). Larger BN doubles the per-stage packed-B + // and SFB cache footprint, which forces `num_stages` down without + // recovering wave utilization, so always pick BN=64 here. + layout.block_m = 128; + layout.block_n = 64; + } + } + layout.cluster_m = 1; + auto config = rebuild_config(layout); + + // The contiguous fallback benchmark explicitly sweeps BLOCK_M/BLOCK_N via + // these overrides; keep this path independent from masked RS heuristics so + // regressions can be attributed to the selected block shape. + if (block_m_override or block_n_override) { + auto layout = config.layout; + if (block_m_override) { + DG_HOST_ASSERT(not use_psum_layout or *block_m_override == 128); + layout.block_m = *block_m_override; + } + if (block_n_override) { + layout.block_n = *block_n_override; + } + DG_HOST_ASSERT((layout.block_m == 64 or layout.block_m == 128 or layout.block_m == 256) and layout.block_n % 16 == 0); + DG_HOST_ASSERT(layout.block_n <= 256); + layout.cluster_m = 1; + config = rebuild_config(layout); + } + 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); + + // Re-derive `num_stages` and `smem_size` to account for: + // (1) the extra packed-FP4 B staging buffer (BLOCK_N * BLOCK_K / 2 bytes per stage) + // — after the A/packed-B barrier merge it is sized at `num_stages` slots + // (same as A/SFA), so it folds into the per-stage cost. + // (2) the enlarged SFB cache: now stores (shape_k_scales, BLOCK_N) floats per block + // so the K loop reads SFB from smem instead of gmem. SFB cache aliases + // smem_d (no separate allocation), so we only subtract the original SFB bytes. + { + const int block_k = config.layout.block_k; + const int block_n = config.layout.block_n; + const int shape_k_scales = ceil_div(static_cast(k), block_k); + const bool uniform_scale_b = (block_k % block_n == 0); + const int sfb_old_bytes = align(shape_k_scales * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + const int sfb_cache_bytes = align(shape_k_scales * block_n * static_cast(sizeof(float)), 16); + const int smem_d_bytes = align(config.layout.block_m * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; + + const int packed_per_stage = block_n * (block_k / 2); + const int smem_a_per_stage = config.storage_config.load_block_m * block_k * + static_cast(c10::elementSize(desc.a_dtype)); + const int smem_b_per_stage = config.storage_config.load_block_n * block_k * + static_cast(c10::elementSize(desc.b_dtype)); + const int smem_sfa_per_stage = align(config.layout.block_m * static_cast(sizeof(float)), 128); + const int original_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage; + const int merged_per_stage = original_per_stage + packed_per_stage; + const int orig_num_stages = config.pipeline_config.num_stages; + const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage + sfb_extra; + + auto fits = [&](int stages) { + return smem_extra + stages * merged_per_stage <= SM90ArchSpec::smem_capacity; + }; + + // Packed-FP4 B halves the per-stage B footprint, so the smem freed by + // the FP4 path can usually accommodate more pipeline stages than the + // FP8 baseline. Mirror the masked path's logic: try to push `num_stages` + // up to `kW4DefaultMaxStages` while it still fits, then fall back if + // the chosen value does not fit (e.g. due to large SFB cache at BN=128). + constexpr int kW4DefaultMaxStages = 8; + int chosen_stages = orig_num_stages; + while (chosen_stages + 1 <= kW4DefaultMaxStages and fits(chosen_stages + 1)) + ++ chosen_stages; + while (chosen_stages >= 3 and not fits(chosen_stages)) + -- chosen_stages; + DG_HOST_ASSERT(chosen_stages >= 3); + config.pipeline_config.num_stages = chosen_stages; + config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; + } + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_mode); + // View packed FP4 B as 1-byte FP8 so TMA loads raw packed bytes (no FP4 unpacking). + const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b_bytes, n, half_k, + config.layout.block_n, config.layout.block_k / 2, + half_k, num_groups, 0); + 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 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 SM90FP8FP4Gemm1D2DRuntime::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()), + .major_sfb = get_major_type_ab(sfb), + .decode_stub = decode_stub, + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .sfb = sfb.data_ptr(), + .grouped_layout = grouped_layout.data_ptr(), + .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 = SM90FP8FP4Gemm1D2DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_contiguous_1d2d_fused", code); + SM90FP8FP4Gemm1D2DRuntime::launch(runtime, args); +} + +static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( + const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& expected_m, + const int& gran_k, + const std::optional& gran_k_a_override, + const std::optional& gran_k_b_override, + const std::string& compiled_dims, + const std::optional& block_m_override, + const std::optional& block_n_override, + const bool& decode_stub, + // INT4-sym (signed [-8, 7]) variant for B. The wire format is shared + // with packed-FP4 (2 nibbles/byte, kPackedFP4 dtype, fp32 SFB), so + // the kernel reuses the same TMA descriptors and SFB layout. Only + // the in-register decode primitive switches to int4_symx4_to_e4m3x4. + const bool& b_is_int4_sym = false, + // DSV4 MTP/speculative-verify hint: caller passes masked_m.max() so + // the host can pick BM matching the hottest group instead of the + // distribution-average expected_m. Fast-path gating (k32 quad_reduce + // / direct_load / compact_sched) still keys on expected_m so existing + // small-M optimizations are preserved. Defaults to expected_m when + // unset. + const std::optional& masked_m_max_hint = std::nullopt, + // active_groups_hint: caller passes count of groups with masked_m > 0 + // (i.e. (masked_m != 0).sum()). Combined with masked_m_max_hint this + // is enough to estimate "工作量分布" and decide fast-path 是否合适： + // * 单热点 / 极少活跃 group → fast-path (BM=32 BN=128) 大胜 + // * 大量活跃 group + 高 max_m → fan-out (BM 阶梯, BN=256) 取胜 + // 不传时退化为旧行为（仅看 max_hint）。 + const std::optional& active_groups_hint = std::nullopt) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + const int gran_k_a_requested = gran_k_a_override.value_or(gran_k); + const int gran_k_b_requested = gran_k_b_override.value_or(gran_k); + DG_HOST_ASSERT(gran_k_a_requested == 128); + DG_HOST_ASSERT(gran_k_b_requested == 32 or gran_k_b_requested == 128); + // INT4-sym path-A is restricted to per-128 fp32 SFB on the device side. + // Reject combinations that would otherwise silently bypass the + // INT4-decode dispatch (e.g. per-32 K-block scales fall into the fused + // decode path, which is not wired for INT4 yet). + DG_HOST_ASSERT(not b_is_int4_sym or gran_k_b_requested == 128); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat or + b.second.scalar_type() == torch::kInt or + b.second.scalar_type() == torch::kBFloat16 or + b.second.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.first.is_contiguous() and b.first.is_contiguous() and d.is_contiguous()); + + const auto [num_groups, m, k] = get_shape<3>(a.first); + const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); + const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(num_groups == num_groups_b and num_groups == num_groups_d); + DG_HOST_ASSERT(masked_m.numel() == num_groups); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0); + DG_HOST_ASSERT(a.second.size(0) == num_groups and a.second.size(1) == m and + a.second.size(2) == ceil_div(k, gran_k_a_requested)); + const int gran_k_b_shape = b.second.scalar_type() == torch::kInt ? + gran_k_b_requested * 4 : gran_k_b_requested; + DG_HOST_ASSERT(b.second.size(0) == num_groups and b.second.size(1) == n and + b.second.size(2) == ceil_div(k, gran_k_b_shape)); + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k_a_requested); + std::optional> recipe_b = std::make_tuple(1, gran_k_b_requested); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, + num_groups, num_groups, false); + DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == gran_k_b_requested); + + auto desc = GemmDesc { + .gemm_type = GemmType::MGroupedMasked, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .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 = expected_m, + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups + }; + auto rebuild_config = [&](Layout layout) { + const auto storage_config = SM90ArchSpec::get_storage_config(desc, layout); + const auto pipeline_config = SM90ArchSpec::get_pipeline_config(desc, layout, storage_config); + DG_HOST_ASSERT(pipeline_config.num_stages >= 3); + const auto launch_config = SM90ArchSpec::get_launch_config(desc, layout); + return GemmConfig{layout, storage_config, pipeline_config, launch_config}; + }; + + const auto layout_candidates = SM90ArchSpec::get_layout_candidates(desc); + DG_HOST_ASSERT(not layout_candidates.empty()); + auto layout = layout_candidates[0]; + auto layout_info = SM90ArchSpec::get_layout_info(desc, layout); + bool found_layout = false; + for (const auto& candidate: layout_candidates) { + if (candidate.cluster_m != 1) { + continue; + } + const auto candidate_info = SM90ArchSpec::get_layout_info(desc, candidate); + if (not found_layout or SM90ArchSpec::compare(candidate_info, layout_info)) { + layout = candidate; + layout_info = candidate_info; + found_layout = true; + } + } + DG_HOST_ASSERT(found_layout); + const auto env_enabled = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr and value[0] != '\0' and value[0] != '0'; + }; + const auto env_disabled = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr and value[0] == '0'; + }; + const auto env_int = [](const char* name, int default_value) { + const char* value = std::getenv(name); + return (value != nullptr and value[0] != '\0') ? std::atoi(value) : default_value; + }; + // BM=64 fast-path 是否启用（函数作用域统一判据，供三处共用：layout 选择 / + // bm32_skew_layout(stages) / bm32_skew_fast_path(device 三件套总闸)）。 + // 两个触发源： + // (1) DG_W4_PATHB_BM64_FASTPATH=1：手动强制（对照实验用）。 + // (2) DG_W4_PATHB_BM64_AUTO=1（默认开）+ 高置信子集 max_hint>=128 且 + // active<=8：自动触发，抓“少数活跃组 + 大 max_m”（MTP verify spec-len + // 命中形态）。详见 layout 选择处 bm64_fastpath 的长注释。 + // (3) DG_W4_PATHB_BM64_FORCE_ALL=1：无条件全部走 BM=64 fast-path + // （用于在线业务峰值吞吐 A/B 对照，绕开 hint 判据；hint 在生产 cuda graph + // 下恒为 None 时，是验证“假设拿到 hint 全命中”的速度上界）。 + // 注意：必须同时让 stages / fast-path 总闸认账，否则 layout.block_m=64 却 + // 走通用慢路径会退化 ~2x（见 DG_W4_PATHB_BM64 非 fastpath 的坑）。 + const bool bm64_fastpath_enabled = + env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 or + env_int("DG_W4_PATHB_BM64_FASTPATH", 0) != 0 or + (env_int("DG_W4_PATHB_BM64_AUTO", 1) != 0 and + masked_m_max_hint.value_or(0) >= 128 and + active_groups_hint.value_or(static_cast(desc.num_groups)) <= 8); + // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages + // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 + // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 + // + // 参数空间：BM ∈ {8, 16, 32, 64, 128}（BM<64 用于 masked small-M）， + // BN ∈ {64, 128, 256}。 + // + // 选择目标：在 (waves <= ceil_div(total_tiles, num_sms)) 的前提下最大化 stages， + // 其次最大化 last_wave 利用率，最后倾向更小的 per-stage（即更小 BN）。 + if (not block_m_override and not block_n_override) { + const int num_sms = desc.num_sms; + const int block_k = layout.block_k; + const int shape_k_scales_b = ceil_div(static_cast(k), gran_k_b); + + auto eval_layout = [&](int bm, int bn) -> std::tuple { + // 返回 (sat_stages, -waves, last_wave_util, -per_stage)，越大越好。 + // stages 在 ~6 之上对 TMA 隐藏几乎饱和，因此用饱和 stage 数比较，避免 + // 小 BN 因 stages=8 击败 wave 利用率更高的候选。 + const int tiles = ceil_div(expected_m, bm) * ceil_div(static_cast(n), bn) * num_groups; + const int waves = ceil_div(tiles, num_sms); + const int last = tiles - (waves - 1) * num_sms; + const int last_util = last <= 0 ? num_sms : last; + const bool uniform_scale_b = (block_k % bn == 0); + const int sfb_old_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + const int sfb_cache_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * bn * static_cast(sizeof(float)), 16); + const int rs_padded_bm = std::max(bm, 64); + const int smem_d_bytes = + align(rs_padded_bm * bn * static_cast(sizeof(nv_bfloat16)), 1024); + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; + const int smem_a_per_stage = rs_padded_bm * block_k * static_cast(c10::elementSize(desc.a_dtype)); + const int smem_sfa_per_stage = + align(rs_padded_bm * static_cast(sizeof(float)), 128); + const int packed_per_stage = bn * (block_k / 2); + const int merged_per_stage = smem_a_per_stage + smem_sfa_per_stage + packed_per_stage; + constexpr int kMaxEvaluatedStages = 10; + constexpr int kBarrierBytes = 16 * kMaxEvaluatedStages * 2; + const int fixed = smem_d_bytes + kBarrierBytes + sfb_extra; + const int max_stages = (SM90ArchSpec::smem_capacity - fixed) / merged_per_stage; + constexpr int kStageSaturation = 6; + const int sat_stages = std::min(std::min(max_stages, kMaxEvaluatedStages), kStageSaturation); + return std::make_tuple(sat_stages, -waves, last_util, -merged_per_stage); + }; + + std::vector> w4_candidates = { + {64, 64}, {64, 128}, {64, 256}, + {128, 64}, {128, 128}, + }; + if (expected_m <= 32) { + w4_candidates.insert(w4_candidates.begin(), {{8, 64}, {16, 64}, {32, 64}}); + } + + std::pair best{layout.block_m, layout.block_n}; + std::tuple best_score{-1, 0, 0, 0}; + bool first = true; + for (const auto& cand : w4_candidates) { + const int bm = cand.first; + const int bn = cand.second; + // 1D2D 内核 unroll 要求 + if (bn > block_k and (bn % (bn - block_k) != 0 and block_k % (bn - block_k) != 0)) + continue; + // masked 路径 multicast 合法性：当前固定 cluster_m=1, cluster_n=1，恒满足 + const auto score = eval_layout(bm, bn); + if (std::get<0>(score) < 3) + continue; + if (first or score > best_score) { + best_score = score; + best = cand; + first = false; + } + } + layout.block_m = best.first; + layout.block_n = best.second; + + // RS masked W4 empirical fallback: + // pick BM close to expected_m to avoid over-computing promotion work, + // and use BN=256 for fewer CTAs while staying within Hopper TMA limits. + // bm_select_m 用 max(expected_m, masked_m_max_hint) —— 当调用方传入 + // hot group 的真实大小时，host 据此选 BM，避免 BM 与 hot group 严重 + // 失配；fast-path gating（k32 quad_reduce / direct_load 等）下面仍按 + // expected_m 判断，保留 small-m 路径的优化。 + const int bm_select_m = std::max( + expected_m, masked_m_max_hint.value_or(expected_m)); + if (desc.gemm_type == GemmType::MGroupedMasked) { + // 公共形状判定：DSV4 MTP/speculative verify 形状下 + // host scheduler 才会做 hard-code 干预，避开通用 shape。 + // A1+A2: 放开到 N∈{4096,7168}, K∈{2048,3072,4096,7168} 以覆盖 wide_n / wide_k 系列。 + // K=3072 引入动机：DSV4 gateup/down_proj 还有一组 n=7168+k=3072 的真实 + // shape，K=3072 满足 BLOCK_K=128 整除 (3072/128=24) + gran_k_b=32 整除 + // (3072/32=96) ⇒ device fast-path (quad_reduce/direct_load/compact_sched) + // 物理可命中。原守护把 k=3072 拒在外面，强制走 path-B 通用 cache_sfb_k32 + // 路径，W4 GB/s 只有 ~600（vs fast-path 命中时 941+ GB/s）。 + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： + // 把守护放宽到 num_groups>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}。 + // 动机：DSV4 down_proj 真实形状是 g24, n=6144, k=7168，原守护把它 + // 挡在 fast-path 外，强制走 path-B 通用 cache_sfb_k32 路径，W4 GB/s + // 只有 627（vs fast-path 命中时 1500-3000+ GB/s），speedup 0.43-0.56x。 + // device fast-path 三件套（quad_reduce / direct_load / compact_sched） + // 实际 BN-agnostic，仅硬约束 BM==32 + gran_k_b==32，冗余专家场景 + // kNumGroups<=36 也可按同一路径处理。 + // 实测：g24+n=6144+k=7168 整张表 0.43-0.56x → 0.71-0.81x； + // g24+n=7168+k=3072 整张表 0.30-0.54x → 0.45-0.79x； + // 默认 dsv4 形状 (g32+n∈{4096,7168}+k∈{2048,4096,7168}) 不变。 + // 开启后所有依赖 dsv4_shape 的路径（BN256 门槛 / small_hot + // / large_bm / cluster_n / bm64_hint）都同步放开，已验证无退化。 + // 设 DG_W4_PATHB_FAST_PATH_RELAX=0 可显式回退到原守护做对照。 + const bool fast_path_relax = + env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0; + const bool dsv4_shape = + (static_cast(desc.num_groups) >= 8 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) + or + (fast_path_relax and + static_cast(desc.num_groups) >= 8 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)); + + if (gran_k_b == 32) { + // path-B：device fast-path（BM=32 + quad_reduce + direct_load） + // 是 gran_k_b=32 下的唯一高效形状。BN 选取根据 hint 大小： + // * hint <= 32（小 hot）：BM=32 BN=128（短 wgmma 流水线 + 高 stage 数） + // * hint > 32（大 hot）：BM=32 BN=256（grid 减半，scheduler 调度密度提升） + // + // 历史教训（skewed masked benchmark 实测）： + // * 试图基于 (active, max_m) 把"path-B 输给 FP8"的 case 切到 + // BM 阶梯 + BN=256 fan-out → device kernel 落入 path-B 通用 + // cache_sfb_k32 路径，比 fast-path 慢 1.5–3x（如 eight_hot_64 + // 190→298us、dense_tail 368→1011us）。 + // * 结论：path-B 必须保持 BM=32 + fast-path 守护命中，否则崩盘。 + // BN=128↔256 两挡都在 device fast-path 守护内（参见 device + // kernel kK32QuadReduce 路径，BN-agnostic：compute_n_0/1 + // 寻址含 m_offset = local_idx * WAVE_BLOCK_M， + // load_k32_quad_scale_b 按 compute_n 单独寻址）。 + // + // hint<=16（纯 fan-out）保留原 BM 阶梯走 BN=256 的 path-B 通用 + // 路径——这部分 device 守护本来就不命中。 + // real_hot_present 判定：caller 是否真的传入 max_hint > 16？ + // * has_value=true：按真实 max_hint 判断（>16 才认为 hot 存在）。 + // * has_value=false（caller 没传 hint）： + // - expected_m > 16：保留旧行为（保守认为 hot 存在）， + // 走 BM=32 fast-path（m>=1024 时命中 fast-path 守护）。 + // - expected_m <= 16：**新行为**，认为 hot 不存在，让 case + // 落到 expected_m candidate 路径（BM∈{8,16,32} 自动选 + + // small_m_simple_sched 命中）。 + // 动机：DSV4 EP 业务真实 shape 大量出现 expected_m=1/2/3 + + // max_hint=None + m=256 的情形。旧逻辑强制 BM=32，但 + // bm32_skew_fast_path 守护要 m>=1024，m=256 时 fast-path 不 + // 命中反而走通用 cache_sfb_k32 路径，padding 浪费 + // (32-expected_m)/32 ≈ 94%。落到 expected_m candidate 后 + // device 端 BLOCK_M<=16 命中 kUseSmallMSimpleSched，物理零浪费。 + const bool real_hot_present = + masked_m_max_hint.has_value() + ? (masked_m_max_hint.value() > 16) + : (expected_m > 16); + // fuse_decode 路径与 fast-path 互斥：env 开启时强制走通用 BM + // 阶梯（device 端 kFuseScaleBDecode 走 path-B 通用 cache_sfb_k32， + // 而 fast-path layout 假设 quad_reduce + direct_load）。 + const bool fuse_decode_hint = + env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; + // BM=64 实验性路径（默认关，env DG_W4_PATHB_BM64=1）： + // 想法：path-B swap_AB 下 wgmma 物理 N = BLOCK_M。BM=64 → wgmma + // 64x64x32，单次 wgmma 内 b_decoded 服务 64 cols 而非 32， + // W4 解码 (prmt+lop3) 成本天然摊薄 2x，弥补 fast-path + // (quad_reduce + direct_load) 关闭带来的 sfb load + fmul + // overhead。同时 grid 减半。 + // 代价：device 端 kK32QuadReduce 假设 BM=32 (compute_n_0/1 只覆盖 + // 32 cols)，BM=64 必须走通用 cache_sfb_k32 + fp32 sfb 路径。 + // 实测结果（DG_W4_PATHB_BM64=1 vs 默认 fast-path）： + // gateup_dense_tail_hot64 356us → 748us (2.1x 退化) + // gateup_mtp_dp4 158us → 275us (1.74x 退化) + // gateup_eight_hot_64 188us → 226us (1.20x 退化) + // gateup_one_hot_32 31us → 47us (1.52x 退化) + // 绝大多数 case 退化 1.2-2.1x，无一受益。 + // 退化根因：cache_sfb_k32 阶段 SMEM build LUT + register fmul + barrier + // 比 fast-path 直读 gmem + quad_reduce 慢得多；BM=64 grid 砍半 + // 反而让单 SM 排队工作量翻倍，stages 不够掩 K-loop（GB/s 从 + // 1818 跌到 452）。decode 复用收益被 sfb pipeline 重建吃光。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + const bool bm64_hint = + dsv4_shape and env_int("DG_W4_PATHB_BM64", 0) != 0; + // BM=64 fast-path 实验路径（默认关，env DG_W4_PATHB_BM64_FASTPATH=1）： + // 与 DG_W4_PATHB_BM64 不同——后者刻意让 device 落回通用 + // cache_sfb_k32 路径；本开关**保持** fast-path 三件套 + // (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched) 开启， + // 仅把 host BM 从 32 抬到 64，验证“weight 解码次数减半 + grid 减半” + // 能否吃掉 max_m>=64 差场景。 + // 依赖：device 端 swap_AB 下 WGMMA N=BLOCK_M=64 指令已存在 + // (FP8MMASelectorRS<64> = MMA_64x64x32_RS_TN)，scale 寻址 BM-agnostic。 + // BM=64 强制 BN=128，避免 A-smem 叠加 BN=256 把 stages 压崩。 + // + // 实测结论（DG_W4_PATHB_BM64_FASTPATH=1 vs 默认 BM=32，ncu 对照， + // g24 N=6144 K=7168 / N=7168 K=3072 skew benchmark）：**全面退化，已证伪**。 + // uniform_17 784us/0.75x → 1391us/0.43x (max_m=17) + // uniform_32 775us/0.77x → 1387us/0.42x (max_m=32) + // dense_tail_hot17 767/0.77x → 1381/0.43x + // uniform_64 1541us/0.38x → 1390us/0.43x (唯一微升) + // one_hot_512 542us/0.58x → 542us/0.58x (完全无变化) + // ncu 四项硬指标 BM32 vs BM64 **逐字节相同**： + // Registers/Thread 168=168, Dyn SMEM 149984=149984, + // Grid (78,1,1)=(78,1,1), Achieved Occupancy 14.0%≈13.9%。 + // 证伪根因（结构性，非实现 bug；正确性 OK diff=0、无 spill）： + // 1. persistent kernel grid 恒 = num_sms，BLOCK_M 不影响 grid。 + // 2. swap_AB 下 WGMMA 形状 / kNumAccum / WAVE_BLOCK_M 全由 + // BLOCK_N 决定，BLOCK_M 加倍不动寄存器、不动 smem。 + // 3. M-tile 数 = ceil(max_m/BM)。DSV4 真实 max_m 多 <=32， + // BM 32→64 时 tile 数不减，但每 tile WGMMA 覆盖 64 行、 + // ~47 行是 padding 空算 → 工作量翻倍、收益为零 → 耗时近翻倍。 + // 4. 仅 max_m=64 时 tile 2→1 抵消 padding 才微升；max_m 更大 + // (one_hot_512 active=1) 瓶颈不在 tile 数，故毫无变化。 + // 结论：当前 DSV4 负载 (max_m 集中 17-64) 下 BM=64 是负优化。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + // 真正受益方向见下方思路 2（active=1 大 max_m 时跨 m_block 复用 + // weight decode），不受 padding 问题影响。 + // BM=64 fast-path **自动触发规则**（默认开，env 可关）： + // 前述 DG_W4_PATHB_BM64_FASTPATH 手动实验证明，BM 32→64 让 + // M-tile 数 ceil(m/32)→ceil(m/64) 减半，仅当“每活跃组平均 + // token 数”够大（每组 m>=64，padding 占比可忽略）才净赚；组小则 + // tile 数不减、padding 翻倍纯亏。判别量本应是 sum_m/active，但 + // host API 只有 max_hint / active（拿不到 sum_m，且不能同步读 + // device 上的 masked_m）。 + // 折中：用高置信子集 `max_hint>=128 且 active<=8`（少数活跃组 + + // 大 max_m，正是 MTP verify 的 spec-len 命中形态）触发，刻意避开 + // max_m=64 / 多活跃组的模糊带（uniform_64 赢、dense_tail_hot64 + // 亏，二者 host 信号 (max_hint=64,active=24) 撞车无法区分）。 + // 实测（g24 N6144/7168 与 g32 N4096，BM64+pair vs 默认 quad）两族 + // shape 一致：命中子集 one_hot_128/384、mtp_384/512_hot、 + // one_hot_512、mtp_768_hot512 等普降 5~26%，被避开的 + // dense_tail/uniform/eight_hot 无一受损。 + // BM=64 下 quad 4 累加器会 spill 爆炸，故命中时**必须**配 + // pair_reduce（见下方 k32_pair_reduce 联动）。 + const bool bm64_fastpath = dsv4_shape and bm64_fastpath_enabled; + if (real_hot_present and not fuse_decode_hint and + not bm64_hint and + env_int("DG_W4_PATHB_FAST_PATH", 1) != 0) { + layout.block_m = bm64_fastpath ? 64 : 32; + // H1: 同时满足 max_m > 32 + active*max_m >= 1024 时走 BN=256， + // grid 减半。门槛公式来自 skewed masked benchmark 实测： + // * eight_hot_64 (8×64=512)、four_hot_64 (4×64=256)、 + // one_hot_384 (1×384=384) 在 BN=256 下 grid 太稀， + // 单 tile WAVE_WGMMA=2 的两次 wgmma 反而拖慢 8-9%。 + // * mtp_dp4 (8×214=1712)、dense_tail_hot64 a32 + // (32×64=2048)、mtp_512_hot384 (8×384=3072) 等 + // 总工作量充足的 case 才能从 BN=256 grid 减半中受益。 + // 需 dsv4_shape 才能命中 device fast-path 守护 + // (N=4096, K∈{4096,2048}, groups=32)，否则 BN=256 落到 + // path-B 通用 cache_sfb_k32 路径反而崩盘。 + const int hint_m = masked_m_max_hint.value_or(0); + // 缺省 active 估计为 num_groups 的一半（保守，避免误升 BN=256） + const int hint_active = active_groups_hint.value_or( + static_cast(desc.num_groups) / 2); + // 总工作量估计（FLOPS 代理）：active × max_m × N × K。 + // 用 int64 防溢出，N/K 单位 elements。 + const int64_t hint_workload = + static_cast(hint_active) * hint_m * + static_cast(desc.n) * + static_cast(desc.k); + // 基础门槛：g32 + N=4096 + K=4096 实测 hint_active*hint_m >= 1024 + // 且 hint_m > 32 时 BN=256 受益。 + const bool bn256_baseline = + hint_m > 32 and + static_cast(hint_active) * hint_m >= 1024; + // 大 N×K 形态扩展（实验性，env DG_W4_PATHB_BN256_BIG=1 默认关）： + // 动机：DSV4 down_proj real shape g24, N=6144, K=7168 下， + // uniform_17/24/32 + dense_tail_hot17 全被 hint_m > 32 + // 卡住走 BN=128 (~941 GB/s)，理论上 grid 4-9 wave 已脱离 + // "grid 太稀"危险区，单 tile work 翻倍应可吞。 + // 阈值原本设计：active>=8 + workload>=1.5e10。 + // 实测结果（DG_W4_PATHB_BN256_BIG=1 vs 默认 BN=128）： + // uniform_17 711us → 797us (-12%, 941→838 GB/s) + // uniform_24 716us → 799us (-12%) + // uniform_32 705us → 790us (-12%) + // dense_tail_hot17 701us → 786us (-12%) + // 退化根因：K=7168 + BN=256 + BM=32 + stages=8 ⇒ SMEM B-tile + // 占用翻倍 (256 cols × 32 K × 8 stages = 64KB)，超出 228KB + // 上限被 JIT 自动降到 stages=4，K-loop 掩蔽崩盘 (GB/s + // 941→842 整齐下跌 11%)。BN=256 grid 减半的收益 << stages + // 减半的带宽损失。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + constexpr int64_t kBN256BigWorkloadThreshold = 15'000'000'000LL; + const bool bn256_big_shape = + env_int("DG_W4_PATHB_BN256_BIG", 0) != 0 and + hint_workload >= kBN256BigWorkloadThreshold and + static_cast(hint_active) >= 8; + const bool bn256_eligible = + dsv4_shape and + masked_m_max_hint.has_value() and + (bn256_baseline or bn256_big_shape) and + env_int("DG_W4_PATHB_BN256", 1) != 0; + // BM=64 fast-path 强制 BN=128：A-smem 已随 BM 翻倍，再叠加 + // BN=256 会把 B-tile + A-tile 占用一起推高，stages 被压崩。 + layout.block_n = (bm64_fastpath or not bn256_eligible) ? 128 : 256; + // H2: cluster_n=2（A 多播）实测**全军退化**，默认关闭。 + // 退化数据（DG_W4_PATHB_CLUSTER_N=1 vs 默认）： + // BN=256 路径 ~9-10% 退化（mtp_512: 201→220, dense_tail_hot64: + // 354→388, mtp_dp4: 158→172） + // BN=128 路径 ~3-12% 退化（four_hot_64: 95→106, + // longtail_128: 157→169） + // 退化根因： + // 1. compact_masked_sched 按 (group, n_block) 调度，cluster=2 + // 把相邻 n_block 绑定到同 cluster，但 masked 下相邻 + // n_block 可能属于不同 group，A tile 复用率被打断。 + // 2. cluster barrier 同步 + cluster launch stagger 对 small + // grid（28-56 个 tile）开销大于多播收益。 + // 3. W4 的 sfb 不参与 multicast，同步开销均摊到两个 CTA。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + if (dsv4_shape and + env_int("DG_W4_PATHB_CLUSTER_N", 0) != 0 and + static_cast(desc.n) % + (static_cast(layout.block_n) * 2) == 0 and + desc.num_sms % 2 == 0) { + layout.cluster_n = 2; + } + } else if (bm64_hint and real_hot_present) { + // BM=64 实验路径：仅 dsv4_shape + DG_W4_PATHB_BM64=1 时进入。 + // 走 path-B 通用 cache_sfb_k32 路径，依赖下面 fast-path 守护 + // 自动关闭 (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched + // 通过 bm32_skew_fast_path 的 BM==32 检查失败而被关闭)。 + layout.block_m = 64; + // BN=128 默认：BM=64 + BN=128 → wgmma 64x64x32 × 2 wgmma per + // K_block，b_decoded 在每次 wgmma 内部服务 64 cols。 + // env DG_W4_PATHB_BM64_BN=256 时尝试更大 grid 减半。 + const int bm64_bn_override = env_int("DG_W4_PATHB_BM64_BN", 0); + layout.block_n = (bm64_bn_override == 256) ? 256 : 128; + } else { + // 纯 fan-out（hint<=16）：base 阶梯，BN=256 + if (bm_select_m <= 8) layout.block_m = 8; + else if (bm_select_m <= 16) layout.block_m = 16; + else if (bm_select_m <= 32) layout.block_m = 32; + else if (bm_select_m <= 64) layout.block_m = 64; + layout.block_n = 256; + } + } else { + // path-A：cooperative prefetch + sfb→smem，BM 阶梯有效。 + if (bm_select_m <= 8) layout.block_m = 8; + else if (bm_select_m <= 16) layout.block_m = 16; + else if (bm_select_m <= 32) layout.block_m = 32; + else if (bm_select_m <= 64) layout.block_m = 64; + layout.block_n = 256; + + // DSV4 + 大 BM：BN=256→128 减小单 tile promote/store 链长， + // 让 grid 上 n-tile 翻倍提升 SM 间负载均衡。 + if (layout.block_m >= 64 and dsv4_shape and + env_int("DG_W4_LARGE_BM_BN128", 1) != 0) { + layout.block_n = 128; + } + // DSV4 + small hot (bm_select∈(32,64])：BM=64 padding 浪费太大， + // 改 BM=32 BN=128 兼顾 hot/small group。 + if (bm_select_m > 32 and bm_select_m <= 64 and dsv4_shape and + env_int("DG_W4_SMALL_HOT_BM32", 1) != 0) { + layout.block_m = 32; + layout.block_n = 128; + } + } + // 历史经验注记： + // * 曾尝试 path-B BM=32→64 升级（hot hint），device kernel 的 + // kK32QuadReduce 在 BM>=64 + swap_AB 下 sfb 寻址 / 4-wgmma fence + // 跟 BM=32 紧耦合，dp4 退化到 0.37x。 + // * 已知天花板（DSV4 skew, groups=32, N=4096）：进一步提升须从 + // device 侧扩展 BM=64 quad_reduce，工作量大且收益不确定。 + } + } + if (not block_m_override and not block_n_override) { + // DG_W4_PATHB_BM64_FORCE_ALL=1 时，把 BLOCK_N 强制锁到 128（与 BM=64 + // fast-path 联动）。动机：BM=64 fast-path 下 final_accum 寄存器量 + // ∝ BLOCK_N/64，BN=256 → 4 份 tile 累加器逼出 168reg/thread、occupancy + // 14%；BN=128 砍半 final_accum，腾出寄存器预算给更多 warp，对大 batch / + // 高 active 场景 occupancy 上限更高。属于在线吞吐对照实验的联动开关， + // 不影响 hint 判据自身路径。 + if (env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 and + bm64_fastpath_enabled and layout.block_m == 64) { + layout.block_n = 128; + } + DG_HOST_ASSERT(layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or + layout.block_m == 64 or layout.block_m == 128); + DG_HOST_ASSERT(layout.block_n == 64 or layout.block_n == 128 or layout.block_n == 256); + } + layout.cluster_m = 1; + auto config = rebuild_config(layout); + + if (block_m_override or block_n_override) { + auto layout = config.layout; + if (block_m_override) { + layout.block_m = *block_m_override; + } + if (block_n_override) { + layout.block_n = *block_n_override; + } + DG_HOST_ASSERT((layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or + layout.block_m == 64 or layout.block_m == 128 or layout.block_m == 256) and layout.block_n % 16 == 0); + DG_HOST_ASSERT(layout.block_n <= 256); + layout.cluster_m = 1; + config = rebuild_config(layout); + } + // Packed FP4 B has half the K bytes of FP8 B. Match PR #287's W4 path: + // TMA writes B with a 64B swizzle and the RS kernel reads it via ldmatrix. + config.storage_config.swizzle_b_mode = config.layout.block_k / 2; + 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 / 2); + + { + const int block_k = config.layout.block_k; + const int block_n = config.layout.block_n; + const int shape_k_scales_b = ceil_div(static_cast(k), gran_k_b); + const bool uniform_scale_b = (block_k % block_n == 0); + const int sfb_old_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + const int sfb_cache_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * block_n * static_cast(sizeof(float)), 16); + const int rs_padded_bm = std::max(config.layout.block_m, 64); + const int base_smem_d_bytes = + align(config.layout.block_m * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int smem_d_bytes = + align(rs_padded_bm * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int smem_d_extra = smem_d_bytes - base_smem_d_bytes; + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; + + const int packed_per_stage = block_n * (block_k / 2); + const int base_smem_a_per_stage = config.storage_config.load_block_m * block_k * + static_cast(c10::elementSize(desc.a_dtype)); + const int base_smem_b_per_stage = config.storage_config.load_block_n * block_k * + static_cast(c10::elementSize(desc.b_dtype)); + const int base_smem_sfa_per_stage = + align(config.layout.block_m * static_cast(sizeof(float)), 128); + const int original_per_stage = + base_smem_a_per_stage + base_smem_b_per_stage + base_smem_sfa_per_stage; + const int smem_a_per_stage = rs_padded_bm * block_k * static_cast(c10::elementSize(desc.a_dtype)); + const int smem_sfa_per_stage = + align(rs_padded_bm * static_cast(sizeof(float)), 128); + const int merged_per_stage = smem_a_per_stage + smem_sfa_per_stage + packed_per_stage; + const int orig_num_stages = config.pipeline_config.num_stages; + const int smem_extra = + config.pipeline_config.smem_size - orig_num_stages * original_per_stage + smem_d_extra + sfb_extra; + + auto fits = [&](int stages) { + return smem_extra + stages * merged_per_stage <= SM90ArchSpec::smem_capacity; + }; + + constexpr int kW4DefaultMaxStages = 8; + int max_fitting = orig_num_stages; + while (max_fitting + 1 <= kW4DefaultMaxStages and fits(max_fitting + 1)) + ++ max_fitting; + int chosen_stages = max_fitting; + while (chosen_stages >= 3 and not fits(chosen_stages)) + -- chosen_stages; + if (gran_k_b == 32 and expected_m <= 16 and + static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) + chosen_stages = std::min(chosen_stages, 6); + const bool bm32_skew_layout = gran_k_b == 32 and + (config.layout.block_m == 32 or + (config.layout.block_m == 64 and bm64_fastpath_enabled)) and + (config.layout.block_n == 128 or config.layout.block_n == 256) and + static_cast(desc.m) >= 1024 and + ((static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) or + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： + // 放开到 g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168} + (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + const int bm32_skew_stage_override = env_int("DG_W4_BM32_SKEW_STAGES", 0); + if (bm32_skew_layout and bm32_skew_stage_override > 0) { + DG_HOST_ASSERT(bm32_skew_stage_override >= 3 and bm32_skew_stage_override <= kW4DefaultMaxStages); + DG_HOST_ASSERT(fits(bm32_skew_stage_override)); + chosen_stages = bm32_skew_stage_override; + } else if (bm32_skew_layout and fits(kW4DefaultMaxStages)) { + chosen_stages = kW4DefaultMaxStages; + } + DG_HOST_ASSERT(chosen_stages >= 3); + config.pipeline_config.num_stages = chosen_stages; + config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; + } + + // R2b-A swap_ab maps original N onto WGMMA M. Use enough math warpgroups + // to cover the 64-row WGMMA-M strips selected by BLOCK_N. + DG_HOST_ASSERT(config.layout.block_n == 64 or config.layout.block_n == 128 or config.layout.block_n == 256); + int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; + const int rs_num_math_threads_override = env_int("DG_W4_RS_MATH_THREADS", 0); + if (rs_num_math_threads_override > 0) { + DG_HOST_ASSERT(rs_num_math_threads_override == 128 or rs_num_math_threads_override == 256); + rs_num_math_threads = rs_num_math_threads_override; + } + config.launch_config.num_math_threads = rs_num_math_threads; + config.launch_config.num_threads = config.launch_config.num_tma_threads + rs_num_math_threads; + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, num_groups, + config.storage_config.swizzle_a_mode); + const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b_bytes, n, half_k, + config.layout.block_n, config.layout.block_k / 2, + half_k, num_groups, + config.storage_config.swizzle_b_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 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 bool bm32_skew_fast_path = gran_k_b == 32 and + (config.layout.block_m == 32 or + (config.layout.block_m == 64 and bm64_fastpath_enabled)) and + (config.layout.block_n == 128 or config.layout.block_n == 256) and + static_cast(desc.m) >= 1024 and + ((static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) or + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）， + // 与 bm32_skew_layout / dsv4_shape 同步生效。device fast-path 三件套 + // (quad_reduce / direct_load / compact_sched) 仅硬约束 BM==32，g<=36 + // 冗余专家场景按同一路径处理， + // 与 N/K 数值无关，物理可命中。 + (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + // Fused-decode 路径（path-B 通用 cache_sfb_k32 + LUT decode）： + // * scale 在 cache 阶段编进 LUT，wgmma 后省一次 fmul（按 quad 算 = 4 次） + // * sfb 4x packed UE8M0：每 4 个 e8m0 打成 1 个 int32，体积 = fp32 的 1/4 + // 与 fast-path（kK32QuadReduce + kScaleBDirectLoad）**互斥**（cuh 强制 + // not kScaleBDirectLoad），开启后强制走通用路径。需要 sfb 是 packed int 布局。 + // 默认关，由 DG_W4_PATHB_FUSE_DECODE=1 显式开启。 + const bool fuse_scale_b_decode = + gran_k_b == 32 and + sfb.scalar_type() == torch::kInt and + get_major_type_ab(sfb) == cute::UMMA::Major::MN and + env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; + const bool scale_b_packed_ue8m0 = fuse_scale_b_decode; + // 思路2 性能探针：把 scale_a 从 promote 内移除（device kScaleAStub 置 scale_a=1 + // 并跳过 SFA 的 TMA load），用于测 fuse_scale_b_decode 单累加器 + scale_a 后移到 + // silu_and_mul 的速度上界。开启后结果数值错误，仅做性能对照。默认关。 + const bool scale_a_stub = env_int("DG_W4_SCALE_A_STUB", 0) != 0; + // 思路2 串行依赖实验：双缓冲预解码（staged_pair_a_regs，issue k 时预取 k+2）， + // 让 decode 与 WGMMA 多重叠一级，代价 +8 reg/线程。仅在 fuse_scale_b_decode 开启 + // 时有意义，用于测"解 issue 串行依赖"能否补回 fuse 路径在大 M 的退化。默认关。 + const bool fuse_predecode_pair = + fuse_scale_b_decode and env_int("DG_W4_FUSE_PREDECODE_PAIR", 0) != 0; + // fuse_scale_b_decode 一旦开启，所有 fast-path / direct-load / e8m0 / bf16 + // 互斥关闭——device 那边 static_assert 会拒绝同时开启。 + const bool scale_b_direct_load = + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and + not fuse_scale_b_decode; + const bool k32_quad_reduce = + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and + not fuse_scale_b_decode; + // 杠杆3 / BM=64 联动：把 4 累加器 quad-reduce（峰值 64 float，在 + // __launch_bounds__(384,1) 的 168reg 硬上限下触发 local spill）退化为已存在的 + // 2 累加器串行 pair-reduce（峰值 32 float）降低寄存器峰值活跃量。仅改单线程内 + // 累加器生命周期、不动线程模型/barrier，数值 bit 级等价。 + // 两种启用： + // (A) BM=32 下**单独**开 pair（DG_W4_K32_PAIR_REDUCE=1）：已证伪——ncu spill + // 223680→134016(-40%) 但 wall-clock 反而 +10~25%（串行累加拉长依赖链， + // spill 不在 BM=32 关键路径上）。默认关。 + // (B) BM=64 fast-path 命中时**必须**配 pair（自动，下面 block_m==64 分支）： + // BM=64 让 kNumAccum 32→quad 4 份 = 128 fp32，spill 爆炸退化 ~2x；pair + // 2 份 = 64 fp32 受控。BM=64 靠 M-tile 数减半换吞吐，pair 是其必要前提。 + // 实测高置信子集（max_hint>=128 且 active<=8）净降 5~26%。 + const bool k32_pair_reduce = + k32_quad_reduce and + ((config.layout.block_m == 64 and bm64_fastpath_enabled) or + env_enabled("DG_W4_K32_PAIR_REDUCE")); + // 杠杆3 续（已证伪，默认关）：pair-reduce 之上把长驻 final_accum 从 fp32 改 bf16 + // 存储（nv_bfloat162）再减半长驻寄存器、继续压 spill（ncu 134016→6480，-97% vs + // quad）。精度无损（W4 diff 0.0001/0.0000）。但 wall-clock 仍慢于 quad baseline + // （随 pair-reduce 一起退化，叠加 bf16 promote 反而更长依赖链）。默认关，需 + // DG_W4_K32_BF16_FINAL_ACCUM=1 显式开启复现。 + const bool k32_bf16_final_accum = + k32_pair_reduce and env_enabled("DG_W4_K32_BF16_FINAL_ACCUM"); + // 杠杆3 续2（已证伪，默认关）：把长驻 final_accum 整体搬进 smem_d（kDirectStore + // + kFinalAccumScratch），final_accum_regs 缩成 1 元素。实测 spill 虽归零，但 + // (1) 寄存器仍顶 168（launch_bounds 硬顶，省下的 reg 没还给 occupancy，Block + // Limit Registers 仍为 1）；(2) DirectStore 逐元素非合并写 gmem 严重拖累大 + // max_m：one_hot_256 0.77x→0.50x、one_hot_512 0.49x→0.29x。净负收益，仅保留 + // 开关供复现，默认关。需 DG_W4_K32_FINAL_ACCUM_SMEM=1 显式开启。 + const bool k32_final_accum_smem = + k32_pair_reduce and env_enabled("DG_W4_K32_FINAL_ACCUM_SMEM"); + // RS baseline for direct E8M0 B scales: keep scale products in split form by default. + // The exponent-adjust path is experimental and can regress some small-M shapes. + const bool scale_b_pow2_promote = + k32_quad_reduce and env_enabled("DG_W4_SCALE_B_POW2_PROMOTE"); + const bool k32_quad_split_promote = + k32_quad_reduce and not env_disabled("DG_W4_K32_QUAD_SPLIT_PROMOTE"); + const bool k32_quad_scale_b_inline = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_INLINE"); + const bool k32_quad_scale_b_prefetch = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_PREFETCH"); + const bool k32_quad_scale_b_vec4 = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_VEC4"); + const bool k32_quad_pair4x2_promote = + k32_quad_reduce and not k32_quad_split_promote and + env_enabled("DG_W4_K32_QUAD_PAIR4X2_PROMOTE"); + // small_m_simple_sched：device 端 kUseSmallMSimpleSched 仅编译期检查 + // BLOCK_M<=16 + GroupedMasked + multicast=1，与 N/K 数值无关。 + // 默认守护 (k>=4096 + n<=4096) 来自历史 g32 + n=4096 dsv4 形状的保守覆盖。 + // 放开到 RELAX 形状集 (g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}) + // 让 DSV4 EP 业务真实 shape (g24 + n∈{6144,7168} + k∈{3072,7168} + expected_m=1/2/3) + // 也能命中 simple_sched，避开通用 masked scheduler 的额外开销。 + const bool small_m_simple_sched = + gran_k_b == 32 and expected_m <= 16 and + ((static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) or + (static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + const bool compact_masked_sched = + bm32_skew_fast_path and not env_disabled("DG_W4_COMPACT_MASKED_SCHED"); + // compact_masked_sched 按 m_max 降序遍历 active group（实验性扩展，默认关）： + // 想法：compact_masked_sched 默认按 group_idx 升序遍历 active group。 + // 当各 group masked_m 不均时，long-tail group 决定 last-wave 收敛节奏。 + // 开启后 inner loop 每步在剩余 active mask 内 O(active) 扫一次找 group_m + // 最大者，让 wave 0 优先吃重 group，wave 末尾留小 group 收尾。 + // 实测结果（DG_W4_PATHB_REORDER_BY_MM=1 vs 默认 fast-path）： + // gateup_dense_tail_hot64 356us → 356us (0%) + // gateup_dense_tail_hot128 399us → 396us (+0.8%) + // gateup_dense_tail_hot214 401us → 400us (-0.2%) + // gateup_mtp_dp4 157us → 157us (0%) + // down_dense_tail_hot64 186us → 185us (+0.5%) + // down_mtp_512_hot384_a8 98us → 100us (-2%) + // down_dense_tail_hot214 206us → 206us (0%) + // gateup_one_hot_32 31us → 31us (+1%, 边缘 case) + // 全军 ±2% 噪声振荡，无任何 case 出现 5-10% 实质收益。 + // 失败根因：compact_masked_sched 已经做了 (group, m_block, n_block) 三维紧凑 + // 展开，dense_tail_hot64 总 tile = 32 group × 2 m_blocks × 32 n_blocks = + // 2048 ≈ 15.5 wave × 132 SM。当 total_tiles >> num_sms 时，wave 边界由 + // ceil(total/132) 决定，与 group 顺序无关；重排只改变"哪些 SM 在 wave 0 + // 拿到 hot group"，last-wave 余数 = (total mod 132) 不变。 + // 边缘 case (one_hot_32 total=32 tile = 1 wave × 32 SM) 的 +1-2% 收益 + // 是 SM 提前进入 idle 的小幅噪声，不构成可推广收益。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + // 仅在 compact_masked_sched 已开启时生效（即 bm32_skew_fast_path）， + // 否则即便 env 开启也不会传到 device。 + const bool reorder_masked_by_max_m = + compact_masked_sched and env_int("DG_W4_PATHB_REORDER_BY_MM", 0) != 0; + // bf16 SFB：体积砍半，**默认开启**。两条 fast-path： + // - path-B (gran_k_b==32 + direct-load)：load_sfb / load_k32_quad_scale_b 直读 gmem。 + // - path-A (gran_k_b==128)：cooperative prefetch 经 smem，smem 也存 bf16。 + // 仍要避开 packed-UE8M0 / fused-decode 等假设 fp32 位级布局的分支。 + // 显式 `DG_W4_SCALE_B_BF16=0` 时回退 fp32；用户传的 sfb 实际是 fp32 也自然回退。 + const bool scale_b_bf16 = + ((scale_b_direct_load and gran_k_b == 32 and not k32_quad_scale_b_prefetch) or + gran_k_b == 128) and + not env_disabled("DG_W4_SCALE_B_BF16") and + sfb.scalar_type() == torch::kBFloat16; + // E8M0 SFB（仅 path-B fast-path）：每元素 1B = fp32 的 8 位指数，体积再砍 2x。 + // 解码 `__uint_as_float(uint32(e) << 23)` 零误差。**默认开启**：当用户传入 + // uint8 sfb 时自动启用；显式 `DG_W4_SCALE_B_E8M0=0` 时回退。 + // 与 bf16 互斥（因 sfb 实际 dtype 已不同），不会同时命中。 + const bool scale_b_e8m0 = + scale_b_direct_load and gran_k_b == 32 and + not k32_quad_scale_b_prefetch and + (not env_disabled("DG_W4_SCALE_B_E8M0") or + env_enabled("DG_W4_SCALE_B_E8M0_ONLY")) and + sfb.scalar_type() == torch::kUInt8; + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat or + (scale_b_bf16 and sfb.scalar_type() == torch::kBFloat16) or + (scale_b_e8m0 and sfb.scalar_type() == torch::kUInt8) or + (fuse_scale_b_decode and sfb.scalar_type() == torch::kInt)); + const SM90FP8FP4Gemm1D2DRSRuntime::Args& rs_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()), + .major_sfb = get_major_type_ab(sfb), + .scale_b_direct_load = scale_b_direct_load, + .scale_b_pow2_promote = scale_b_pow2_promote, + .k32_quad_reduce = k32_quad_reduce, + .k32_pair_reduce = k32_pair_reduce, + .k32_bf16_final_accum = k32_bf16_final_accum, + .k32_final_accum_smem = k32_final_accum_smem, + .k32_quad_split_promote = k32_quad_split_promote, + .k32_quad_scale_b_inline = k32_quad_scale_b_inline, + .k32_quad_scale_b_prefetch = k32_quad_scale_b_prefetch, + .k32_quad_scale_b_vec4 = k32_quad_scale_b_vec4, + .k32_quad_pair4x2_promote = k32_quad_pair4x2_promote, + .small_m_simple_sched = small_m_simple_sched, + .compact_masked_sched = compact_masked_sched, + .fuse_scale_b_decode = fuse_scale_b_decode, + .scale_a_stub = scale_a_stub, + .fuse_predecode_pair = fuse_predecode_pair, + .scale_b_packed_ue8m0 = scale_b_packed_ue8m0, + .scale_b_gran_k = static_cast(gran_k_b), + .b_is_int4_sym = b_is_int4_sym, + .scale_b_bf16 = scale_b_bf16, + .scale_b_e8m0 = scale_b_e8m0, + .reorder_masked_by_max_m = reorder_masked_by_max_m, + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .sfb = sfb.data_ptr(), + .grouped_layout = masked_m.data_ptr(), + .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 = SM90FP8FP4Gemm1D2DRSRuntime::generate(rs_args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_fused", code); + SM90FP8FP4Gemm1D2DRSRuntime::launch(runtime, rs_args); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp new file mode 100644 index 0000000000..65653763be --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -0,0 +1,277 @@ +#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/layout.hpp" +#include "epilogue.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +// JIT runtime for the RS-mode SM90 FP8xFP4 1D2D kernel. +// +// S2 stage: identical wire-format to `SM90FP8FP4Gemm1D2DRuntime`, only the +// included header and instantiated kernel name differ. This gives the new +// kernel its own JIT artifact and entry point so subsequent steps (S3+) can +// safely diverge the device code without affecting the SS path or the +// contiguous host route. +class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + cute::UMMA::Major major_sfb; + bool scale_b_direct_load; + bool scale_b_pow2_promote; + bool k32_quad_reduce; + // 杠杆3：把 quad-reduce（4 累加器，峰值 64 float→spill）退化为已存在的 + // 2 累加器串行 pair-reduce（峰值 32 float），消除寄存器 spill。保持 + // pingpong + scale_b_direct_load，仅改单线程内累加器生命周期，数值等价。 + // 由 DG_W4_K32_PAIR_REDUCE=1 开启，默认关。 + bool k32_pair_reduce; + // 杠杆3 续：pair-reduce 下进一步把长驻的 final_accum 从 fp32 改 bf16 + // 存储（nv_bfloat162），每线程长驻寄存器减半，继续压低残留 spill。 + // 代价是 promote 累加用 bf16，大 K 有精度风险，需 diff 验证。 + // 仅 pair-reduce 生效，由 DG_W4_K32_BF16_FINAL_ACCUM=1 开启，默认关。 + bool k32_bf16_final_accum; + // 杠杆3 续2：把长驻的 final_accum 整体搬进 smem_d（kDirectStore + + // kFinalAccumScratch），final_accum_regs 数组缩成 1 个元素，长驻寄存器 + // 再砍一大块。DirectStore 路径自包含、逐元素写 gmem、不走 STSM+TMA、无 + // store-barrier，故不会 hang。仅 pair-reduce 生效，由 + // DG_W4_K32_FINAL_ACCUM_SMEM=1 开启，默认关。 + bool k32_final_accum_smem; + bool k32_quad_split_promote; + bool k32_quad_scale_b_inline; + bool k32_quad_scale_b_prefetch; + bool k32_quad_scale_b_vec4; + bool k32_quad_pair4x2_promote; + bool small_m_simple_sched; + bool compact_masked_sched; + // Fused-decode：在 cache_sfb_k32 阶段把 e8m0 scale 编进 LUT，wgmma 后省一次 fmul。 + // 与 scale_b_direct_load / k32_quad_reduce / e8m0/bf16/pow2/compact_sched 互斥 + // （host 决策侧保证），路径走 path-B 通用 cache_sfb_k32 + LUT decode。 + bool fuse_scale_b_decode; + // 思路2 性能探针（DG_W4_SCALE_A_STUB）：把 scale_a 从 promote 内彻底移除 + // （device 端 kScaleAStub 置 scale_a=1 并跳过 SFA 的 TMA load），用于验证 + // "fuse_scale_b_decode 单累加器 + scale_a 后移到 silu_and_mul" 完整方案的 + // 速度上界。注意：开启后结果数值错误（scale_a 未参与），仅做性能对照， + // 不可用于生产。配合 fuse_scale_b_decode 一起测才有意义。 + bool scale_a_stub; + // 思路2 串行依赖实验（DG_W4_FUSE_PREDECODE_PAIR）：fuse_scale_b_decode 默认 + // 1 级预取（decode k+1 重叠 WGMMA k）下，WGMMA k+1 仍卡 decode k+1。本开关 + // 改走双缓冲预解码（staged_pair_a_regs[2][4]，issue k 时预取 k+2），让 decode + // 与 WGMMA 多重叠一级，代价 +8 reg/线程。仅在 fuse_scale_b_decode 开启时有效， + // 用于测"解 issue 串行依赖"能否补回 fuse 路径在大 M 的退化。默认关。 + bool fuse_predecode_pair; + // 配合 fuse_scale_b_decode：sfb 物理布局是 [groups, K/32/4, N]（MN-major + 4 个 + // e8m0 打包成 1 个 int32），体积 = fp32 的 1/4。需 sfb.scalar_type() == kInt。 + bool scale_b_packed_ue8m0; + uint32_t scale_b_gran_k; + // INT4-sym (signed [-8, 7] packed two nibbles/byte) variant for B. + // Path-A: per-128 fp32 SFB, no fused-decode. See kernel header. + bool b_is_int4_sym; + // bf16 SFB（path-A k128 / path-B fast-path）：体积砍半，scale 用 __nv_bfloat16。 + bool scale_b_bf16; + // E8M0 SFB（path-B fast-path 专用）：每元素 1B = fp32 的 8 位指数。 + // 解码 `__uint_as_float(uint32(e) << 23)` 零误差。仅 pow2 scale 适用。 + bool scale_b_e8m0; + // compact_masked_sched 按 m_max 降序遍历 active group（实验性扩展）： + // 仅在 compact_masked_sched 开启时生效，让 wave 0 优先吃重 group， + // 减少 last-wave imbalance。host 默认关。 + bool reorder_masked_by_max_m; + void *gmem_b_ptr; + void *gmem_d_ptr; + void *sfb; + void *grouped_layout; + 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_fp8_fp4_gemm_1d2d_rs_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", + to_string(args.major_sfb), + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', 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), + get_default_epilogue_type(std::nullopt), + "false", + args.scale_a_stub ? "true" : "false", + "false", + "false", + "false", + args.scale_b_pow2_promote ? "true" : "false", + "false", + "false", + args.scale_b_direct_load ? "true" : "false", + "false", + "false", + "false", + args.k32_final_accum_smem ? "true" : "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + args.k32_bf16_final_accum ? "true" : "false", + "false", + "false", + "false", + "false", + "false", + args.k32_quad_reduce ? "true" : "false", + args.k32_pair_reduce ? "true" : "false", + (args.k32_quad_reduce and not args.k32_pair_reduce) ? "true" : "false", + args.k32_quad_split_promote ? "true" : "false", + args.k32_quad_scale_b_inline ? "true" : "false", + args.k32_quad_scale_b_prefetch ? "true" : "false", + args.k32_quad_scale_b_vec4 ? "true" : "false", + args.k32_quad_pair4x2_promote ? "true" : "false", + "false", + args.k32_final_accum_smem ? "true" : "false", + "false", + "false", + args.small_m_simple_sched ? "true" : "false", + args.fuse_scale_b_decode ? "true" : "false", + "false", + "false", + "false", + "false", + args.fuse_predecode_pair ? "true" : "false", + "false", + "false", + "false", + "false", + 0, + args.scale_b_packed_ue8m0 ? "true" : "false", + args.scale_b_gran_k, + 1, + 1, + 0, + args.b_is_int4_sym ? "true" : "false", + args.scale_b_bf16 ? "true" : "false", + args.scale_b_e8m0 ? "true" : "false", + args.reorder_masked_by_max_m ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.gmem_b_ptr, args.sfb, args.grouped_layout, args.gmem_d_ptr, + 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)); + } +}; + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 5d1f17b5b8..ef022e6c07 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -99,11 +99,23 @@ static void __instantiate_kernel() {{ } }; -static std::tuple preprocess_sf(const torch::Tensor& sf) { +static std::tuple preprocess_sf(const torch::Tensor& sf, + const bool& require_float = true) { // NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA const auto dim = sf.dim(); DG_HOST_ASSERT(dim == 2 or dim == 3); - DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); + if (require_float) { + DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); + } else { + // 放过 bf16 / E8M0 (UInt8) sfb（path-B fast-path，体积砍半 / 砍 4 倍）。 + DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat or + sf.scalar_type() == torch::kInt or + sf.scalar_type() == torch::kBFloat16 or + sf.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(sf.element_size() == sizeof(float) or + sf.scalar_type() == torch::kBFloat16 or + sf.scalar_type() == torch::kUInt8); + } const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf; const auto [num_groups, mn, sf_k] = get_shape<3>(batched_sf); @@ -112,12 +124,16 @@ static std::tuple preprocess_sf(const to } static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { - const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf); + const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf, false); // The last kernel already gives a column-major TMA aligned layout if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn) return (dim == 2) ? batched_sf.squeeze(0) : batched_sf; + // bf16 / E8M0 (UInt8) sfb 必须由调用方按 MN-major + tma_aligned_mn 直接构造， + // 不能落入下面的 fp32-only TransposeFP32Runtime。 + DG_HOST_ASSERT(sf.scalar_type() != torch::kBFloat16 and sf.scalar_type() != torch::kUInt8); + const auto out = torch::empty_strided({num_groups, mn, sf_k}, {tma_aligned_mn * sf_k, 1, tma_aligned_mn}, batched_sf.options()); diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 928472d35a..17b08e4831 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,10 +1,12 @@ #pragma once +#include + #include #include -#include "math.hpp" -#include "exception.hpp" +#include "../utils/math.hpp" +#include "../utils/exception.hpp" #include "../jit/device_runtime.hpp" namespace deep_gemm { @@ -90,12 +92,20 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, // Always do shape checks const auto sf_dtype = sf.scalar_type(); - DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt); + // BF16 / E8M0 (UInt8) SFB fast path 也允许；其它情况维持 fp32/int 限制。 + DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt or + sf_dtype == torch::kBFloat16 or sf_dtype == torch::kUInt8); DG_HOST_ASSERT(sf.dim() == static_cast(num_groups.has_value()) + 2); if (num_groups.has_value()) DG_HOST_ASSERT(sf.size(-3) == num_groups.value()); DG_HOST_ASSERT(sf.size(-2) == ceil_div(mn, gran_mn)); - DG_HOST_ASSERT(sf.size(-1) == ceil_div(k, gran_k * (sf_dtype == torch::kFloat ? 1 : 4))); + const char* scale_b_lut_env = std::getenv("DG_W4_SCALE_B_LUT"); + const bool scale_b_lut = sf_dtype == torch::kInt and + scale_b_lut_env != nullptr and scale_b_lut_env[0] != '\0' and scale_b_lut_env[0] != '0'; + const int expected_sf_k = scale_b_lut ? + ceil_div(k, gran_k) * 2 : + ceil_div(k, gran_k * (sf_dtype == torch::kInt ? 4 : 1)); + DG_HOST_ASSERT(sf.size(-1) == expected_sf_k); // TMA stride checks: TMA aligned and MN-major if (tma_stride_check) { diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a9542e2f44..8c4cad5cdc 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -76,6 +76,25 @@ # TODO: remove these later fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked + try: + fp8_fp4_gemm_nt_sm90_fused_wgmma = _C.fp8_fp4_gemm_nt_sm90_fused_wgmma + except AttributeError: + pass + try: + m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma + ) + except AttributeError: + pass + try: + m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma + ) + m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma + ) + except AttributeError: + pass except ImportError: # Expected behavior for CUDA runtime version before 12.1 pass diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh new file mode 100644 index 0000000000..0aebdf7dca --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh @@ -0,0 +1,396 @@ +#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 +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + + +__device__ __forceinline__ uint8_t sm90_fp8_fp4_fused_e2m1_to_e4m3_bits(uint8_t code) { + // E2M1 values {0, 0.5, 1, 1.5, 2, 3, 4, 6} map exactly to E4M3. + constexpr uint8_t kE2M1ToE4M3[8] = {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c}; + const uint8_t value_idx = code & 0x07u; + const uint8_t sign = (value_idx != 0u && (code & 0x08u) != 0u) ? 0x80u : 0x00u; + return kE2M1ToE4M3[value_idx] | sign; +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void +sm90_fp8_fp4_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, int8_t* gmem_b_ptr, + cd_dtype_t* gmem_d_ptr, int* grouped_layout, + cute::TmaDescriptor* tensor_map_buffer, + uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a_base, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, + const __grid_constant__ cute::TmaDescriptor tensor_map_cd) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) + // Scaling checks + DG_STATIC_ASSERT(kNumTMAThreads == 128 and kNumMathThreads % 128 == 0, "Invalid Threads"); + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, "Invalid C/D data dtype"); + DG_STATIC_ASSERT(kGemmType == GemmType::Normal or kGemmType == GemmType::MGroupedContiguous or + kGemmType == GemmType::MGroupedContiguousWithPsumLayout, + "SM90 FP8xFP4 fused only supports normal and m-grouped contiguous GEMM"); + + // Types + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size"); + + // Overwrite shape constants if the compiler gives + 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; + + // Shared memory + static constexpr uint32_t SMEM_TENSOR_MAP_SIZE = (kGemmType == GemmType::KGroupedContiguous ? sizeof(cute::TmaDescriptor) * 2 : 0); + static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(cd_dtype_t); + static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * sizeof(float); + static constexpr uint32_t ALIGNED_SMEM_SFB_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFB_SIZE_PER_STAGE, 128u); + DG_STATIC_ASSERT(SMEM_SFA_SIZE_PER_STAGE % 128 == 0, "Invalid TMA alignment"); + + // Configs + const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane_idx = threadIdx.x % 32; + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a_base); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_sfb); + cute::prefetch_tma_descriptor(&tensor_map_cd); + } + __syncwarp(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + + // Tensor maps on shared and global memory + (void)gmem_a_ptr; + (void)tensor_map_buffer; + (void)gmem_d_ptr; + + // Data on shared memory + auto smem_d = reinterpret_cast(smem_buffer + SMEM_TENSOR_MAP_SIZE); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + (SMEM_TENSOR_MAP_SIZE + 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_TENSOR_MAP_SIZE + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE)); + }); + constexpr auto SMEM_SF_OFFSET = SMEM_TENSOR_MAP_SIZE + 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 * SMEM_SFA_SIZE_PER_STAGE)); + }); + auto smem_sfb = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * ALIGNED_SMEM_SFB_SIZE_PER_STAGE)); + }); + + // Barriers on shared memory + constexpr auto SMEM_BARRIER_OFFSET = SMEM_SF_OFFSET + kNumStages * (SMEM_SFA_SIZE_PER_STAGE + ALIGNED_SMEM_SFB_SIZE_PER_STAGE); + auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_BARRIER_OFFSET + i * static_cast(sizeof(Barrier)))); + }); + auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_BARRIER_OFFSET + (kNumStages + i) * static_cast(sizeof(Barrier)))); + }); + + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // Initialize barriers + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Pipeline unroll control + constexpr uint32_t kNumPipelineUnrolls = (kGemmType == GemmType::KGroupedContiguous ? 0 : kNumStages); + + // Register reconfigurations (more math registers are needed with unrolling) + constexpr uint32_t kNumTMARegisters = (kNumPipelineUnrolls == 0 ? 40 : 24); + constexpr uint32_t kNumMathRegisters = (kNumPipelineUnrolls == 0 ? 232 : 240); + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // TMA and MMA pipeline + const auto get_pipeline = [=](const uint32_t& iter_idx) -> cute::tuple { + return {iter_idx % kNumStages, (iter_idx / kNumStages) & 1}; // Pipeline stage and phase + }; + uint32_t iter_idx = 0; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + 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"); + + const uint32_t num_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + const uint32_t m_idx = m_block_idx * BLOCK_M; + const uint32_t n_idx = n_block_idx * BLOCK_N; + + #pragma unroll kNumPipelineUnrolls + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; ++ k_block_idx) { + // Wait consumer release + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Issue TMA + auto& full_barrier = *full_barriers[stage_idx]; + const uint32_t k_idx = k_block_idx * BLOCK_K; + const uint32_t current_group_idx = get_current_group_idx(); + const uint32_t sf_k_idx_a = k_block_idx; + const uint32_t sf_k_idx_b = current_group_idx * math::ceil_div(shape_k, BLOCK_K) + k_block_idx; + const auto tensor_map_a_ptr = &tensor_map_a_base; + tma::copy(&tensor_map_sfa, &full_barrier, smem_sfa[stage_idx], m_idx, sf_k_idx_a, num_tma_multicast_a); + tma::copy(&tensor_map_sfb, &full_barrier, smem_sfb[stage_idx], n_idx, sf_k_idx_b, num_tma_multicast_b); + tma::copy(tensor_map_a_ptr, &full_barrier, smem_a[stage_idx], k_idx, m_idx, num_tma_multicast_a); + full_barrier.arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + SMEM_SFB_SIZE_PER_STAGE); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + #pragma unroll + for (uint32_t s = 0; s < kNumStages; ++ s) { + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + empty_barriers[stage_idx]->wait(phase ^ 1); + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto r_0 = warp_idx * 16 + row_idx, r_1 = r_0 + 8; + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Accumulation for WGMMA or CUDA promotion + DG_STATIC_ASSERT(BLOCK_M == WGMMA::M * (BLOCK_M <= 64 ? 1 : 2), "Invalid block sizes"); + const uint32_t current_shape_k = shape_k; + const uint32_t current_group_idx = get_current_group_idx(); + const uint32_t num_k_blocks = math::ceil_div(current_shape_k, BLOCK_K); + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0}; + float2 scales_b[WGMMA::kNumAccum / 4]; + + // Empty barrier arrival + auto empty_barrier_arrive = [&](uint32_t s) { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? empty_barriers[s]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta) : void(); + } + }; + + #pragma unroll kNumPipelineUnrolls + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; ++ k_block_idx) { + // Wait TMA arrivals + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + full_barriers[stage_idx]->wait(phase); + + // Decode one packed byte into two FP8 values. SM90 cannot use + // Blackwell's FP4 TMA/UMMA path, so keep B loads coalesced here. + for (uint32_t idx = threadIdx.x; idx < BLOCK_N * (BLOCK_K / 2); idx += kNumMathThreads) { + const uint32_t tile_n = idx / (BLOCK_K / 2); + const uint32_t packed_k = idx % (BLOCK_K / 2); + const uint32_t tile_k = packed_k * 2; + const uint32_t global_n = n_block_idx * BLOCK_N + tile_n; + const uint32_t global_k = k_block_idx * BLOCK_K + tile_k; + constexpr bool kFullStaticNK = SHAPE_N != 0 and SHAPE_K != 0 and + SHAPE_N % BLOCK_N == 0 and SHAPE_K % BLOCK_K == 0; + uint8_t packed; + const uint32_t b_group_offset = current_group_idx * shape_n * (shape_k / 2); + if constexpr (kFullStaticNK) { + packed = static_cast(gmem_b_ptr[b_group_offset + global_n * (shape_k / 2) + global_k / 2]); + } else { + packed = 0; + if (global_n < shape_n and global_k < current_shape_k) + packed = static_cast(gmem_b_ptr[b_group_offset + global_n * (shape_k / 2) + global_k / 2]); + } + + const uint32_t n_group = tile_n / 8; + const uint32_t n_in_group = tile_n % 8; + const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); + const uint32_t smem_idx = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K + swizzled_k; + auto smem_b_bytes = reinterpret_cast(smem_b[stage_idx]); + const uint8_t lo = sm90_fp8_fp4_fused_e2m1_to_e4m3_bits(packed & 0x0fu); + const uint8_t hi = sm90_fp8_fp4_fused_e2m1_to_e4m3_bits((packed >> 4) & 0x0fu); + if constexpr (kFullStaticNK) { + *reinterpret_cast(smem_b_bytes + smem_idx) = static_cast(lo) | (static_cast(hi) << 8); + } else { + if (global_k + 1 < current_shape_k) { + *reinterpret_cast(smem_b_bytes + smem_idx) = static_cast(lo) | (static_cast(hi) << 8); + } else { + smem_b_bytes[smem_idx] = lo; + } + } + } + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + + // Read A scales + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + auto scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + r_0); + auto scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + r_1); + + // Read B scales + #pragma unroll + for (int i = 0; i < WGMMA::kNumAccum / 4; ++i) + scales_b[i] = ptx::ld_shared(reinterpret_cast(smem_sfb[stage_idx] + i * 8 + col_idx * 2)); + + // Commit WGMMA instructions + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc(smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc(smem_b[stage_idx] + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + // Notify barrier arrival + empty_barrier_arrive(stage_idx); + + // Promote with scales + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float &scale_b_0 = scales_b[i].x; + const float &scale_b_1 = scales_b[i].y; + final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + } + } + + if constexpr (kGemmType == GemmType::Normal) { + // Flush previous stores + if (warp_idx % 4 == 0 and cute::elect_one_sync()) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(128, math_wg_idx); + + // Store to D shared memory + const auto smem_d_0 = reinterpret_cast(smem_d + r_0 * BLOCK_N + col_idx * 2); + const auto smem_d_1 = reinterpret_cast(smem_d + r_1 * BLOCK_N + col_idx * 2); + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + ptx::st_shared(smem_d_0 + i * 4, {final_accum[i * 4 + 0], final_accum[i * 4 + 1]}); + ptx::st_shared(smem_d_1 + i * 4, {final_accum[i * 4 + 2], final_accum[i * 4 + 3]}); + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(128, math_wg_idx); + + // Use TMA reduce-add to accumulate C into D for the normal FP32 path. + if (warp_idx % 4 == 0 and cute::elect_one_sync()) { + cute::SM90_TMA_REDUCE_ADD_2D::copy( + &tensor_map_cd, smem_d_0, n_block_idx * BLOCK_N, + m_block_idx * BLOCK_M + r_0); + cute::tma_store_arrive(); + } + } else { + const uint32_t row_0 = m_block_idx * BLOCK_M + r_0; + const uint32_t row_1 = m_block_idx * BLOCK_M + r_1; + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t col = n_block_idx * BLOCK_N + i * 8 + col_idx * 2; + if (row_0 < shape_m and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = cd_dtype_t(final_accum[i * 4 + 0]); + if (row_0 < shape_m and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = cd_dtype_t(final_accum[i * 4 + 1]); + if (row_1 < shape_m and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = cd_dtype_t(final_accum[i * 4 + 2]); + if (row_1 < shape_m and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = cd_dtype_t(final_accum[i * 4 + 3]); + } + } + __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/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh new file mode 100644 index 0000000000..dd0a0eaf83 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh @@ -0,0 +1,713 @@ +#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 +#include + +namespace deep_gemm { + +template +CUTLASS_DEVICE void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) { + if (num_former_iters == kNumFormerIters) { + func(cute::Int{}); + return; + } + + if constexpr (kNumFormerIters + kGap <= kEnd) + dispatch_num_former_iters(num_former_iters, func); +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void +sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, + nv_bfloat16* gmem_d_ptr, + 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__) + // Scaling checks + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT( + math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or + (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); + + // Types + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M, "Invalid block size"); + + // Overwrite shape constants if the compiler gives + 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; + + // Shared memory + static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0); + 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); + // Packed FP4 B is loaded by TMA into a separate buffer; each row is BLOCK_K / 2 bytes. + static constexpr uint32_t BLOCK_K_PACKED = BLOCK_K / 2; + static constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = BLOCK_N * BLOCK_K_PACKED; + 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 shape_k_scales = math::ceil_div(shape_k, BLOCK_K); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); + // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough + // smem_d capacity, so they fall back to a separate SFB region. + const uint32_t smem_sfb_bytes = math::align(shape_k_scales * BLOCK_N * sizeof(float), 16u); + constexpr uint32_t COMPILED_SHAPE_K_SCALES = SHAPE_K == 0 ? 0 : math::constexpr_ceil_div(SHAPE_K, BLOCK_K); + constexpr uint32_t COMPILED_SMEM_SFB_BYTES = + math::constexpr_align(COMPILED_SHAPE_K_SCALES * BLOCK_N * static_cast(sizeof(float)), 16u); + constexpr bool kUseSeparateSFB = SHAPE_K != 0 and COMPILED_SMEM_SFB_BYTES > SMEM_D_SIZE; + constexpr uint32_t SMEM_SFB_SIZE = kUseSeparateSFB ? COMPILED_SMEM_SFB_BYTES : 0; + + // NOTES: Make sure we have enough shared memory for WGMMA padding + static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); + DG_STATIC_ASSERT(WGMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, "Memory Out of bound for WGMMA"); + + // Configs + 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(); + + // Prefetch TMA descriptors at the very beginning + 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(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + auto fp4_to_e4m3_bits = [](uint32_t code) -> uint32_t { + // E2M1 mag {0..7} -> E4M3 magnitude bytes {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c}. + // Use hardware __byte_perm as a single-cycle 8-entry byte LUT: + // LUT_LO[mag] = {0x00, 0x30, 0x38, 0x3c}, LUT_HI[mag-4] = {0x40, 0x44, 0x48, 0x4c} + // __byte_perm(LUT_LO, LUT_HI, mag) returns 4 copies of the byte at index `mag` + // (mag is in 0..7, fits in the lower nibble, used as the byte selector). + // Sign: bit 3 of code shifted to MSB. We deliberately allow the FP4 "negative zero" + // (mag=0, sign=1) to map to E4M3 0x80 (which is -0, not NaN). WGMMA treats -0 as 0, + // so dropping the explicit mag!=0 mask saves ~4 ops per nibble (~128 per uint4 of + // packed FP4) without affecting numerical results. + constexpr uint32_t LUT_LO = 0x3c383000u; + constexpr uint32_t LUT_HI = 0x4c484440u; + const uint32_t mag = code & 0x07u; + const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; + const uint32_t sign = (code & 0x08u) << 4; // 0 or 0x80 + return mag_byte | sign; + }; + auto fp4_pair_to_e4m3_pair = [&](uint32_t packed) { + const uint32_t lo = fp4_to_e4m3_bits(packed & 0x0fu); + const uint32_t hi = fp4_to_e4m3_bits((packed >> 4) & 0x0fu); + return lo | (hi << 8); + }; + + // Data on shared memory + 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_B_PACKED_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + auto smem_b_packed = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_B_PACKED_OFFSET + i * SMEM_B_PACKED_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumStages * SMEM_B_PACKED_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); + }); + constexpr uint32_t SMEM_SFB_OFFSET = SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE; + // Prefer aliasing SFB onto smem_d. For small-M tiles smem_d is too small, + // so allocate a separate SFB cache to enable BLOCK_M=64 masked kernels. + auto smem_sfb = reinterpret_cast(smem_buffer + (kUseSeparateSFB ? SMEM_SFB_OFFSET : 0)); + if constexpr (not kUseSeparateSFB) { + DG_TRAP_ONLY_DEVICE_ASSERT(smem_sfb_bytes <= SMEM_D_SIZE); + } + + // Fill barriers. + // After the A/packed-B barrier merge there is only one set of full/empty barriers. + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SFB_OFFSET + SMEM_SFB_SIZE); + 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; }); + + // Initialize barriers + DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast"); + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Register reconfigurations + constexpr uint32_t kNumTMARegisters = 40; + constexpr uint32_t kNumMathRegisters = kNumMathThreads == 128 ? 248 : 232; + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // `gmem_b_ptr` is no longer used: B is now loaded by TMA into `smem_b_packed`. + (void)gmem_b_ptr; + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedMasked) { + return scheduler.current_group_idx; + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // Pipeline and TMA phases (single shared pipeline for A/SFA/packed-B after the barrier merge) + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used. + // We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32`. + if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + 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)) { + // Wait consumer release for the (now merged) A/SFA/packed-B slot + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Issue TMA A + constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); + const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); + + 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, batch_idx); + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx), + num_tma_multicast_a); + + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier + const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; + tma::copy(&tensor_map_b, &full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), + k_idx_packed, + scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b, batch_idx); + + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + + SMEM_B_PACKED_SIZE_PER_STAGE); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; ++ i) { + empty_barriers[stage_idx]->wait(phase ^ 1); + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto r_0 = warp_idx * 16 + row_idx, 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); + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const uint32_t current_group_idx = get_current_group_idx(); + + // Cooperatively prefetch the SFB tile for this block from gmem to smem. + // Layout in smem: [shape_k_scales, BLOCK_N] (k outer, n inner). + // Out-of-bound n is filled with 1.0f to keep `n_idx >= shape_n` neutral. + // + // Optimization: use cp.async to copy gmem->smem directly (no register + // round-trip). For MN-major (n innermost in gmem) we can issue 16-byte + // (float4) cp.async per thread, cutting #instructions by 4x. K-major + // and the OOB tail use scalar 4-byte cp.async / st.shared. + { + const uint32_t n_block_base = n_block_idx * BLOCK_N; + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = shape_k_scales * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales; + // Issue cp.async (16B per thread) for fully in-bounds vectors; + // fall back to scalar st.shared with 1.0f-fill for the tail. + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_idx = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + k_idx * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + // Tail: at least one element is OOB; use scalar st with 1.0f fill. + float4 vals; + vals.x = (n_idx + 0 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 0] : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 1] : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 2] : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 3] : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + // K-major: sfb is strided along n; cannot easily vectorize across n. + // Use scalar 4B cp.async for in-bounds, st.shared with 1.0f for OOB. + const uint32_t total = shape_k_scales * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_idx = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales + + n_idx * shape_k_scales + k_idx; + ptx::cp_async_4(smem_dst, gmem_src); + } + } + } + // Commit and wait for all cp.async issued above; pair with the + // existing NamedBarrier so the smem cache is visible to all + // math warps before they enter the K loop. + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + + auto load_sfb = [&](uint32_t n_idx, uint32_t k_block_idx) { + // SFB has been staged into smem above; out-of-bound `n_idx` already + // resolves to 1.0f because we wrote 1.0f for those slots. + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + return ptx::ld_shared(smem_sfb + k_block_idx * BLOCK_N + n_off); + }; + + // Decide the number of scales B to load + DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0); + uint32_t num_former_iters = BLOCK_N / 8; + if constexpr (not kMustUseUniformedScaleB) { + num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; + } + + // Accumulation for WGMMA or CUDA promotion + constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; + + // Pick threads whose WGMMA results are to be stored in shared memory + DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for `BLOCK_M < 64`"); + constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + const bool do_wgmma_store = BLOCK_M >= WGMMA::M or warp_idx < kNumWGMMAStoreThreads / 32; + + // Empty barrier arrival + 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(); + } + }; + + // Skip useless computations + const bool is_cta_computation_valid = scheduler.is_computation_valid(m_block_idx, 0); + if (is_cta_computation_valid) { + // The compiler must know the dynamic variable `num_former_iters`'s real value + constexpr bool kShouldOptimize = BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; + constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; + + // Decode helper: wait for the (merged) full[s] barrier, then + // decode packed FP4 B from `smem_b_packed[s]` into the swizzled + // `smem_b[s]`. Caller must subsequently issue + // NamedBarrier::sync(...,7) to publish the decoded data to the + // wgmma async-proxy. The packed slot is released together with + // A/SFA via the single empty_barrier_arrive() after wgmma_wait. + auto wait_and_decode = [&](uint32_t s, uint32_t p) { + full_barriers[s]->wait(p); + auto smem_b_packed_bytes = smem_b_packed[s]; + auto smem_b_bytes = reinterpret_cast(smem_b[s]); + constexpr uint32_t kVecPackedBytes = 16; // 16 packed bytes -> 32 e4m3 bytes + constexpr uint32_t kVecsPerRow = BLOCK_K_PACKED / kVecPackedBytes; + constexpr uint32_t kNumVecs = BLOCK_N * kVecsPerRow; + DG_STATIC_ASSERT(BLOCK_K_PACKED % kVecPackedBytes == 0, + "Packed K must be multiple of 16-byte vector width"); + DG_STATIC_ASSERT(BLOCK_K == 128, + "Swizzle assumes BLOCK_K == 128 so 32B store stays in-range"); + if constexpr (not kDecodeStub) { + for (uint32_t idx = threadIdx.x; idx < kNumVecs; idx += kNumMathThreads) { + const uint32_t tile_n = idx / kVecsPerRow; + const uint32_t vec_k = idx % kVecsPerRow; + const uint32_t tile_k = vec_k * (kVecPackedBytes * 2); // 32-byte step in K + const uint4 packed16 = *reinterpret_cast( + smem_b_packed_bytes + tile_n * BLOCK_K_PACKED + vec_k * kVecPackedBytes); + uint64_t decoded[4] = {0, 0, 0, 0}; + auto decode_u32 = [&](uint32_t packed_word, uint64_t& out) { + #pragma unroll + for (uint32_t b = 0; b < 4; ++ b) { + const uint32_t pair = fp4_pair_to_e4m3_pair((packed_word >> (b * 8)) & 0xffu); + out |= static_cast(pair) << (b * 16); + } + }; + decode_u32(packed16.x, decoded[0]); + decode_u32(packed16.y, decoded[1]); + decode_u32(packed16.z, decoded[2]); + decode_u32(packed16.w, decoded[3]); + const uint32_t n_group = tile_n / 8; + const uint32_t n_in_group = tile_n % 8; + const uint32_t row_base = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K; + // Each segment writes 16 bytes (= 2 u64). Use a + // single st.shared.v2.u64 instead of two + // st.shared.u64, halving the store instruction + // count for the decoded B tile. + { + const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); + ptx::st_shared_v2_u64(smem_b_bytes + row_base + swizzled_k, + decoded[0], decoded[1]); + } + { + const uint32_t swizzled_k = (tile_k + 16) ^ (n_in_group * 16); + ptx::st_shared_v2_u64(smem_b_bytes + row_base + swizzled_k, + decoded[2], decoded[3]); + } + } + } + }; + + + // Dispatch `num_former_iters` and launch MMAs with decode/wgmma + // pipeline overlap: decode runs one stage ahead of wgmma so that + // each iter's wait+decode for stage k+1 hides under iter k's + // wgmma async work. + dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) { + // Lead pointers: track the stage that decode is currently + // working on (one ahead of wgmma's `stage_idx`). + uint32_t lead_stage = stage_idx, lead_phase = phase; + auto advance_lead = [&]() { + lead_stage = lead_stage == kNumStages - 1 ? 0 : lead_stage + 1; + lead_phase ^= lead_stage == 0; + }; + + // Prologue: decode the first stage so iter 0's wgmma can + // consume it without waiting. + wait_and_decode(stage_idx, phase); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + advance_lead(); // lead now points one stage ahead of wgmma + + #pragma unroll 8 + 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); + + // smem_b[stage_idx] is already decoded (prologue or the + // previous iter's hide-decode) and made visible by the + // matching NamedBarrier::sync below. + // smem_a[stage_idx] / smem_sfa[stage_idx] are guaranteed + // ready by the wait_and_decode that targeted stage_idx. + + float scale_b_0_regs[WGMMA::kNumAccum / 4]; + float scale_b_1_regs[WGMMA::kNumAccum / 4]; + if (do_wgmma_store) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t n_idx = n_block_idx * BLOCK_N + i * 8 + col_idx * 2; + scale_b_0_regs[i] = load_sfb(n_idx, k_block_idx); + scale_b_1_regs[i] = load_sfb(n_idx + 1, k_block_idx); + } + } + + const bool has_next = (k_block_idx + 1) < num_total_k_blocks; + + // TODO: remove some useless computation for unaligned Ms + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto m_offset = local_idx * WAVE_BLOCK_M; + + // Read A scales + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + auto scale_a_0 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_0 + m_offset) : 0; + auto scale_a_1 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_1 + m_offset) : 0; + + // Commit WGMMA instructions + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + + // Hide decode of next stage under this wgmma's async + // work. Run only on the last wave so all warps follow + // the same wait/decode/sync sequence per K iter. + const bool is_last_wave = (local_idx == BLOCK_M / WAVE_BLOCK_M - 1); + if (is_last_wave and has_next) { + wait_and_decode(lead_stage, lead_phase); + } + + ptx::warpgroup_wait<0>(); + + // Notify barrier arrival at the last warpgroup wave + if (is_last_wave) + empty_barrier_arrive(); + + // Skip promotion for the unfilled parts + if (not do_wgmma_store) + continue; + + // Promote with scales + // NOTES: making it as predicates is very important for performance, comparing to two loops + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float scale_b_0 = scale_b_0_regs[i]; + const float scale_b_1 = scale_b_1_regs[i]; + 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]; + } + } + + // Publish next iter's decoded smem_b to the wgmma + // async-proxy. The packed-B slot is released together + // with A/SFA via the unified empty_barrier_arrive() above. + if (has_next) { + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + advance_lead(); + } + } + }); + } else { + #pragma unroll + 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(); + } + } + + // Psum layout can have a final partial M tile. TMA store writes a + // full BLOCK_M tile and may go out of the tensor-map bounds, so use + // a guarded scalar store for this layout. + if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + const uint32_t psum_store_end = + scheduler.current_group_idx + 1 < kNumGroups ? + math::align(scheduler.current_psum_m, BLOCK_M) : scheduler.current_psum_m; + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + const uint32_t row_0 = m_block_idx * BLOCK_M + m_offset + r_0; + const uint32_t row_1 = m_block_idx * BLOCK_M + m_offset + r_1; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t base_col = epilogue_type_t::template apply_index_n<8>( + n_block_idx * BLOCK_N + i * 8); + const uint32_t col = base_col + col_idx * 2; + const bool row_0_valid = (row_0 >= scheduler.last_psum_m and row_0 < psum_store_end); + const bool row_1_valid = (row_1 >= scheduler.last_psum_m and row_1 < psum_store_end); + if (row_0_valid and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = __float2bfloat16_rn(shifted_accum[i * 4 + 0]); + if (row_0_valid and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = __float2bfloat16_rn(shifted_accum[i * 4 + 1]); + if (row_1_valid and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = __float2bfloat16_rn(shifted_accum[i * 4 + 2]); + if (row_1_valid and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = __float2bfloat16_rn(shifted_accum[i * 4 + 3]); + } + } + __syncwarp(); + continue; + } + + // TMA checks + 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"); + + // Skip WGMMA store for the unfilled parts + if (not do_wgmma_store) + continue; + + // Wait last TMA store to be finished + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + // Write back to shared memory using STSM and issue TMA stores + DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); + #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) { + // Swizzle or padding into the correct address + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + // Calculate the swizzling atom offset and in-atom offset + constexpr uint32_t kNumBankGroupBytes = 16; + auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8); + + // Calculate the index of the bank group to be written in the atom + auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes); + + // Reshape the atom in another view and swizzle + // - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)` + // - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)` + 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); + + // Add back into the base pointer + // NOTES: think twice before modifying this, as changes may affect the number of instructions + smem_ptr = reinterpret_cast(smem_d) + // Base pointer + warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp offset + m_offset * kSwizzleDMode + // Wave offset + atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom offset (constants) + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset + } else { + // No swizzling, just padding + smem_ptr = reinterpret_cast(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8); + } + + // NOTES: only 16 lanes' addresses are used + 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); + + // Use TMA store to write back to global memory + // TODO: compatible with FP32 output + constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked; + DG_STATIC_ASSERT(kNumWGMMAStoreThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks"); + 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 = epilogue_type_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); + auto m_idx = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); + if constexpr (kGemmType == GemmType::Batched) { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, + n_idx, m_idx, scheduler.current_group_idx); + } else { + 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 support sm_90a"); +#endif +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh new file mode 100644 index 0000000000..ab68ef4013 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh @@ -0,0 +1,3236 @@ +#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 +#include + +namespace deep_gemm { + +// SM90 FP8 x FP4 GEMM, RS-mode variant. +// +// S2 stage: this kernel is a verbatim copy of `sm90_fp8_fp4_gemm_1d2d_impl`, +// renamed to `sm90_fp8_fp4_gemm_1d2d_rs_impl` so it gets its own JIT artifact. +// Functionally identical to the SS variant; later steps will replace the +// inner WGMMA loop with RS-mode (register A) WGMMA + register-resident FP4 +// dequant (using `FP8MMASelectorRS`, ldmatrix, prmt+lop3) without touching +// the SS variant or the contiguous code path. + +// ============================================================================= +// S3 prep (M1): RS-mode utility helpers. +// These helpers are intentionally not yet called by `sm90_fp8_fp4_gemm_1d2d_rs_impl`; +// they are introduced as a self-contained, side-effect-free building block so +// the second-round kernel rewrite can plug them in without further header +// surgery. Numerical equivalence with the SS variant is preserved this round. +// ============================================================================= +namespace fp4_rs_detail { + +template +struct FinalAccumStorage { + using type = float; +}; + +template <> +struct FinalAccumStorage { + using type = nv_bfloat162; +}; + +template +CUTLASS_DEVICE void final_accum_init(dtype_t* base, uint32_t idx) { + if constexpr (kBF16FinalAccum) { + base[idx] = __float22bfloat162_rn({0.0f, 0.0f}); + } else { + base[idx] = 0.0f; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair(dtype_t* base, uint32_t pair_idx, + float scale_0, float value_0, + float scale_1, float value_1) { + if constexpr (kBF16PromoteMath) { + const nv_bfloat162 scale = __float22bfloat162_rn({scale_0, scale_1}); + const nv_bfloat162 value = __float22bfloat162_rn({value_0, value_1}); + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + const nv_bfloat162 out = __hfma2(scale, value, dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = out; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(out); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0, value_0, dst.x); + dst.y = __fmaf_rn(scale_1, value_1, dst.y); + } else { + dst.x += scale_0 * value_0; + dst.y += scale_1 * value_1; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + if constexpr (kFmaPromote) { + base[idx + 0] = __fmaf_rn(scale_0, value_0, base[idx + 0]); + base[idx + 1] = __fmaf_rn(scale_1, value_1, base[idx + 1]); + } else { + base[idx + 0] += scale_0 * value_0; + base[idx + 1] += scale_1 * value_1; + } + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair2(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b) { + if constexpr (kBF16PromoteMath) { + const nv_bfloat162 scale_a = __float22bfloat162_rn({scale_0_a, scale_1_a}); + const nv_bfloat162 value_a = __float22bfloat162_rn({value_0_a, value_1_a}); + const nv_bfloat162 scale_b = __float22bfloat162_rn({scale_0_b, scale_1_b}); + const nv_bfloat162 value_b = __float22bfloat162_rn({value_0_b, value_1_b}); + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(scale_a, value_a, dst); + dst = __hfma2(scale_b, value_b, dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0_a, value_0_a, dst.x); + dst.y = __fmaf_rn(scale_1_a, value_1_a, dst.y); + dst.x = __fmaf_rn(scale_0_b, value_0_b, dst.x); + dst.y = __fmaf_rn(scale_1_b, value_1_b, dst.y); + } else { + dst.x += scale_0_a * value_0_a; + dst.y += scale_1_a * value_1_a; + dst.x += scale_0_b * value_0_b; + dst.y += scale_1_b * value_1_b; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(scale_0_a, value_0_a, dst_0); + dst_1 = __fmaf_rn(scale_1_a, value_1_a, dst_1); + dst_0 = __fmaf_rn(scale_0_b, value_0_b, dst_0); + dst_1 = __fmaf_rn(scale_1_b, value_1_b, dst_1); + } else { + dst_0 += scale_0_a * value_0_a; + dst_1 += scale_1_a * value_1_a; + dst_0 += scale_0_b * value_0_b; + dst_1 += scale_1_b * value_1_b; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b, + float scale_0_c, float value_0_c, + float scale_1_c, float value_1_c, + float scale_0_d, float value_0_d, + float scale_1_d, float value_1_d) { + if constexpr (kBF16PromoteMath) { + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(__float22bfloat162_rn({scale_0_a, scale_1_a}), + __float22bfloat162_rn({value_0_a, value_1_a}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_b, scale_1_b}), + __float22bfloat162_rn({value_0_b, value_1_b}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_c, scale_1_c}), + __float22bfloat162_rn({value_0_c, value_1_c}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_d, scale_1_d}), + __float22bfloat162_rn({value_0_d, value_1_d}), dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0_a, value_0_a, dst.x); + dst.y = __fmaf_rn(scale_1_a, value_1_a, dst.y); + dst.x = __fmaf_rn(scale_0_b, value_0_b, dst.x); + dst.y = __fmaf_rn(scale_1_b, value_1_b, dst.y); + dst.x = __fmaf_rn(scale_0_c, value_0_c, dst.x); + dst.y = __fmaf_rn(scale_1_c, value_1_c, dst.y); + dst.x = __fmaf_rn(scale_0_d, value_0_d, dst.x); + dst.y = __fmaf_rn(scale_1_d, value_1_d, dst.y); + } else { + dst.x += scale_0_a * value_0_a; + dst.y += scale_1_a * value_1_a; + dst.x += scale_0_b * value_0_b; + dst.y += scale_1_b * value_1_b; + dst.x += scale_0_c * value_0_c; + dst.y += scale_1_c * value_1_c; + dst.x += scale_0_d * value_0_d; + dst.y += scale_1_d * value_1_d; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(scale_0_a, value_0_a, dst_0); + dst_1 = __fmaf_rn(scale_1_a, value_1_a, dst_1); + dst_0 = __fmaf_rn(scale_0_b, value_0_b, dst_0); + dst_1 = __fmaf_rn(scale_1_b, value_1_b, dst_1); + dst_0 = __fmaf_rn(scale_0_c, value_0_c, dst_0); + dst_1 = __fmaf_rn(scale_1_c, value_1_c, dst_1); + dst_0 = __fmaf_rn(scale_0_d, value_0_d, dst_0); + dst_1 = __fmaf_rn(scale_1_d, value_1_d, dst_1); + } else { + dst_0 += scale_0_a * value_0_a; + dst_1 += scale_1_a * value_1_a; + dst_0 += scale_0_b * value_0_b; + dst_1 += scale_1_b * value_1_b; + dst_0 += scale_0_c * value_0_c; + dst_1 += scale_1_c * value_1_c; + dst_0 += scale_0_d * value_0_d; + dst_1 += scale_1_d * value_1_d; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4x2(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b, + float scale_0_c, float value_0_c, + float scale_1_c, float value_1_c, + float scale_0_d, float value_0_d, + float scale_1_d, float value_1_d, + float scale_2_a, float value_2_a, + float scale_3_a, float value_3_a, + float scale_2_b, float value_2_b, + float scale_3_b, float value_3_b, + float scale_2_c, float value_2_c, + float scale_3_c, float value_3_c, + float scale_2_d, float value_2_d, + float scale_3_d, float value_3_d) { + final_accum_promote_pair4( + base, pair_idx, + scale_0_a, value_0_a, scale_1_a, value_1_a, + scale_0_b, value_0_b, scale_1_b, value_1_b, + scale_0_c, value_0_c, scale_1_c, value_1_c, + scale_0_d, value_0_d, scale_1_d, value_1_d); + final_accum_promote_pair4( + base, pair_idx + 1, + scale_2_a, value_2_a, scale_3_a, value_3_a, + scale_2_b, value_2_b, scale_3_b, value_3_b, + scale_2_c, value_2_c, scale_3_c, value_3_c, + scale_2_d, value_2_d, scale_3_d, value_3_d); +} + +CUTLASS_DEVICE void keep_float_live(float value) { + asm volatile("" :: "f"(value) : "memory"); +} + +struct FP8MMAF16AccumM64N8K32RS { + static constexpr uint32_t M = 64; + static constexpr uint32_t N = 8; + static constexpr uint32_t K = 32; + static constexpr uint32_t kNumAccum = 2; + + template + CUTLASS_DEVICE static void wgmma(uint32_t const* a, GmmaDescriptor const& desc, + uint32_t* d, bool scale_d) { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %7, 0;\n" + " wgmma.mma_async.sync.aligned.m64n8k32.f16.e4m3.e4m3 " + " {%0, %1}, {%2, %3, %4, %5}, %6, p, 1, 1;\n" + "}\n" + : "+r"(d[0]), "+r"(d[1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "l"(desc.desc_), "r"(static_cast(scale_d))); + } +}; + +CUTLASS_DEVICE void unpack_f16_accum_m64n8(uint32_t const* src, float* dst) { + #pragma unroll + for (uint32_t i = 0; i < FP8MMAF16AccumM64N8K32RS::kNumAccum; ++ i) { + const half2 value_h = *reinterpret_cast(src + i); + const float2 value_f = __half22float2(value_h); + dst[i * 2 + 0] = value_f.x; + dst[i * 2 + 1] = value_f.y; + } +} + +CUTLASS_DEVICE float scale_float_by_pow2(float value, float pow2_scale) { + uint32_t value_bits = *reinterpret_cast(&value); + const uint32_t scale_bits = *reinterpret_cast(&pow2_scale); + const int32_t exp_shift = static_cast((scale_bits >> 23) & 0xffu) - 127; + value_bits = static_cast(static_cast(value_bits) + (exp_shift << 23)); + return *reinterpret_cast(&value_bits); +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4_split_scale(dtype_t* base, uint32_t pair_idx, + float scale_a_0, float scale_a_1, + float scale_b_a, float scale_b_b, + float scale_b_c, float scale_b_d, + float value_0_a, float value_1_a, + float value_0_b, float value_1_b, + float value_0_c, float value_1_c, + float value_0_d, float value_1_d) { + auto prod_0 = [&](float scale_b) { + return kScaleBPow2Promote ? scale_float_by_pow2(scale_a_0, scale_b) : scale_a_0 * scale_b; + }; + auto prod_1 = [&](float scale_b) { + return kScaleBPow2Promote ? scale_float_by_pow2(scale_a_1, scale_b) : scale_a_1 * scale_b; + }; + if constexpr (kBF16PromoteMath) { + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_a), prod_1(scale_b_a)}), + __float22bfloat162_rn({value_0_a, value_1_a}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_b), prod_1(scale_b_b)}), + __float22bfloat162_rn({value_0_b, value_1_b}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_c), prod_1(scale_b_c)}), + __float22bfloat162_rn({value_0_c, value_1_c}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_d), prod_1(scale_b_d)}), + __float22bfloat162_rn({value_0_d, value_1_d}), dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(prod_0(scale_b_a), value_0_a, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_a), value_1_a, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_b), value_0_b, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_b), value_1_b, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_c), value_0_c, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_c), value_1_c, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_d), value_0_d, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_d), value_1_d, dst.y); + } else { + dst.x += prod_0(scale_b_a) * value_0_a; + dst.y += prod_1(scale_b_a) * value_1_a; + dst.x += prod_0(scale_b_b) * value_0_b; + dst.y += prod_1(scale_b_b) * value_1_b; + dst.x += prod_0(scale_b_c) * value_0_c; + dst.y += prod_1(scale_b_c) * value_1_c; + dst.x += prod_0(scale_b_d) * value_0_d; + dst.y += prod_1(scale_b_d) * value_1_d; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(prod_0(scale_b_a), value_0_a, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_a), value_1_a, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_b), value_0_b, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_b), value_1_b, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_c), value_0_c, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_c), value_1_c, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_d), value_0_d, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_d), value_1_d, dst_1); + } else { + dst_0 += prod_0(scale_b_a) * value_0_a; + dst_1 += prod_1(scale_b_a) * value_1_a; + dst_0 += prod_0(scale_b_b) * value_0_b; + dst_1 += prod_1(scale_b_b) * value_1_b; + dst_0 += prod_0(scale_b_c) * value_0_c; + dst_1 += prod_1(scale_b_c) * value_1_c; + dst_0 += prod_0(scale_b_d) * value_0_d; + dst_1 += prod_1(scale_b_d) * value_1_d; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE float final_accum_load_scalar(const dtype_t* base, uint32_t idx) { + if constexpr (kBF16FinalAccum) { + const float2 value = __bfloat1622float2(base[idx / 2]); + return (idx % 2 == 0) ? value.x : value.y; + } else { + return base[idx]; + } +} + +template +CUTLASS_DEVICE nv_bfloat162 final_accum_load_pair_bf16(const dtype_t* base, uint32_t pair_idx) { + if constexpr (kBF16FinalAccum) { + return base[pair_idx]; + } else { + return __float22bfloat162_rn({base[pair_idx * 2 + 0], base[pair_idx * 2 + 1]}); + } +} + +// `ldmatrix.sync.aligned.x4.m8n8.trans.shared.b16`: load 4 8x8 b16 matrices +// from swizzled smem into 4 b32 registers per lane, transposed. The transpose +// option is the key to feeding RS-mode WGMMA's A operand layout when the +// source tile is laid out K-major in smem (which is the only layout the +// dequant step can produce cheaply from packed FP4 input). +// +// Note: WGMMA RS-mode A still consumes 8-bit (E4M3) pairs, so we phrase the +// granularity as b16 ldmatrix and let the upper layer treat each b32 register +// as 4 packed E4M3 lanes. +struct SM90_U32x4_LDSM_T { + CUTLASS_DEVICE static void + copy(uint32_t& dst_0, uint32_t& dst_1, uint32_t& dst_2, uint32_t& dst_3, void* smem_src) { + asm volatile( + "ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst_0), "=r"(dst_1), "=r"(dst_2), "=r"(dst_3) + : "l"(__cvta_generic_to_shared(smem_src))); + } +}; + +template +CUTLASS_DEVICE uint32_t permute_col(const uint32_t row, const uint32_t col) { + constexpr uint32_t strd = 128 / (COL_BYTES < 128u ? COL_BYTES : 128u); + return ((col / NV) ^ (row % 8 / strd)) * NV; +} + +template +struct SM90_U32x2_STSM_T { + CUTLASS_DEVICE static void + copy(dtype_t src_0, dtype_t src_1, void* smem_dst) { + DG_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src[2] = {*reinterpret_cast(&src_0), *reinterpret_cast(&src_1)}; + asm volatile("stmatrix.sync.aligned.x2.m8n8.shared.b16.trans [%0], {%1, %2};\n" + :: "l"(__cvta_generic_to_shared(smem_dst)), "r"(src[0]), "r"(src[1])); + } +}; + +// Decode a single FP4 (E2M1) code (low 4 bits) to its E4M3 byte +// representation. Same byte-LUT formulation as the scalar path inside the SS +// kernel, hoisted to namespace scope so it can also be called from +// register-level dequant routines. +// mag {0..7} -> {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c} +// sign bit copied from FP4 bit 3 to E4M3 bit 7 (FP4 -0 maps to E4M3 -0; +// WGMMA treats -0 as 0, so this is numerically safe and saves a branch). +CUTLASS_DEVICE uint32_t fp4_to_e4m3_byte(uint32_t code) { + constexpr uint32_t LUT_LO = 0x3c383000u; + constexpr uint32_t LUT_HI = 0x4c484440u; + const uint32_t mag = code & 0x07u; + const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; + const uint32_t sign = (code & 0x08u) << 4; // 0 or 0x80 + return mag_byte | sign; +} + +// Decode 8 packed FP4 codes (one 32-bit word holding 8 nibbles) into 8 E4M3 +// bytes packed into a single 64-bit value (low byte = code 0). +// +// This is the register-resident analogue of the SS variant's +// `fp4_pair_to_e4m3_pair` chained over 4 bytes. It is the building block the +// second-round kernel will use after `ldmatrix.x4.trans` to materialise A +// operand registers for RS-mode WGMMA without round-tripping through smem. +// +// Implementation kept straightforward (LUT per nibble). PTX-level +// vectorisation via `prmt`/`lop3` is left for a follow-up micro-optimisation +// in round 2 once the rest of the data path is verified correct. +CUTLASS_DEVICE uint64_t fp4x8_to_e4m3x8(uint32_t packed) { + uint64_t out = 0; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t nib = (packed >> (i * 4)) & 0x0fu; + out |= static_cast(fp4_to_e4m3_byte(nib)) << (i * 8); + } + return out; +} + +CUTLASS_DEVICE void fast_fp4_to_e4m3_convert(uint32_t outputs[2], uint32_t input) { + const uint64_t decoded = fp4x8_to_e4m3x8(input); + outputs[0] = static_cast(decoded); + outputs[1] = static_cast(decoded >> 32); +} + +CUTLASS_DEVICE uint32_t fp4x4_to_e4m3x4(uint32_t packed) { + constexpr uint32_t pos0 = 0x3c383000u; + constexpr uint32_t pos1 = 0x4c484440u; + const uint32_t lut_idx = packed & 0x7777u; + const uint32_t sign_shifted = packed << 4; + uint32_t mag_bytes, sign_bytes; + asm volatile( + "{\n" + " prmt .b32 %0, %3, %4, %2;\n" + " prmt .b32 %1, %5, %6, 0xd9c8;\n" + "}\n" + : "=r"(mag_bytes), "=r"(sign_bytes) + : "r"(lut_idx), "r"(pos0), "r"(pos1), "r"(sign_shifted), "r"(packed)); + uint32_t out; + asm volatile( + "{\n" + " lop3.b32 %0, %1, %2, 0x80808080, 0xf8;\n" + "}\n" + : "=r"(out) + : "r"(mag_bytes), "r"(sign_bytes)); + return out; +} + +CUTLASS_DEVICE int32_t pow2_scale_to_exp_shift(float scale) { + const uint32_t scale_bits = *reinterpret_cast(&scale); + return static_cast((scale_bits >> 23) & 0xffu) - 127; +} + +struct ScaledE4M3Lut { + uint32_t lo; + uint32_t hi; +}; + +CUTLASS_DEVICE ScaledE4M3Lut make_scaled_e4m3_lut(uint32_t exp_offset) { + // Random FP4 per-32 scales are overwhelmingly ceil-pow2(max(abs(x))/6) + // in {2^-1, 1}, which maps to exp_offset {5, 6}. Fast-path those to + // avoid the dynamic IMAD chain in the WGMMA issue loop. + if (exp_offset == 5u) + return {0x34302800u, 0x44403c38u}; + if (exp_offset == 6u) + return {0x3c383000u, 0x4c484440u}; + const uint32_t exp_offset_buffer1 = + exp_offset * 0x08080800u + (exp_offset ? 0xfffffc00u : 0u); + const uint32_t exp_offset_buffer2 = exp_offset * 0x08080808u; + constexpr uint32_t mantissa_lo = 0x0c080400u; + constexpr uint32_t mantissa_hi = 0x1c181410u; + return {mantissa_lo + exp_offset_buffer1, mantissa_hi + exp_offset_buffer2}; +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_lut(uint32_t packed, ScaledE4M3Lut lut) { + const uint32_t lut_idx = packed & 0x7777u; + const uint32_t sign_shifted = packed << 4; + uint32_t mantissa_bytes, sign_bytes; + asm volatile( + "{\n" + " prmt .b32 %0, %3, %4, %2;\n" + " prmt .b32 %1, %5, %6, 0xd9c8;\n" + "}\n" + : "=r"(mantissa_bytes), "=r"(sign_bytes) + : "r"(lut_idx), "r"(lut.lo), "r"(lut.hi), "r"(sign_shifted), "r"(packed)); + return (sign_bytes & 0x80808080u) | mantissa_bytes; +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_offset(uint32_t packed, uint32_t exp_offset) { + return fp4x4_to_scaled_e4m3x4_lut(packed, make_scaled_e4m3_lut(exp_offset)); +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_humming(uint32_t packed_nibbles, uint32_t exp_offset) { + const uint32_t packed = + (packed_nibbles & 0x000fu) | + ((packed_nibbles & 0x00f0u) << 4) | + ((packed_nibbles & 0x0f00u) << 8) | + ((packed_nibbles & 0xf000u) << 12); + + const uint32_t exp_offset_buffer1 = exp_offset * 0x08080800u + (exp_offset ? 0xfffffc00u : 0u); + const uint32_t exp_offset_buffer2 = exp_offset * 0x08080808u; + const uint32_t exp_offsets0 = __byte_perm(exp_offset_buffer1, exp_offset_buffer2, packed); + const uint32_t exp_offsets1 = __byte_perm(exp_offset_buffer1, exp_offset_buffer2, packed >> 16); + + uint32_t scaled_mag; + asm volatile( + "{\n" + " lop3.b32 %0, %1, 0x80808080, %2, 0xf8;\n" + "}\n" + : "=r"(scaled_mag) + : "r"(packed << 4), "r"((packed & 0x07070707u) << 2)); + return scaled_mag + __byte_perm(exp_offsets0, exp_offsets1, 0x6420); +} + +} // namespace fp4_rs_detail + +template +CUTLASS_DEVICE void dispatch_num_former_iters_rs(uint32_t num_former_iters, const func_t& func) { + if (num_former_iters == kNumFormerIters) { + func(cute::Int{}); + return; + } + + if constexpr (kNumFormerIters + kGap <= kEnd) + dispatch_num_former_iters_rs(num_former_iters, func); +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, kLaunchBoundsMinBlocks) void +sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, + nv_bfloat16* gmem_d_ptr, + 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__) + // Scaling checks + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT(kScaleBGranK == 32 or kScaleBGranK == 128, + "Only support per-32 or per-128-channel FP4 scaling"); + DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or (kFuseScaleBDecode and kScaleBGranK == 32), + "Packed UE8M0 SFB is only supported by fused per-32 scale_b decode"); + DG_STATIC_ASSERT(kScaleBGranK == 128 or not kOverlapPromote, + "DG_W4_OVERLAP_PROMOTE does not support per-32 FP4 scaling"); + DG_STATIC_ASSERT(kScaleBGranK == 128 or kScaleKGroup == 1, + "DG_W4_SCALE_K_GROUP does not support per-32 FP4 scaling"); + DG_STATIC_ASSERT(kScaleKGroup == 1 or kScaleKGroup == 2 or kScaleKGroup == 4, + "DG_W4_SCALE_K_GROUP only supports 1/2/4"); + DG_STATIC_ASSERT(not kBIsInt4Sym, "RS-mode FP8xFP4 kernel does not support INT4-sym B"); + DG_STATIC_ASSERT(not (kScaleBBF16 and kScaleBE8M0), "Scale-B cannot be both BF16 and E8M0"); + DG_STATIC_ASSERT((not kScaleBBF16 and not kScaleBE8M0) or (kScaleBDirectLoad and kScaleBGranK == 32), + "Compressed Scale-B dtypes are only supported by direct-load per-32 path"); + DG_STATIC_ASSERT(kFuseScaleBDecodeAssumeExp == 0 or kFuseScaleBDecodeAssumeExp == 5 or + kFuseScaleBDecodeAssumeExp == 6, + "DG_W4_FUSE_SCALE_B_DECODE_ASSUME_EXP only supports 0/5/6"); + DG_STATIC_ASSERT( + math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or + (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); + // Types + using WGMMA = typename mma::sm90::FP8MMASelectorRS::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M <= 256, "Invalid RS-mode MMA N size"); + DG_STATIC_ASSERT(BLOCK_N % WGMMA::M == 0, "RS-mode swap_ab requires BLOCK_N to be tiled by 64-row WGMMA M"); + + // Overwrite shape constants if the compiler gives + 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; + + // Shared memory + static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0); + static constexpr uint32_t SMEM_D_ROWS = BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M; + static constexpr uint32_t SMEM_D_SIZE = math::constexpr_align(SMEM_D_ROWS * BLOCK_N * static_cast(sizeof(__nv_bfloat16)), 1024u); + static constexpr uint32_t SMEM_A_TMA_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = + (BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M) * BLOCK_K * sizeof(__nv_fp8_e4m3); + // Packed FP4 B is loaded by TMA into a separate buffer; each row is BLOCK_K / 2 bytes. + static constexpr uint32_t BLOCK_K_PACKED = BLOCK_K / 2; + static constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = BLOCK_N * BLOCK_K_PACKED; + static constexpr uint32_t kNumRSMathWGs = kNumMathThreads / 128; + static constexpr uint32_t SMEM_SFA_TMA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = + (BLOCK_M < WGMMA::M ? WGMMA::M : 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 SCALE_B_ELEMENT_SIZE = + kScaleBE8M0 ? static_cast(sizeof(uint8_t)) : + (kScaleBBF16 ? static_cast(sizeof(nv_bfloat16)) : + static_cast(sizeof(float))); + const uint32_t shape_k_scales_a = math::ceil_div(shape_k, BLOCK_K); + const uint32_t shape_k_scales_b = math::ceil_div(shape_k, kScaleBGranK); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / SCALE_B_ELEMENT_SIZE); + // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough + // smem_d capacity, so they fall back to a separate SFB region. + const uint32_t smem_sfb_bytes = kScaleBGranK == 32 ? + math::align((BLOCK_K / kScaleBGranK) * + (kFuseScaleBDecode ? math::ceil_div(BLOCK_N, 4u) : BLOCK_N) * + sizeof(float), 16u) : + math::align(shape_k_scales_b * BLOCK_N * sizeof(float), 16u); + // NOTES: Make sure we have enough shared memory for WGMMA padding + static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); + DG_STATIC_ASSERT(WGMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE, "Memory Out of bound for WGMMA"); + + // Configs + 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(); + constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); + constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; + constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; + constexpr bool kParallelNWavesEnabled = + kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; + DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, + "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); + constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + constexpr uint32_t kEmptyBarrierMathWarps = + kParallelNWavesEnabled ? kBaseWGMMAStoreThreads / 32 : kNumMathThreads / 32; + + // Prefetch TMA descriptors at the very beginning + 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(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + // Data on shared memory + 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); + }); + constexpr uint32_t SMEM_B_PACKED_OFFSET = SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE; + auto smem_b_packed = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_B_PACKED_OFFSET + i * SMEM_B_PACKED_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumStages * SMEM_B_PACKED_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); + }); + constexpr uint32_t SMEM_SFB_OFFSET = SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE; + // Prefer aliasing SFB onto smem_d. Small-M tiles usually have too little + // smem_d, and SHAPE_K can be runtime-only, so choose the separate SFB region + // dynamically to match the host-side smem_size calculation. + const bool use_separate_sfb = smem_sfb_bytes > SMEM_D_SIZE; + auto smem_sfb = reinterpret_cast(smem_buffer + (use_separate_sfb ? SMEM_SFB_OFFSET : 0)); + auto smem_sfb_exp = reinterpret_cast(smem_sfb); + + // Fill barriers. + // After the A/packed-B barrier merge there is only one set of full/empty barriers. + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SFB_OFFSET + (use_separate_sfb ? smem_sfb_bytes : 0)); + 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; }); + + // Initialize barriers + DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast"); + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kEmptyBarrierMathWarps); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Register reconfigurations + constexpr uint32_t kNumTMARegisters = 40; + constexpr uint32_t kDefaultMathRegisters = kNumMathThreads == 128 ? 248 : 232; + constexpr uint32_t kNumMathRegisters = kMathRegCap == 0 ? kDefaultMathRegisters : kMathRegCap; + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // `gmem_b_ptr` is no longer used: B is now loaded by TMA into `smem_b_packed`. + (void)gmem_b_ptr; + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + uint32_t simple_sched_linear_idx = blockIdx.x; + constexpr bool kUseSmallMSimpleSched = + kSmallMSimpleSched and kGemmType == GemmType::MGroupedMasked and BLOCK_M <= 8 and kNumTMAMulticast == 1; + auto get_next_block = [&]() { + if constexpr (kUseSmallMSimpleSched) { + const uint32_t n_blocks = math::ceil_div(shape_n, BLOCK_N); + const uint32_t total_blocks = n_blocks * kNumGroups; + while (simple_sched_linear_idx < total_blocks) { + scheduler.current_group_idx = simple_sched_linear_idx / n_blocks; + n_block_idx = simple_sched_linear_idx - scheduler.current_group_idx * n_blocks; + m_block_idx = 0; + simple_sched_linear_idx += gridDim.x; + if (scheduler.is_computation_valid(m_block_idx, 0)) + return true; + } + return false; + } else { + return scheduler.get_next_block(m_block_idx, n_block_idx); + } + }; + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedMasked) { + return scheduler.current_group_idx; + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // Pipeline and TMA phases (single shared pipeline for A/SFA/packed-B after the barrier merge) + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used. + // We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32`. + if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (get_next_block()) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + 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)) { + // Wait consumer release for the (now merged) A/SFA/packed-B slot + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Issue TMA A + constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); + const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); + + 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.template get_global_idx(shape_m, BLOCK_M, m_block_idx), + num_tma_multicast_a, batch_idx); + if constexpr (not kScaleAStub) { + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales_a, 1, k_block_idx), + num_tma_multicast_a); + } + + if constexpr (not kWeightStub) { + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier. + const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; + tma::copy(&tensor_map_b, &full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), + k_idx_packed, + scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b, batch_idx); + } + + constexpr uint32_t kExpectedTxBytes = SMEM_A_TMA_SIZE_PER_STAGE + + (kWeightStub ? 0 : SMEM_B_PACKED_SIZE_PER_STAGE) + + (kScaleAStub ? 0 : SMEM_SFA_TMA_SIZE_PER_STAGE); + full_barrier.arrive_and_expect_tx(kExpectedTxBytes); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; ++ i) { + empty_barriers[stage_idx]->wait(phase ^ 1); + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto warp_in_wg = warp_idx % 4; + + auto a_desc = mma::sm90::make_smem_desc(smem_a[0], 1); + const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0); + + constexpr uint32_t kLdmatrixVecBytes = 16 / sizeof(__nv_fp8_e4m3); + + // Persistently schedule over blocks + while (get_next_block()) { + const uint32_t current_group_idx = get_current_group_idx(); + + if constexpr (kScaleBGranK == 128) { + // Cooperatively prefetch the SFB tile for this block from gmem to smem. + // Layout in smem: [shape_k_scales_b, BLOCK_N] (k outer, n inner). + // Out-of-bound n is filled with 1.0f to keep `n_idx >= shape_n` neutral. + // + // Optimization: use cp.async to copy gmem->smem directly (no register + // round-trip). For MN-major (n innermost in gmem) we can issue 16-byte + // (float4) cp.async per thread, cutting #instructions by 4x. K-major + // and the OOB tail use scalar 4-byte cp.async / st.shared. + const uint32_t n_block_base = n_block_idx * BLOCK_N; + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = shape_k_scales_b * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + // Issue cp.async (16B per thread) for fully in-bounds vectors; + // fall back to scalar st.shared with 1.0f-fill for the tail. + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_idx = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + k_idx * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + // Tail: at least one element is OOB; use scalar st with 1.0f fill. + float4 vals; + const float* sfb_row = sfb_base + k_idx * aligned_shape_n_sfb + n_idx; + vals.x = (n_idx + 0 < shape_n) ? *(sfb_row + 0) : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? *(sfb_row + 1) : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? *(sfb_row + 2) : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? *(sfb_row + 3) : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + // K-major: sfb is strided along n; cannot easily vectorize across n. + // Use scalar 4B cp.async for in-bounds, st.shared with 1.0f for OOB. + const uint32_t total = shape_k_scales_b * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_idx = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_idx; + ptx::cp_async_4(smem_dst, gmem_src); + } + } + } + // Commit and wait for all cp.async issued above; pair with the + // existing NamedBarrier so the smem cache is visible to all + // math warps before they enter the K loop. + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + + auto cache_sfb_k32 = [&](uint32_t k_block_idx) { + if constexpr (kScaleBGranK == 32 and not kScaleBStub and not kFuseScaleBDecodeStub and not kScaleBDirectLoad) { + const uint32_t n_block_base = n_block_idx * BLOCK_N; + const uint32_t scale_k_base = k_block_idx * (BLOCK_K / kScaleBGranK); + constexpr uint32_t kScaleRows = BLOCK_K / kScaleBGranK; + if constexpr (kFuseScaleBDecode) { + DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or kMajorSFB == cute::UMMA::Major::MN, + "Packed UE8M0 SFB path expects MN-major transformed scale layout"); + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for packed K/32 SFB exp cache"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = kScaleRows * kVecsPerRow; + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_off = i / kVecsPerRow; + const uint32_t n_off = (i % kVecsPerRow) * kVec; + const uint32_t n_idx = n_block_base + n_off; + uint32_t packed_offsets = 0; + #pragma unroll + for (uint32_t j = 0; j < kVec; ++ j) { + uint32_t exp_offset = 6; + if constexpr (kScaleBPackedUE8M0) { + if (n_idx + j < shape_n) { + const uint32_t packed_shape_k_scales_b = math::ceil_div(shape_k_scales_b, 4u); + const uint32_t packed_k_idx = (scale_k_base + k_off) / 4u; + const uint32_t byte_idx = (scale_k_base + k_off) % 4u; + const uint32_t* sfb_packed = reinterpret_cast(sfb); + const uint32_t packed_scale = sfb_packed[ + current_group_idx * aligned_shape_n_sfb * packed_shape_k_scales_b + + packed_k_idx * aligned_shape_n_sfb + n_idx + j]; + const uint32_t e8m0_exp = (packed_scale >> (byte_idx * 8)) & 0xffu; + exp_offset = (e8m0_exp > 121u) ? (e8m0_exp - 121u) : 0u; + } + } else { + float val = 1.0f; + if (n_idx + j < shape_n) { + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + const float* ptr = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx + j; + val = *ptr; + } else { + const float* ptr = sfb + + current_group_idx * shape_n * shape_k_scales_b + + (n_idx + j) * shape_k_scales_b + scale_k_base + k_off; + val = *ptr; + } + } + const int32_t exp_shift = fp4_rs_detail::pow2_scale_to_exp_shift(val); + exp_offset = static_cast(exp_shift + 6) & 0xffu; + } + packed_offsets |= exp_offset << (j * 8); + } + ptx::st_shared(smem_sfb_exp + k_off * kVecsPerRow + n_off / kVec, packed_offsets); + } + } else if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized K/32 SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = kScaleRows * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_off = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_off * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + float4 vals; + const float* sfb_row = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx; + vals.x = (n_idx + 0 < shape_n) ? *(sfb_row + 0) : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? *(sfb_row + 1) : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? *(sfb_row + 2) : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? *(sfb_row + 3) : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + const uint32_t total = kScaleRows * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_off = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_off * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + scale_k_base + k_off; + ptx::cp_async_4(smem_dst, gmem_src); + } + } + } + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + }; + + auto load_sfb = [&](uint32_t n_idx, uint32_t k_block_idx) { + if (n_idx >= shape_n) + return 1.0f; + if constexpr (kScaleBDirectLoad and kScaleBGranK == 32) { + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const uint32_t offset = + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b + + k_block_idx * aligned_shape_n_sfb + n_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } + } else { + const uint32_t offset = current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_block_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } + } + } else if constexpr (kScaleBGranK == 32) { + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + const uint32_t k_off = k_block_idx % (BLOCK_K / kScaleBGranK); + return ptx::ld_shared(smem_sfb + k_off * BLOCK_N + n_off); + } else { + // SFB has been staged into smem above; out-of-bound `n_idx` already + // resolves to 1.0f because we wrote 1.0f for those slots. + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + return ptx::ld_shared(smem_sfb + k_block_idx * BLOCK_N + n_off); + } + }; + + auto load_sfb_exp_offset = [&](uint32_t n_idx, uint32_t k_block_idx) { + if (n_idx >= shape_n) + return uint32_t(6); + constexpr uint32_t kVec = 4; + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + const uint32_t k_off = k_block_idx % (BLOCK_K / kScaleBGranK); + const uint32_t packed_offsets = ptx::ld_shared(smem_sfb_exp + k_off * kVecsPerRow + n_off / kVec); + return (packed_offsets >> ((n_off % kVec) * 8)) & 0xffu; + }; + + // Decide the number of scales B to load + DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0); + uint32_t num_former_iters = BLOCK_N / 8; + if constexpr (not kMustUseUniformedScaleB) { + num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; + } + + // Accumulation for WGMMA or CUDA promotion + constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); + constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; + constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; + constexpr bool kParallelNWavesEnabled = + kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; + constexpr bool kUseScaleKGroup = (kScaleKGroup > 1 and (WAVE_WGMMA == 1 or WAVE_WGMMA == 2)); + constexpr bool kUseScaleKGroupExact = + (kScaleKGroupExact and kUseScaleKGroup and kScaleKGroup == 2 and not kFusedPromote); + DG_STATIC_ASSERT(kNumMathThreads % 128 == 0, "RS-mode math threads must be whole warpgroups"); + DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, + "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); + const uint32_t wave_group_idx = kParallelNWavesEnabled ? math_wg_idx / kWGsPerNWave : 0; + const uint32_t wave_mwg_idx = kParallelNWavesEnabled ? math_wg_idx % kWGsPerNWave : math_wg_idx; + const uint32_t wave_warp_idx = wave_mwg_idx * 4 + warp_in_wg; + const uint32_t r_0 = wave_warp_idx * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + using final_accum_t = typename fp4_rs_detail::FinalAccumStorage::type; + constexpr uint32_t kNumAccumSets = kUseScaleKGroup ? WAVE_WGMMA : 1; + constexpr uint32_t kFinalAccumStride = + kBF16FinalAccum ? WGMMA::kNumAccum / 2 : WGMMA::kNumAccum; + constexpr uint32_t kNumFinalAccumRegs = kFinalAccumStride * WAVE_WGMMA; + constexpr bool kUseFinalAccumScratch = + kFinalAccumScratch and kDirectStore and kGemmType == GemmType::MGroupedMasked and + BLOCK_M < WGMMA::M and not kUseScaleKGroup; + constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + constexpr uint32_t kNumWGMMAStoreThreads = + kParallelNWavesEnabled ? kBaseWGMMAStoreThreads * WAVE_WGMMA : kBaseWGMMAStoreThreads; + constexpr uint32_t kFinalAccumScratchBytes = + kNumMathThreads * kNumFinalAccumRegs * sizeof(final_accum_t); + float accum_storage[WGMMA::kNumAccum * kNumAccumSets]; + final_accum_t final_accum_regs[kUseFinalAccumScratch ? 1 : kNumFinalAccumRegs]; + final_accum_t* final_accum = final_accum_regs; + if constexpr (kUseFinalAccumScratch) { + DG_STATIC_ASSERT(kFinalAccumScratchBytes <= SMEM_D_SIZE, + "DG_W4_FINAL_ACCUM_SCRATCH needs more smem_d scratch space"); + final_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * kNumFinalAccumRegs; + } + #pragma unroll + for (uint32_t i = 0; i < kNumFinalAccumRegs; ++ i) + fp4_rs_detail::final_accum_init(final_accum, i); + + // Pick threads whose WGMMA results are to be stored in shared memory + DG_STATIC_ASSERT(BLOCK_N >= 64, "RS-mode swap_ab requires at least one 64-row compute tile"); + const bool do_wgmma_store = warp_idx < kNumWGMMAStoreThreads / 32; + + // Empty barrier arrival + auto empty_barrier_arrive_stage = [&](uint32_t target_stage) { + if constexpr (kParallelNWavesEnabled) + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 2); + if constexpr (kNumTMAMulticast == 1) { + (lane_idx == 0 and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? + empty_barriers[target_stage]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + (lane_idx < kNumTMAMulticast and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? + empty_barriers[target_stage]->arrive(target_cta) : void(); + } + }; + auto empty_barrier_arrive = [&]() { + empty_barrier_arrive_stage(stage_idx); + }; + + // Skip useless computations + const bool is_cta_computation_valid = scheduler.is_computation_valid(m_block_idx, 0); + if (is_cta_computation_valid) { + // The compiler must know the dynamic variable `num_former_iters`'s real value + constexpr bool kShouldOptimize = BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; + constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; + + auto wait_stage = [&](uint32_t s, uint32_t p) { + full_barriers[s]->wait(p); + }; + + + // Dispatch `num_former_iters` and launch MMAs. + dispatch_num_former_iters_rs<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) { + constexpr uint32_t kAccumScratchBytes = kNumWGMMAStoreThreads * WGMMA::kNumAccum * sizeof(float); + if constexpr (kOverlapPromote and WAVE_WGMMA == 1 and kNumStages >= 2 and + not kWGMMAStub and kAccumScratchBytes <= SMEM_D_SIZE) { + auto accum = accum_storage; + auto smem_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * WGMMA::kNumAccum; + float scale_0_0_regs[WGMMA::kNumAccum / 4]; + float scale_1_0_regs[WGMMA::kNumAccum / 4]; + float scale_0_1_regs[WGMMA::kNumAccum / 4]; + float scale_1_1_regs[WGMMA::kNumAccum / 4]; + bool prev_valid = false; + + auto snapshot_accum = [&]() { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + smem_accum[i] = accum[i]; + asm volatile("" ::: "memory"); + }; + + auto promote_snapshot = [&]() { + if constexpr (not kPromoteStub) { + asm volatile("" ::: "memory"); + auto shifted_accum = final_accum; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_regs[i], smem_accum[i * 4 + 0], + scale_1_0_regs[i], smem_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_regs[i], smem_accum[i * 4 + 2], + scale_1_1_regs[i], smem_accum[i * 4 + 3]); + } + } + }; + + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks;) { + const uint32_t cur_stage = stage_idx; + const uint32_t cur_phase = phase; + const auto a_desc_base_lo = a_desc_lo + cur_stage * (SMEM_A_SIZE_PER_STAGE / 16); + wait_stage(cur_stage, cur_phase); + + if (not do_wgmma_store) { + empty_barrier_arrive_stage(cur_stage); + advance_pipeline(k_block_idx); + continue; + } + + constexpr uint32_t local_idx = 0; + constexpr uint32_t m_offset = 0; + const uint32_t lane_row = lane_idx / 4; + const uint32_t lane_col_pair = lane_idx % 4; + const uint32_t packed_shift = (lane_col_pair & 1u) * 16u; + const uint32_t lane_pair_col = (lane_col_pair & 2u) * 2u; + uint32_t rs_row_offset[4]; + uint32_t rs_raw_col[4]; + uint32_t rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + lane_row; + const uint32_t addr_tid_g = (warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + rs_raw_col[mat] = addr_t_col / 2 + lane_pair_col; + rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + auto load_a_regs = [&](uint32_t k, uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[cur_stage] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[cur_stage] + rs_row_offset[mat] + swizzled_col)); + } + const uint32_t packed_shifted = packed_word >> packed_shift; + if constexpr (kDecodeStub) { + a_regs[mat] = 0x38383838u; + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } + }; + + uint32_t a_regs[4]; + load_a_regs(0, a_regs); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, k); + asm volatile("" ::: "memory"); + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + + if (prev_valid) { + promote_snapshot(); + } + + const uint32_t compute_n_0 = n_block_idx * BLOCK_N + r_0; + const uint32_t compute_n_1 = n_block_idx * BLOCK_N + r_1; + float scale_b_0 = 1.0f; + float scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + scale_b_0 = load_sfb(compute_n_0, k_block_idx); + scale_b_1 = load_sfb(compute_n_1, k_block_idx); + } + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[cur_stage] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[cur_stage] + m_idx + 1); + } + auto scale_product = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else { + return scale_a * scale_b; + } + }; + scale_0_0_regs[i] = scale_product(scale_a_0, scale_b_0); + scale_1_0_regs[i] = scale_product(scale_a_1, scale_b_0); + scale_0_1_regs[i] = scale_product(scale_a_0, scale_b_1); + scale_1_1_regs[i] = scale_product(scale_a_1, scale_b_1); + } + } + + ptx::warpgroup_wait<0>(); + empty_barrier_arrive_stage(cur_stage); + snapshot_accum(); + prev_valid = true; + advance_pipeline(k_block_idx); + } + + if (prev_valid) { + promote_snapshot(); + } + } else { + constexpr uint32_t kScaleSumStride = WGMMA::kNumAccum / 4; + float scale_0_0_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_1_0_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_0_1_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_1_1_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float accum_first_storage[kUseScaleKGroupExact ? WAVE_WGMMA * WGMMA::kNumAccum : 1]; + #pragma unroll 8 + 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); + wait_stage(stage_idx, phase); + cache_sfb_k32(k_block_idx); + + // Small-N/BM=32 experiments may launch more math WGs than + // the current B tile has 64-row RS slices. Inactive WGs must + // still release the pipeline slot, but must not read B rows + // beyond BLOCK_N. + if (not do_wgmma_store) { + empty_barrier_arrive(); + continue; + } + + // TODO: remove some useless computation for unaligned Ms + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + auto accum = accum_storage + (kUseScaleKGroup ? local_idx * WGMMA::kNumAccum : 0); + auto m_offset = local_idx * WAVE_BLOCK_M; + const uint32_t scale_group_pos = + kUseScaleKGroup ? (k_block_idx % kScaleKGroup) : 0; + const bool should_promote_group = + (not kUseScaleKGroup) or (scale_group_pos == kScaleKGroup - 1) or + (k_block_idx + 1 == num_total_k_blocks); + + // Read scales before `warpgroup_arrive` so the next CTA cannot pollute shared memory. + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + const uint32_t compute_n_0 = n_block_idx * BLOCK_N + m_offset + r_0; + const uint32_t compute_n_1 = n_block_idx * BLOCK_N + m_offset + r_1; + float scale_b_0 = 0.0f; + float scale_b_1 = 0.0f; + if constexpr (kScaleBGranK == 128) { + if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + if constexpr (kScaleBStub) { + scale_b_0 = 1.0f; + scale_b_1 = 1.0f; + } else { + scale_b_0 = load_sfb(compute_n_0, k_block_idx); + scale_b_1 = load_sfb(compute_n_1, k_block_idx); + } + } + } + float scale_a_0_regs[kLateScaleA ? 1 : WGMMA::kNumAccum / 4]; + float scale_a_1_regs[kLateScaleA ? 1 : WGMMA::kNumAccum / 4]; + if constexpr (kLateScaleA) { + // Keep SFA live in shared memory and load it in the promotion loop. + } else if constexpr (kScaleAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + scale_a_0_regs[i] = 1.0f; + scale_a_1_regs[i] = 1.0f; + } + } else if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0_regs[i] = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1_regs[i] = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } + + const uint32_t lane_row = lane_idx / 4; + const uint32_t lane_col_pair = lane_idx % 4; + const uint32_t packed_shift = (lane_col_pair & 1u) * 16u; + const uint32_t lane_pair_col = (lane_col_pair & 2u) * 2u; + uint32_t rs_row_offset[4]; + uint32_t rs_raw_col[4]; + uint32_t rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + lane_row; + const uint32_t addr_tid_g = (warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + rs_raw_col[mat] = addr_t_col / 2 + lane_pair_col; + rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + auto load_a_regs = [&](uint32_t k, uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + } + const uint32_t packed_shifted = packed_word >> packed_shift; + if constexpr (kDecodeStub) { + // Four E4M3 1.0 values. Keeps the LDS/address path but removes FP4 decode. + a_regs[mat] = 0x38383838u; + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } + }; + auto load_a_regs_scaled_b = [&](uint32_t k, + const uint32_t exp_offsets[2], + const fp4_rs_detail::ScaledE4M3Lut scaled_luts[2], + uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + auto load_packed_word = [&](uint32_t mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + } + return packed_word; + }; + if constexpr (kDecodeStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + static_cast(load_packed_word(mat)); + a_regs[mat] = 0x38383838u; + } + } else if constexpr (kBLoadStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = load_packed_word(mat); + } else if constexpr (kFuseScaleBDecodeAssumeExp == 6) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } else if constexpr (kFuseScaleBDecodeAssumeExp == 5) { + constexpr fp4_rs_detail::ScaledE4M3Lut kExp5Lut{0x34302800u, 0x44403c38u}; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, kExp5Lut); + } + } else { + if constexpr (kFuseScaleBDecodeFastCommon) { + const bool uniform_exp = exp_offsets[0] == exp_offsets[1]; + if (uniform_exp and exp_offsets[0] == 6u) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } else if (uniform_exp and exp_offsets[0] == 5u) { + constexpr fp4_rs_detail::ScaledE4M3Lut kExp5Lut{0x34302800u, 0x44403c38u}; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, kExp5Lut); + } + } else { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + if constexpr (kFuseScaleBHummingDecode) { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_humming( + packed_shifted, exp_offsets[mat & 1u]); + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, scaled_luts[mat & 1u]); + } + } + } + } else { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + if constexpr (kFuseScaleBHummingDecode) { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_humming( + packed_shifted, exp_offsets[mat & 1u]); + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, scaled_luts[mat & 1u]); + } + } + } + } + }; + + // Decode the next RS-A fragment immediately after issuing the + // current WGMMA so decode/address work can overlap the async MMA. + uint32_t a_regs[4]; + if constexpr (not kFuseScaleBDecode) + load_a_regs(0, a_regs); + if constexpr (kScaleBGranK == 32) { + DG_STATIC_ASSERT(WGMMA::K == 32, + "per-32 FP4 scale path assumes each WGMMA K slice is one scale group"); + DG_STATIC_ASSERT(not kUseScaleKGroup, + "per-32 FP4 scale path does not support DG_W4_SCALE_K_GROUP"); + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + if constexpr (kFuseScaleBDecode) { + if constexpr (not kWGMMAStub) { + uint32_t scale_b_exp_offsets[BLOCK_K / WGMMA::K][2]; + fp4_rs_detail::ScaledE4M3Lut scale_b_luts[kFuseScaleBOnDemandLut ? 1 : (BLOCK_K / WGMMA::K)][2]; + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + const uint32_t n_idx = + n_block_idx * BLOCK_N + rs_row_offset[pair] / BLOCK_K_PACKED; + if constexpr (kScaleBStub) + scale_b_exp_offsets[kk][pair] = 6u; + else + scale_b_exp_offsets[kk][pair] = + load_sfb_exp_offset(n_idx, scale_b_k_idx); + } + } + if constexpr (not kFuseScaleBOnDemandLut and not kDecodeStub and not kBLoadStub) { + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + scale_b_luts[kk][pair] = + fp4_rs_detail::make_scaled_e4m3_lut(scale_b_exp_offsets[kk][pair]); + } + } + } + uint32_t staged_a_regs[BLOCK_K / WGMMA::K][4]; + uint32_t staged_pair_a_regs[2][4]; + uint32_t* staged_smem_a_regs = reinterpret_cast(smem_d) + + math::ceil_div(smem_sfb_bytes, 128u) * 32u + + (warp_idx * 32 + lane_idx) * (BLOCK_K / WGMMA::K) * 4u; + uint32_t* ws_decoded_regs = reinterpret_cast(smem_d) + + math::ceil_div(smem_sfb_bytes, 128u) * 32u; + auto load_a_regs_scaled_b_for_k = [&](uint32_t kk, uint32_t regs[4]) { + if constexpr (kFuseScaleBOnDemandLut) { + fp4_rs_detail::ScaledE4M3Lut on_demand_luts[2]; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) + on_demand_luts[pair] = + fp4_rs_detail::make_scaled_e4m3_lut(scale_b_exp_offsets[kk][pair]); + load_a_regs_scaled_b(kk, scale_b_exp_offsets[kk], on_demand_luts, regs); + } else { + load_a_regs_scaled_b(kk, scale_b_exp_offsets[kk], scale_b_luts[kk], regs); + } + }; + if constexpr (kFuseScaleBSlicePromote) { + DG_STATIC_ASSERT(not kFuseScaleBPredecode and not kFuseScaleBPredecodePair and + not kFuseScaleBSharedStage and not kFuseScaleBWSDecode, + "slice-promote is only for the direct fused decode path"); + DG_STATIC_ASSERT(not kPromoteFromSmem, + "slice-promote does not use promote-from-smem"); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + load_a_regs_scaled_b_for_k(k, a_regs); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + if constexpr (not kWGMMAStub) { + #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 + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + 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 constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (not kScaleBEarlyProduct) { + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, accum[i * 4 + 0], + scale_a_1, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, accum[i * 4 + 2], + scale_a_1, accum[i * 4 + 3]); + } + } + } + } else { + if constexpr (kFuseScaleBWSDecode) { + DG_STATIC_ASSERT(BLOCK_N == 256 and kNumMathThreads == 256, + "DG_W4_FUSE_SCALE_B_WS_DECODE only targets BN256/256-thread experiments"); + DG_STATIC_ASSERT(SMEM_D_SIZE >= 1024 + kNumMathThreads * + (BLOCK_K / WGMMA::K) * 4 * sizeof(uint32_t), + "smem_d is too small for WS decoded FP4 staging"); + if (warp_idx == 0) { + #pragma unroll + for (uint32_t target_tid = lane_idx; target_tid < kNumMathThreads; target_tid += 32) { + const uint32_t target_warp_idx = target_tid / 32; + const uint32_t target_lane_idx = target_tid % 32; + const uint32_t target_math_wg_idx = target_tid / 128; + const uint32_t target_warp_in_wg = target_warp_idx % 4; + const uint32_t target_wave_mwg_idx = + kParallelNWavesEnabled ? target_math_wg_idx % kWGsPerNWave : target_math_wg_idx; + const uint32_t target_lane_row = target_lane_idx / 4; + const uint32_t target_lane_col_pair = target_lane_idx % 4; + const uint32_t target_packed_shift = (target_lane_col_pair & 1u) * 16u; + const uint32_t target_lane_pair_col = (target_lane_col_pair & 2u) * 2u; + uint32_t target_rs_row_offset[4]; + uint32_t target_rs_raw_col[4]; + uint32_t target_rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + target_lane_row; + const uint32_t addr_tid_g = (target_warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = + target_wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + target_rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + target_rs_raw_col[mat] = addr_t_col / 2 + target_lane_pair_col; + target_rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + uint32_t target_exp_offsets[2]; + fp4_rs_detail::ScaledE4M3Lut target_luts[2]; + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + const uint32_t n_idx = + n_block_idx * BLOCK_N + + target_rs_row_offset[pair] / BLOCK_K_PACKED; + if constexpr (kScaleBStub) + target_exp_offsets[pair] = 6u; + else + target_exp_offsets[pair] = + load_sfb_exp_offset(n_idx, scale_b_k_idx); + target_luts[pair] = + fp4_rs_detail::make_scaled_e4m3_lut(target_exp_offsets[pair]); + } + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + const uint32_t swizzled_col = + ((kk ^ target_rs_swizzle_xor[mat]) << 4) + + target_rs_raw_col[mat]; + const uint32_t packed_word = ptx::ld_shared( + reinterpret_cast( + smem_b_packed[stage_idx] + + target_rs_row_offset[mat] + swizzled_col)); + const uint32_t packed_shifted = packed_word >> target_packed_shift; + uint32_t decoded = 0x38383838u; + if constexpr (not kDecodeStub) { + decoded = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, target_luts[mat & 1u]); + } + ptx::st_shared(ws_decoded_regs + + target_tid * (BLOCK_K / WGMMA::K) * 4u + + kk * 4u + mat, + decoded); + } + } + } + } + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 1); + } else if constexpr (kFuseScaleBSharedStage) { + DG_STATIC_ASSERT(BLOCK_N == 256, + "DG_W4_FUSE_SCALE_B_SHARED_STAGE only targets BN256 experiments"); + DG_STATIC_ASSERT(SMEM_D_SIZE >= 1024 + kNumMathThreads * + (BLOCK_K / WGMMA::K) * 4 * sizeof(uint32_t), + "smem_d is too small for shared decoded FP4 staging"); + uint32_t tmp_a_regs[4]; + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + load_a_regs_scaled_b_for_k(kk, tmp_a_regs); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + ptx::st_shared(staged_smem_a_regs + kk * 4 + mat, tmp_a_regs[mat]); + } + } + } else if constexpr (kFuseScaleBPredecode) { + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + load_a_regs_scaled_b_for_k(kk, staged_a_regs[kk]); + } + } else if constexpr (kFuseScaleBPredecodePair) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_FUSE_SCALE_B_PREDECODE_PAIR assumes four K=32 slices"); + load_a_regs_scaled_b_for_k(0, staged_pair_a_regs[0]); + load_a_regs_scaled_b_for_k(1, staged_pair_a_regs[1]); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + if constexpr (not kFuseScaleBPredecode and not kFuseScaleBPredecodePair and + not kFuseScaleBSharedStage and not kFuseScaleBWSDecode) { + load_a_regs_scaled_b_for_k(0, a_regs); + } + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + if constexpr (kFuseScaleBWSDecode) { + uint32_t ws_a_regs[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + ws_a_regs[mat] = ptx::ld_shared( + ws_decoded_regs + threadIdx.x * (BLOCK_K / WGMMA::K) * 4u + + k * 4u + mat); + } + WGMMA::wgmma(ws_a_regs, a_desc, accum, static_cast(k)); + } else if constexpr (kFuseScaleBSharedStage) { + uint32_t smem_a_regs[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) + smem_a_regs[mat] = ptx::ld_shared(staged_smem_a_regs + k * 4 + mat); + WGMMA::wgmma(smem_a_regs, a_desc, accum, static_cast(k)); + } else if constexpr (kFuseScaleBPredecode) + WGMMA::wgmma(staged_a_regs[k], a_desc, accum, static_cast(k)); + else if constexpr (kFuseScaleBPredecodePair) + WGMMA::wgmma(staged_pair_a_regs[k & 1u], a_desc, accum, + static_cast(k)); + else + WGMMA::wgmma(a_regs, a_desc, accum, static_cast(k)); + asm volatile("" ::: "memory"); + if constexpr (kFuseScaleBPredecodePair) { + if (k + 2 < BLOCK_K / WGMMA::K) { + load_a_regs_scaled_b_for_k(k + 2, staged_pair_a_regs[k & 1u]); + } + } else if constexpr (not kFuseScaleBPredecode) { + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) { + load_a_regs_scaled_b_for_k(k + 1, a_regs); + } + } + } + } + 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 constexpr (kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + + if constexpr (not kPromoteStub) { + float* smem_promote_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * WGMMA::kNumAccum; + if constexpr (kPromoteFromSmem) { + DG_STATIC_ASSERT(BLOCK_N == 256 and kNumMathThreads == 256, + "DG_W4_PROMOTE_FROM_SMEM only targets BN256/256-thread experiments"); + DG_STATIC_ASSERT(kNumMathThreads * WGMMA::kNumAccum * sizeof(float) <= SMEM_D_SIZE, + "smem_d is too small for promote-from-smem accum scratch"); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::st_shared(smem_promote_accum + i, accum[i]); + asm volatile("" ::: "memory"); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 0) : accum[i * 4 + 0], + scale_a_1, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 1) : accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 2) : accum[i * 4 + 2], + scale_a_1, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 3) : accum[i * 4 + 3]); + } + } + } else if constexpr (kFuseScaleBDecodeStub) { + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, static_cast(k)); + asm volatile("" ::: "memory"); + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + 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>(); + } else { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, accum[i * 4 + 0], + scale_a_1, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, accum[i * 4 + 2], + scale_a_1, accum[i * 4 + 3]); + } + } + } else { + if constexpr (kK32Pingpong) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_K32_PINGPONG currently assumes four K=32 slices"); + float accum_alt[WGMMA::kNumAccum]; + auto promote_k32_slice = [&](float* slice_accum, uint32_t k_slice) { + if constexpr (not kPromoteStub) { + float slice_scale_b_0 = 1.0f; + float slice_scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + const float prod_0_0 = scale_prod(scale_a_0, slice_scale_b_0); + const float prod_1_0 = scale_prod(scale_a_1, slice_scale_b_0); + const float prod_0_1 = scale_prod(scale_a_0, slice_scale_b_1); + const float prod_1_1 = scale_prod(scale_a_1, slice_scale_b_1); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod_0_0 + slice_accum[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod_1_0 + slice_accum[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod_0_1 + slice_accum[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod_1_1 + slice_accum[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + 1.0f, slice_accum[i * 4 + 0], + 1.0f, slice_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + 1.0f, slice_accum[i * 4 + 2], + 1.0f, slice_accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + prod_0_0, slice_accum[i * 4 + 0], + prod_1_0, slice_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + prod_0_1, slice_accum[i * 4 + 2], + prod_1_1, slice_accum[i * 4 + 3]); + } + } + } + }; + auto promote_k32_pair = [&](float* slice_accum_0, float* slice_accum_1, + uint32_t k_slice_0, uint32_t k_slice_1) { + if constexpr (not kPromoteStub) { + float slice0_scale_b_0 = 1.0f; + float slice0_scale_b_1 = 1.0f; + float slice1_scale_b_0 = 1.0f; + float slice1_scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + const uint32_t scale_b_k_idx_0 = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice_0; + const uint32_t scale_b_k_idx_1 = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice_1; + slice0_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx_0); + slice0_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx_0); + slice1_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx_1); + slice1_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx_1); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice0_scale_b_0); + fp4_rs_detail::keep_float_live(slice0_scale_b_1); + fp4_rs_detail::keep_float_live(slice1_scale_b_0); + fp4_rs_detail::keep_float_live(slice1_scale_b_1); + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + const float prod0_0_0 = scale_prod(scale_a_0, slice0_scale_b_0); + const float prod0_1_0 = scale_prod(scale_a_1, slice0_scale_b_0); + const float prod0_0_1 = scale_prod(scale_a_0, slice0_scale_b_1); + const float prod0_1_1 = scale_prod(scale_a_1, slice0_scale_b_1); + const float prod1_0_0 = scale_prod(scale_a_0, slice1_scale_b_0); + const float prod1_1_0 = scale_prod(scale_a_1, slice1_scale_b_0); + const float prod1_0_1 = scale_prod(scale_a_0, slice1_scale_b_1); + const float prod1_1_1 = scale_prod(scale_a_1, slice1_scale_b_1); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod0_0_0 + slice_accum_0[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod0_1_0 + slice_accum_0[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod0_0_1 + slice_accum_0[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod0_1_1 + slice_accum_0[i * 4 + 3]); + fp4_rs_detail::keep_float_live(prod1_0_0 + slice_accum_1[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod1_1_0 + slice_accum_1[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod1_0_1 + slice_accum_1[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod1_1_1 + slice_accum_1[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 0, + 1.0f, slice_accum_0[i * 4 + 0], + 1.0f, slice_accum_0[i * 4 + 1], + 1.0f, slice_accum_1[i * 4 + 0], + 1.0f, slice_accum_1[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 1, + 1.0f, slice_accum_0[i * 4 + 2], + 1.0f, slice_accum_0[i * 4 + 3], + 1.0f, slice_accum_1[i * 4 + 2], + 1.0f, slice_accum_1[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 0, + prod0_0_0, slice_accum_0[i * 4 + 0], + prod0_1_0, slice_accum_0[i * 4 + 1], + prod1_0_0, slice_accum_1[i * 4 + 0], + prod1_1_0, slice_accum_1[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 1, + prod0_0_1, slice_accum_0[i * 4 + 2], + prod0_1_1, slice_accum_0[i * 4 + 3], + prod1_0_1, slice_accum_1[i * 4 + 2], + prod1_1_1, slice_accum_1[i * 4 + 3]); + } + } + } + }; + auto load_k32_quad_scale_b = [&](float* scale_b_0, float* scale_b_1) { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = 1.0f; + scale_b_1[kk] = 1.0f; + } + if constexpr (not kScaleBStub) { + if constexpr ((kK32QuadScaleBInline or kK32QuadScaleBVec4) and + kScaleBDirectLoad and kScaleBGranK == 32) { + const uint32_t scale_b_k_base = + k_block_idx * (BLOCK_K / kScaleBGranK); + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b + + scale_b_k_base * aligned_shape_n_sfb; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = compute_n_0 < shape_n ? + *(sfb_base + kk * aligned_shape_n_sfb + compute_n_0) : 1.0f; + scale_b_1[kk] = compute_n_1 < shape_n ? + *(sfb_base + kk * aligned_shape_n_sfb + compute_n_1) : 1.0f; + } + } else { + const float* sfb_base = sfb + + current_group_idx * shape_n * shape_k_scales_b + scale_b_k_base; + if constexpr (kK32QuadScaleBVec4) { + if (compute_n_0 < shape_n) { + const float4 v0 = *reinterpret_cast( + sfb_base + compute_n_0 * shape_k_scales_b); + scale_b_0[0] = v0.x; + scale_b_0[1] = v0.y; + scale_b_0[2] = v0.z; + scale_b_0[3] = v0.w; + } + if (compute_n_1 < shape_n) { + const float4 v1 = *reinterpret_cast( + sfb_base + compute_n_1 * shape_k_scales_b); + scale_b_1[0] = v1.x; + scale_b_1[1] = v1.y; + scale_b_1[2] = v1.z; + scale_b_1[3] = v1.w; + } + } else { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = compute_n_0 < shape_n ? + *(sfb_base + compute_n_0 * shape_k_scales_b + kk) : 1.0f; + scale_b_1[kk] = compute_n_1 < shape_n ? + *(sfb_base + compute_n_1 * shape_k_scales_b + kk) : 1.0f; + } + } + } + } else { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + scale_b_0[kk] = load_sfb(compute_n_0, scale_b_k_idx); + scale_b_1[kk] = load_sfb(compute_n_1, scale_b_k_idx); + } + } + } + }; + auto promote_k32_quad = [&](float* accum_0, float* accum_1, + float* accum_2, float* accum_3, + float* scale_b_0, float* scale_b_1) { + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + fp4_rs_detail::keep_float_live(scale_b_0[kk]); + fp4_rs_detail::keep_float_live(scale_b_1[kk]); + } + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[0]) + accum_0[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[0]) + accum_0[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[0]) + accum_0[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[0]) + accum_0[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[1]) + accum_1[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[1]) + accum_1[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[1]) + accum_1[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[1]) + accum_1[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[2]) + accum_2[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[2]) + accum_2[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[2]) + accum_2[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[2]) + accum_2[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[3]) + accum_3[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[3]) + accum_3[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[3]) + accum_3[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[3]) + accum_3[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 0, + 1.0f, accum_0[i * 4 + 0], 1.0f, accum_0[i * 4 + 1], + 1.0f, accum_1[i * 4 + 0], 1.0f, accum_1[i * 4 + 1], + 1.0f, accum_2[i * 4 + 0], 1.0f, accum_2[i * 4 + 1], + 1.0f, accum_3[i * 4 + 0], 1.0f, accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 1, + 1.0f, accum_0[i * 4 + 2], 1.0f, accum_0[i * 4 + 3], + 1.0f, accum_1[i * 4 + 2], 1.0f, accum_1[i * 4 + 3], + 1.0f, accum_2[i * 4 + 2], 1.0f, accum_2[i * 4 + 3], + 1.0f, accum_3[i * 4 + 2], 1.0f, accum_3[i * 4 + 3]); + } else { + if constexpr (kK32QuadPersistentScaleProduct) { + float prod_0_0[4], prod_1_0[4]; + float prod_0_1[4], prod_1_1[4]; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + prod_0_0[kk] = scale_prod(scale_a_0, scale_b_0[kk]); + prod_1_0[kk] = scale_prod(scale_a_1, scale_b_0[kk]); + prod_0_1[kk] = scale_prod(scale_a_0, scale_b_1[kk]); + prod_1_1[kk] = scale_prod(scale_a_1, scale_b_1[kk]); + } + fp4_rs_detail::final_accum_promote_pair4< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2 + 0, + prod_0_0[0], accum_0[i * 4 + 0], prod_1_0[0], accum_0[i * 4 + 1], + prod_0_0[1], accum_1[i * 4 + 0], prod_1_0[1], accum_1[i * 4 + 1], + prod_0_0[2], accum_2[i * 4 + 0], prod_1_0[2], accum_2[i * 4 + 1], + prod_0_0[3], accum_3[i * 4 + 0], prod_1_0[3], accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2 + 1, + prod_0_1[0], accum_0[i * 4 + 2], prod_1_1[0], accum_0[i * 4 + 3], + prod_0_1[1], accum_1[i * 4 + 2], prod_1_1[1], accum_1[i * 4 + 3], + prod_0_1[2], accum_2[i * 4 + 2], prod_1_1[2], accum_2[i * 4 + 3], + prod_0_1[3], accum_3[i * 4 + 2], prod_1_1[3], accum_3[i * 4 + 3]); + } else if constexpr (kK32QuadSplitPromote or kK32QuadShortProductPromote) { + fp4_rs_detail::final_accum_promote_pair4_split_scale< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath, kScaleBPow2Promote>( + shifted_accum, i * 2 + 0, scale_a_0, scale_a_1, + scale_b_0[0], scale_b_0[1], scale_b_0[2], scale_b_0[3], + accum_0[i * 4 + 0], accum_0[i * 4 + 1], + accum_1[i * 4 + 0], accum_1[i * 4 + 1], + accum_2[i * 4 + 0], accum_2[i * 4 + 1], + accum_3[i * 4 + 0], accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4_split_scale< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath, kScaleBPow2Promote>( + shifted_accum, i * 2 + 1, scale_a_0, scale_a_1, + scale_b_1[0], scale_b_1[1], scale_b_1[2], scale_b_1[3], + accum_0[i * 4 + 2], accum_0[i * 4 + 3], + accum_1[i * 4 + 2], accum_1[i * 4 + 3], + accum_2[i * 4 + 2], accum_2[i * 4 + 3], + accum_3[i * 4 + 2], accum_3[i * 4 + 3]); + } else { + if constexpr (kK32QuadPair4x2Promote) { + fp4_rs_detail::final_accum_promote_pair4x2< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2, + scale_prod(scale_a_0, scale_b_0[0]), accum_0[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[0]), accum_0[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[1]), accum_1[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[1]), accum_1[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[2]), accum_2[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[2]), accum_2[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[3]), accum_3[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[3]), accum_3[i * 4 + 1], + scale_prod(scale_a_0, scale_b_1[0]), accum_0[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[0]), accum_0[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[1]), accum_1[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[1]), accum_1[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[2]), accum_2[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[2]), accum_2[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[3]), accum_3[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[3]), accum_3[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 0, + scale_prod(scale_a_0, scale_b_0[0]), accum_0[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[0]), accum_0[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[1]), accum_1[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[1]), accum_1[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[2]), accum_2[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[2]), accum_2[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[3]), accum_3[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[3]), accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 1, + scale_prod(scale_a_0, scale_b_1[0]), accum_0[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[0]), accum_0[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[1]), accum_1[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[1]), accum_1[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[2]), accum_2[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[2]), accum_2[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[3]), accum_3[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[3]), accum_3[i * 4 + 3]); + } + } + } + } + } + }; + + if constexpr (kK32QuadReduce) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_K32_QUAD_REDUCE assumes four K=32 slices"); + float scale_b_0[4]; + float scale_b_1[4]; + if constexpr (kWGMMAF16Accum) { + DG_STATIC_ASSERT(BLOCK_M == 8, + "DG_W4_WGMMA_F16_ACCUM currently targets BM8 small-M path"); + DG_STATIC_ASSERT(WGMMA::kNumAccum == 4, + "FP16 accumulator unpack assumes m64n8 f32 accumulator layout"); + using WGMMAF16 = fp4_rs_detail::FP8MMAF16AccumM64N8K32RS; + uint32_t accum_h0[WGMMAF16::kNumAccum]; + uint32_t accum_h1[WGMMAF16::kNumAccum]; + uint32_t accum_h2[WGMMAF16::kNumAccum]; + uint32_t accum_h3[WGMMAF16::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + accum_h0[i] = 0; + accum_h1[i] = 0; + accum_h2[i] = 0; + accum_h3[i] = 0; + } + if constexpr (not kWGMMAStub) { + auto fence_u32 = [](uint32_t& value) { + asm volatile("" : "+r"(value) :: "memory"); + }; + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + fence_u32(accum_h0[i]); + fence_u32(accum_h1[i]); + fence_u32(accum_h2[i]); + fence_u32(accum_h3[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo; + WGMMAF16::wgmma(a_regs, a_desc, accum_h0, false); + asm volatile("" ::: "memory"); + load_a_regs(1, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h1, false); + asm volatile("" ::: "memory"); + load_a_regs(2, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 2 * WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h2, false); + asm volatile("" ::: "memory"); + load_a_regs(3, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 3 * WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h3, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + fence_u32(accum_h0[i]); + fence_u32(accum_h1[i]); + fence_u32(accum_h2[i]); + fence_u32(accum_h3[i]); + } + if constexpr (kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + ptx::warpgroup_wait<0>(); + } + if constexpr (kK32QuadScaleBPrefetch) { + // Already loaded while WGMMA was in flight. + } else { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + float accum_f0[WGMMA::kNumAccum]; + float accum_f1[WGMMA::kNumAccum]; + float accum_f2[WGMMA::kNumAccum]; + float accum_f3[WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum_f0[i] = 0.0f; + accum_f1[i] = 0.0f; + accum_f2[i] = 0.0f; + accum_f3[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h0, accum_f0); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h1, accum_f1); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h2, accum_f2); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h3, accum_f3); + } + promote_k32_quad(accum_f0, accum_f1, accum_f2, accum_f3, scale_b_0, scale_b_1); + } else { + float accum_2[WGMMA::kNumAccum]; + float accum_3[WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum[i] = 0.0f; + accum_alt[i] = 0.0f; + accum_2[i] = 0.0f; + accum_3[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + ptx::warpgroup_fence_operand(accum_2[i]); + ptx::warpgroup_fence_operand(accum_3[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + load_a_regs(1, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_alt, false); + asm volatile("" ::: "memory"); + load_a_regs(2, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 2 * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_2, false); + asm volatile("" ::: "memory"); + load_a_regs(3, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 3 * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_3, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + ptx::warpgroup_fence_operand(accum_2[i]); + ptx::warpgroup_fence_operand(accum_3[i]); + } + if constexpr (kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + ptx::warpgroup_wait<0>(); + } + if constexpr (not kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + promote_k32_quad(accum, accum_alt, accum_2, accum_3, scale_b_0, scale_b_1); + } + } else { + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; k += 2) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum[i] = 0.0f; + accum_alt[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + } + load_a_regs(k + 1, a_regs); + if constexpr (not kWGMMAStub) { + a_desc.reg32_[0] = a_desc_base_lo + (k + 1) * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_alt, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + } + ptx::warpgroup_wait<0>(); + } + + if constexpr (kK32PairReduce) { + promote_k32_pair(accum, accum_alt, k, k + 1); + } else { + promote_k32_slice(accum, k); + promote_k32_slice(accum_alt, k + 1); + } + + if constexpr (BLOCK_K / WGMMA::K > 2) { + if (k + 2 < BLOCK_K / WGMMA::K) + load_a_regs(k + 2, a_regs); + } + } + } + } else { + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + float slice_scale_b_0 = 1.0f; + float slice_scale_b_1 = 1.0f; + if constexpr ((kScaleBEarlyLoad or kScaleBEarlyProduct) and not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + float early_prod_0_0[WGMMA::kNumAccum / 4]; + float early_prod_1_0[WGMMA::kNumAccum / 4]; + float early_prod_0_1[WGMMA::kNumAccum / 4]; + float early_prod_1_1[WGMMA::kNumAccum / 4]; + if constexpr (kScaleBEarlyProduct and not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + early_prod_0_0[i] = kScaleBMulStub ? scale_a_0 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_0) : + scale_a_0 * slice_scale_b_0); + early_prod_1_0[i] = kScaleBMulStub ? scale_a_1 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_0) : + scale_a_1 * slice_scale_b_0); + early_prod_0_1[i] = kScaleBMulStub ? scale_a_0 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_1) : + scale_a_0 * slice_scale_b_1); + early_prod_1_1[i] = kScaleBMulStub ? scale_a_1 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_1) : + scale_a_1 * slice_scale_b_1); + } + } + if constexpr (not kWGMMAStub) { + #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 + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + 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 constexpr (not kPromoteStub) { + if constexpr (not kScaleBEarlyLoad and not kScaleBEarlyProduct and not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBEarlyProduct and not kPromoteMulStub) { + // Product was computed before WGMMA wait to test latency hiding. + } else if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + const float prod_0_0 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_0 : + (kScaleBEarlyProduct ? early_prod_0_0[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_0) : + scale_a_0 * slice_scale_b_0)); + const float prod_1_0 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_1 : + (kScaleBEarlyProduct ? early_prod_1_0[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_0) : + scale_a_1 * slice_scale_b_0)); + const float prod_0_1 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_0 : + (kScaleBEarlyProduct ? early_prod_0_1[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_1) : + scale_a_0 * slice_scale_b_1)); + const float prod_1_1 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_1 : + (kScaleBEarlyProduct ? early_prod_1_1[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_1) : + scale_a_1 * slice_scale_b_1)); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod_0_0 + accum[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod_1_0 + accum[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod_0_1 + accum[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod_1_1 + accum[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + 1.0f, accum[i * 4 + 0], + 1.0f, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + 1.0f, accum[i * 4 + 2], + 1.0f, accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + prod_0_0, accum[i * 4 + 0], + prod_1_0, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + prod_0_1, accum[i * 4 + 2], + prod_1_1, accum[i * 4 + 3]); + } + } + } + + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + } + } + + const bool is_last_wave = (local_idx == WAVE_WGMMA - 1); + if (kParallelNWavesEnabled or is_last_wave) + empty_barrier_arrive(); + continue; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + } + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + if constexpr (not kWGMMAStub) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + const bool scale_d = kUseScaleKGroup ? (scale_group_pos != 0 or k != 0) : static_cast(k); + WGMMA::wgmma(a_regs, a_desc, accum, scale_d); + asm volatile("" ::: "memory"); + } + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + float scale_0_0_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_1_0_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_0_1_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_1_1_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + if constexpr (not kPromoteStub and not kFusedPromote) { + if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + const uint32_t scale_sum_offset = local_idx * kScaleSumStride; + if constexpr (kUseScaleKGroup) { + if (scale_group_pos == 0) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + scale_0_0_sum[scale_sum_offset + i] = 0.0f; + scale_1_0_sum[scale_sum_offset + i] = 0.0f; + scale_0_1_sum[scale_sum_offset + i] = 0.0f; + scale_1_1_sum[scale_sum_offset + i] = 0.0f; + } + } + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float prod_0_0 = scale_a_0_regs[i] * scale_b_0; + const float prod_1_0 = scale_a_1_regs[i] * scale_b_0; + const float prod_0_1 = scale_a_0_regs[i] * scale_b_1; + const float prod_1_1 = scale_a_1_regs[i] * scale_b_1; + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos == 0) { + scale_0_0_sum[scale_sum_offset + i] = prod_0_0; + scale_1_0_sum[scale_sum_offset + i] = prod_1_0; + scale_0_1_sum[scale_sum_offset + i] = prod_0_1; + scale_1_1_sum[scale_sum_offset + i] = prod_1_1; + } + if (should_promote_group) { + scale_0_0_regs[i] = prod_0_0; + scale_1_0_regs[i] = prod_1_0; + scale_0_1_regs[i] = prod_0_1; + scale_1_1_regs[i] = prod_1_1; + } + } else if constexpr (kUseScaleKGroup) { + scale_0_0_sum[scale_sum_offset + i] += prod_0_0; + scale_1_0_sum[scale_sum_offset + i] += prod_1_0; + scale_0_1_sum[scale_sum_offset + i] += prod_0_1; + scale_1_1_sum[scale_sum_offset + i] += prod_1_1; + if (should_promote_group) { + const float inv_group = 1.0f / static_cast(scale_group_pos + 1); + scale_0_0_regs[i] = scale_0_0_sum[scale_sum_offset + i] * inv_group; + scale_1_0_regs[i] = scale_1_0_sum[scale_sum_offset + i] * inv_group; + scale_0_1_regs[i] = scale_0_1_sum[scale_sum_offset + i] * inv_group; + scale_1_1_regs[i] = scale_1_1_sum[scale_sum_offset + i] * inv_group; + } + } else { + scale_0_0_regs[i] = prod_0_0; + scale_1_0_regs[i] = prod_1_0; + scale_0_1_regs[i] = prod_0_1; + scale_1_1_regs[i] = prod_1_1; + } + } + } + } + if constexpr (not kWGMMAStub) { + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + } + + const bool is_last_wave = (local_idx == WAVE_WGMMA - 1); + if constexpr (not kWGMMAStub) { + ptx::warpgroup_wait<0>(); + } else { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos == 0 and not should_promote_group) { + auto accum_first = accum_first_storage + local_idx * WGMMA::kNumAccum; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum_first[i] = accum[i]; + } + } + + // Notify barrier arrival at the last warpgroup wave + if ((kParallelNWavesEnabled or is_last_wave) and + (not kLateScaleA or not do_wgmma_store or not should_promote_group)) + empty_barrier_arrive(); + + // Skip promotion for the unfilled parts + if (not do_wgmma_store or not should_promote_group) + continue; + + // Promote with scales + // NOTES: making it as predicates is very important for performance, comparing to two loops + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + if constexpr (kFusedPromote) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0 * scale_b_0, accum[i * 4 + 0], + scale_a_1 * scale_b_0, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0 * scale_b_1, accum[i * 4 + 2], + scale_a_1 * scale_b_1, accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_regs[i], accum[i * 4 + 0], + scale_1_0_regs[i], accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_regs[i], accum[i * 4 + 2], + scale_1_1_regs[i], accum[i * 4 + 3]); + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos != 0) { + const uint32_t scale_sum_offset = local_idx * kScaleSumStride; + auto accum_first = accum_first_storage + local_idx * WGMMA::kNumAccum; + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_sum[scale_sum_offset + i] - scale_0_0_regs[i], + accum_first[i * 4 + 0], + scale_1_0_sum[scale_sum_offset + i] - scale_1_0_regs[i], + accum_first[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_sum[scale_sum_offset + i] - scale_0_1_regs[i], + accum_first[i * 4 + 2], + scale_1_1_sum[scale_sum_offset + i] - scale_1_1_regs[i], + accum_first[i * 4 + 3]); + } + } + } + } + } + if constexpr (kLateScaleA) { + if (kParallelNWavesEnabled or is_last_wave) + empty_barrier_arrive(); + } + } + } + } + }); + } else { + #pragma unroll + 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(); + } + } + + if constexpr (kStoreStub) { + __syncwarp(); + continue; + } + + if constexpr (kDirectStore and kGemmType == GemmType::MGroupedMasked and BLOCK_M < WGMMA::M) { + const uint32_t global_m_base = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t n_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kDirectStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (uint32_t i = 0; i < kDirectStoreIters; ++ i) { + const uint32_t local_m_0 = i * 8 + col_idx * 2; + const uint32_t local_m_1 = local_m_0 + 1; + const uint32_t col_0 = epilogue_type_t::template apply_index_n<1>( + n_block_idx * BLOCK_N + n_offset + r_0); + const uint32_t col_1 = epilogue_type_t::template apply_index_n<1>( + n_block_idx * BLOCK_N + n_offset + r_1); + const bool row_0_valid = scheduler.is_computation_valid(m_block_idx, local_m_0); + const bool row_1_valid = scheduler.is_computation_valid(m_block_idx, local_m_1); + auto direct_store = [&](bool row_valid, uint32_t local_m, uint32_t col, uint32_t accum_idx) { + if (row_valid and col < shape_n) { + const nv_bfloat16 out = __float2bfloat16_rn( + fp4_rs_detail::final_accum_load_scalar(shifted_accum, accum_idx)); + if constexpr (kTMAStoreStub) { + const uint32_t sink = static_cast(*reinterpret_cast(&out)); + asm volatile("" :: "r"(sink) : "memory"); + } else { + gmem_d_ptr[(global_m_base + local_m) * shape_n + col] = out; + } + } + }; + direct_store(row_0_valid, local_m_0, col_0, i * 4 + 0); + direct_store(row_1_valid, local_m_1, col_0, i * 4 + 1); + direct_store(row_0_valid, local_m_0, col_1, i * 4 + 2); + direct_store(row_1_valid, local_m_1, col_1, i * 4 + 3); + } + } + __syncwarp(); + continue; + } + + // Psum layout can have a final partial M tile. TMA store writes a + // full BLOCK_M tile and may go out of the tensor-map bounds, so use + // a guarded scalar store for this layout. + if constexpr (kGemmType == GemmType::MGroupedMasked and BLOCK_M < WGMMA::M) { + // Full small-BM tiles can use the normal STSM+TMA store path below. + // Guarded scalar copy-back is only needed for partial masked tiles. + if (not scheduler.is_computation_valid(m_block_idx, BLOCK_M - 1)) { + const uint32_t global_m_base = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (auto i = 0; i < kStoreIters; ++ i) { + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumBankGroupBytes = 16; + const uint32_t row = i * 8 + lane_idx % 8; + uint32_t col = warp_in_wg * 2 + lane_idx / 8; + col ^= row % (kSwizzleDMode / 16); + const uint32_t n_atom_idx = m_offset / WGMMA::M + wave_mwg_idx; + smem_ptr = reinterpret_cast(smem_d) + + n_atom_idx * SMEM_D_ROWS * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + col * kNumBankGroupBytes; + } else { + const uint32_t row = i * 8 + lane_idx % 8; + const uint32_t col = m_offset + wave_mwg_idx * WGMMA::M + warp_in_wg * 2 + lane_idx / 8; + smem_ptr = reinterpret_cast(smem_d + row * BLOCK_N + col); + } + + if constexpr (not kSTSMStub) { + const nv_bfloat162 out_0 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 0); + const nv_bfloat162 out_1 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 1); + if constexpr (kSTSMConvertOnly) { + const uint32_t sink_0 = *reinterpret_cast(&out_0); + const uint32_t sink_1 = *reinterpret_cast(&out_1); + asm volatile("" :: "r"(sink_0), "r"(sink_1) : "memory"); + } else { + fp4_rs_detail::SM90_U32x2_STSM_T::copy(out_0, out_1, smem_ptr); + } + } + } + } + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + if constexpr (not kTMAStoreStub) { + const uint32_t group_m = __ldg(grouped_layout + current_group_idx); + const uint32_t valid_rows = group_m - m_block_idx * BLOCK_M; + const uint32_t store_elems = valid_rows * BLOCK_N; + for (uint32_t idx = threadIdx.x; idx < store_elems; idx += kNumWGMMAStoreThreads) { + const uint32_t row = idx / BLOCK_N; + const uint32_t col = idx % BLOCK_N; + const uint32_t global_col = epilogue_type_t::template apply_index_n<1>(n_block_idx * BLOCK_N + col); + if (global_col < shape_n) { + nv_bfloat16* smem_src = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode / kNumElemBytes; + const uint32_t n_atom_idx = col / TMA_D_BLOCK_N; + const uint32_t in_atom_col = col % TMA_D_BLOCK_N; + uint32_t bank_col = in_atom_col / (kNumBankGroupBytes / kNumElemBytes); + bank_col ^= row % (kSwizzleDMode / kNumBankGroupBytes); + const uint32_t bank_offset = in_atom_col % (kNumBankGroupBytes / kNumElemBytes); + smem_src = reinterpret_cast( + reinterpret_cast(smem_d) + + n_atom_idx * SMEM_D_ROWS * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + bank_col * kNumBankGroupBytes + + bank_offset * kNumElemBytes); + } else { + smem_src = smem_d + row * BLOCK_N + col; + } + gmem_d_ptr[(global_m_base + row) * shape_n + global_col] = *smem_src; + } + } + } + __syncwarp(); + continue; + } + } + + if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + const uint32_t psum_store_end = + kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (scheduler.current_group_idx + 1 < kNumGroups ? + math::align(scheduler.current_psum_m, BLOCK_M) : scheduler.current_psum_m) : + shape_m; + constexpr uint32_t kStoreWaves = + BLOCK_M < WAVE_BLOCK_M ? 1 : BLOCK_M / WAVE_BLOCK_M; + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < kStoreWaves; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr bool kStoreWithGroupOffset = kGemmType == GemmType::MGroupedMasked; + const uint32_t row_0 = scheduler.template get_global_idx( + shape_m, BLOCK_M, m_block_idx) + m_offset + r_0; + const uint32_t row_1 = scheduler.template get_global_idx( + shape_m, BLOCK_M, m_block_idx) + m_offset + r_1; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t base_col = epilogue_type_t::template apply_index_n<8>( + n_block_idx * BLOCK_N + i * 8); + const uint32_t col = base_col + col_idx * 2; + const bool row_0_valid = kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (row_0 >= scheduler.last_psum_m and row_0 < psum_store_end) : + scheduler.is_computation_valid(m_block_idx, m_offset + r_0); + const bool row_1_valid = kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (row_1 >= scheduler.last_psum_m and row_1 < psum_store_end) : + scheduler.is_computation_valid(m_block_idx, m_offset + r_1); + if (row_0_valid and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 0)); + if (row_0_valid and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 1)); + if (row_1_valid and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 2)); + if (row_1_valid and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 3)); + } + } + __syncwarp(); + continue; + } + + // TMA checks + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes); + 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"); + + // Skip WGMMA store for the unfilled parts + if (not do_wgmma_store) + continue; + + // Wait last TMA store to be finished + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + // Write back to shared memory using STSM and issue TMA stores + DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (auto i = 0; i < kStoreIters; ++ i) { + // RS swap_ab accumulates an original-N x original-M tile. + // Store with stmatrix.trans so smem_d remains original-M x original-N for TMA store. + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumBankGroupBytes = 16; + const uint32_t row = i * 8 + lane_idx % 8; + uint32_t col = warp_in_wg * 2 + lane_idx / 8; + col ^= row % (kSwizzleDMode / 16); + const uint32_t n_atom_idx = m_offset / WGMMA::M + wave_mwg_idx; + smem_ptr = reinterpret_cast(smem_d) + + n_atom_idx * BLOCK_M * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + col * kNumBankGroupBytes; + } else { + // No swizzling, just padding + const uint32_t row = i * 8 + lane_idx % 8; + const uint32_t col = m_offset + wave_mwg_idx * WGMMA::M + warp_in_wg * 2 + lane_idx / 8; + smem_ptr = reinterpret_cast(smem_d + row * BLOCK_N + col); + } + + // NOTES: only 16 lanes' addresses are used + if constexpr (not kSTSMStub) { + const nv_bfloat162 out_0 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 0); + const nv_bfloat162 out_1 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 1); + if constexpr (kSTSMConvertOnly) { + const uint32_t sink_0 = *reinterpret_cast(&out_0); + const uint32_t sink_1 = *reinterpret_cast(&out_1); + asm volatile("" :: "r"(sink_0), "r"(sink_1) : "memory"); + } else { + fp4_rs_detail::SM90_U32x2_STSM_T::copy(out_0, out_1, smem_ptr); + } + } + } + } + if constexpr (not kTMAStoreStub) { + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + } + + // Use TMA store to write back to global memory + // TODO: compatible with FP32 output + constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked; + DG_STATIC_ASSERT(kNumWGMMAStoreThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks"); + if constexpr (not kTMAStoreStub) { + 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 = epilogue_type_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); + auto m_idx = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + if constexpr (kGemmType == GemmType::Batched) { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, + n_idx, m_idx, scheduler.current_group_idx); + } else { + 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 support sm_90a"); +#endif +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/mma/sm90.cuh b/deep_gemm/include/deep_gemm/mma/sm90.cuh index 2c061940de..ad66688501 100644 --- a/deep_gemm/include/deep_gemm/mma/sm90.cuh +++ b/deep_gemm/include/deep_gemm/mma/sm90.cuh @@ -74,6 +74,76 @@ struct FP8MMASelector { using type = decltype(select_type()); }; +// FP8 RS-mode WGMMA wrapper: A is in registers (4 × u32 per lane for K=32), +// B is in shared memory via a GMMA descriptor. Used by the W4A8/W4-FP8 path +// where packed FP4 weights are dequantized in registers before WGMMA, avoiding +// the smem write-back round-trip. +template +struct FP8MMARS { + template + CUTLASS_DEVICE static void call_fma_impl(uint32_t* a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(a[0], a[1], a[2], a[3], desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(uint32_t* a, uint64_t const& desc_b, float* d, bool scale_d) { + call_fma_impl(a, desc_b, d, scale_d, cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 32; + static constexpr int kNumAccum = M * N / 128; +}; + +// Mirror of `FP8MMASelector` but selects RS-mode (A in regs) variants. +// Kept as a separate selector to preserve binary compatibility with all +// existing callers of `FP8MMASelector`. +template +struct FP8MMASelectorRS { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 168) return MMA_64x168x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 176) return MMA_64x176x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 184) return MMA_64x184x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 200) return MMA_64x200x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 208) return MMA_64x208x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 216) return MMA_64x216x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 224) return MMA_64x224x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 232) return MMA_64x232x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 240) return MMA_64x240x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 248) return MMA_64x248x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 256) return MMA_64x256x32_F32E4M3E4M3_RS_TN(); + } + + static constexpr auto select_type() { + return FP8MMARS(); + } + + using type = decltype(select_type()); +}; + template struct BF16MMA { template diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index c3e03bec73..398d131a16 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -138,6 +138,42 @@ CUTLASS_DEVICE void st_shared(const __int128_t* ptr, __int128_t val) { asm volatile("st.shared.b128 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "q"(val)); } +CUTLASS_DEVICE void st_shared(const float4* ptr, float4 val) { + asm volatile("st.shared.v4.f32 [%0], {%1, %2, %3, %4};" :: + "l"(__cvta_generic_to_shared(ptr)), + "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)); +} + +// 16-byte vector store as two 64-bit values (st.shared.v2.u64). Useful when +// the data is naturally produced as two 64-bit lanes; saves one instruction +// versus issuing 2 separate `st.shared.u64`. +CUTLASS_DEVICE void st_shared_v2_u64(void* ptr, uint64_t a, uint64_t b) { + asm volatile("st.shared.v2.u64 [%0], {%1, %2};" :: + "l"(__cvta_generic_to_shared(ptr)), "l"(a), "l"(b)); +} + +/// Async copy from global to shared (cp.async) +// 16-byte cache-global async copy +CUTLASS_DEVICE void cp_async_16(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "l"(__cvta_generic_to_shared(smem_dst)), "l"(gmem_src)); +} + +// 4-byte cache-all async copy (for scalar / strided fallback) +CUTLASS_DEVICE void cp_async_4(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" :: + "l"(__cvta_generic_to_shared(smem_dst)), "l"(gmem_src)); +} + +CUTLASS_DEVICE void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template +CUTLASS_DEVICE void cp_async_wait_group() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); +} + CUTLASS_DEVICE void st_shared_bulk(void* smem_ptr, const uint32_t& num_bytes) { // `size` must be 64-bit before PTX ISA 9.0 asm volatile("st.bulk.weak.shared::cta [%0], %1, 0;" :: diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index f1582ed560..3153d41cd0 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -140,4 +140,85 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128 x_dequantized = _dequantize_from_fp4_e2m1(unpacked) group_idx = torch.arange(n, device=packed.device) // gran_k x_restored = x_dequantized * sf[:, group_idx] - return x_restored \ No newline at end of file + return x_restored + + +# --------------------------------------------------------------------------- +# INT4 symmetric (signed [-8, 7]) helpers used by sm90 INT4-A8 path. +# 4-bit two's-complement codes: 0..7 -> 0..7, 8..15 -> -8..-1. +# Two nibbles are packed into one int8 byte, low nibble first. +# +# All 16 INT4 sym values are exactly representable in FP8 E4M3, so we provide +# a lossless INT4 -> packed E4M3 byte converter that lets us drive the existing +# FP8xFP8 fused kernel with INT4-quantised B for end-to-end accuracy testing +# while a dedicated INT4-A8 fused kernel is being landed. +# --------------------------------------------------------------------------- + +def _int4_sym_decode(nibble: torch.Tensor) -> torch.Tensor: + """Map low-4-bit two's-complement nibble to signed int in [-8, 7].""" + nib = (nibble & 0x0F).to(torch.int32) + return torch.where(nib >= 8, nib - 16, nib) + + +def per_token_cast_to_int4_sym(x: torch.Tensor, gran_k: int = 128 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-(row, K-block) symmetric INT4 quantisation. + + Returns (packed, sf): + packed: int8 tensor of shape (m, n // 2), two nibbles per byte (low first). + sf: fp32 tensor of shape (m, n // gran_k), absmax / 7. + """ + assert x.dim() == 2 + m, n = x.shape + assert n % 2 == 0 + padded_n = align(n, gran_k) + x_padded = torch.zeros((m, padded_n), dtype=x.dtype, device=x.device) + x_padded[:, :n] = x + x_view = x_padded.view(m, padded_n // gran_k, gran_k) + x_amax = x_view.abs().float().amax(dim=2).clamp_min(1e-4) + sf = x_amax / 7.0 + x_scaled = x_view.float() * (1.0 / sf.unsqueeze(2)) + # Round-half-to-even and clamp to int4 sym range. + q = torch.round(x_scaled).clamp_(-8, 7).to(torch.int8).view(m, padded_n) + # Pack two nibbles per byte; mask with 0x0F first to drop the sign extension. + q2 = q.view(m, padded_n // 2, 2) + packed = ((q2[:, :, 0] & 0x0F) | ((q2[:, :, 1] & 0x0F) << 4)).to(torch.int8) + return packed[:, : n // 2].contiguous(), sf + + +def cast_back_from_int4_sym(packed: torch.Tensor, sf: torch.Tensor, + gran_k: int = 128) -> torch.Tensor: + """Dequantise INT4-sym packed weights back to fp32.""" + m, n2 = packed.shape + n = n2 * 2 + lo = _int4_sym_decode(packed) + hi = _int4_sym_decode(packed >> 4) + unpacked = torch.empty((m, n), dtype=torch.float, device=packed.device) + unpacked[:, 0::2] = lo.float() + unpacked[:, 1::2] = hi.float() + group_idx = torch.arange(n, device=packed.device) // gran_k + return unpacked * sf[:, group_idx] + + +# Lossless 16-entry LUT: INT4 sym (4-bit two's-complement) -> FP8 E4M3 byte. +# Verified: every signed int in [-8, 7] is exactly representable in E4M3, so +# casting q.float().to(torch.float8_e4m3fn) round-trips identically. +def int4_sym_to_e4m3_bytes(packed_int4: torch.Tensor) -> torch.Tensor: + """Convert packed INT4-sym (2 nibbles / byte) to unpacked E4M3 bytes. + + Args: + packed_int4: int8 tensor (..., n // 2). Low nibble is element 2*i, + high nibble is element 2*i+1. + Returns: + torch.float8_e4m3fn tensor (..., n) with byte-identical encoding to + a fresh fp32 -> fp8_e4m3 cast of the dequantised integer values. + """ + assert packed_int4.dtype == torch.int8 + *prefix, n2 = packed_int4.shape + n = n2 * 2 + lo = _int4_sym_decode(packed_int4) + hi = _int4_sym_decode(packed_int4 >> 4) + unpacked = torch.empty((*prefix, n), dtype=torch.int32, device=packed_int4.device) + unpacked[..., 0::2] = lo + unpacked[..., 1::2] = hi + return unpacked.float().to(torch.float8_e4m3fn) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py new file mode 100644 index 0000000000..9f2f687d89 --- /dev/null +++ b/tests/test_sm90_fp8_fp4.py @@ -0,0 +1,921 @@ +import sys +import time +import os +from pathlib import Path + +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 ( + cast_back_from_fp4, + per_token_cast_to_fp4, + per_token_cast_to_fp8, +) + + +def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int = 128) -> torch.Tensor: + group_idx = torch.arange(x.size(-1), device=x.device) // gran_k + return x.float() * sf[..., group_idx] + + +def _require_sm90() -> None: + assert torch.cuda.is_available() + major, _ = torch.cuda.get_device_capability() + if major != 9: + raise RuntimeError(f"This benchmark is intended for SM90, got sm_{major}x") + + +def _time_cuda(fn, warmup: int = 3, iters: int = 10) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 + + +def _effective_bytes( + groups: int, + m_per_group: int, + n: int, + k: int, + a_gran_k: int, + *, + fp8_b: bool, + b_gran_k: int | None = None, +) -> int: + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + logical_m = groups * m_per_group + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _effective_masked_bytes( + groups: int, + masked_m_values: list[int], + n: int, + k: int, + a_gran_k: int, + *, + fp8_b: bool, + b_gran_k: int | None = None, +) -> int: + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + logical_m = sum(masked_m_values) + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _build_grouped_layout(groups: int, m_per_group: int): + m = groups * m_per_group + group_starts = [group_id * m_per_group for group_id in range(groups)] + group_ends = [(group_id + 1) * m_per_group for group_id in range(groups)] + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) + return m, group_starts, group_ends, grouped_layout + + +def _benchmark_case(groups: int, m_per_group: int, n: int, k: int, gran_k: int = 128) -> dict[str, float | int]: + m, group_starts, group_ends, grouped_layout = _build_grouped_layout(groups, m_per_group) + a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for group_id in range(groups): + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=gran_k + ) + b_w4 = (b_fp4, b_sf) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + b_dequant = cast_back_from_fp4(b_w4[0][group_id], b_w4[1][group_id], gran_k=gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=gran_k + ) + start = group_starts[group_id] + end = group_ends[group_id] + if start != end: + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + use_psum_layout=False, + ) + + run_fp8() + fp8_diff = calc_diff(d_fp8, ref) + fp8_elapsed = _time_cuda(run_fp8) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b_w4, + d_w4, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=False, + ) + + run_w4() + w4_diff = calc_diff(d_w4, ref) + w4_elapsed = _time_cuda(run_w4) + + fp8_threshold = 0.05 + assert w4_diff < 0.015 + assert fp8_diff < fp8_threshold + + w4_bytes = _effective_bytes(groups, m_per_group, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_bytes(groups, m_per_group, n, k, gran_k, fp8_b=True) + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + + +def _print_markdown_table(rows: list[dict[str, float | int]]) -> None: + print("groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for row in rows: + prefix = f"{row['groups']} | {row['m_per_group']} | {row['n']} | {row['k']} | " + print( + prefix + + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " + f"{row['speedup']:.2f}x" + ) + + +def _masked_benchmark_case( + groups: int, + m_per_group: int, + n: int, + k: int, + a_gran_k: int = 128, + b_gran_k: int = 32, +) -> dict[str, float | int]: + sm90_masked_w4 = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if sm90_masked_w4 is None: + raise RuntimeError( + "SM90 FP8xFP4 masked fused kernel is not exposed yet. " + "Do not call generic m_grouped_fp8_fp4_gemm_nt_masked on SM90; " + "it is currently routed to the SM100 FP8xFP4 masked path." + ) + + max_m = 128 + masked_m = torch.full((groups,), m_per_group, device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = 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, k // a_gran_k), device="cuda", dtype=torch.float) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) + b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k + ) + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + a = (a_data, a_sf) + b_w4 = (b_fp4, b_sf) + + assert a[1].shape == (groups, max_m, k // a_gran_k) + assert b_w4[1].shape == (groups, n, b_sf_k) + if b_gran_k == 128 and not use_packed_b_sf: + assert b_w4[1].dtype == torch.float + assert b_w4[1].shape == (groups, n, k // 128) + + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + valid_m = int(masked_m[group_id].item()) + b_dequant = cast_back_from_fp4( + b_w4[0][group_id], b_w4[1][group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=a_gran_k + ) + if valid_m > 0: + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + sm90_masked_w4( + a, + b_w4, + d_w4, + masked_m, + m_per_group, + gran_k=a_gran_k, + gran_k_a=a_gran_k, + gran_k_b=b_gran_k, + ) + + run_w4() + w4_diff = max( + calc_diff(d_w4[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + w4_elapsed = _time_cuda(run_w4) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, + b_fp8, + d_fp8, + masked_m, + m_per_group, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + ) + + run_fp8() + fp8_diff = max( + calc_diff(d_fp8[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + fp8_elapsed = _time_cuda(run_fp8) + + #assert w4_diff < 0.015 + #assert fp8_diff < 0.05 + + w4_bytes = _effective_bytes(groups, m_per_group, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_bytes(groups, m_per_group, n, k, a_gran_k, fp8_b=True) + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + + +def _masked_skew_benchmark_case( + name: str, + masked_m_values: list[int], + expected_m: int, + n: int, + k: int, + max_m: int = 1024, + a_gran_k: int = 128, + b_gran_k: int = 32, + pass_hints: bool = True, +) -> dict[str, float | int | str]: + sm90_masked_w4 = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if sm90_masked_w4 is None: + raise RuntimeError("SM90 FP8xFP4 masked fused kernel is not exposed yet.") + + groups = len(masked_m_values) + assert groups > 0 + assert max(masked_m_values) <= max_m + masked_m_max_hint = max(masked_m_values) if pass_hints else None + active_groups_hint = ( + sum(1 for v in masked_m_values if v > 0) if pass_hints else None + ) + masked_m = torch.tensor(masked_m_values, device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = 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, k // a_gran_k), device="cuda", dtype=torch.float) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) + b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) + block_m_override = int(os.getenv("DG_W4_BLOCK_M_OVERRIDE", "0")) or None + block_n_override = int(os.getenv("DG_W4_BLOCK_N_OVERRIDE", "0")) or None + bm32_skew_fast_path = ( + b_gran_k == 32 + and masked_m_max_hint is not None + and masked_m_max_hint > 16 + and os.getenv("DG_W4_PATHB_FUSE_DECODE", "0") == "0" + and os.getenv("DG_W4_PATHB_FAST_PATH", "1") != "0" + and os.getenv("DG_W4_PATHB_BM64", "0") == "0" + and (block_m_override is None or block_m_override == 32) + and (block_n_override is None or block_n_override in (128, 256)) + and max_m >= 1024 + and groups == 32 + and n in (4096, 7168) + and k in (2048, 3072, 4096, 7168) + ) + scale_b_direct_load = b_gran_k == 32 and ( + expected_m <= 16 or bm32_skew_fast_path + ) + scale_b_dtype_fast_path = ( + scale_b_direct_load + and os.getenv("DG_W4_K32_QUAD_SCALE_B_PREFETCH", "0") == "0" + and os.getenv("DG_W4_SCALE_B_POW2_PROMOTE", "0") == "0" + ) + # bf16 sfb：path-A (gran_k_b=128) 与 path-B fast-path (gran_k_b=32) 都支持， + # 体积砍半。按 MN-major + tma_aligned_mn=align(N, 8) 直接构造，避开 host 端 + # fp32-only 的 transpose 路径。**默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 + use_bf16_b_sf = ( + ((b_gran_k == 32 and scale_b_dtype_fast_path) or b_gran_k == 128) + and not use_packed_b_sf + and bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) + ) + # E8M0 sfb（uint8）：仅 path-B (gran_k_b=32)。每元素 = fp32 pow2 scale 的指数位。 + # per_token_cast_to_fp4(..., use_ue8m0=True) 已保证 sf 是严格 pow2，因此抽指数无损。 + # **默认开启**，DG_W4_SCALE_B_E8M0=0 时回退。优先级高于 bf16（互斥）。 + use_e8m0_b_sf = ( + b_gran_k == 32 + and scale_b_dtype_fast_path + and not use_packed_b_sf + and bool(int(os.getenv("DG_W4_SCALE_B_E8M0", "1"))) + ) + if use_e8m0_b_sf: + use_bf16_b_sf = False # e8m0 优先 + if use_e8m0_b_sf: + # uint8: tma_aligned_mn = ceil(N, 16)，要求 N % 16 == 0。 + assert n % 16 == 0 + tma_aligned_n = (n + 15) // 16 * 16 + b_sf = torch.empty_strided( + (groups, n, b_sf_k), + (tma_aligned_n * b_sf_k, 1, tma_aligned_n), + device="cuda", + dtype=torch.uint8, + ) + b_sf_fp32 = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) + elif use_bf16_b_sf: + # bf16: tma_aligned_mn = ceil(N, 8)，要求 N % 8 == 0。 + assert n % 8 == 0 + tma_aligned_n = (n + 7) // 8 * 8 + b_sf = torch.empty_strided( + (groups, n, b_sf_k), + (tma_aligned_n * b_sf_k, 1, tma_aligned_n), + device="cuda", + dtype=torch.bfloat16, + ) + b_sf_fp32 = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) + else: + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + b_sf_fp32 = b_sf + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k + ) + b_fp4[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + if use_e8m0_b_sf: + # b_sf_fp32 已是严格 pow2（per_token_cast_to_fp4 use_ue8m0=True）。 + # 抽 fp32 指数位 = ((bits >> 23) & 0xff)：sign=0、mantissa=0 时 fp32 == 2^(e-127)。 + b_sf_bits = b_sf_fp32.view(torch.int32) + b_sf_e8m0 = ((b_sf_bits >> 23) & 0xff).to(torch.uint8) + b_sf.copy_(b_sf_e8m0) + elif use_bf16_b_sf: + # 写入 MN-major bf16 buffer：fp32 → bf16 round-down，再按目标 stride copy。 + b_sf.copy_(b_sf_fp32.to(torch.bfloat16)) + a = (a_data, a_sf) + b_w4 = (b_fp4, b_sf) + + # caller hint：把 hot group 的真实大小传给 host，host 据此选 BM。 + # path-A 通过 masked_m_max_hint 接收 hint；FP8 baseline 没有 hint API， + # 但它的 expected_m 可以直接传 max（FP8 路径无 expected_m<=8 fast-path 副作用）， + # 这样两条路径都按 hot 调度，speedup 反映的是算法差距而非 caller-side gap。 + gemm_expected_m = expected_m + # pass_hints=False 时 max_hint=None：两边都用 expected_m（mirror 业务真实 + # caller 不传 hint 的状态，host 走没有 hint 的 small-m candidate 路径）。 + fp8_expected_m = ( + max(masked_m_max_hint, expected_m) + if masked_m_max_hint is not None + else expected_m + ) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + # ref 反量化必须用 fp32 scale（cast_back_from_fp4 直接做 sf 乘法，不会解码 e8m0/bf16）。 + # use_e8m0_b_sf / use_bf16_b_sf 时 b_w4[1] 是 uint8/bf16 编码，需走 b_sf_fp32 兜底。 + b_sf_for_ref = b_sf_fp32 if (use_e8m0_b_sf or use_bf16_b_sf) else b_w4[1] + for group_id, valid_m in enumerate(masked_m_values): + b_dequant = cast_back_from_fp4( + b_w4[0][group_id], b_sf_for_ref[group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=a_gran_k + ) + if valid_m > 0: + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + sm90_masked_w4( + a, + b_w4, + d_w4, + masked_m, + gemm_expected_m, + gran_k=a_gran_k, + gran_k_a=a_gran_k, + gran_k_b=b_gran_k, + masked_m_max_hint=masked_m_max_hint, + active_groups_hint=active_groups_hint, + ) + + run_w4() + w4_diff = max( + calc_diff(d_w4[group_id, :valid_m], ref[group_id, :valid_m]) if valid_m > 0 else 0.0 + for group_id, valid_m in enumerate(masked_m_values) + ) + w4_elapsed = _time_cuda(run_w4) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, + b_fp8, + d_fp8, + masked_m, + fp8_expected_m, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + ) + + run_fp8() + fp8_diff = max( + calc_diff(d_fp8[group_id, :valid_m], ref[group_id, :valid_m]) if valid_m > 0 else 0.0 + for group_id, valid_m in enumerate(masked_m_values) + ) + fp8_elapsed = _time_cuda(run_fp8) + + w4_bytes = _effective_masked_bytes(groups, masked_m_values, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_masked_bytes(groups, masked_m_values, n, k, a_gran_k, fp8_b=True) + return { + "case": name, + "groups": groups, + "expected_m": expected_m, + "masked_m_hint": masked_m_max_hint, + "b_gran_k": b_gran_k, + "max_m": max_m, + "sum_m": sum(masked_m_values), + "masked_max": max(masked_m_values), + "active_groups": sum(1 for value in masked_m_values if value > 0), + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + + +def _print_skew_table(rows: list[dict[str, float | int | str]]) -> None: + print("case | groups | b_gran_k | exp_m_avg | hint | sum_m | max_m | active | n | k | " + "W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for row in rows: + print( + f"{row['case']} | {row['groups']} | {row['b_gran_k']} | {row['expected_m']} | " + f"{row['masked_m_hint']} | {row['sum_m']} | " + f"{row['masked_max']} | {row['active_groups']} | {row['n']} | {row['k']} | " + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " + f"{row['speedup']:.2f}x" + ) + + +def _accuracy_case( + groups: int, + m_per_group: int, + n: int, + k: int, + gran_k: int = 128, + *, + block_m: int = 128, + block_n: int = 128, +) -> tuple[float, float]: + m, group_starts, group_ends, grouped_layout = _build_grouped_layout(groups, m_per_group) + a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for group_id in range(groups): + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=gran_k + ) + b_w4 = (b_fp4, b_sf) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + b_dequant = cast_back_from_fp4(b_w4[0][group_id], b_w4[1][group_id], gran_k=gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=gran_k + ) + start = group_starts[group_id] + end = group_ends[group_id] + if start != end: + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b_w4, + d_w4, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=False, + block_m_override=block_m, + block_n_override=block_n, + ) + + d_fp8 = torch.empty_like(ref) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + use_psum_layout=False, + ) + + return calc_diff(d_w4, ref), calc_diff(d_fp8, ref) + + +def test_sm90_fp8_fp4_contiguous() -> None: + _require_sm90() + torch.manual_seed(0) + + rows = [] + for groups in (8, 16, 24, 32): + for m_per_group in (128, 256, 512, 1024): + rows.append(_benchmark_case(groups, m_per_group, n=4096, k=7168)) + _print_markdown_table(rows) + + +def test_sm90_fp8_fp4_masked() -> None: + _require_sm90() + torch.manual_seed(2) + + print("direct E8M0 B scale case: b.second shape = [groups, N, K/32]") + + rows = [] + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048)) + _print_markdown_table(rows) + +def test_sm90_fp8_fp4_masked_direct_fp32_scale() -> None: + _require_sm90() + torch.manual_seed(3) + + print("direct FP32 B scale case: b.second shape = [groups, N, K/128]") + rows = [] + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168, b_gran_k=128)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048, b_gran_k=128)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168, b_gran_k=128)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048, b_gran_k=128)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168, b_gran_k=128)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048, b_gran_k=128)) + _print_markdown_table(rows) + + +def test_sm90_fp8_fp4_masked_skew_cases() -> None: + _require_sm90() + torch.manual_seed(4) + + def skew_values( + total: int, hot: int, active: int = 8, groups: int = 32 + ) -> list[int]: + assert active >= 1 + assert active <= groups + assert total >= hot + values = [0] * groups + values[0] = hot + remaining = total - hot + for idx in range(1, active): + share = (remaining + active - idx - 1) // (active - idx) + values[idx] = share + remaining -= share + assert sum(values) == total + return values + + def values_from_active(active_values: list[int], groups: int = 32) -> list[int]: + assert len(active_values) <= groups + assert all(value >= 0 for value in active_values) + return active_values + [0] * (groups - len(active_values)) + + def values_from_repeated(value: int, active: int, groups: int = 32) -> list[int]: + assert active <= groups + return values_from_active([value] * active, groups=groups) + + def long_tail_values( + hot: int, tail_start: int, active: int, groups: int = 32 + ) -> list[int]: + values = [hot] + next_value = tail_start + for _ in range(active - 1): + values.append(max(1, next_value)) + next_value = max(1, next_value // 2) + return values_from_active(values, groups=groups) + + print("skewed masked case: shape mirrors DSV4 MTP verify; " + "groups=32, b_gran_k=32 walks path-B (k32 fast path) " + "with b.second shape [groups, N, K/32]") + rows = [] + shapes = [ + ("gateup", 4096, 4096), + ("down", 4096, 2048), + # Non-DSV4 dimensions keep the same group/masked pattern but stress + # different N/K ratios that can expose scheduler or scale-load cliffs. + ("wide_n", 7168, 2048), + ("wide_k", 4096, 7168), + ] + distributions = [ + # Uniform small-M cases validate that hint-aware BM32 selection does not + # regress the common non-skew path. + ("uniform_1", [1] * 32, 1), + ("uniform_2", [2] * 32, 2), + ("uniform_4", [4] * 32, 4), + ("uniform_8", [8] * 32, 8), + ("uniform_16", [16] * 32, 16), + ("uniform_32", [32] * 32, 32), + # Around the BM16 -> BM32 hot threshold. + ("one_hot_15", values_from_active([15]), 1), + ("one_hot_16", values_from_active([16]), 1), + ("one_hot_17", values_from_active([17]), 1), + ("one_hot_24", values_from_active([24]), 1), + ("one_hot_32", values_from_active([32]), 1), + ("one_hot_48", values_from_active([48]), 2), + ("one_hot_64", values_from_active([64]), 2), + ("one_hot_128", values_from_active([128]), 4), + ("one_hot_214", values_from_active([214]), 1), + ("one_hot_384", values_from_active([384]), 12), + # Original MTP verify distributions copied from observed DP logs. + ("mtp_dp2", skew_values(total=144, hot=50), 7), + ("mtp_dp0", skew_values(total=195, hot=160), 7), + ("mtp_dp4", skew_values(total=290, hot=214), 7), + # Same total token count but different active/hotness patterns. + ("mtp_144_hot96_a4", skew_values(total=144, hot=96, active=4), 7), + ("mtp_144_hot96_a8", skew_values(total=144, hot=96, active=8), 7), + ("mtp_195_hot96_a16", skew_values(total=195, hot=96, active=16), 7), + ("mtp_290_hot160_a16", skew_values(total=290, hot=160, active=16), 7), + ("mtp_384_hot256_a8", skew_values(total=384, hot=256, active=8), 12), + ("mtp_512_hot384_a8", skew_values(total=512, hot=384, active=8), 16), + # Multi-hot cases mimic router concentration on a few experts rather + # than a single dominant expert. + ("two_hot_64_64", values_from_active([64, 64]), 4), + ("two_hot_128_64", values_from_active([128, 64]), 6), + ("two_hot_160_96", values_from_active([160, 96]), 8), + ("four_hot_32", values_from_repeated(32, active=4), 4), + ("four_hot_64", values_from_repeated(64, active=4), 8), + ("eight_hot_32", values_from_repeated(32, active=8), 8), + ("eight_hot_64", values_from_repeated(64, active=8), 16), + # Long tails stress compact masked scheduling and active-group scanning. + ("longtail_128_a8", long_tail_values(hot=128, tail_start=64, active=8), 8), + ("longtail_214_a8", long_tail_values(hot=214, tail_start=48, active=8), 7), + ("longtail_256_a16", long_tail_values(hot=256, tail_start=64, active=16), 12), + # Dense active but skewed cases can happen when all experts receive a + # few tokens and one or two experts still become hot. + ( + "dense_tail_hot64", + values_from_active([64, 32, 16, 8] + [4] * 28), + 8, + ), + ( + "dense_tail_hot128", + values_from_active([128, 64, 32, 16] + [4] * 28), + 8, + ), + ( + "dense_tail_hot214", + values_from_active([214, 64, 32, 16] + [4] * 28), + 8, + ), + ] + # for shape_name, n, k in shapes: + # for dist_name, masked_m_values, expected_m in distributions: + # rows.append( + # _masked_skew_benchmark_case( + # f"{shape_name}_{dist_name}", + # masked_m_values, + # expected_m=expected_m, + # n=n, + # k=k, + # max_m=1024, + # b_gran_k=32, + # ) + # ) + + group24_shapes = [ + ("g24_m4096_n6144_k7168", 6144, 7168, 4096), + ("g24_m4096_n7168_k3072", 7168, 3072, 4096), + ] + group24_distributions = [ + # Reproduce the warmup cliff that showed up as: + # m=17, max_m=4096, n=6144, k=7168, num_groups=24. + ("uniform_1", values_from_repeated(1, active=24, groups=24), 1), + ("uniform_8", values_from_repeated(8, active=24, groups=24), 8), + ("uniform_16", values_from_repeated(16, active=24, groups=24), 16), + ("uniform_17", values_from_repeated(17, active=24, groups=24), 17), + ("uniform_24", values_from_repeated(24, active=24, groups=24), 24), + ("uniform_32", values_from_repeated(32, active=24, groups=24), 32), + ("uniform_64", values_from_repeated(64, active=24, groups=24), 64), + # Around the E8M0 direct-load threshold and BM16/BM32 transition. + ("one_hot_16", values_from_active([16], groups=24), 16), + ("one_hot_17", values_from_active([17], groups=24), 17), + ("one_hot_32", values_from_active([32], groups=24), 32), + ("one_hot_64", values_from_active([64], groups=24), 64), + ("one_hot_128", values_from_active([128], groups=24), 128), + ("one_hot_256", values_from_active([256], groups=24), 256), + ("one_hot_512", values_from_active([512], groups=24), 512), + ("two_hot_17", values_from_active([17, 17], groups=24), 17), + ("two_hot_64_64", values_from_active([64, 64], groups=24), 64), + ("two_hot_128_64", values_from_active([128, 64], groups=24), 128), + ("four_hot_17", values_from_repeated(17, active=4, groups=24), 17), + ("four_hot_32", values_from_repeated(32, active=4, groups=24), 32), + ("four_hot_64", values_from_repeated(64, active=4, groups=24), 64), + ("eight_hot_17", values_from_repeated(17, active=8, groups=24), 17), + ("eight_hot_32", values_from_repeated(32, active=8, groups=24), 32), + ("eight_hot_64", values_from_repeated(64, active=8, groups=24), 64), + # MTP-like skew for 24 local groups, including the m=17 average. + ("mtp_408_hot160_a8", skew_values(408, 160, active=8, groups=24), 17), + ("mtp_408_hot256_a8", skew_values(408, 256, active=8, groups=24), 17), + ("mtp_576_hot384_a8", skew_values(576, 384, active=8, groups=24), 24), + ("mtp_768_hot512_a8", skew_values(768, 512, active=8, groups=24), 32), + # Long tails and dense active tails stress compact scheduling when + # active_groups is neither tiny nor fully uniform. + ( + "longtail_128_a8", + long_tail_values(128, tail_start=64, active=8, groups=24), + 17, + ), + ( + "longtail_256_a12", + long_tail_values(256, tail_start=96, active=12, groups=24), + 24, + ), + ( + "dense_tail_hot17", + values_from_active([17, 16, 8, 4] + [1] * 20, groups=24), + 17, + ), + ( + "dense_tail_hot64", + values_from_active([64, 32, 16, 8] + [4] * 20, groups=24), + 17, + ), + ( + "dense_tail_hot128", + values_from_active([128, 64, 32, 16] + [4] * 20, groups=24), + 24, + ), + ( + "dense_tail_hot256", + values_from_active([256, 128, 64, 32] + [8] * 20, groups=24), + 32, + ), + ] + for shape_name, n, k, max_m in group24_shapes: + for dist_name, masked_m_values, expected_m in group24_distributions: + rows.append( + _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, + k=k, + max_m=max_m, + b_gran_k=32, + ) + ) + + # DSV4 EP 真实业务 shape mirror：m=256 + expected_m∈{1,2,3} + max_hint=None + # （caller 没传 hint，business dispatch_output 不携带 max_hint/active_hint）。 + # 三组 shape：gateup(g24, n=6144, k=7168) / down(g24, n=7168, k=3072) + # 业务比例：expected_m=2 占 71%, expected_m=3 占 18%, expected_m=1 占 11%。 + # 物理 m=256 但 caller 不知道每个 group 的真实 masked_m，masked_m_values 设 + # expected_m * groups 模拟 uniform 分布，用 pass_hints=False 让 host 走没有 + # hint 的 small-m candidate + simple_sched 路径。 + business_shapes = [ + ("biz_gateup_g24_m256_n6144_k7168", 6144, 7168, 256), + ("biz_down_g24_m256_n7168_k3072", 7168, 3072, 256), + ] + business_distributions = [ + ("expected_m_1", values_from_repeated(1, active=24, groups=24), 1), + ("expected_m_2", values_from_repeated(2, active=24, groups=24), 2), + ("expected_m_3", values_from_repeated(3, active=24, groups=24), 3), + ] + for shape_name, n, k, max_m in business_shapes: + for dist_name, masked_m_values, expected_m in business_distributions: + rows.append( + _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, + k=k, + max_m=max_m, + b_gran_k=32, + pass_hints=False, + ) + ) + _print_skew_table(rows) + + # print("\nworst W4/FP8 speedup cases") + # _print_skew_table(sorted(rows, key=lambda row: float(row["speedup"]))[:12]) + + +if __name__ == "__main__": + start_time = time.time() + # if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): + # test_sm90_fp8_fp4_contiguous() + if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + test_sm90_fp8_fp4_masked_skew_cases() + elif os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): + test_sm90_fp8_fp4_masked_direct_fp32_scale() + else: + test_sm90_fp8_fp4_masked() + print(f"done in {time.time() - start_time:.2f}s") diff --git a/tests/test_sm90_int4_a8.py b/tests/test_sm90_int4_a8.py new file mode 100644 index 0000000000..a9d654ed4f --- /dev/null +++ b/tests/test_sm90_int4_a8.py @@ -0,0 +1,524 @@ +"""Accuracy + benchmark for INT4-A8 (B = INT4 sym, A = FP8 e4m3) masked GEMM +on SM90, exercising the dedicated INT4-sym device path. + +Pipeline: + - B is symmetric INT4 (signed [-8, 7]) packed two nibbles per byte. The + wire format is byte-identical to packed FP4 (kPackedFP4 == kInt8). + - A is FP8 e4m3 with per-(row, K-block=128) fp32 scales. + - SFB is per-(row, K-block=128) fp32 (path-A). + - The kernel decodes B nibbles in registers via int4_symx4_to_e4m3x4 + instead of fp4x4_to_e4m3x4. This is selected via the new + `b_is_int4_sym=True` argument on + `m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma`. + +Verification +------------ +INT4-sym 的 16 个码点在 E4M3 中**无损可表示**，因此 INT4 kernel 与 fp32 ref +的差距应当极小（与 bf16 cast + fp32 累加噪声同量级，即 cos > 0.999）。这是 +本测试的硬正确性指标 —— **cos_abs**。 + +作为对照，我们还跑一条 FP8(b_dequant) baseline：把 INT4 反量化回 fp32 后 +再 re-quantize 成 fp8 喂进常规 FP8 kernel。这条路径自身要走一次 fp8 取整 +（3-bit mantissa），相对 fp32 ref 会落到 ~cos 0.98 的 FP8 误差墙；INT4 与 +它的差距正好反映这层 re-quant 噪声 —— **cos_eq**，仅做信息性记录，不参与 +pass/fail 判定。 + + cos_abs = cos(INT4-kernel out, fp32 INT4-dequant ref) > 0.99 [硬指标] + cos_eq = cos(INT4-kernel out, FP8(b_dequant)-kernel) info only + +性能列与 test_sm90_fp8_fp4.py 一致（GB/s 用 effective bytes 模型）。 +""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import deep_gemm +from deep_gemm.utils.math import ( + cast_back_from_int4_sym, + per_token_cast_to_fp8, + per_token_cast_to_int4_sym, +) + + +COS_ABS_THRESHOLD = 0.99 # INT4-sym is lossless in E4M3, must hit ~1.0 +COS_EQ_INFO_ONLY = True # cos vs FP8(b_dequant) baseline is informative, + # not a pass/fail signal: the FP8 baseline is the + # one with re-quant noise. + + +def _require_sm90() -> None: + assert torch.cuda.is_available() + major, _ = torch.cuda.get_device_capability() + if major != 9: + raise RuntimeError(f"This test is intended for SM90, got sm_{major}x") + + +def _resolve_kernel(): + fn = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if fn is None: + raise RuntimeError( + "deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma not found. " + "Rebuild the C++ extension after the INT4-sym wiring." + ) + return fn + + +def _time_cuda(fn, warmup: int = 3, iters: int = 10) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 + + +def _cos_sim(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.float().reshape(-1) + b = b.float().reshape(-1) + return float(F.cosine_similarity(a, b, dim=0)) + + +def _max_abs_err(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max()) + + +def _build_int4_a8_inputs(groups: int, m_per_group: int, max_m: int, + n: int, k: int, *, gran_k: int = 128): + """Build one INT4-A8 problem with two parallel B representations. + + INT4 path: + - b_int4: packed signed [-8, 7] nibbles, dtype int8, shape (G, n, k/2) + - b_int4_sf: per-(n, K-block=128) fp32 scale, shape (G, n, k/128) + + FP8 baseline path (for equivalence check): + - b_fp8: round-tripped fp8_e4m3, shape (G, n, k) + - b_fp8_sf: per-(n, K-block=128) fp32 scale, shape (G, n, k/128) + Built by dequantising INT4 to fp32 and re-quantising via + per_token_cast_to_fp8 (use_ue8m0=False). + """ + masked_m = torch.full((groups,), m_per_group, device="cuda", dtype=torch.int32) + + # ---- A: bf16 -> per-(row, 128-block) fp32-scale fp8_e4m3 ---- + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + a_fp8 = torch.empty_like(a_ref, dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + a_fp8[g], a_sf[g] = per_token_cast_to_fp8(a_ref[g], use_ue8m0=False, gran_k=gran_k) + + # ---- B: bf16 -> per-(row, 128-block) INT4 sym ---- + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + b_int4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_int4_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_int4[g], b_int4_sf[g] = per_token_cast_to_int4_sym(b_ref[g], gran_k=gran_k) + + # ---- fp32 dequant references (for ground-truth matmul) ---- + group_idx_a = torch.arange(k, device="cuda") // gran_k + a_dequant = a_fp8.float() * a_sf[..., group_idx_a] + b_dequant = torch.empty((groups, n, k), device="cuda", dtype=torch.float) + for g in range(groups): + b_dequant[g] = cast_back_from_int4_sym(b_int4[g], b_int4_sf[g], gran_k=gran_k) + + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for g in range(groups): + valid_m = int(masked_m[g].item()) + if valid_m == 0: + continue + ref[g, :valid_m] = (a_dequant[g, :valid_m] @ b_dequant[g].t()).to(torch.bfloat16) + + # ---- FP8(b_dequant) baseline path (equivalence reference) ---- + b_fp8 = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_fp8[g], b_fp8_sf[g] = per_token_cast_to_fp8( + b_dequant[g].to(torch.bfloat16), use_ue8m0=False, gran_k=gran_k, + ) + + return dict( + masked_m=masked_m, + a_fp8=a_fp8, a_sf=a_sf, + b_int4=b_int4, b_int4_sf=b_int4_sf, + b_fp8=b_fp8, b_fp8_sf=b_fp8_sf, + ref=ref, + ) + + +def _run_int4_kernel(fn, case, m_per_group, gran_k=128): + """Drive the FP4 masked entry with INT4-sym packed B + b_is_int4_sym=True.""" + a = (case["a_fp8"], case["a_sf"]) + # The kernel's B-first tensor is dtype-checked as kPackedFP4 (== kInt8); + # our int8-packed INT4 fits the wire format directly. + b = (case["b_int4"], case["b_int4_sf"]) + d = torch.empty_like(case["ref"]) + fn( + a, b, d, case["masked_m"], m_per_group, + gran_k=gran_k, gran_k_a=gran_k, gran_k_b=gran_k, + b_is_int4_sym=True, + ) + return d + + +def _run_fp8_baseline(case, m_per_group, gran_k=128): + """Equivalence reference: FP8 fused kernel on (A_fp8, fp8(b_dequant)).""" + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_fp8"], case["b_fp8_sf"]) + d = torch.empty_like(case["ref"]) + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, b, d, case["masked_m"], m_per_group, + recipe_a=(1, gran_k), recipe_b=(1, gran_k), + ) + return d + + +def _per_group_cos_min(d, target, m_per_group): + return min( + _cos_sim(d[g, :m_per_group], target[g, :m_per_group]) + for g in range(d.shape[0]) + ) + + +def _per_group_mae_max(d, target, m_per_group): + return max( + _max_abs_err(d[g, :m_per_group], target[g, :m_per_group]) + for g in range(d.shape[0]) + ) + + +def _effective_masked_bytes(groups, m_per_group, n, k, a_gran_k, *, + fp8_b: bool, b_gran_k=None): + """Bandwidth model identical to test_sm90_fp8_fp4._effective_masked_bytes + (each masked group treats its valid m rows as the logical_m payload).""" + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + logical_m = groups * m_per_group + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _accuracy_case(fn, groups, m_per_group, n, k, *, gran_k=128, seed=0): + torch.manual_seed(seed) + max_m = max(m_per_group, 128) + case = _build_int4_a8_inputs(groups, m_per_group, max_m, n, k, gran_k=gran_k) + + d_int4 = _run_int4_kernel(fn, case, m_per_group, gran_k=gran_k) + d_fp8 = _run_fp8_baseline(case, m_per_group, gran_k=gran_k) + + cos_eq = _per_group_cos_min(d_int4, d_fp8, m_per_group) + cos_abs = _per_group_cos_min(d_int4, case["ref"], m_per_group) + mae_abs = _per_group_mae_max(d_int4, case["ref"], m_per_group) + int4_us = _time_cuda(lambda: _run_int4_kernel(fn, case, m_per_group, gran_k=gran_k)) + fp8_us = _time_cuda(lambda: _run_fp8_baseline(case, m_per_group, gran_k=gran_k)) + + int4_bytes = _effective_masked_bytes(groups, m_per_group, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_masked_bytes(groups, m_per_group, n, k, gran_k, fp8_b=True) + return dict( + groups=groups, m_per_group=m_per_group, n=n, k=k, + cos_eq=cos_eq, cos_abs=cos_abs, max_abs_err=mae_abs, + int4_us=int4_us * 1e6, fp8_us=fp8_us * 1e6, + int4_gbps=int4_bytes / int4_us / 1e9, + fp8_gbps=fp8_bytes / fp8_us / 1e9, + speedup=fp8_us / int4_us, + ) + + +def test_int4_a8_masked_accuracy() -> None: + _require_sm90() + fn = _resolve_kernel() + print("INT4-A8 masked accuracy + perf vs FP8 (masked m_grouped)") + print(f" cos_abs = cos(INT4-kernel, fp32-ref) > {COS_ABS_THRESHOLD:.3f} -- pass/fail") + print(f" cos_eq = cos(INT4-kernel, FP8(b_dequant)-kernel) -- info only") + print() + print("groups | m/group | n | k | cos_abs | cos_eq | " + "INT4 us | INT4 GB/s | FP8 us | FP8 GB/s | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + + pass_count = 0 + fail_count = 0 + failed_rows = [] + for groups in (8, 16, 32): + for m_per_group in (1, 4, 8, 16, 32): + for (n, k) in [(4096, 7168), (7168, 2048), (4096, 4096)]: + row = _accuracy_case(fn, groups, m_per_group, n, k) + ok_abs = row["cos_abs"] > COS_ABS_THRESHOLD + ok = ok_abs + pass_count += int(ok) + fail_count += int(not ok) + marker = "" if ok else f" ## FAIL (cos_abs<={COS_ABS_THRESHOLD})" + print( + f"{row['groups']} | {row['m_per_group']} | {row['n']} | " + f"{row['k']} | {row['cos_abs']:.4f} | {row['cos_eq']:.4f} | " + f"{row['int4_us']:.0f} | {row['int4_gbps']:.0f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | " + f"{row['speedup']:.2f}x{marker}" + ) + if not ok: + failed_rows.append(row) + print(f"\n{pass_count} passed, {fail_count} failed") + if fail_count > 0: + details = "; ".join( + f"g={r['groups']} m={r['m_per_group']} n={r['n']} k={r['k']} " + f"cos_abs={r['cos_abs']:.4f}" + for r in failed_rows[:5] + ) + raise AssertionError( + f"{fail_count} INT4-sym cases failed cos_abs > {COS_ABS_THRESHOLD}: " + + details + (" ..." if len(failed_rows) > 5 else "") + ) + + +# --------------------------------------------------------------------------- +# Skewed (uneven) masked-m benchmark — mirrors test_sm90_fp8_fp4 layout +# --------------------------------------------------------------------------- +# 在不均匀 mask 模式下度量 INT4 vs FP8 的 speedup。每条 case 用一组 per-group +# 的 masked_m_values（长度 == 32 个 group 槽，未激活槽 mask=0）。INT4 路径仍 +# 走 b_is_int4_sym=True；FP8 baseline 用 cast_back_from_int4_sym 反量化后 +# re-quant 到 fp8。bandwidth model 用每 group 的有效 m。 +def _build_int4_a8_skew_inputs(masked_m_values, max_m, n, k, *, gran_k=128): + groups = len(masked_m_values) + masked_m = torch.tensor(masked_m_values, device="cuda", dtype=torch.int32) + + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + a_fp8 = torch.empty_like(a_ref, dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + a_fp8[g], a_sf[g] = per_token_cast_to_fp8(a_ref[g], use_ue8m0=False, gran_k=gran_k) + + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + b_int4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + # path-A bf16 sfb：与 fp8_fp4 测试同款，按 MN-major + tma_aligned_mn=align(N,8) 直构造。 + # **默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 + use_bf16_b_sf = bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) + if use_bf16_b_sf: + assert n % 8 == 0 + tma_aligned_n = (n + 7) // 8 * 8 + b_int4_sf = torch.empty_strided( + (groups, n, k // gran_k), + (tma_aligned_n * (k // gran_k), 1, tma_aligned_n), + device="cuda", + dtype=torch.bfloat16, + ) + b_int4_sf_fp32 = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + else: + b_int4_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + b_int4_sf_fp32 = b_int4_sf + for g in range(groups): + b_int4[g], b_int4_sf_fp32[g] = per_token_cast_to_int4_sym(b_ref[g], gran_k=gran_k) + if use_bf16_b_sf: + b_int4_sf.copy_(b_int4_sf_fp32.to(torch.bfloat16)) + + group_idx_a = torch.arange(k, device="cuda") // gran_k + a_dequant = a_fp8.float() * a_sf[..., group_idx_a] + b_dequant = torch.empty((groups, n, k), device="cuda", dtype=torch.float) + for g in range(groups): + b_dequant[g] = cast_back_from_int4_sym(b_int4[g], b_int4_sf_fp32[g], gran_k=gran_k) + + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for g, valid_m in enumerate(masked_m_values): + if valid_m > 0: + ref[g, :valid_m] = (a_dequant[g, :valid_m] @ b_dequant[g].t()).to(torch.bfloat16) + + b_fp8 = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_fp8[g], b_fp8_sf[g] = per_token_cast_to_fp8( + b_dequant[g].to(torch.bfloat16), use_ue8m0=False, gran_k=gran_k, + ) + + return dict( + masked_m=masked_m, + masked_m_values=masked_m_values, + a_fp8=a_fp8, a_sf=a_sf, + b_int4=b_int4, b_int4_sf=b_int4_sf, + b_fp8=b_fp8, b_fp8_sf=b_fp8_sf, + ref=ref, + ) + + +def _run_int4_kernel_skew(fn, case, expected_m, gran_k=128, masked_m_max_hint=None): + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_int4"], case["b_int4_sf"]) + d = torch.empty_like(case["ref"]) + kwargs = dict( + gran_k=gran_k, gran_k_a=gran_k, gran_k_b=gran_k, + b_is_int4_sym=True, + ) + if masked_m_max_hint is not None: + kwargs["masked_m_max_hint"] = masked_m_max_hint + fn(a, b, d, case["masked_m"], expected_m, **kwargs) + return d + + +def _run_fp8_baseline_skew(case, expected_m, gran_k=128): + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_fp8"], case["b_fp8_sf"]) + d = torch.empty_like(case["ref"]) + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, b, d, case["masked_m"], expected_m, + recipe_a=(1, gran_k), recipe_b=(1, gran_k), + ) + return d + + +def _per_group_cos_min_skew(d, target, masked_m_values): + cos = [] + for g, valid_m in enumerate(masked_m_values): + if valid_m > 0: + cos.append(_cos_sim(d[g, :valid_m], target[g, :valid_m])) + return min(cos) if cos else 1.0 + + +def _effective_masked_bytes_skew(masked_m_values, n, k, a_gran_k, *, + fp8_b: bool, b_gran_k=None): + """Bandwidth model identical to test_sm90_fp8_fp4 skew path: each active + group contributes its valid_m rows to A/D and its full B+SFB tensor.""" + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + sum_m = sum(masked_m_values) + active_groups = sum(1 for v in masked_m_values if v > 0) + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = sum_m * k + sum_m * a_scale_k * 4 + b_data_bytes = active_groups * n * (k if fp8_b else k // 2) + b_scale_bytes = active_groups * n * b_scale_k * 4 + d_bytes = sum_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _masked_skew_benchmark_case(name, masked_m_values, expected_m, n, k, *, + max_m=1024, gran_k=128): + fn = _resolve_kernel() + assert max(masked_m_values) <= max_m + case = _build_int4_a8_skew_inputs(masked_m_values, max_m, n, k, gran_k=gran_k) + + # 透传 masked_m_max_hint：caller 告知 hot group 大小，host 据此选 BM。 + # expected_m 保持原 distribution 平均值（与生产语义一致）。FP8 baseline + # 没有 hint API 但 expected_m 接受任意值，让它走 max 以反映"两条路径都 + # 按 hot 调度"下的真实算法差距。 + masked_m_max_hint = max(masked_m_values) + gemm_expected_m = expected_m + fp8_expected_m = max(masked_m_max_hint, expected_m) + + d_int4 = _run_int4_kernel_skew(fn, case, gemm_expected_m, gran_k=gran_k, + masked_m_max_hint=masked_m_max_hint) + d_fp8 = _run_fp8_baseline_skew(case, fp8_expected_m, gran_k=gran_k) + + cos_abs = _per_group_cos_min_skew(d_int4, case["ref"], masked_m_values) + cos_eq = _per_group_cos_min_skew(d_int4, d_fp8, masked_m_values) + int4_us = _time_cuda(lambda: _run_int4_kernel_skew( + fn, case, gemm_expected_m, gran_k=gran_k, + masked_m_max_hint=masked_m_max_hint)) + fp8_us = _time_cuda(lambda: _run_fp8_baseline_skew(case, fp8_expected_m, gran_k=gran_k)) + + int4_bytes = _effective_masked_bytes_skew(masked_m_values, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_masked_bytes_skew(masked_m_values, n, k, gran_k, fp8_b=True) + return dict( + case=name, + groups=len(masked_m_values), + expected_m=expected_m, + masked_m_hint=masked_m_max_hint, + sum_m=sum(masked_m_values), + masked_max=max(masked_m_values), + active_groups=sum(1 for v in masked_m_values if v > 0), + n=n, k=k, + cos_abs=cos_abs, cos_eq=cos_eq, + int4_us=int4_us * 1e6, fp8_us=fp8_us * 1e6, + int4_gbps=int4_bytes / int4_us / 1e9, + fp8_gbps=fp8_bytes / fp8_us / 1e9, + speedup=fp8_us / int4_us, + ) + + +def _print_skew_table(rows): + print("case | groups | exp_m_avg | hint | sum_m | max_m | active | n | k | " + "cos_abs | cos_eq | INT4 us | INT4 GB/s | FP8 us | FP8 GB/s | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for r in rows: + print( + f"{r['case']} | {r['groups']} | {r['expected_m']} | {r['masked_m_hint']} | {r['sum_m']} | " + f"{r['masked_max']} | {r['active_groups']} | {r['n']} | {r['k']} | " + f"{r['cos_abs']:.4f} | {r['cos_eq']:.4f} | " + f"{r['int4_us']:.0f} | {r['int4_gbps']:.0f} | " + f"{r['fp8_us']:.0f} | {r['fp8_gbps']:.0f} | " + f"{r['speedup']:.2f}x" + ) + + +def test_int4_a8_masked_skew_cases() -> None: + _require_sm90() + torch.manual_seed(4) + + def skew_values(total: int, hot: int, active: int = 8) -> list[int]: + # 32 个槽位，前 active 个非零（第 0 个为 hot，其余分摊 total-hot）， + # 剩余 32-active 个槽 mask=0。 + assert active >= 1 and total >= hot + values = [0] * 32 + values[0] = hot + remaining = total - hot + for idx in range(1, active): + share = (remaining + active - idx - 1) // (active - idx) + values[idx] = share + remaining -= share + assert sum(values) == total + return values + + print("INT4-A8 masked SKEW perf vs FP8 (uneven masked_m distribution)") + print(f" cos_abs > {COS_ABS_THRESHOLD:.3f} -- pass/fail") + print() + rows = [] + shapes = [ + ("gateup", 4096, 4096), + ("down", 4096, 2048), + ] + distributions = [ + ("uniform_1", [1] * 32, 1), + ("uniform_8", [8] * 32, 8), + ("mtp_dp2", skew_values(total=144, hot=50), 7), + ("mtp_dp0", skew_values(total=195, hot=160), 7), + ("mtp_dp4", skew_values(total=290, hot=214), 7), + ("one_hot_214", [214] + [0] * 31, 1), + ] + pass_count = 0 + fail_count = 0 + for shape_name, n, k in shapes: + for dist_name, masked_m_values, expected_m in distributions: + row = _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, k=k, max_m=1024, + ) + ok = row["cos_abs"] > COS_ABS_THRESHOLD + pass_count += int(ok) + fail_count += int(not ok) + rows.append(row) + _print_skew_table(rows) + print(f"\n{pass_count} passed, {fail_count} failed") + if fail_count > 0: + raise AssertionError(f"{fail_count} skew cases failed cos_abs > {COS_ABS_THRESHOLD}") + + +if __name__ == "__main__": + start_time = time.time() + _require_sm90() + if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + test_int4_a8_masked_skew_cases() + else: + test_int4_a8_masked_accuracy() + print(f"\ndone in {time.time() - start_time:.2f}s")