SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32

csrc/jit_kernels/heuristics/runtime.hpp

@@ -49,9 +49,9 @@ class HeuristicsRuntime {
     struct ContiguousMKAlignment { int max_block_m, min_block_m, step; };
 
     static ContiguousMKAlignment get_contiguous_mk_alignment(const int& arch_major) {
-        // SM120: warp layout is kMWarps(4) * MMA_M(16), so BLOCK_M is a multiple of 64
+        // SM120: BLOCK_M in {32, 64, 128}; 32 uses the kNWarps=4 warp layout (decode)
         if (arch_major == 12)
-            return {128, 64, 64};
+            return {128, 32, 32};
         // SM100: 224 down to 32 by 32 (sgl tuning)
         if (arch_major == 10)
             return {224, 32, 32};
@@ -70,6 +70,9 @@ class HeuristicsRuntime {
                 per_group_m = (per_group_m + num_groups.value() - 1) / num_groups.value();
             for (; block_m > spec.min_block_m and block_m - spec.step >= per_group_m; block_m -= spec.step);
         }
+        // SM120 supports no 96-row tile (per-expert m > 64 measures best at 128)
+        if (device_runtime->get_arch_major() == 12 and block_m == 96)
+            block_m = 128;
         return block_m;
     }
 };

csrc/jit_kernels/heuristics/sm120.hpp

