[SM90] Add native block32 FP8 GEMM support (dsv41)

csrc/apis/gemm.hpp

@@ -199,12 +199,17 @@ static void fp8_fp4_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
 
         if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
             DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100");
+            // SM90 kernels read one activation scale per row.
+            const int gran_m = recipe.has_value() ? std::get<0>(recipe.value()) : std::get<0>(recipe_a.value());
             const int gran_n = recipe.has_value() ? std::get<1>(recipe.value()) : std::get<0>(recipe_b.value());
+            DG_HOST_ASSERT(gran_m == 1);
+            DG_HOST_ASSERT(gran_k_a == gran_k_b);
+            DG_HOST_ASSERT((gran_n == 32 and gran_k_a == 32) or (gran_k_a == 128 and (gran_n == 1 or gran_n == 128)));
             if (gran_n == 1) {
                 sm90_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
             } else {
                 const auto major_sfb = get_major_type_ab(sfb);
-                sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims);
+                sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims, std::nullopt, gran_k_a);
             }
         } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
             sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, gran_k_a, gran_k_b,
@@ -319,6 +324,7 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
 
     // Dispatch implementation
     if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
+        DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128);
         const auto major_sfb = get_major_type_ab(sfb);
         sm90_m_grouped_fp8_gemm_contiguous_1d2d(a.first, sfa, b.first, sfb, d, grouped_layout,
                                                 num_groups, m, n, k, major_a, major_b, major_sfb,
@@ -393,6 +399,7 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair<torch::Tensor, torc
 
     // Dispatch implementation
     if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
+        DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128);
         const auto major_sfb = get_major_type_ab(sfb);
         sm90_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m,
                                             num_groups, m, n, k, expected_m, major_a, major_b, major_sfb, compiled_dims);

csrc/apis/layout.hpp

@@ -34,12 +34,12 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
     // Pre-transform checks
     check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups);
 
-    // (FP32, 1, 128) on SM90: transform to TMA-aligned and MN-major
-    if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
+    // FP32 per-token scales on SM90: transform to TMA-aligned and MN-major
+    if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 128 or gran_k == 32) and (arch_major == 9 or disable_ue8m0_cast))
         return get_mn_major_tma_aligned_tensor(sf);
 
-    // (FP32, 128, 128) on SM90: no need to transform, check SFB requirements
-    if (sf.scalar_type() == torch::kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
+    // FP32 block scales on SM90: no need to transform, check SFB requirements
+    if (sf.scalar_type() == torch::kFloat and ((gran_mn == 128 and gran_k == 128) or (gran_mn == 32 and gran_k == 32)) and (arch_major == 9 or disable_ue8m0_cast))
         return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, false, true, torch::kFloat);
 
     // (FP32, x, gran_k) on SM100/SM120: transform to (INT, 1, gran_k), TMA-aligned and MN-major

csrc/jit_kernels/heuristics/sm90.hpp

