feat: support SM90 FP8 × FP4 grouped GEMM

csrc/apis/gemm.hpp

@@ -5,6 +5,7 @@
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
+#include "../jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp"
 #include "../jit_kernels/impls/sm90_bf16_gemm.hpp"
 #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
@@ -140,6 +141,27 @@ static void fp8_fp4_gemm_tt(const std::pair<torch::Tensor, torch::Tensor>& a,
                     d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast);
 }
 
+static void fp8_fp4_gemm_nt_sm90_fused_wgmma(const std::pair<torch::Tensor, torch::Tensor>& a,
+                                             const std::pair<torch::Tensor, torch::Tensor>& b,
+                                             const torch::Tensor& d,
+                                             const std::optional<torch::Tensor>& c,
+                                             const int& gran_k,
+                                             const std::string& compiled_dims) {
+    const auto [m, k] = get_shape<2>(a.first);
+    const auto [n, half_k] = get_shape<2>(b.first);
+    DG_HOST_ASSERT(half_k * 2 == k);
+
+    std::optional<std::tuple<int, int, int>> recipe = std::nullopt;
+    std::optional<std::tuple<int, int>> recipe_a = std::make_tuple(1, gran_k);
+    std::optional<std::tuple<int, int>> recipe_b = std::make_tuple(1, gran_k);
+    const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout(
+        a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, false);
+    DG_HOST_ASSERT(gran_k_a == gran_k and gran_k_b == gran_k);
+
+    sm90_fp8_fp4_gemm_1d1d_fused(
+        {a.first, sfa}, {b.first, sfb}, d, c, gran_k, compiled_dims);
+}
+
 static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor, torch::Tensor>& a,
                                                  const std::pair<torch::Tensor, torch::Tensor>& b,
                                                  const torch::Tensor& d,
@@ -624,6 +646,37 @@ static void register_apis(pybind11::module_& m) {
           py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt,
           py::arg("compiled_dims") = "mn",
           py::arg("disable_ue8m0_cast") = false);
+    m.def("fp8_fp4_gemm_nt_sm90_fused_wgmma", &fp8_fp4_gemm_nt_sm90_fused_wgmma,
+          py::arg("a"), py::arg("b"), py::arg("d"),
+          py::arg("c") = std::nullopt,
+          py::arg("gran_k") = 128,
+          py::arg("compiled_dims") = "nk");
+    m.def("m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma",
+          &sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused,
+          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"),
+          py::arg("gran_k") = 128,
+          py::arg("compiled_dims") = "nk",
+          py::arg("use_psum_layout") = false,
+          py::arg("expected_m_for_psum_layout") = std::nullopt,
+          py::arg("block_m_override") = std::nullopt,
+          py::arg("block_n_override") = std::nullopt,
+          py::arg("decode_stub") = false);
+    m.def("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma",
+          &sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused,
+          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"),
+          py::arg("expected_m"),
+          py::arg("gran_k") = 128,
+          py::arg("gran_k_a") = std::nullopt,
+          py::arg("gran_k_b") = std::nullopt,
+          py::arg("compiled_dims") = "nk",
+          py::arg("block_m_override") = std::nullopt,
+          py::arg("block_n_override") = std::nullopt,
+          py::arg("decode_stub") = false,
+          py::arg("b_is_int4_sym") = false,
+          py::arg("masked_m_max_hint") = std::nullopt,
+          py::arg("active_groups_hint") = std::nullopt);
+    m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") =
+        m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma");
     m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous,
           py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"),
           py::arg("recipe") = std::nullopt,

csrc/apis/layout.hpp

@@ -36,8 +36,20 @@ 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))
+    // SM90: BF16 SFB supports gran_k {32, 128}; expects MN-major TMA-aligned strides.
+    if (sf.scalar_type() == torch::kBFloat16 and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9)
+        return get_mn_major_tma_aligned_tensor(sf);
+    // SM90: UInt8 E8M0 SFB supports gran_k=32; expects MN-major TMA-aligned strides.
+    if (sf.scalar_type() == torch::kUInt8 and gran_mn == 1 and gran_k == 32 and arch_major == 9)
+        return get_mn_major_tma_aligned_tensor(sf);
+
+    // (FP32, 1, 32/128) on SM90: transform to TMA-aligned and MN-major
+    if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and
+        (arch_major == 9 or disable_ue8m0_cast))
+        return get_mn_major_tma_aligned_tensor(sf);
+
+    // (INT packed UE8M0, 1, 32/128) on SM90: transform to TMA-aligned and MN-major.
+    if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9)
         return get_mn_major_tma_aligned_tensor(sf);
 
     // (FP32, 128, 128) on SM90: no need to transform, check SFB requirements