Support SM120

.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
+deep_gemm/*.so

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<int>());
-    } 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::tuple<torch::Tensor, st
     const auto arch_major = device_runtime->get_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::tuple<torch::Tensor, st
         sm100_fp4_paged_mqa_logits(q_fp, q_sf.value(), kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, schedule_meta,
                                    logits_dtype, batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
                                    is_varlen, aligned_max_context_len, block_table_stride, num_sms, split_kv);
+    } else if (not is_fp4 and arch_major == 12) {
+        sm120_fp8_paged_mqa_logits(q_fp, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table,
+                                   logits_dtype, batch_size, next_n, num_heads, head_dim, block_kv,
+                                   is_context_lens_2d, aligned_max_context_len, block_table_stride);
     } else if (not is_fp4 and (arch_major == 9 or arch_major == 10)) {
         smxx_fp8_paged_mqa_logits(q_fp, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, schedule_meta,
                                   logits_dtype, batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,

csrc/apis/einsum.hpp

@@ -14,6 +14,7 @@
 #include "../jit_kernels/impls/sm100_bmk_bnk_mn.hpp"
 #include "../jit_kernels/impls/sm90_bf16_gemm.hpp"
 #include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
+#include "../jit_kernels/impls/sm120_fp8_einsum.hpp"
 #include "../jit_kernels/impls/smxx_cublaslt.hpp"
 #endif
 
@@ -186,6 +187,13 @@ static void fp8_einsum(const std::string& expr,
     // Some hardcoded Einstein sum kernels
     const auto arch_major = device_runtime->get_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});

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");
     }

csrc/indexing/main.cu

@@ -20,6 +20,7 @@
 // Hyperconnection kernels
 #include <deep_gemm/impls/sm90_tf32_hc_prenorm_gemm.cuh>
 #include <deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh>
+#include <deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh>
 
 // Layout kernels
 #include <deep_gemm/impls/smxx_layout.cuh>

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,

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";

csrc/jit_kernels/heuristics/sm120.hpp

@@ -0,0 +1,243 @@
+#pragma once
+
+#include <cute/arch/mma_sm120.hpp>
+#include <deep_gemm/common/types.cuh>
+
+#include "common.hpp"
+#include "utils.hpp"
+#include "../../utils/exception.hpp"
+
+namespace deep_gemm {
+
+struct SM120ArchSpec {
+
+    static constexpr int smem_capacity = 232448;
+
+    static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
+        // Block M candidates
+        std::vector<int> 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<int> 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<Layout> 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<int>(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<int>(sizeof(float)), 128);
+        const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ?
+            0 : align(layout.block_n * static_cast<int>(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<int>(ceil_div(desc.k, layout.block_k) * static_cast<int>(sizeof(float)) * use_uniform_sfb, 8);
+
+        // Extra tensormap for 1D1D kernels
+        const int smem_tensormap =
+            desc.gemm_type == GemmType::KGroupedContiguous ? 4 * static_cast<int>(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<float>(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<int64_t>::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