@@ -14,8 +14,17 @@ namespace deep_gemm {
 
 struct SM90ArchSpec {
     static constexpr int smem_capacity = 232448;
-
     static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
+        constexpr int short_k_stages = 8;
+        // Block K is fixed by the MMA operand width.
+        const int block_k = 128 * 8 / get_num_element_bits(desc.get_mma_kind());
+        const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32;
+        const int n_tile_multiple = block32 ? desc.max_gran_k : 16;
+        const int max_1d2d_block_n = block32 ? block_k : 192;
+        const int short_k = short_k_stages * block_k;
+        const int max_block_m = desc.k <= short_k ? 2 * block_k : block_k;
+        const bool allow_small_m = not block32 or desc.k <= 2 * short_k;
+
         // Block M candidates
         std::vector<int> block_m_candidates;
         if (desc.gemm_type == GemmType::Normal or
@@ -24,8 +33,8 @@ struct SM90ArchSpec {
             // TODO: check 256's performance
             block_m_candidates = {64, 128};
             // NOTES: smaller block M can avoid TMA L2 OOB bound
-            if (desc.m <= 16) block_m_candidates.push_back(16);
-            if (desc.m <= 32) block_m_candidates.push_back(32);
+            if (desc.m <= 16 and allow_small_m) block_m_candidates.push_back(16);
+            if (desc.m <= 32 and allow_small_m) block_m_candidates.push_back(32);
 
             // BF16 output GEMM supports 256
             if (desc.cd_dtype != torch::kFloat)
@@ -39,7 +48,7 @@ struct SM90ArchSpec {
 
         // Block N candidates
         std::vector<int> block_n_candidates;
-        int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of());
+        int step = std::lcm(n_tile_multiple, 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) {
@@ -51,22 +60,19 @@ struct SM90ArchSpec {
         // Register spills
         int end = 256;
         if (desc.kernel_type == KernelType::Kernel1D2D)
-            end = 192;
+            end = max_1d2d_block_n;
         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 * 8 / get_num_element_bits(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);
+            (desc.gemm_type == GemmType::Batched) or block32;
 
         // Enumerate all candidates
         std::vector<Layout> candidates;
@@ -82,6 +88,10 @@ struct SM90ArchSpec {
 
                 for (int block_m: block_m_candidates) {
                     for (int block_n: block_n_candidates) {
+                        // Long K favors overlapping wide single-wave tiles; short K can
+                        // amortize B loads across two M waves without spilling.
+                        if (block32 and block_m > max_block_m)
+                            continue;
                         // 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;
@@ -149,6 +159,11 @@ struct SM90ArchSpec {
 
     static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) {
         constexpr int kNumMaxStages = 16;
+        const bool block32 = desc.kernel_type == KernelType::Kernel1D2D and desc.max_gran_k == 32;
+        const int scale_gran_k = block32 ? desc.max_gran_k : layout.block_k;
+        const int scales_per_stage = layout.block_k / scale_gran_k;
+        const int scale_b_rows = block32 ? layout.block_n / scale_gran_k :
+            (layout.block_k % layout.block_n == 0 ? 1 : 2);
 
         // TODO: consider swap AB
         // C/D for TMA stores
@@ -163,14 +178,13 @@ struct SM90ArchSpec {
 
         // 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);
+            0 : align(layout.block_m * scales_per_stage * 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);
+            0 : align<int>(ceil_div(desc.k, scale_gran_k) * scale_b_rows * static_cast<int>(sizeof(float)), 8);
 
         // Extra tensormap for 1D1D kernels
         const int smem_tensormap =

csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp

@@ -27,13 +27,19 @@ class SM90FP8Gemm1D2DRuntime final {
         CUtensorMap tensor_map_b;
         CUtensorMap tensor_map_d;
         CUtensorMap tensor_map_sfa;
+        int sf_granularity = 128;
     };
 
     static void compile_and_launch(const std::string& tag, const Args& args) {
         // A known N keeps wide single-CTA SM90 tiles free of register spills.
         const bool specialize_n = args.gemm_desc.gemm_type == GemmType::Normal &&
             args.gemm_desc.compiled_dims.empty() && args.gemm_config.layout.block_m == 128 &&
             args.gemm_config.layout.block_n > 128 && args.gemm_config.layout.get_cluster_size() == 1;
+        // Specialize fixed N/K to simplify scale indexing.
+        // Keep M dynamic to avoid specialization for each token count.
+        auto compiled_dims = args.gemm_desc.compiled_dims;
+        if (args.sf_granularity == 32 and compiled_dims.empty())
+            compiled_dims = "nk";
         const auto kernel = jit->compile(tag, std::format(R"(
 #include <deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh>
 
@@ -51,14 +57,14 @@ static void __instantiate_kernel() {{
         {}, {},
         {}, {},
         {},
-        {}
+        {}, {}
     >);
 }};
 )",
         to_string(args.major_sfb),
-        get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
-        get_compiled_dim(args.gemm_desc.n, 'n', specialize_n ? "n" : args.gemm_desc.compiled_dims),
-        get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
+        get_compiled_dim(args.gemm_desc.m, 'm', compiled_dims),
+        get_compiled_dim(args.gemm_desc.n, 'n', specialize_n ? "n" : compiled_dims),
+        get_compiled_dim(args.gemm_desc.k, 'k', compiled_dims),
         args.gemm_desc.num_groups,
         args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
         args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
@@ -67,7 +73,7 @@ static void __instantiate_kernel() {{
         args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
         args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type),
         to_string(args.gemm_desc.cd_dtype),
-        get_default_epilogue_type(args.epilogue_type)));
+        get_default_epilogue_type(args.epilogue_type), args.sf_granularity));
 
         // Launch
         jit->launch(
@@ -87,8 +93,11 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
                                const int& m, const int& n, const int& k,
                                const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
                                const std::string& compiled_dims,
-                               const std::optional<std::string>& epilogue_type = std::nullopt) {
+                               const std::optional<std::string>& epilogue_type = std::nullopt,
+                               const int sf_granularity = 128) {
     DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
+    DG_HOST_ASSERT(sf_granularity == 128 or sf_granularity == 32);
+    DG_HOST_ASSERT(sf_granularity != 32 or (n % 8 == 0 and k % 32 == 0));
     DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
 
     const auto desc = GemmDesc {
@@ -101,7 +110,8 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
         .with_accumulation = c.has_value(),
         .num_sms = runtime->get_num_sms(),
         .tc_util = runtime->get_tc_util(),
-        .compiled_dims = compiled_dims
+        .compiled_dims = compiled_dims,
+        .max_gran_k = sf_granularity
     };
     const auto config = get_best_config<SM90ArchSpec>(desc);
 
@@ -124,7 +134,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
                                                static_cast<int>(d.stride(-2)), 1,
                                                config.storage_config.swizzle_cd_mode);
     const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
-                                                 config.layout.block_m, config.layout.block_k, 1, 0);
+                                                 config.layout.block_m, sf_granularity, 1, 0, 0, false, config.layout.block_k / sf_granularity);
 
     // Compile and launch
     SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_fp8_gemm_1d2d", {
@@ -144,6 +154,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
         .tensor_map_b = tensor_map_b,
         .tensor_map_d = tensor_map_d,
         .tensor_map_sfa = tensor_map_sfa,
+        .sf_granularity = sf_granularity,
     });
 }