Add SM90 BF16 grouped GEMM with gathered activation rows

README.md

@@ -77,6 +77,8 @@ To perform a basic non-grouped FP8 GEMM, call the `fp8_gemm_{nt, nn, tn, tt}` fu
 
 Unlike traditional grouped GEMMs in CUTLASS, DeepGEMM groups only the M-axis, while N and K must remain fixed. This design is tailored for scenarios where experts in an MoE model share the same shape. For training forward passes or inference prefilling, where each expert may process a varying number of tokens, we concatenate these tokens into a single tensor, referred to as the "contiguous" layout. Note that each expert segment must be aligned to the GEMM M block size (`get_mk_alignment_for_contiguous_layout()`).  For more information, please refer to the `m_grouped_fp8_gemm_{nt, nn}_contiguous` function documentation.
 
+On SM90, `m_grouped_bf16_gemm_nt_contiguous_gathered(a, b, d, grouped_layout, a_row_indices)` gathers the rows of `a` inside the kernel, which equals `m_grouped_bf16_gemm_nt_contiguous(a.index_select(0, a_row_indices), b, d, grouped_layout)` without materializing the gathered activations. It requires K-major BF16 inputs with `K % 64 == 0`, an M alignment of 128, and contiguous int64 row indices that address valid rows (including padding rows); see `test_m_grouped_gemm_contiguous_gathered` in `tests/test_bf16.py` for an example.
+
 We also provide a K-axis-grouped API for MoE weight backward (with M and N must remain fixed), please refer to `k_grouped_fp8_gemm_tn_contiguous` for more information.
 
 #### Grouped GEMMs (masked layout)

csrc/apis/gemm.hpp

@@ -632,6 +632,44 @@ static void m_grouped_bf16_gemm_nn_contiguous(const torch::Tensor& a, const torc
                                       d, grouped_layout, compiled_dims, use_psum_layout, ensure_zero_padding, std::nullopt, epilogue_class);
 }
 
+static void m_grouped_bf16_gemm_nt_contiguous_gathered(const torch::Tensor& a, const torch::Tensor& b,
+                                                       const torch::Tensor& d, const torch::Tensor& grouped_layout,
+                                                       const torch::Tensor& a_row_indices) {
+    // Shape must be `[M, K] @ [G, N, K].mT`, where row `i` of the GEMM input is `a[a_row_indices[i]]`
+    // NOTES: padding rows must also reference valid source rows
+    DG_HOST_ASSERT(jit->device.get_arch_major() == 9 and "Gathered A requires SM90");
+    DG_HOST_ASSERT(a.is_cuda() and b.device() == a.device() and d.device() == a.device() and
+                   grouped_layout.device() == a.device() and a_row_indices.device() == a.device());
+    DG_HOST_ASSERT(grouped_layout.is_contiguous() and a_row_indices.is_contiguous());
+
+    // Type and shape checks
+    const auto [num_rows, k] = get_shape<2>(a);
+    const auto [num_groups, n, k_] = get_shape<3>(b);
+    const auto [m, n_] = get_shape<2>(d);
+    DG_HOST_ASSERT(n == n_ and k == k_);
+    DG_HOST_ASSERT(n > 0 and k > 0 and k % 64 == 0 and num_groups > 0);
+    DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16);
+    DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16);
+    DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
+    const auto [num_layout_rows] = get_shape<1>(grouped_layout);
+    const auto [num_row_indices] = get_shape<1>(a_row_indices);
+    DG_HOST_ASSERT(num_layout_rows == m and num_row_indices == m);
+    DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and a_row_indices.scalar_type() == torch::kInt64);
+
+    // Layout checks: K-major A with 16-byte aligned rows (a padded row stride is allowed), contiguous B, N-major D
+    DG_HOST_ASSERT(heuristics_runtime->get_mk_alignment_for_contiguous_layout() == 128);
+    DG_HOST_ASSERT(a.stride(1) == 1 and a.stride(0) % 8 == 0 and reinterpret_cast<uintptr_t>(a.data_ptr()) % 16 == 0);
+    DG_HOST_ASSERT(b.is_contiguous());
+    check_major_type_cd(d);
+
+    // Do nothing if empty
+    if (m == 0)
+        return;
+    DG_HOST_ASSERT(num_rows > 0);
+
+    sm90_m_grouped_bf16_gemm_contiguous_gathered(a, b, d, grouped_layout, a_row_indices, num_groups, m, n, k);
+}
+
 static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::Tensor& b,
                                           const torch::Tensor& d, const torch::Tensor& masked_m,
                                           const int& expected_m, const std::string& compiled_dims,
