diff --git a/.gitignore b/.gitignore index d0cdf6ca42..dd11d39075 100644 --- a/.gitignore +++ b/.gitignore @@ -17,8 +17,11 @@ deep_gemm/include/cutlass /.clang* /.cache +# Local Codex skills and agent configuration +/.codex/ + # Generated stub files stubs/ # Symlinks to compiled extensions -deep_gemm/*.so \ No newline at end of file +deep_gemm/*.so diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 505b0c0973..fc82df3e11 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -207,13 +207,13 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_ const auto arch_major = device_runtime->get_arch_major(); if (is_varlen) { const auto& indices_tensor = indices.value(); - DG_HOST_ASSERT(arch_major == 10 and next_n == 1 and (block_kv == 64 or block_kv == 32)); + DG_HOST_ASSERT((arch_major == 10 or arch_major == 12) and next_n == 1 and (block_kv == 64 or block_kv == 32)); DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size); DG_HOST_ASSERT(indices_tensor.is_contiguous()); DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt); smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, true, indices_tensor.data_ptr()); - } else if (arch_major == 9 or arch_major == 10) { - DG_HOST_ASSERT(block_kv == 64 or (arch_major == 10 and block_kv == 32)); + } else if (arch_major == 9 or arch_major == 10 or arch_major == 12) { + DG_HOST_ASSERT(block_kv == 64 or ((arch_major == 10 or arch_major == 12) and block_kv == 32)); smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, false, nullptr); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); @@ -335,7 +335,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleget_arch_major(); const auto indices_tensor = indices.value_or(torch::Tensor()); if (is_varlen) { - DG_HOST_ASSERT(arch_major == 10 and next_n == 1); + DG_HOST_ASSERT((arch_major == 10 or arch_major == 12) and next_n == 1); DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size); DG_HOST_ASSERT(indices_tensor.is_contiguous()); DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt); @@ -368,6 +368,10 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleget_arch_major(); if (expr == "bhr,hdr->bhd") { + if (arch_major == 12) { + DG_HOST_ASSERT(not c.has_value()); + DG_HOST_ASSERT(recipe == std::make_tuple(1, 128, 128)); + sm120_fp8_bhr_hdr_bhd_scalar(a.first, a.second, b.first, b.second, d); + return; + } + // Permute dims to satisfy the order of (batch_size, m, n, k) // (batch_size, m, n, k): (h, b, d, r) const auto perm_a = a.first.permute({1, 0, 2}); diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 1a13984d24..97c4b36fd2 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -5,6 +5,7 @@ #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp" #include "../jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp" +#include "../jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp" #endif namespace deep_gemm::hyperconnection { @@ -52,6 +53,8 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, sm90_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else if (arch_major == 10) { sm100_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); + } else if (arch_major == 12) { + sm120_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } diff --git a/csrc/indexing/main.cu b/csrc/indexing/main.cu index a42b66f9e6..24c188f674 100644 --- a/csrc/indexing/main.cu +++ b/csrc/indexing/main.cu @@ -20,6 +20,7 @@ // Hyperconnection kernels #include #include +#include // Layout kernels #include diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp index 7d85a5f556..e1d1f20114 100644 --- a/csrc/jit/compiler.hpp +++ b/csrc/jit/compiler.hpp @@ -202,10 +202,11 @@ class NVCCCompiler final: public Compiler { // The override the compiler flags // Only NVCC >= 12.9 supports arch-specific family suffix const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9); - flags = fmt::format("{} -I{} --gpu-architecture=sm_{} " + const auto arch_flag = fmt::format("-gencode=arch=compute_{},code=sm_{}", arch, arch); + flags = fmt::format("{} -I{} {} " "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " "-O3 --expt-relaxed-constexpr --expt-extended-lambda", - flags, library_include_path.c_str(), arch); + flags, library_include_path.c_str(), arch_flag); } void compile(const std::string &code, const std::filesystem::path& dir_path, diff --git a/csrc/jit/device_runtime.hpp b/csrc/jit/device_runtime.hpp index 2321aded7a..60a5a11257 100644 --- a/csrc/jit/device_runtime.hpp +++ b/csrc/jit/device_runtime.hpp @@ -88,6 +88,11 @@ class DeviceRuntime { std::string get_arch(const bool& number_only = false, const bool& support_arch_family = false) { const auto [major, minor] = get_arch_pair(); + if (major == 12) { + if (number_only) + return "120"; + return support_arch_family ? "120f" : "120"; + } if (major == 10 and minor != 1) { if (number_only) return "100"; diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp new file mode 100644 index 0000000000..a204b61b65 --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -0,0 +1,243 @@ +#pragma once + +#include +#include + +#include "common.hpp" +#include "utils.hpp" +#include "../../utils/exception.hpp" + +namespace deep_gemm { + +struct SM120ArchSpec { + + static constexpr int smem_capacity = 232448; + + static std::vector get_layout_candidates(const GemmDesc& desc) { + // Block M candidates + std::vector block_m_candidates; + if (desc.gemm_type == GemmType::Normal or + desc.gemm_type == GemmType::Batched or + desc.gemm_type == GemmType::KGroupedContiguous) { + block_m_candidates = {64, 128}; + if (desc.m <= 16) block_m_candidates.push_back(16); + if (desc.m <= 32) block_m_candidates.push_back(32); + + // BF16 output GEMM supports 256 + if (desc.cd_dtype != torch::kFloat) + block_m_candidates.push_back(256); + } else if (desc.gemm_type == GemmType::MGroupedContiguous or + desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) { + block_m_candidates = std::vector{heuristics_runtime->get_mk_alignment_for_contiguous_layout()}; + } else if (desc.gemm_type == GemmType::MGroupedMasked) { + block_m_candidates = {64, 128}; + } + + // Block N candidates + std::vector block_n_candidates; + int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of()); + int start = step; + // Avoid bank conflicts for 1D1D kernel FP32 output + if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { + DG_HOST_ASSERT(desc.major_a == cute::UMMA::Major::K); + DG_HOST_ASSERT(desc.major_b == cute::UMMA::Major::K); + start = 24; + block_n_candidates.push_back(16); + } + // Register spills + int end = 256; + if (desc.kernel_type == KernelType::Kernel1D2D) + end = 192; + if (desc.kernel_type == KernelType::Kernel1D1D) + end = 160; + // Enumerate + for (int i = start; i <= end; i += step) + block_n_candidates.push_back(i); + + // Block K is always in a fixed manner + const int block_k = 128 / get_element_size(desc.get_mma_kind()); + + // Disable multicast for performance + const bool disable_multicast = + // The number of k-groups is large (a heuristic) + (desc.gemm_type == GemmType::KGroupedContiguous and desc.num_groups > 4) or + // Not supported + (desc.gemm_type == GemmType::Batched); + + // Enumerate all candidates + std::vector candidates; + for (int cluster_m = 1; cluster_m <= (disable_multicast ? 1 : 2); ++ cluster_m) { + for (int cluster_n = 1; cluster_n <= (disable_multicast ? 1 : 2); ++ cluster_n) { + // We only support cluster 2 + if (cluster_m * cluster_n > 2) + continue; + + // SM count must be divisible + if (desc.num_sms % (cluster_m * cluster_n) != 0) + continue; + + for (int block_m: block_m_candidates) { + for (int block_n: block_n_candidates) { + // 1D2D kernel unroll requirement + if (desc.kernel_type == KernelType::Kernel1D2D and block_n > block_k and (block_n % (block_n - block_k) != 0 and block_k % (block_n - block_k) != 0)) + continue; + + // Multicast legality for masked layout + // TODO: add some comments about it + if ((desc.gemm_type == GemmType::MGroupedMasked or desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) and + ceil_div(desc.n, block_n) % (cluster_m * cluster_n) != 0) + continue; + + // The block sizes cannot be too large (for enough registers), so at least one dim less than 128 + if (block_m > 128 and block_n > 128) + continue; + + // Calculate swizzling + const auto layout = Layout{0, block_m, block_n, block_k, cluster_m, cluster_n}; + const auto storage_config = get_storage_config(desc, layout); + + // Make sure swizzling is large enough (32B's performance is low) + if (storage_config.swizzle_a_mode % 64 != 0 or storage_config.swizzle_b_mode % 64 != 0) + continue; + + // To hide TMA latency, the stage count should be at least 3; for small matrices, at least 4 + int num_stages = get_pipeline_config(desc, layout, storage_config).num_stages; + if (num_stages < 3 or (block_m * block_n < 128 * 192 and num_stages < 4)) + continue; + + candidates.push_back(layout); + } + } + } + } + + DG_HOST_ASSERT(not candidates.empty()); + return candidates; + } + + static StorageConfig get_storage_config(const GemmDesc& desc, const Layout& layout) { + constexpr int mma_m = 64; + + // Load/store block sizes (w/o consideration of swizzling atoms, w/ consideration of loop atoms) + // TODO: support swap AB + DG_HOST_ASSERT(layout.swap_ab == 0); + const auto load_block_m = layout.block_m; + const auto load_block_n = layout.block_n; + // 1D1D kernel will do single warp-group stores + const auto store_block_m = desc.kernel_type == KernelType::Kernel1D1D ? mma_m : layout.block_m; + const auto store_block_n = layout.block_n; + + // Decide swizzling by the inner dim + const auto swizzle_mode_a = get_swizzle_mode( + desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m, c10::elementSize(desc.a_dtype)); + const auto swizzle_mode_b = get_swizzle_mode( + desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n, c10::elementSize(desc.b_dtype)); + // We only enable swizzling for non-FP32 outputs + const auto swizzle_mode_cd = desc.cd_dtype != torch::kFloat ? + get_swizzle_mode(store_block_n, c10::elementSize(desc.cd_dtype)) : 0; + + return { + load_block_m, load_block_n, + store_block_m, store_block_n, + swizzle_mode_a, swizzle_mode_b, swizzle_mode_cd + }; + } + + static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) { + constexpr int kNumMaxStages = 16; + + // TODO: consider swap AB + // C/D for TMA stores + // NOTES: 1024 is for TMA swizzling alignment requirement + const int smem_cd = + align(layout.block_m * layout.block_n * static_cast(c10::elementSize(desc.cd_dtype)), 1024); + const int smem_barriers = kNumMaxStages * 8 * 2; + + // Calculate A/B per stages + const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype); + const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype); + + // Calculate SF A/B per stages + const int smem_sfa_per_stage = desc.kernel_type == KernelType::KernelNoSF ? + 0 : align(layout.block_m * static_cast(sizeof(float)), 128); + const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ? + 0 : align(layout.block_n * static_cast(sizeof(float)), 128); + + // Extra SFB sizes for 1D2D kernels + const int use_uniform_sfb = layout.block_k % layout.block_n == 0 ? 1 : 2; + const int smem_extra_sfb = desc.kernel_type != KernelType::Kernel1D2D ? + 0 : align(ceil_div(desc.k, layout.block_k) * static_cast(sizeof(float)) * use_uniform_sfb, 8); + + // Extra tensormap for 1D1D kernels + const int smem_tensormap = + desc.gemm_type == GemmType::KGroupedContiguous ? 4 * static_cast(sizeof(CUtensorMap)) : 0; + + // Calculate stages + const int smem_extra = smem_cd + smem_barriers + smem_extra_sfb + smem_tensormap; + const int smem_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage + smem_sfb_per_stage; + const int num_stages = std::min( + (smem_capacity - smem_extra) / smem_per_stage, + kNumMaxStages); + return { + smem_extra + num_stages * smem_per_stage, + num_stages + }; + } + + static LaunchConfig get_launch_config(const GemmDesc& desc, const Layout& layout) { + const int num_tma_threads = 128; + const int num_math_threads = layout.block_m <= 64 ? 128 : 256; + return { + desc.num_sms, + layout.get_cluster_size(), + num_tma_threads + num_math_threads, + num_tma_threads, num_math_threads, + 0, 0 // Meaningless for SM120 + }; + } + + static LayoutInfo get_layout_info(const GemmDesc& desc, const Layout& layout) { + const auto num_blocks = + ceil_div(desc.get_expected_m(), layout.block_m) * + ceil_div(desc.get_expected_n(), layout.block_n) * + desc.get_expected_num_groups(); + const auto num_waves = ceil_div(num_blocks, desc.num_sms); + const auto num_last_blocks = num_blocks % desc.num_sms; + const auto last_wave_util = num_last_blocks == 0 ? desc.num_sms : num_last_blocks; + + // Utils + const int l2_bandwidth_per_cycle = std::min(64. * desc.num_sms, 8e6 / (1.3e3)); // B/cycle + const int l1_bandwidth_per_cycle = 128 * desc.num_sms; // B/cycle + const int mma_m = 64; + const int elem_size_ab = c10::elementSize(desc.a_dtype); + const int elem_size_cd = c10::elementSize(desc.cd_dtype); + DG_HOST_ASSERT(desc.a_dtype == desc.b_dtype); + + // Data movement per block + int64_t expected_k = desc.get_expected_k(); + int64_t num_bytes_l2_ab = expected_k * (layout.block_m / layout.cluster_n + layout.block_n / layout.cluster_m) * elem_size_ab; + int64_t num_bytes_l1_ab = expected_k * (layout.block_m + layout.block_n) * elem_size_ab; + int64_t num_bytes_l1_tc = expected_k * (std::max(mma_m, layout.block_m) + layout.block_n) * elem_size_ab + + layout.block_m * layout.block_n * elem_size_cd; + int64_t num_bytes_l1_l2_cd = layout.block_m * layout.block_n * elem_size_cd * (desc.with_accumulation ? 2 : 1); + + // HBM bandwidth and total compute (Tensor/CUDA cores) are constant across configs + // We only model L1/L2 cycles as they are the primary variables between configs + int64_t num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bandwidth_per_cycle; + int64_t num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + num_bytes_l1_l2_cd) * num_blocks / l1_bandwidth_per_cycle; + float wave_efficiency = static_cast(num_blocks) / (num_waves * desc.num_sms); + int64_t num_cycles = std::max(num_l1_cycles, num_l2_cycles) / wave_efficiency; + + // Disable multicasting if only one wave exists + if (layout.cluster_n * layout.cluster_m > 1 and num_waves <= 1) + num_cycles = std::numeric_limits::max(); + + return {num_waves, last_wave_util, num_cycles, layout}; + } + + static bool compare(const LayoutInfo& a, const LayoutInfo& b) { + return a.num_cycles < b.num_cycles; + } +}; + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp new file mode 100644 index 0000000000..967221a20d --- /dev/null +++ b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp @@ -0,0 +1,149 @@ +#pragma once + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +class SM120FP8BhrHdrBhdScalarRuntime final: public LaunchRuntime { +public: + struct Args { + int shape_b; + int shape_h; + int shape_d; + int shape_r; + + int64_t a_stride_b; + int64_t a_stride_h; + int64_t a_stride_r; + int64_t sfa_stride_b; + int64_t sfa_stride_h; + int64_t sfa_stride_r; + int64_t b_stride_h; + int64_t b_stride_d; + int64_t b_stride_r; + int64_t sfb_stride_h; + int64_t sfb_stride_d; + int64_t sfb_stride_r; + int64_t d_stride_b; + int64_t d_stride_h; + int64_t d_stride_d; + + void* a; + float* sfa; + void* b; + float* sfb; + void* d; + at::ScalarType output_dtype; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + constexpr int output_tile_d = 128; + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp8_bhr_hdr_bhd_scalar< + {}, {} + >); +}}; +)", to_string(args.output_dtype), output_tile_d); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.shape_b, + args.shape_h, + args.shape_d, + args.shape_r, + args.a_stride_b, + args.a_stride_h, + args.a_stride_r, + args.sfa_stride_b, + args.sfa_stride_h, + args.sfa_stride_r, + args.b_stride_h, + args.b_stride_d, + args.b_stride_r, + args.sfb_stride_h, + args.sfb_stride_d, + args.sfb_stride_r, + args.d_stride_b, + args.d_stride_h, + args.d_stride_d, + args.a, + args.sfa, + args.b, + args.sfb, + args.d + )); + } +}; + +static void sm120_fp8_bhr_hdr_bhd_scalar(const torch::Tensor& a, + const torch::Tensor& sfa, + const torch::Tensor& b, + const torch::Tensor& sfb, + const torch::Tensor& d) { + constexpr int gran_mn = 128; + constexpr int gran_k = 128; + constexpr int output_tile_d = 128; + + const auto [shape_b, shape_h, shape_r] = get_shape<3>(a); + const auto [shape_h_, shape_d, shape_r_] = get_shape<3>(b); + const auto [shape_b_, shape_h__, shape_d_] = get_shape<3>(d); + + DG_HOST_ASSERT(shape_b == shape_b_ and shape_h == shape_h_ and shape_h == shape_h__); + DG_HOST_ASSERT(shape_r == shape_r_ and shape_d == shape_d_); + DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(a.stride(2) == 1 and b.stride(2) == 1 and d.stride(2) == 1); + + check_sf_layout(sfa, shape_h, shape_r, 1, gran_k, shape_b); + check_sf_layout(sfb, shape_d, shape_r, gran_mn, gran_k, shape_h); + + const int num_d_tiles = ceil_div(shape_d, output_tile_d); + const int grid_dim = shape_b * shape_h * num_d_tiles; + const SM120FP8BhrHdrBhdScalarRuntime::Args args = { + .shape_b = shape_b, + .shape_h = shape_h, + .shape_d = shape_d, + .shape_r = shape_r, + .a_stride_b = a.stride(0), + .a_stride_h = a.stride(1), + .a_stride_r = a.stride(2), + .sfa_stride_b = sfa.stride(0), + .sfa_stride_h = sfa.stride(1), + .sfa_stride_r = sfa.stride(2), + .b_stride_h = b.stride(0), + .b_stride_d = b.stride(1), + .b_stride_r = b.stride(2), + .sfb_stride_h = sfb.stride(0), + .sfb_stride_d = sfb.stride(1), + .sfb_stride_r = sfb.stride(2), + .d_stride_b = d.stride(0), + .d_stride_h = d.stride(1), + .d_stride_d = d.stride(2), + .a = a.data_ptr(), + .sfa = sfa.data_ptr(), + .b = b.data_ptr(), + .sfb = sfb.data_ptr(), + .d = d.data_ptr(), + .output_dtype = d.scalar_type(), + .launch_args = LaunchArgs(grid_dim, output_tile_d) + }; + const auto code = SM120FP8BhrHdrBhdScalarRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_bhr_hdr_bhd_scalar", code); + SM120FP8BhrHdrBhdScalarRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp new file mode 100644 index 0000000000..4a2a1bc0a5 --- /dev/null +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -0,0 +1,154 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "../heuristics/sm120.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +class SM120BF16HCPrenormGemmRuntime final: public LaunchRuntime { +public: + struct Args { + int m, n, k; + int block_m, block_n, block_k; + int num_splits; + int swizzle_cd_mode; + int num_stages; + int num_math_threads, num_tma_threads; + + LaunchArgs launch_args; + + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + float* sqr_sum; + }; + + 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(&sm120_tf32_hc_prenorm_gemm_impl< + {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {} + >); +}}; +)", + args.n, args.k, + args.block_m, args.block_n, args.block_k, + args.num_splits, + args.swizzle_cd_mode, + args.num_stages, + args.num_math_threads, args.num_tma_threads); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.m, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.sqr_sum)); + } +}; + +static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& d, + const torch::Tensor& sqr_sum, + const int& m, const int& n, const int& k, + const int& num_splits) { + + // SM120 HC currently uses the same small-N tile shape as the SM90 path. + constexpr int block_m = 64; + constexpr int block_k = 64; + constexpr int num_math_threads = 128; + constexpr int num_tma_threads = 128; + constexpr int num_threads = num_math_threads + num_tma_threads; + + const int block_n = align(n, 16); + DG_HOST_ASSERT(n <= block_n); + // Only support small N for now + DG_HOST_ASSERT(n <= 32 and n % 8 == 0); + DG_HOST_ASSERT(k % block_k == 0); + DG_HOST_ASSERT(num_splits >= 1 and num_splits <= k / block_k); + + const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float)); + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, + block_m, block_k, + static_cast(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1, + get_swizzle_mode(block_k, a.element_size()), 0, + true); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k, + block_n, block_k, + static_cast(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1, + get_swizzle_mode(block_k, b.element_size()), 0, + true); + const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n, + block_m, block_n, + static_cast(d.stride(-2)), 1, + swizzle_cd_mode) + : make_tma_3d_desc(d, n, m, num_splits, + block_n, block_m, 1, + static_cast(d.stride(-2)), + static_cast(d.stride(-3)), + swizzle_cd_mode); + + // Calculate stages + int num_stages = 12, smem_size = 0; + while (num_stages > 0) { + const int smem_a_per_stage = block_m * block_k * static_cast(sizeof(nv_bfloat16)); + const int smem_b_per_stage = block_n * block_k * static_cast(sizeof(float)); + const int smem_cd = block_m * swizzle_cd_mode; + const int smem_barriers = num_stages * 2 * 8; + smem_size = (smem_a_per_stage + smem_b_per_stage) * num_stages + + smem_cd + smem_barriers; + + const int smem_capacity = device_runtime->get_prop()->sharedMemPerBlockOptin; + if (smem_size <= smem_capacity) + break; + -- num_stages; + } + DG_HOST_ASSERT(num_stages > 0); + + if (get_env("DG_JIT_DEBUG", 0)) { + printf("M: %d, N: %d, K: %d -> " + "block M: %d, block N: %d, block K: %d, split K: %d" + "stages: %d, shared memory: %d, swizzle CD: %d\n", + m, n, k, block_m, block_n, block_k, num_splits, + num_stages, smem_size, swizzle_cd_mode); + } + + smem_size = device_runtime->get_prop()->sharedMemPerBlockOptin; + + const SM120BF16HCPrenormGemmRuntime::Args args = { + .m = m, .n = n, .k = k, + .block_m = block_m, .block_n = block_n, .block_k = block_k, + .num_splits = num_splits, + .swizzle_cd_mode = swizzle_cd_mode, + .num_stages = num_stages, + .num_math_threads = num_math_threads, + .num_tma_threads = num_tma_threads, + .launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .sqr_sum = sqr_sum.data_ptr() + }; + + const auto code = SM120BF16HCPrenormGemmRuntime::generate(args); + const auto runtime = compiler->build("sm120_tf32_hc_prenorm_gemm", code); + SM120BF16HCPrenormGemmRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index 2a3288ee0c..98d28555ab 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -166,6 +166,211 @@ static void __instantiate_kernel() {{ } }; +class SM120FP8PagedMQALogitsScalarRuntime final: public LaunchRuntime { +public: + struct Args { + int batch_size; + int next_n; + int num_heads; + int head_dim; + int block_kv; + bool is_context_lens_2d; + int block_table_stride; + int logits_stride; + int tokens_per_block; + int token_groups; + bool cache_q; + + int64_t q_batch_stride; + int64_t q_next_stride; + int64_t q_head_stride; + int64_t kv_block_stride; + int64_t kv_token_stride; + int64_t kv_scale_block_stride; + int64_t weights_row_stride; + + void* q; + void* kv_cache; + float* kv_cache_scales; + float* weights; + int* context_lens; + void* logits; + int* block_table; + at::ScalarType logits_dtype; + + LaunchArgs launch_args; + }; + + 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(&sm120_fp8_paged_mqa_logits_scalar< + {}, {}, {}, + {}, {}, + {} + >); +}}; +)", args.num_heads, args.head_dim, args.block_kv, + args.is_context_lens_2d ? "true" : "false", args.tokens_per_block, + to_string(args.logits_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.batch_size, + args.next_n, + args.logits_stride, + args.block_table_stride, + args.q_batch_stride, + args.q_next_stride, + args.q_head_stride, + args.kv_block_stride, + args.kv_token_stride, + args.kv_scale_block_stride, + args.weights_row_stride, + args.q, + args.kv_cache, + args.kv_cache_scales, + args.weights, + args.context_lens, + args.logits, + args.block_table + )); + } +}; + +class SM120FP8PagedMQALogitsTiledRuntime final: public LaunchRuntime { +public: + using Args = SM120FP8PagedMQALogitsScalarRuntime::Args; + + 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(&sm120_fp8_paged_mqa_logits_mma_tiled< + {}, + {}, + {}, + {}, + {} + >); +}}; +)", args.block_kv, args.is_context_lens_2d ? "true" : "false", + args.token_groups, + args.cache_q ? "true" : "false", + to_string(args.logits_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.batch_size, + args.next_n, + args.logits_stride, + args.block_table_stride, + args.q_batch_stride, + args.q_next_stride, + args.q_head_stride, + args.kv_block_stride, + args.kv_token_stride, + args.kv_scale_block_stride, + args.weights_row_stride, + args.q, + args.kv_cache, + args.kv_cache_scales, + args.weights, + args.context_lens, + args.logits, + args.block_table + )); + } +}; + +static void sm120_fp8_paged_mqa_logits(const torch::Tensor& q, + const torch::Tensor& kv_cache, + const torch::Tensor& kv_cache_scales, + const torch::Tensor& weights, + const torch::Tensor& context_lens, + const torch::Tensor& logits, + const torch::Tensor& block_table, + const at::ScalarType& logits_dtype, + const int& batch_size, const int& next_n, + const int& num_heads, const int& head_dim, + const int& block_kv, + const bool& is_context_lens_2d, + const int& logits_stride, + const int& block_table_stride) { + constexpr int tokens_per_block = 128; + DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); + DG_HOST_ASSERT(block_kv == 32 or block_kv == 64); + + const SM120FP8PagedMQALogitsScalarRuntime::Args args = { + .batch_size = batch_size, + .next_n = next_n, + .num_heads = num_heads, + .head_dim = head_dim, + .block_kv = block_kv, + .is_context_lens_2d = is_context_lens_2d, + .block_table_stride = block_table_stride, + .logits_stride = logits_stride, + .tokens_per_block = tokens_per_block, + .token_groups = 1, + .cache_q = false, + .q_batch_stride = q.stride(0), + .q_next_stride = q.stride(1), + .q_head_stride = q.stride(2), + .kv_block_stride = kv_cache.stride(0), + .kv_token_stride = kv_cache.stride(1), + .kv_scale_block_stride = kv_cache_scales.stride(0), + .weights_row_stride = weights.stride(0), + .q = q.data_ptr(), + .kv_cache = kv_cache.data_ptr(), + .kv_cache_scales = kv_cache_scales.data_ptr(), + .weights = weights.data_ptr(), + .context_lens = context_lens.data_ptr(), + .logits = logits.data_ptr(), + .block_table = block_table.data_ptr(), + .logits_dtype = logits_dtype, + .launch_args = LaunchArgs(std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_block)), + tokens_per_block) + }; + const bool use_tiled = get_env("DG_SM120_PAGED_MQA_TILED", 0) != 0; + if (use_tiled and num_heads == 64 and head_dim == 128 and block_kv == 64 and logits_dtype == torch::kFloat32) { + const int token_groups = get_env("DG_SM120_PAGED_MQA_TILED_GROUPS", 8); + DG_HOST_ASSERT(token_groups == 1 or token_groups == 2 or token_groups == 4 or token_groups == 8); + constexpr int tiled_token_tile = 8; + constexpr int tiled_num_threads = 128; + const int tokens_per_tiled_cta = tiled_token_tile * token_groups; + const int tiled_smem_size = 64 * tokens_per_tiled_cta * sizeof(float); + auto tiled_args = args; + tiled_args.token_groups = token_groups; + const bool should_cache_q_by_default = token_groups != 1; + tiled_args.cache_q = get_env( + "DG_SM120_PAGED_MQA_TILED_CACHE_Q", + should_cache_q_by_default ? 1 : 0) != 0; + DG_HOST_ASSERT(not tiled_args.cache_q or token_groups == 2 or token_groups == 4 or token_groups == 8); + tiled_args.launch_args = LaunchArgs( + std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_tiled_cta)), + tiled_num_threads, + tiled_smem_size); + const auto code = SM120FP8PagedMQALogitsTiledRuntime::generate(tiled_args); + const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_mma_tiled", code); + SM120FP8PagedMQALogitsTiledRuntime::launch(runtime, tiled_args); + return; + } + + const auto code = SM120FP8PagedMQALogitsScalarRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_scalar", code); + SM120FP8PagedMQALogitsScalarRuntime::launch(runtime, args); +} + static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, const torch::Tensor& kv_cache, const torch::Tensor& kv_cache_scales, diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh new file mode 100644 index 0000000000..e24008bbd6 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh @@ -0,0 +1,126 @@ +#pragma once + +#include + +#include + +#include + +namespace deep_gemm { + +CUTLASS_DEVICE float sm120_fp8_einsum_dot_scalar( + const __nv_fp8_e4m3* a, + const __nv_fp8_e4m3* b, + const float* sfa, + const float* sfb, + const uint32_t shape_r, + const uint64_t sfa_stride_r, + const uint64_t sfb_stride_r) { + float accum = 0.0f; + for (uint32_t r_idx = 0; r_idx < shape_r; ++r_idx) { + const uint32_t r_scale_idx = r_idx / 128; + const float scale = sfa[r_scale_idx * sfa_stride_r] * sfb[r_scale_idx * sfb_stride_r]; + const float a_value = static_cast(a[r_idx]); + const float b_value = static_cast(b[r_idx]); + accum = fmaf(a_value, b_value * scale, accum); + } + return accum; +} + +CUTLASS_DEVICE float sm120_fp8_einsum_dot_fp8x4( + const __nv_fp8_e4m3* a, + const __nv_fp8_e4m3* b, + const float* sfa, + const float* sfb, + const uint32_t shape_r, + const uint64_t sfa_stride_r, + const uint64_t sfb_stride_r) { + float accum = 0.0f; + const uint32_t num_full_scale_blocks = shape_r / 128; + for (uint32_t scale_idx = 0; scale_idx < num_full_scale_blocks; ++scale_idx) { + const float scale = sfa[scale_idx * sfa_stride_r] * sfb[scale_idx * sfb_stride_r]; + const uint32_t r_start = scale_idx * 128; + #pragma unroll + for (uint32_t r_offset = 0; r_offset < 128; r_offset += 4) { + const uint32_t r_idx = r_start + r_offset; + const auto a_values = static_cast( + *reinterpret_cast(a + r_idx)); + const auto b_values = static_cast( + *reinterpret_cast(b + r_idx)); + accum = fmaf(a_values.x, b_values.x * scale, accum); + accum = fmaf(a_values.y, b_values.y * scale, accum); + accum = fmaf(a_values.z, b_values.z * scale, accum); + accum = fmaf(a_values.w, b_values.w * scale, accum); + } + } + const uint32_t tail_start = num_full_scale_blocks * 128; + if (tail_start >= shape_r) + return accum; + + const float tail_scale = sfa[num_full_scale_blocks * sfa_stride_r] * + sfb[num_full_scale_blocks * sfb_stride_r]; + for (uint32_t r_idx = tail_start; r_idx < shape_r; r_idx += 4) { + const auto a_values = static_cast( + *reinterpret_cast(a + r_idx)); + const auto b_values = static_cast( + *reinterpret_cast(b + r_idx)); + accum = fmaf(a_values.x, b_values.x * tail_scale, accum); + accum = fmaf(a_values.y, b_values.y * tail_scale, accum); + accum = fmaf(a_values.z, b_values.z * tail_scale, accum); + accum = fmaf(a_values.w, b_values.w * tail_scale, accum); + } + return accum; +} + +template +CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_scalar( + const uint32_t shape_b, + const uint32_t shape_h, + const uint32_t shape_d, + const uint32_t shape_r, + const uint64_t a_stride_b, + const uint64_t a_stride_h, + const uint64_t a_stride_r, + const uint64_t sfa_stride_b, + const uint64_t sfa_stride_h, + const uint64_t sfa_stride_r, + const uint64_t b_stride_h, + const uint64_t b_stride_d, + const uint64_t b_stride_r, + const uint64_t sfb_stride_h, + const uint64_t sfb_stride_d, + const uint64_t sfb_stride_r, + const uint64_t d_stride_b, + const uint64_t d_stride_h, + const uint64_t d_stride_d, + const __nv_fp8_e4m3* a, + const float* sfa, + const __nv_fp8_e4m3* b, + const float* sfb, + output_t* d) { + const uint32_t num_d_tiles = (shape_d + kOutputTileD - 1) / kOutputTileD; + const uint32_t tile_idx = blockIdx.x; + const uint32_t d_tile_idx = tile_idx % num_d_tiles; + const uint32_t h_idx = (tile_idx / num_d_tiles) % shape_h; + const uint32_t b_idx = tile_idx / (num_d_tiles * shape_h); + const uint32_t d_idx = d_tile_idx * kOutputTileD + threadIdx.x; + + if (b_idx >= shape_b or d_idx >= shape_d) + return; + + const auto a_base = b_idx * a_stride_b + h_idx * a_stride_h; + const auto b_base = h_idx * b_stride_h + d_idx * b_stride_d; + const auto sfa_base = b_idx * sfa_stride_b + h_idx * sfa_stride_h; + const auto sfb_base = h_idx * sfb_stride_h + (d_idx / 128) * sfb_stride_d; + const bool can_vectorize = (shape_r % 4 == 0) and (a_stride_r == 1) and (b_stride_r == 1); + const float accum = can_vectorize ? + sm120_fp8_einsum_dot_fp8x4(a + a_base, b + b_base, sfa + sfa_base, sfb + sfb_base, + shape_r, sfa_stride_r, sfb_stride_r) : + sm120_fp8_einsum_dot_scalar(a + a_base, b + b_base, sfa + sfa_base, sfb + sfb_base, + shape_r, sfa_stride_r, sfb_stride_r); + + const auto d_offset = b_idx * d_stride_b + h_idx * d_stride_h + d_idx * d_stride_d; + d[d_offset] = static_cast(accum); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh new file mode 100644 index 0000000000..56c6becfff --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -0,0 +1,395 @@ +#pragma once + +#include + +#include +#include + +#include + +#include + +namespace deep_gemm { + +template +CUTLASS_DEVICE float sm120_fp8_dot_scaled( + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv, + const float kv_scale) { + float score = 0.0f; + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += 4) { + const auto q_values = static_cast( + *reinterpret_cast(q + dim_idx)); + const auto kv_values = static_cast( + *reinterpret_cast(kv + dim_idx)); + score = fmaf(q_values.x, kv_values.x * kv_scale, score); + score = fmaf(q_values.y, kv_values.y * kv_scale, score); + score = fmaf(q_values.z, kv_values.z * kv_scale, score); + score = fmaf(q_values.w, kv_values.w * kv_scale, score); + } + return score; +} + +template +CUTLASS_GLOBAL __launch_bounds__(kTokensPerBlock, 1) +void sm120_fp8_paged_mqa_logits_scalar( + const uint32_t batch_size, + const uint32_t next_n, + const uint32_t logits_stride, + const uint32_t block_table_stride, + const uint64_t q_batch_stride, + const uint64_t q_next_stride, + const uint64_t q_head_stride, + const uint64_t kv_block_stride, + const uint64_t kv_token_stride, + const uint64_t kv_scale_block_stride, + const uint64_t weights_row_stride, + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv_cache, + const float* kv_cache_scales, + const float* weights, + const uint32_t* context_lens, + logits_dtype_t* logits, + const uint32_t* block_table) { + const auto row = blockIdx.x; + const auto token_idx = blockIdx.y * kTokensPerBlock + threadIdx.x; + if (row >= batch_size * next_n or token_idx >= logits_stride) + return; + + const auto batch_idx = row / next_n; + const auto next_idx = row - batch_idx * next_n; + const auto context_lens_idx = kIsContextLens2D ? row : batch_idx; + const auto context_len = context_lens[context_lens_idx]; + const auto output_offset = row * static_cast(logits_stride) + token_idx; + + if (token_idx >= context_len) { + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + return; + } + + const auto logical_block_idx = token_idx / BLOCK_KV; + const auto token_in_block = token_idx - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto q_base = q + batch_idx * q_batch_stride + next_idx * q_next_stride; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + const auto kv_scale = kv_cache_scales[kv_block_idx * kv_scale_block_stride + token_in_block]; + const auto weight_base = weights + row * weights_row_stride; + + float total = 0.0f; + #pragma unroll + for (uint32_t head_idx = 0; head_idx < kNumHeads; ++head_idx) { + const auto q_head = q_base + head_idx * q_head_stride; + const auto score = sm120_fp8_dot_scaled(q_head, kv_base, kv_scale); + total += fmaxf(score, 0.0f) * weight_base[head_idx]; + } + + logits[output_offset] = static_cast(total); +} + +template +CUTLASS_GLOBAL __launch_bounds__(128, 1) +void sm120_fp8_paged_mqa_logits_mma_tiled( + const uint32_t batch_size, + const uint32_t next_n, + const uint32_t logits_stride, + const uint32_t block_table_stride, + const uint64_t q_batch_stride, + const uint64_t q_next_stride, + const uint64_t q_head_stride, + const uint64_t kv_block_stride, + const uint64_t kv_token_stride, + const uint64_t kv_scale_block_stride, + const uint64_t weights_row_stride, + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv_cache, + const float* kv_cache_scales, + const float* weights, + const uint32_t* context_lens, + logits_dtype_t* logits, + const uint32_t* block_table) { + using Element = cute::float_e4m3_t; + using Accumulator = float; + + constexpr uint32_t kNumHeads = 64; + constexpr uint32_t kHeadDim = 128; + constexpr uint32_t kHeadsPerWarp = 16; + constexpr uint32_t kTokensPerTile = 8; + constexpr uint32_t kTokensPerCta = kTokensPerTile * kTokenGroups; + constexpr uint32_t kMmaK = 32; + constexpr uint32_t kHeadTiles = kNumHeads / kHeadsPerWarp; + + extern __shared__ float shared_scores[]; + + const auto row = blockIdx.x; + const auto token_tile_start = blockIdx.y * kTokensPerCta; + if (row >= batch_size * next_n) + return; + + const auto batch_idx = row / next_n; + const auto next_idx = row - batch_idx * next_n; + const auto context_lens_idx = kIsContextLens2D ? row : batch_idx; + const auto context_len = context_lens[context_lens_idx]; + + if (token_tile_start >= context_len) { + if (threadIdx.x < kTokensPerCta) { + const auto token_idx = token_tile_start + threadIdx.x; + if (token_idx < logits_stride) { + const auto output_offset = row * static_cast(logits_stride) + token_idx; + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + } + } + return; + } + + const auto lane_idx = threadIdx.x % 32; + const auto warp_idx = threadIdx.x / 32; + const auto q_base = q + batch_idx * q_batch_stride + next_idx * q_next_stride; + + auto tiled_mma = cute::make_tiled_mma( + cute::SM120_16x8x32_TN{}); + auto thread_mma = tiled_mma.get_slice(lane_idx); + + if constexpr (kCacheQ and (kTokenGroups == 2 or kTokenGroups == 4 or kTokenGroups == 8)) { + auto make_score_tile = [&](auto token_group) { + constexpr uint32_t kGroup = decltype(token_group)::value; + return cute::make_tensor( + cute::make_smem_ptr( + shared_scores + + warp_idx * kHeadsPerWarp * kTokensPerCta + + kGroup * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + }; + + auto score_tile0 = make_score_tile(cute::Int<0>{}); + auto score_tile1 = make_score_tile(cute::Int<1>{}); + auto tCgC0 = thread_mma.partition_C(score_tile0); + auto tCgC1 = thread_mma.partition_C(score_tile1); + auto tCrC0 = thread_mma.make_fragment_C(tCgC0); + auto tCrC1 = thread_mma.make_fragment_C(tCgC1); + cute::clear(tCrC0); + cute::clear(tCrC1); + + auto accumulate_group = [&](const uint32_t dim_idx, const auto& tArA, auto token_group, auto& tCrC) { + constexpr uint32_t kGroup = decltype(token_group)::value; + const auto group_token_start = token_tile_start + kGroup * kTokensPerTile; + if (group_token_start >= context_len or group_token_start >= logits_stride) + return; + + const auto logical_block_idx = group_token_start / BLOCK_KV; + const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + const auto kv_tile = reinterpret_cast(kv_base + dim_idx); + auto kv_tensor = cute::make_tensor( + cute::make_gmem_ptr(kv_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(kv_token_stride, cute::Int<1>{})); + auto tBgB = thread_mma.partition_B(kv_tensor); + auto tBrB = thread_mma.partition_fragment_B(kv_tensor); + cute::copy(tBgB, tBrB); + cute::gemm(tiled_mma, tArA, tBrB, tCrC); + }; + + if constexpr (kTokenGroups == 2) { + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + } else if constexpr (kTokenGroups == 4) { + auto score_tile2 = make_score_tile(cute::Int<2>{}); + auto score_tile3 = make_score_tile(cute::Int<3>{}); + auto tCgC2 = thread_mma.partition_C(score_tile2); + auto tCgC3 = thread_mma.partition_C(score_tile3); + auto tCrC2 = thread_mma.make_fragment_C(tCgC2); + auto tCrC3 = thread_mma.make_fragment_C(tCgC3); + cute::clear(tCrC2); + cute::clear(tCrC3); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + accumulate_group(dim_idx, tArA, cute::Int<2>{}, tCrC2); + accumulate_group(dim_idx, tArA, cute::Int<3>{}, tCrC3); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + cute::copy(tCrC2, tCgC2); + cute::copy(tCrC3, tCgC3); + } else { + auto score_tile2 = make_score_tile(cute::Int<2>{}); + auto score_tile3 = make_score_tile(cute::Int<3>{}); + auto score_tile4 = make_score_tile(cute::Int<4>{}); + auto score_tile5 = make_score_tile(cute::Int<5>{}); + auto score_tile6 = make_score_tile(cute::Int<6>{}); + auto score_tile7 = make_score_tile(cute::Int<7>{}); + auto tCgC2 = thread_mma.partition_C(score_tile2); + auto tCgC3 = thread_mma.partition_C(score_tile3); + auto tCgC4 = thread_mma.partition_C(score_tile4); + auto tCgC5 = thread_mma.partition_C(score_tile5); + auto tCgC6 = thread_mma.partition_C(score_tile6); + auto tCgC7 = thread_mma.partition_C(score_tile7); + auto tCrC2 = thread_mma.make_fragment_C(tCgC2); + auto tCrC3 = thread_mma.make_fragment_C(tCgC3); + auto tCrC4 = thread_mma.make_fragment_C(tCgC4); + auto tCrC5 = thread_mma.make_fragment_C(tCgC5); + auto tCrC6 = thread_mma.make_fragment_C(tCgC6); + auto tCrC7 = thread_mma.make_fragment_C(tCgC7); + cute::clear(tCrC2); + cute::clear(tCrC3); + cute::clear(tCrC4); + cute::clear(tCrC5); + cute::clear(tCrC6); + cute::clear(tCrC7); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + accumulate_group(dim_idx, tArA, cute::Int<2>{}, tCrC2); + accumulate_group(dim_idx, tArA, cute::Int<3>{}, tCrC3); + accumulate_group(dim_idx, tArA, cute::Int<4>{}, tCrC4); + accumulate_group(dim_idx, tArA, cute::Int<5>{}, tCrC5); + accumulate_group(dim_idx, tArA, cute::Int<6>{}, tCrC6); + accumulate_group(dim_idx, tArA, cute::Int<7>{}, tCrC7); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + cute::copy(tCrC2, tCgC2); + cute::copy(tCrC3, tCgC3); + cute::copy(tCrC4, tCgC4); + cute::copy(tCrC5, tCgC5); + cute::copy(tCrC6, tCgC6); + cute::copy(tCrC7, tCgC7); + } + } else { + #pragma unroll + for (uint32_t token_group = 0; token_group < kTokenGroups; ++token_group) { + const auto group_token_start = token_tile_start + token_group * kTokensPerTile; + if (group_token_start >= context_len or group_token_start >= logits_stride) + continue; + + const auto logical_block_idx = group_token_start / BLOCK_KV; + const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + + auto score_tile = cute::make_tensor( + cute::make_smem_ptr( + shared_scores + + warp_idx * kHeadsPerWarp * kTokensPerCta + + token_group * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto tCgC = thread_mma.partition_C(score_tile); + auto tCrC = thread_mma.make_fragment_C(tCgC); + cute::clear(tCrC); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + const auto kv_tile = reinterpret_cast(kv_base + dim_idx); + + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto kv_tensor = cute::make_tensor( + cute::make_gmem_ptr(kv_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(kv_token_stride, cute::Int<1>{})); + + auto tAgA = thread_mma.partition_A(q_tensor); + auto tBgB = thread_mma.partition_B(kv_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + auto tBrB = thread_mma.partition_fragment_B(kv_tensor); + + cute::copy(tAgA, tArA); + cute::copy(tBgB, tBrB); + cute::gemm(tiled_mma, tArA, tBrB, tCrC); + } + + cute::copy(tCrC, tCgC); + } + } + __syncthreads(); + + const auto token_idx = token_tile_start + threadIdx.x; + if (threadIdx.x >= kTokensPerCta or token_idx >= logits_stride) + return; + + const auto output_offset = row * static_cast(logits_stride) + token_idx; + if (token_idx >= context_len) { + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + return; + } + + const auto logical_block_idx = token_idx / BLOCK_KV; + const auto token_in_block = token_idx - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_scale = kv_cache_scales[kv_block_idx * kv_scale_block_stride + token_in_block]; + const auto weight_base = weights + row * weights_row_stride; + + float total = 0.0f; + #pragma unroll + for (uint32_t head_tile = 0; head_tile < kHeadTiles; ++head_tile) { + #pragma unroll + for (uint32_t head = 0; head < kHeadsPerWarp; ++head) { + const auto head_idx = head_tile * kHeadsPerWarp + head; + const auto score = shared_scores[ + head_tile * kHeadsPerWarp * kTokensPerCta + + head * kTokensPerCta + threadIdx.x] * kv_scale; + total += fmaxf(score, 0.0f) * weight_base[head_idx]; + } + } + + logits[output_offset] = static_cast(total); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh new file mode 100644 index 0000000000..d7d659daf1 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh @@ -0,0 +1,288 @@ +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace deep_gemm { + +template +CUTLASS_DEVICE +uint32_t get_sm120_swizzled_bank_group_idx(const uint32_t& offset, const uint32_t& lane_idx) { + constexpr uint32_t kGroupsInSwizzleRange = kSwizzleMode / kSwizzleBase; + + const auto bank_group_idx = offset + lane_idx * kGroupsInSwizzleRange; + + constexpr uint32_t kNumBankGroups = 128 / kSwizzleBase; + constexpr bool kHasShortcut = kGroupsInSwizzleRange == kNumBankGroups; + auto row = kHasShortcut ? (offset / kNumBankGroups + lane_idx) : (bank_group_idx / kNumBankGroups); + auto col = kHasShortcut ? (offset) : (bank_group_idx % kNumBankGroups); + col ^= row % kGroupsInSwizzleRange; + + return (row * kNumBankGroups + col) % kGroupsInSwizzleRange; +} + +template +CUTLASS_GLOBAL void __launch_bounds__(kNumMathThreads + kNumTMAThreads, 1) +sm120_tf32_hc_prenorm_gemm_impl(const uint32_t shape_m, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, + const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_d, + float* sqr_sum) { + +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1200)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + DG_STATIC_ASSERT(BLOCK_K == 64, "Invalid block K"); + DG_STATIC_ASSERT(kSwizzleAMode == 128, "Invalid swizzle A mode"); + DG_STATIC_ASSERT(kSwizzleBMode == 128, "Invalid swizzle B mode"); + + DG_STATIC_ASSERT(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + DG_STATIC_ASSERT(kNumMathThreads == 128, "Invalid MMA threads"); + + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + + if (warp_idx == 0 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_d); + } + + auto smem_cd = reinterpret_cast(smem_buffer); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE)); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE)); + }); + + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); + + if (warp_idx == 1 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(128); + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + constexpr uint32_t kNumKBlocks = math::constexpr_ceil_div(SHAPE_K, BLOCK_K); + constexpr uint32_t kNumKBlocksPerSplit = kNumKBlocks / kNumSplits; + constexpr uint32_t kRemainKBlocks = kNumKBlocks % kNumSplits; + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t k_offset = (k_split_idx * kNumKBlocksPerSplit + cute::min(k_split_idx, kRemainKBlocks)) * BLOCK_K; + const uint32_t m_offset = shape_m * k_split_idx; + const uint32_t num_total_stages = kNumKBlocksPerSplit + (k_split_idx < kRemainKBlocks); + + constexpr uint32_t kNumTMARegisters = 40; + + cudaGridDependencySynchronize(); + + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cutlass::arch::warpgroup_reg_dealloc(); + for (uint32_t s = 0; s < num_total_stages; ++ s) { + const auto stage_idx = s % kNumStages; + empty_barriers[stage_idx]->wait(((s / kNumStages) & 1) ^ 1); + + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t k_idx = k_offset + s * BLOCK_K; + + tma::copy(&tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx); + tma::copy(&tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], k_idx, 0); + + constexpr uint32_t kNumArrivalBytes = SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE; + full_barriers[stage_idx]->arrive_and_expect_tx(kNumArrivalBytes); + } + + for (uint32_t s = num_total_stages; s < num_total_stages + kNumStages; ++ s) { + const auto stage_idx = s % kNumStages; + empty_barriers[stage_idx]->wait(((s / kNumStages) & 1) ^ 1); + } + } else if (warp_idx < kNumMathThreads / 32) { + + DG_STATIC_ASSERT(BLOCK_M == 64, "Invalid block M"); + DG_STATIC_ASSERT(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "Invalid block K"); + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + constexpr uint32_t WGMMA_N = BLOCK_N; + + using MMASelector = mma::sm120::TF32MMASelector; + float accum[MMASelector::N_atoms * MMASelector::kNumAccumPerAtom] = {0}; + + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + float sqr_sum_acc_0 = 0; + float sqr_sum_acc_1 = 0; + + #pragma unroll kNumStages < 8 ? kNumStages : kNumStages / 2 + for (uint32_t s = 0; s < num_total_stages; ++ s) { + const auto& stage_idx = s % kNumStages; + full_barriers[stage_idx]->wait((s / kNumStages) & 1); + + // Using MMA atom dimensions mapped to registers + constexpr uint32_t kNumRegPerMMA = MMASelector::M * MMASelector::K / 32; // 4 floats + constexpr uint32_t kNumMMAPerBlockK = BLOCK_K / MMASelector::K; + + float a[kNumRegPerMMA * kNumMMAPerBlockK]; + DG_STATIC_ASSERT(kSwizzleAMode == 128, "Invalid swizzle A mode"); + + uint32_t row = warp_idx * 16 + lane_idx / 4; + + #pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++ i) { + uint32_t bank_group_idx = (row ^ i) % 8; + nv_bfloat16* a_bf16_smem_ptr_upper = smem_a[stage_idx] + row * BLOCK_K + bank_group_idx * kNumElemsPerBankGroup; + nv_bfloat16* a_bf16_smem_ptr_lower = smem_a[stage_idx] + (row + 8) * BLOCK_K + bank_group_idx * kNumElemsPerBankGroup; + + uint32_t elem_offset = lane_idx % 4; + + nv_bfloat16 a_bf16[kNumRegPerMMA]; + a_bf16[0] = a_bf16_smem_ptr_upper[elem_offset]; + a_bf16[1] = a_bf16_smem_ptr_lower[elem_offset]; + a_bf16[2] = a_bf16_smem_ptr_upper[elem_offset + 4]; + a_bf16[3] = a_bf16_smem_ptr_lower[elem_offset + 4]; + + auto a_bf16x2_ptr = reinterpret_cast(a_bf16); + auto a_float2_ptr = reinterpret_cast(a); + + float2 a_float2_0 = __bfloat1622float2(a_bf16x2_ptr[0]); + float2 a_float2_1 = __bfloat1622float2(a_bf16x2_ptr[1]); + + a_float2_ptr[i * 2 + 0] = a_float2_0; + a_float2_ptr[i * 2 + 1] = a_float2_1; + + sqr_sum_acc_0 += a_float2_0.x * a_float2_0.x + a_float2_1.x * a_float2_1.x; + sqr_sum_acc_1 += a_float2_0.y * a_float2_0.y + a_float2_1.y * a_float2_1.y; + } + + __syncwarp(); + if (s > 0) + empty_barriers[(s - 1) % kNumStages]->arrive(); + + constexpr int kNumElemsInSwizzleRange = 128 / sizeof(float); + constexpr uint32_t kNumAtomsInSwizzleRange = kNumElemsInSwizzleRange / MMASelector::K; + DG_STATIC_ASSERT(BLOCK_K % kNumElemsInSwizzleRange == 0, "Invalid block K"); + + #pragma unroll + for (int i = 0; i < BLOCK_K / kNumElemsInSwizzleRange; i++) { + #pragma unroll + for (int k = 0; k < kNumAtomsInSwizzleRange; k++) { + + float* a_step = a + (i * kNumAtomsInSwizzleRange + k) * kNumRegPerMMA; + + #pragma unroll + for (int n = 0; n < MMASelector::N_atoms; n++) { + + uint32_t atom_n = lane_idx / 4; + uint32_t atom_k = lane_idx % 4; + + uint32_t global_n = n * 8 + atom_n; + uint32_t global_k = (i * kNumAtomsInSwizzleRange + k) * 8 + atom_k; + + uint32_t b_atom_idx = global_k / kNumElemsInSwizzleRange; + uint32_t b_atom_k = global_k % kNumElemsInSwizzleRange; + uint32_t atom_linear_idx_0 = global_n * kNumElemsInSwizzleRange + b_atom_k; + uint32_t atom_linear_idx_1 = atom_linear_idx_0 + 4; + + uint32_t swizzle_xor = ((atom_linear_idx_0 >> 5) & 7) << 2; + uint32_t atom_swizzled_idx_0 = atom_linear_idx_0 ^ swizzle_xor; + uint32_t atom_swizzled_idx_1 = atom_linear_idx_1 ^ swizzle_xor; + uint32_t atom_base_idx = b_atom_idx * BLOCK_N * kNumElemsInSwizzleRange; + + float b0 = smem_b[stage_idx][atom_base_idx + atom_swizzled_idx_0]; + float b1 = smem_b[stage_idx][atom_base_idx + atom_swizzled_idx_1]; + + MMASelector::type::fma( + accum[n * 4 + 0], accum[n * 4 + 1], accum[n * 4 + 2], accum[n * 4 + 3], + a_step[0], a_step[1], a_step[2], a_step[3], + b0, b1, + accum[n * 4 + 0], accum[n * 4 + 1], accum[n * 4 + 2], accum[n * 4 + 3] + ); + } + } + } + } + + const auto& reduced_sum_0 = math::warp_reduce_sum<4>(sqr_sum_acc_0); + const auto& reduced_sum_1 = math::warp_reduce_sum<4>(sqr_sum_acc_1); + + const auto& m_idx = m_block_idx * BLOCK_M + (warp_idx * BLOCK_M_PER_WARP + lane_idx / 4); + if (lane_idx % 4 == 0) { + if (m_idx < shape_m) + sqr_sum[m_offset + m_idx] = reduced_sum_0; + if (m_idx + 8 < shape_m) + sqr_sum[m_offset + m_idx + 8] = reduced_sum_1; + } + + __syncwarp(); + empty_barriers[(num_total_stages-1) % kNumStages]->arrive(); + + uint32_t is_odd_pair = lane_idx / 2 % 2; + uint32_t row_idx = lane_idx / 4; + uint32_t reordered_pair_idx = is_odd_pair * 8 + row_idx; + + auto shifted_smem_ptr = reinterpret_cast(smem_cd) + + (warp_idx * BLOCK_M_PER_WARP + row_idx) * kSwizzleCDMode + + lane_idx % 2 * 8; + + #pragma unroll + for (uint32_t i = 0; i < (kSwizzleCDMode / sizeof(float)) / 4; i += 2) { + uint32_t bank_group_idx = get_sm120_swizzled_bank_group_idx(i + is_odd_pair, reordered_pair_idx); + auto smem_ptr = shifted_smem_ptr + bank_group_idx * kNumBankGroupBytes; + + auto values = reinterpret_cast(accum + i * 2); + ptx::st_shared(smem_ptr, values[0], values[1]); + ptx::st_shared(smem_ptr + 8 * kSwizzleCDMode, values[2], values[3]); + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(128, 1); + + if (warp_idx == 0 and cute::elect_one_sync()) { + if constexpr (kNumSplits == 1) { + cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_cd, 0, m_block_idx * BLOCK_M); + } else { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_cd, 0, m_block_idx * BLOCK_M, k_split_idx); + } + cute::tma_store_arrive(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_120"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/mma/sm120.cuh b/deep_gemm/include/deep_gemm/mma/sm120.cuh new file mode 100644 index 0000000000..91eec0efa0 --- /dev/null +++ b/deep_gemm/include/deep_gemm/mma/sm120.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include +#include // Required for __float_as_uint + +namespace deep_gemm::mma::sm120 { + +CUTLASS_DEVICE void mma_m16n8k8_f32_tf32accum( + float& d0, float& d1, float& d2, float& d3, + float a0, float a1, float a2, float a3, + float b0, float b1, + float c0, float c1, float c2, float c3) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + // tf32 multiplicands expect .b32 registers ("r") + : "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)), "r"(__float_as_uint(a2)), "r"(__float_as_uint(a3)), + "r"(__float_as_uint(b0)), "r"(__float_as_uint(b1)), + // f32 accumulators expect .f32 registers ("f") + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +} + +template +struct TF32MMASync { + static constexpr int MMA_M = 16; + static constexpr int MMA_N = 8; + static constexpr int MMA_K = 8; + static constexpr int kNumAccum = M * N / 32; + + static_assert(M == 16 and N == 8 and K == 8, "SM120 TF32 mma.sync atom is 16x8x8"); + + CUTLASS_DEVICE static void fma( + float& d0, float& d1, float& d2, float& d3, + float a0, float a1, float a2, float a3, + float b0, float b1, + float c0, float c1, float c2, float c3) { + mma_m16n8k8_f32_tf32accum(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3); + } +}; + +template +struct TF32MMASelector { + static constexpr auto select_type() { + static_assert(N == 8 or N == 16 or N == 32, "SM120 TF32 hc_prenorm supports N <= 32"); + return TF32MMASync<16, 8, 8>{}; + } + + using type = decltype(select_type()); + + static constexpr int M = 16; + static constexpr int K = 8; + static constexpr int kNumAccumPerAtom = 4; + + static constexpr int MMA_N_per_atom() { return 8; } + static constexpr int N_atoms = N / MMA_N_per_atom(); +}; + +} // namespace deep_gemm::mma::sm120 diff --git a/tests/test_sm120_kernels.py b/tests/test_sm120_kernels.py new file mode 100644 index 0000000000..f7f95d0a97 --- /dev/null +++ b/tests/test_sm120_kernels.py @@ -0,0 +1,380 @@ +import torch +import pytest + +import deep_gemm +from deep_gemm.testing import calc_diff, get_arch_major, test_filter as _test_filter +from deep_gemm.utils.math import ceil_div, per_block_cast_to_fp8, per_token_cast_to_fp8 + + +def _cast_kv_cache_to_fp8(kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + num_blocks, block_kv, num_heads, head_dim = kv_cache.shape + assert num_heads == 1 + + scale = kv_cache.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) / 448.0 + kv_fp8 = (kv_cache * scale.reciprocal()).to(torch.float8_e4m3fn) + kv_simulated = kv_fp8.float() * scale + + fused = torch.empty( + (num_blocks, block_kv * (head_dim + 4)), + device=kv_cache.device, + dtype=torch.uint8, + ) + fused[:, : block_kv * head_dim] = kv_fp8.view(num_blocks, block_kv * head_dim).view(torch.uint8) + fused[:, block_kv * head_dim :] = scale.view(num_blocks, block_kv).view(torch.uint8) + return fused.view(num_blocks, block_kv, num_heads, head_dim + 4), kv_simulated + + +def _paged_mqa_reference( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + max_context_len: int, +) -> torch.Tensor: + batch_size, next_n, num_heads, head_dim = q.shape + _, block_kv, _, _ = kv_cache.shape + logits = torch.full( + (batch_size * next_n, max_context_len), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + for batch_idx in range(batch_size): + for next_idx in range(next_n): + row = batch_idx * next_n + next_idx + context_len = int(context_lens[batch_idx, next_idx].item()) + for token_idx in range(context_len): + block_idx = int(block_table[batch_idx, token_idx // block_kv].item()) + kv = kv_cache[block_idx, token_idx % block_kv, 0].float() + score = (q[batch_idx, next_idx].float() * kv).sum(dim=1) + logits[row, token_idx] = (score.relu() * weights[row]).sum() + + return logits + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_hc_prenorm_gemm_kernel_path() -> None: + torch.manual_seed(0) + + m, n, k = 5, 8, 256 + a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + b = torch.randn((n, k), dtype=torch.float32, device="cuda") + for num_splits in (None, 3): + d = torch.empty((m, n), dtype=torch.float32, device="cuda") if num_splits is None else \ + torch.empty((num_splits, m, n), dtype=torch.float32, device="cuda") + sqr_sum = torch.empty((m,), dtype=torch.float32, device="cuda") if num_splits is None else \ + torch.empty((num_splits, m), dtype=torch.float32, device="cuda") + + deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=num_splits) + + final_d = d if num_splits is None else d.sum(dim=0) + final_sqr_sum = sqr_sum if num_splits is None else sqr_sum.sum(dim=0) + ref_d = a.float() @ b.T + ref_sqr_sum = a.float().square().sum(dim=-1) + diff = max(calc_diff(final_d, ref_d), calc_diff(final_sqr_sum, ref_sqr_sum)) + assert diff < 1e-6, f"{num_splits=}, {diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_hc_prenorm_gemm_larger_tiles() -> None: + torch.manual_seed(6) + + for m, n, k, num_splits in ( + (70, 16, 512, 2), + (70, 32, 512, 4), + (129, 32, 1024, 4), + ): + a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + b = torch.randn((n, k), dtype=torch.float32, device="cuda") + d = torch.empty((num_splits, m, n), dtype=torch.float32, device="cuda") + sqr_sum = torch.empty((num_splits, m), dtype=torch.float32, device="cuda") + + deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=num_splits) + + ref_d = a.float() @ b.T + ref_sqr_sum = a.float().square().sum(dim=-1) + diff = max(calc_diff(d.sum(dim=0), ref_d), calc_diff(sqr_sum.sum(dim=0), ref_sqr_sum)) + assert diff < 1e-6, f"{m=}, {n=}, {k=}, {num_splits=}, {diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_scalar_path() -> None: + torch.manual_seed(1) + + batch_size, next_n, num_heads, head_dim = 2, 2, 32, 32 + block_kv = 64 + max_context_len = 96 + num_blocks = 4 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[17, 23], [31, 47]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0], [1]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens.view(-1, 1) + for logits_dtype in (torch.float32, torch.bfloat16): + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=logits_dtype, + ) + + assert logits.dtype == logits_dtype + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + threshold = 1e-6 if logits_dtype == torch.float32 else 2e-6 + assert diff < threshold, f"{logits_dtype=}, {diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_v4_shape_scalar_path() -> None: + torch.manual_seed(4) + + batch_size, next_n, num_heads, head_dim = 1, 1, 64, 128 + block_kv = 64 + max_context_len = 65 + num_blocks = 2 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[max_context_len]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0, 1]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + ) + + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + assert diff < 1e-6, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +@pytest.mark.parametrize( + "token_groups, cache_q", + ((1, False), (2, False), (4, False), (8, False), (2, True), (4, True), (8, True)), +) +def test_sm120_fp8_paged_mqa_logits_tiled_path(monkeypatch, token_groups: int, cache_q: bool) -> None: + torch.manual_seed(5) + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED", "1") + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED_GROUPS", str(token_groups)) + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED_CACHE_Q", "1" if cache_q else "0") + + batch_size, next_n, num_heads, head_dim = 2, 1, 64, 128 + block_kv = 64 + max_context_len = 73 + num_blocks = 4 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[max_context_len], [37]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0, 1], [2, 3]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + ) + + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + assert diff < 1e-6, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_varlen_scalar_path() -> None: + torch.manual_seed(2) + + raw_batch_size, num_heads, head_dim = 3, 32, 32 + block_kv = 32 + max_context_len = 80 + num_blocks = 8 + tokens_per_seq = torch.tensor([1, 2, 1], device="cuda", dtype=torch.int32) + indices = torch.arange(raw_batch_size, device="cuda", dtype=torch.int32) + indices = indices.repeat_interleave(tokens_per_seq) + batch_size = int(tokens_per_seq.sum().item()) + next_n = 1 + + base_context_lens = torch.tensor([11, 19, 7], device="cuda", dtype=torch.int32) + offsets = torch.cat([ + torch.arange(int(token_count.item()), device="cuda", dtype=torch.int32) + for token_count in tokens_per_seq + ]) + context_lens = (base_context_lens.repeat_interleave(tokens_per_seq) + offsets).view(-1, 1) + + q = torch.randn( + (batch_size, next_n, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + + base_block_table = torch.tensor( + [ + [0, 1, 2], + [3, 4, 5], + [6, 7, 0], + ], + device="cuda", + dtype=torch.int32, + ) + block_table = base_block_table.repeat_interleave(tokens_per_seq, dim=0) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + indices=indices, + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + indices=indices, + ) + + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + assert diff < 1e-6, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_einsum_scalar_path() -> None: + torch.manual_seed(3) + + batch_size, num_heads, rank, output_dim = 3, 4, 256, 128 + x = torch.randn((batch_size, num_heads, rank), device="cuda", dtype=torch.bfloat16) + y = torch.randn((num_heads, output_dim, rank), device="cuda", dtype=torch.bfloat16) + expected = torch.einsum("bhr,hdr->bhd", x, y) + + x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, rank), use_ue8m0=False) + x_in = x_fp8.view(batch_size, num_heads, rank), x_sf.view(batch_size, num_heads, ceil_div(rank, 128)) + y_fp8 = torch.empty_like(y, dtype=torch.float8_e4m3fn) + y_sf = torch.empty( + (num_heads, ceil_div(output_dim, 128), ceil_div(rank, 128)), + device="cuda", + dtype=torch.float32, + ) + for head_idx in range(num_heads): + y_fp8[head_idx], y_sf[head_idx] = per_block_cast_to_fp8(y[head_idx], use_ue8m0=False) + + actual = torch.empty((batch_size, num_heads, output_dim), device="cuda", dtype=torch.bfloat16) + deep_gemm.fp8_einsum("bhr,hdr->bhd", x_in, (y_fp8, y_sf), actual) + + diff = calc_diff(actual, expected) + assert diff < 1e-3, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_einsum_v4_shape_scalar_path() -> None: + torch.manual_seed(5) + + batch_size, num_heads, rank, output_dim = 5, 64, 128, 128 + x = torch.randn((batch_size, num_heads, rank), device="cuda", dtype=torch.bfloat16) + y = torch.randn((num_heads, output_dim, rank), device="cuda", dtype=torch.bfloat16) + expected = torch.einsum("bhr,hdr->bhd", x, y) + + x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, rank), use_ue8m0=False) + x_in = x_fp8.view(batch_size, num_heads, rank), x_sf.view(batch_size, num_heads, 1) + y_fp8 = torch.empty_like(y, dtype=torch.float8_e4m3fn) + y_sf = torch.empty((num_heads, 1, 1), device="cuda", dtype=torch.float32) + for head_idx in range(num_heads): + y_fp8[head_idx], y_sf[head_idx] = per_block_cast_to_fp8(y[head_idx], use_ue8m0=False) + + actual = torch.empty((batch_size, num_heads, output_dim), device="cuda", dtype=torch.bfloat16) + deep_gemm.fp8_einsum("bhr,hdr->bhd", x_in, (y_fp8, y_sf), actual) + + diff = calc_diff(actual, expected) + assert diff < 1e-3, f"{diff=}" diff --git a/third-party/cutlass b/third-party/cutlass index f3fde58372..da5e086dab 160000 --- a/third-party/cutlass +++ b/third-party/cutlass @@ -1 +1 @@ -Subproject commit f3fde58372d33e9a5650ba7b80fc48b3b49d40c8 +Subproject commit da5e086dab31d63815acafdac9a9c5893b1c69e2