@@ -13,7 +13,7 @@ namespace deep_gemm {
 struct SM120ArchSpec {
     static constexpr int smem_capacity = 101376;  // 99KB
 
-    static constexpr int kMinBlockM = 64;   // kMWarps(4) * MMA_M(16), both FP8 and BF16 with kNWarps=2
+    static constexpr int kMinBlockM = 64;   // kNWarps=2 floor; BLOCK_M=32 uses kNWarps=4 (m-grouped decode)
 
     static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
         const int elem_size = get_element_size(desc.get_mma_kind());
@@ -25,7 +25,15 @@ struct SM120ArchSpec {
         const bool is_small_n = (n_for_tile > 0 and n_for_tile <= 32);
 
         std::vector<int> block_m_candidates;
-        if (runtime_align <= kMinBlockM)
+        const bool is_m_grouped = is_m_grouped_contiguous(desc.gemm_type)
+                                  or desc.gemm_type == GemmType::MGroupedMasked;
+        if (runtime_align == 32 and is_m_grouped) {
+            // Decode alignment: BLOCK_M=32 via the kNWarps=4 warp layout. No valid warp
+            // layout exists for BLOCK_M=32 with N <= 32 tiles — reject explicitly.
+            DG_HOST_ASSERT(not is_small_n and "m-grouped with N <= 32 requires alignment >= 64");
+            block_m_candidates.push_back(32);
+        }
+        else if (runtime_align <= kMinBlockM)
             block_m_candidates.push_back(64);
         else {
             if (is_small_n)
@@ -48,7 +56,8 @@ struct SM120ArchSpec {
         // but only beneficial for large M (>= 2048) and non-mixed dtypes.
         const bool is_mixed = (desc.a_dtype != desc.b_dtype);
         std::vector<int> block_k_candidates;
-        if (!is_mixed and expected_m >= 2048)
+        // BLOCK_M=32 decode keeps BK=128: enables the padding-skip cp.async A path (SW128)
+        if (!is_mixed and expected_m >= 2048 and not (is_m_grouped and runtime_align == 32))
             block_k_candidates.push_back(64 / elem_size);
         block_k_candidates.push_back(128 / elem_size);
 
@@ -75,10 +84,16 @@ struct SM120ArchSpec {
         for (int block_m : block_m_candidates) {
         // For grouped contiguous GEMM, BLOCK_M must divide runtime_align
         // to prevent tiles from straddling expert boundaries.
-        if (is_m_grouped_contiguous(desc.gemm_type) and block_m > runtime_align)
+        if (is_m_grouped_contiguous(desc.gemm_type)
+            and (block_m > runtime_align or runtime_align % block_m != 0))
             continue;
         for (int block_k : block_k_candidates) {
         for (int block_n : block_n_candidates) {
+            // BLOCK_M < 64 (kNWarps=4): BLOCK_N=64 only. Even kNTilesPerWarp (ldmatrix.x2)
+            // requires BLOCK_N % 64; measured vs BN=128: never worse, -1.6% on thin-N decode,
+            // and doubles SM fill at tiny M.
+            if (block_m < 64 and block_n != 64)
+                continue;
             if (!is_small_n and (block_n > 128 or block_n > mn_major_b_max_n))
                 continue;
 
@@ -145,7 +160,11 @@ struct SM120ArchSpec {
 
         const int cd_size = c10::elementSize(desc.cd_dtype);
         // cd_n_contiguous gates the TMA-store epilogue (off for AB-swap transposed output).
-        const auto swizzle_mode_cd = (desc.cd_n_contiguous and layout.block_n * cd_size >= 128) ? 128 : 0;
+        // M-grouped 1D1D uses the direct-store epilogue so the kernel can skip M-padding row
+        // stores (padding rows of D are unspecified); faster at decode AND prefill (fewer bytes).
+        const bool grouped_skip_padding_store = desc.kernel_type == KernelType::Kernel1D1D
+            and (desc.gemm_type == GemmType::MGroupedContiguous or desc.gemm_type == GemmType::MGroupedMasked);
+        const auto swizzle_mode_cd = (not grouped_skip_padding_store and desc.cd_n_contiguous and layout.block_n * cd_size >= 128) ? 128 : 0;
 
         // Sub-tile epilogue: reduce SMEM_D by storing smaller M sub-tiles.
         // Try store_block_m = 64 (sub-tile) and see if it gains pipeline stages.

csrc/jit_kernels/impls/sm120_bf16_gemm.hpp

@@ -34,10 +34,10 @@ class SM120BF16GemmRuntime final: public LaunchRuntime<SM120BF16GemmRuntime> {
     };
 
     static std::string generate_impl(const Args& args) {
-        // kNWarps: BM=64 needs kNWarps=2 (kMWarps=4); BM=128 uses default
-        // kNWarps=kNumMathWarps (all warps along M, original layout).
+        // kNWarps: BM=32 needs kNWarps=4 (kMWarps=2), BM=64 needs kNWarps=2 (kMWarps=4);
+        // BM=128 uses default kNWarps=kNumMathWarps (all warps along M, original layout).
         const uint32_t block_m = args.gemm_config.layout.block_m;
-        const uint32_t k_n_warps = (block_m % 64 == 0 and block_m < 128) ? 2u : 1u;
+        const uint32_t k_n_warps = (block_m == 32) ? 4u : ((block_m % 64 == 0 and block_m < 128) ? 2u : 1u);
         return fmt::format(R"(
 #include <deep_gemm/impls/sm120_bf16_gemm.cuh>
 

csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp

@@ -15,6 +15,14 @@
 
 namespace deep_gemm {
 
+// Skip D stores of M-padding rows in m-grouped layouts, leaving them untouched (their contents
+// are unspecified). Requires the direct-store epilogue (swizzle_cd == 0) and no split-K.
+static bool should_skip_padding_store(const GemmDesc& desc, const GemmConfig& config) {
+    return (desc.gemm_type == GemmType::MGroupedContiguous or desc.gemm_type == GemmType::MGroupedMasked)
+        and config.storage_config.swizzle_cd_mode == 0
+        and config.split_k_factor == 1;
+}
+
 class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM120FP8FP4Gemm1D1DRuntime> {
 public:
     struct Args {
@@ -31,6 +39,7 @@ class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM120FP8FP4Gemm1D1D
         int stride_cd_m;
         int stride_cd_n;
         int stride_cd_batch;
+        bool a_cpasync;  // A is contiguous and K % BLOCK_K == 0: padding-skip cp.async A path allowed
 
         void* gmem_d;
         void* gmem_c;
@@ -72,6 +81,8 @@ static void __instantiate_kernel() {{
         {},
         {},
         {},
+        {},
+        {},
         {}
     >);
 }};
@@ -96,7 +107,9 @@ static void __instantiate_kernel() {{
         (args.gemm_desc.major_b == cute::UMMA::Major::K) ? "true" : "false",
         args.k_grouped_constant_stride ? "true" : "false",
         args.gemm_config.storage_config.store_block_m,
-        args.gemm_config.split_k_factor);
+        args.gemm_config.split_k_factor,
+        should_skip_padding_store(args.gemm_desc, args.gemm_config) ? "true" : "false",
+        args.a_cpasync ? "true" : "false");
     }
 
     static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -481,9 +494,10 @@ static void sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a,
         .stride_cd_m = n,
         .stride_cd_n = 0,
         .stride_cd_batch = 0,
+        .a_cpasync = a.is_contiguous() and (k % config.layout.block_k == 0),
         .gmem_d = d.data_ptr(),
         .gmem_c = nullptr,
-        .gmem_a_ptr = nullptr,
+        .gmem_a_ptr = a.data_ptr(),
         .gmem_b_ptr = nullptr,
         .gmem_workspace = nullptr,
         .grouped_layout = grouped_layout.data_ptr(),
@@ -571,9 +585,10 @@ static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con
         .stride_cd_m = n,
         .stride_cd_n = 0,
         .stride_cd_batch = 0,
+        .a_cpasync = a.is_contiguous() and (k % config.layout.block_k == 0),
         .gmem_d = d.data_ptr(),
         .gmem_c = nullptr,
-        .gmem_a_ptr = nullptr,
+        .gmem_a_ptr = a.data_ptr(),
         .gmem_b_ptr = nullptr,
         .gmem_workspace = nullptr,
         .grouped_layout = masked_m.data_ptr(),

deep_gemm/include/deep_gemm/common/sm120_utils.cuh

@@ -252,6 +252,32 @@ __device__ __forceinline__ uint32_t load_sf(const char* smem_sf, int idx) {
     return *reinterpret_cast<const uint32_t*>(smem_sf + idx * sizeof(int32_t));
 }
 
+
+// Warp-cooperative cp.async.cg row loader with TMA-equivalent SW128 swizzle.
+// Loads rows [0, num_rows) of kRowBytes each; rows beyond keep stale SMEM.
+template <uint32_t kRowBytes, uint32_t kSwizzleMode>
+CUTLASS_DEVICE void cpasync_load_rows(char* smem_dst, const char* gmem_src,
+                                      const uint32_t& num_rows, const uint32_t& gmem_row_stride,
+                                      const uint32_t& lane_idx) {
+    static_assert(kSwizzleMode == 128 and kRowBytes == 128, "Only the SW128 layout is verified");
+    constexpr uint32_t kChunks = kRowBytes / 16;
+    const uint32_t total = num_rows * kChunks;
+    for (uint32_t item = lane_idx; item < total; item += 32) {
+        const uint32_t row = item / kChunks;
+        const uint32_t offset = CuTeSwizzle<kSwizzleMode>::apply(item * 16);
+        const auto dst = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst + offset));
+        const char* src = gmem_src + row * static_cast<uint64_t>(gmem_row_stride) + (item % kChunks) * 16;
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src) : "memory");
+    }
+}
+
+// The mbarrier receives one arrival per calling thread once all its prior cp.asyncs land.
+// .noinc: the arrival is part of the barrier's fixed expected count (init +1 per lane).
+CUTLASS_DEVICE void cpasync_mbarrier_arrive(void* mbar) {
+    const auto addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar));
+    asm volatile("cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n" :: "r"(addr) : "memory");
+}
+
 } // namespace deep_gemm::sm120
 
 #endif