@@ -966,6 +1004,8 @@ static void register_apis(pybind11::module_& m) {
           py::arg("use_psum_layout") = false,
           py::arg("ensure_zero_padding") = true,
           py::arg("epilogue") = nullptr);
+    m.def("m_grouped_bf16_gemm_nt_contiguous_gathered", &m_grouped_bf16_gemm_nt_contiguous_gathered,
+          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("a_row_indices"));
     m.def("m_grouped_bf16_gemm_nt_masked", &m_grouped_bf16_gemm_nt_masked,
           py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"),
           py::arg("expected_m"), py::arg("compiled_dims") = "nk",

csrc/jit_kernels/impls/sm90_bf16_gemm.hpp

@@ -21,6 +21,10 @@ class SM90BF16GemmRuntime final {
         CUtensorMap tensor_map_a;
         CUtensorMap tensor_map_b;
         CUtensorMap tensor_map_cd;
+        void* gathered_a = nullptr;
+        void* a_row_indices = nullptr;
+        uint64_t gathered_a_stride = 0;
+        uint32_t gathered_a_rows = 0;
     };
 
     static void compile_and_launch(const std::string& tag, const Args& args) {
@@ -41,7 +45,7 @@ static void __instantiate_kernel() {{
         {}, {},
         {},
         {}, {},
-        {}
+        {}, {}
     >);
 }};
 )",
@@ -61,15 +65,16 @@ 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), args.gemm_desc.with_accumulation,
-        to_string(args.gemm_desc.cd_dtype)));
+        to_string(args.gemm_desc.cd_dtype), args.a_row_indices != nullptr));
 
         // Launch
         jit->launch(
             kernel, args.options,
             args.grouped_layout,
             args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
             args.tensor_map_a, args.tensor_map_b,
-            args.tensor_map_cd
+            args.tensor_map_cd,
+            args.gathered_a, args.a_row_indices, args.gathered_a_stride, args.gathered_a_rows
         );
     }
 };
@@ -200,6 +205,70 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
     });
 }
 
+// Gathered A is loaded with `cp.async` by each CTA, so TMA multicast clusters do not apply
+struct SM90GatheredArchSpec: SM90ArchSpec {
+    static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
+        auto layouts = SM90ArchSpec::get_layout_candidates(desc);
+        std::erase_if(layouts, [](const Layout& layout) { return layout.get_cluster_size() != 1; });
+        return layouts;
+    }
+};
+
+static void sm90_m_grouped_bf16_gemm_contiguous_gathered(const torch::Tensor& a,
+                                                         const torch::Tensor& b,
+                                                         const torch::Tensor& d,
+                                                         const torch::Tensor& m_indices,
+                                                         const torch::Tensor& a_row_indices,
+                                                         const int& num_groups, const int& m, const int& n, const int& k) {
+    const auto desc = GemmDesc {
+        .gemm_type = GemmType::MGroupedContiguous,
+        .kernel_type = KernelType::KernelNoSF,
+        .m = m, .n = n, .k = k, .num_groups = num_groups,
+        .a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
+        .cd_dtype = d.scalar_type(),
+        .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K,
+        .with_accumulation = false,
+        .num_sms = runtime->get_num_sms(),
+        .tc_util = runtime->get_tc_util(),
+        .compiled_dims = "nk",
+        .expected_m = m, .expected_n = n, .expected_k = k,
+        .expected_num_groups = 1
+    };
+    const auto config = get_best_config<SM90GatheredArchSpec>(desc);
+
+    // A is gathered by the kernel itself, so only B and D use TMA
+    const auto tensor_map_b = make_tma_b_desc(desc.major_b, b, n, k,
+                                              config.storage_config.load_block_n,
+                                              config.layout.block_k,
+                                              static_cast<int>(b.stride(get_non_contiguous_dim(desc.major_b))), num_groups,
+                                              config.storage_config.swizzle_b_mode);
+    const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
+                                                config.storage_config.store_block_m,
+                                                config.storage_config.store_block_n,
+                                                static_cast<int>(d.stride(-2)), 1,
+                                                config.storage_config.swizzle_cd_mode);
+
+    // Compile and launch
+    SM90BF16GemmRuntime::compile_and_launch("sm90_m_grouped_bf16_gemm_contiguous_gathered", {
+        .gemm_desc = desc,
+        .gemm_config = config,
+        .options = {
+            .num_smem_bytes = config.pipeline_config.smem_size,
+            .grid_dim = dim3(config.launch_config.num_sms, 1, 1),
+            .block_dim = dim3(config.launch_config.num_threads, 1, 1),
+            .cluster_dim = dim3(1, 1, 1),
+        },
+        .grouped_layout = m_indices.data_ptr(),
+        .tensor_map_a = {},
+        .tensor_map_b = tensor_map_b,
+        .tensor_map_cd = tensor_map_cd,
+        .gathered_a = a.data_ptr(),
+        .a_row_indices = a_row_indices.data_ptr(),
+        .gathered_a_stride = static_cast<uint64_t>(a.stride(0)),
+        .gathered_a_rows = static_cast<uint32_t>(a.size(0))
+    });
+}
+
 static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
                                             const torch::Tensor& b,
                                             const torch::Tensor& d,

