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_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp new file mode 100644 index 0000000000..7399ddc4a5 --- /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_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp new file mode 100644 index 0000000000..3da5b907cf --- /dev/null +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -0,0 +1,153 @@ +#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) { + + // Note: may need to change these later + 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); + + 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/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..a9dcbd9c61 --- /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..0ca91bd073 --- /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