deep_gemm/__init__.py

@@ -60,6 +60,7 @@
     bf16_gemm_tn, bf16_gemm_tt,
     m_grouped_bf16_gemm_nt_contiguous,
     m_grouped_bf16_gemm_nn_contiguous,
+    m_grouped_bf16_gemm_nt_contiguous_gathered,
     m_grouped_bf16_gemm_nt_masked,
     k_grouped_bf16_gemm_tn_contiguous,
     # MegaGate kernels

deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh

@@ -34,13 +34,15 @@ template <cute::UMMA::Major kMajorA, cute::UMMA::Major kMajorB,
           uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           uint32_t kNumSMs,
           GemmType kGemmType, bool kWithAccumulation,
-          typename cd_dtype_t>
+          typename cd_dtype_t, bool kGatherA = false>
 CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void
 sm90_bf16_gemm_impl(int* grouped_layout,
                     uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
                     const __grid_constant__ cute::TmaDescriptor tensor_map_a,
                     const __grid_constant__ cute::TmaDescriptor tensor_map_b,
-                    const __grid_constant__ cute::TmaDescriptor tensor_map_cd) {
+                    const __grid_constant__ cute::TmaDescriptor tensor_map_cd,
+                    const nv_bfloat16* gathered_a, const int64_t* a_row_indices,
+                    uint64_t gathered_a_stride, uint32_t gathered_a_rows) {
 #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
     // Enlarge `BLOCK_K` for some cases
     // NOTES: this is for reducing the `warpgroup_wait<0>()` overhead
@@ -83,7 +85,8 @@ sm90_bf16_gemm_impl(int* grouped_layout,
 
     // Prefetch TMA descriptors at the very beginning
     if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
-        cute::prefetch_tma_descriptor(&tensor_map_a);
+        if constexpr (not kGatherA)
+            cute::prefetch_tma_descriptor(&tensor_map_a);
         cute::prefetch_tma_descriptor(&tensor_map_b);
         cute::prefetch_tma_descriptor(&tensor_map_cd);
     }
@@ -112,7 +115,7 @@ sm90_bf16_gemm_impl(int* grouped_layout,
     if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) {
         #pragma unroll
         for (uint32_t i = 0; i < kNumStages; ++ i) {
-            full_barriers[i]->init(1);
+            full_barriers[i]->init(kGatherA ? kNumTMAThreads + 1 : 1);
             empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
         }
 
@@ -124,7 +127,8 @@ sm90_bf16_gemm_impl(int* grouped_layout,
     (kNumTMAMulticast > 1) ? comm::cluster_sync_with_relaxed_arrive() : __syncthreads();
 
     // Register reconfigurations
-    constexpr uint32_t kNumTMARegisters = 48;
+    // NOTES: the gathering loader warps keep per-row source offsets in registers; 56 still leaves 224 for both math warpgroups
+    constexpr uint32_t kNumTMARegisters = kGatherA ? 56 : 48;
     constexpr uint32_t kNumMathRegisters = kNumMathThreads == 128 ? 248 : 224;
 
     // Wait for primary kernel completion
@@ -150,7 +154,61 @@ sm90_bf16_gemm_impl(int* grouped_layout,
 
         // NOTES: only one thread (or warp) will be used
         // We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32`
-        if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) {
+        if constexpr (kGatherA) {
+            DG_STATIC_ASSERT(kMajorA == cute::UMMA::Major::K and kMajorB == cute::UMMA::Major::K,
+                             "Gathered A requires K-major BF16 inputs");
+            DG_STATIC_ASSERT(kGemmType == GemmType::MGroupedContiguous and kNumTMAMulticast == 1,
+                             "Gathered A supports contiguous M groups without multicast");
+            DG_STATIC_ASSERT(BLOCK_K == 64 and kSwizzleAMode == 128 and BLOCK_M == 128,
+                             "Gathered A currently supports 128x64 A tiles with 128-byte swizzle");
+            constexpr uint32_t kVectorsPerRow = BLOCK_K * sizeof(nv_bfloat16) / 16;
+            constexpr uint32_t kVectorsPerThread = BLOCK_M * kVectorsPerRow / kNumTMAThreads;
+            const uint32_t loader_idx = threadIdx.x - kNumMathThreads;
+            const bool issue_b = warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync();
+            while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+                const auto m_idx = m_block_idx * BLOCK_M;
+                const auto n_idx = scheduler.template get_global_idx<true, sched::IndexType::MN>(
+                    shape_n, BLOCK_N, n_block_idx, m_block_idx);
+                uint64_t source_offsets[kVectorsPerThread];
+                bool valid_rows[kVectorsPerThread];
+                #pragma unroll
+                for (uint32_t i = 0; i < kVectorsPerThread; ++ i) {
+                    const auto vector_idx = loader_idx + i * kNumTMAThreads;
+                    const auto row = vector_idx / kVectorsPerRow;
+                    valid_rows[i] = m_idx + row < shape_m;
+                    const auto source = valid_rows[i] ? a_row_indices[m_idx + row] : 0;
+                    DG_TRAP_ONLY_DEVICE_ASSERT(not valid_rows[i] or (source >= 0 and source < gathered_a_rows));
+                    source_offsets[i] = static_cast<uint64_t>(source) * gathered_a_stride;
+                }
+                const auto num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K);
+                for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
+                    empty_barriers[stage_idx]->wait(phase ^ 1);
+                    auto& full_barrier = *full_barriers[stage_idx];
+                    if (issue_b) {
+                        tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode, cutlass::bfloat16_t>(
+                            &tensor_map_b, &full_barrier, smem_b[stage_idx],
+                            k_block_idx * BLOCK_K, n_idx);
+                        full_barrier.arrive_and_expect_tx(SMEM_B_SIZE_PER_STAGE);
+                    }
+                    #pragma unroll
+                    for (uint32_t i = 0; i < kVectorsPerThread; ++ i) {
+                        const auto vector_idx = loader_idx + i * kNumTMAThreads;
+                        const auto row = vector_idx / kVectorsPerRow;
+                        const auto column = vector_idx % kVectorsPerRow;
+                        // Match the existing TMA descriptor's 128-byte XOR swizzle.
+                        const auto destination = reinterpret_cast<uint4*>(smem_a[stage_idx]) +
+                            row * kVectorsPerRow + (column ^ (row % 8));
+                        const auto source = gathered_a + source_offsets[i] +
+                            k_block_idx * BLOCK_K + column * 8;
+                        asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16, %2;" ::
+                            "r"(static_cast<uint32_t>(__cvta_generic_to_shared(destination))),
+                            "l"(source), "r"(valid_rows[i] ? 16 : 0) : "memory");
+                    }
+                    // The barrier counts every loader plus the B transaction issuer.
+                    cutlass::arch::cpasync_barrier_arrive_noinc(reinterpret_cast<uint64_t*>(&full_barrier));
+                }
+            }
+        } else if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) {
             DG_STATIC_ASSERT(kNumTMAThreads >= 128, "Need at least 128 threads for TMA warp-group");
 
             // Persistently schedule over blocks