Performance: Larger BlockTile optimizations enable 1470+ TFLOPS FP8 performance on the H800-SXM platform

2025-03-25 10:44:57 +08:00
parent 3b3783d06c
commit 46eb0d08fb
5 changed files with 1008 additions and 184 deletions
@@ -98,6 +98,98 @@ struct Scheduler {
        return true;
    }
 };
+template <GemmType kGemmType,
+          uint32_t SHAPE_M, uint32_t BLOCK_M, uint32_t BLOCK_N,
+          uint32_t kNumGroups, uint32_t kNumTMAMulticast,
+          uint32_t kNumMBlocks = ceil_div(SHAPE_M, BLOCK_M),
+          uint32_t kNumMBlocksPerGroup = 16>
+struct SchedulerLargeBlockTile {
+    int current_iter = -1;
+    uint32_t num_aligned_n_blocks;
+
+    // For normal GEMM
+    // Maybe not used in the masked grouped GEMM
+    uint32_t num_blocks;
+
+    // For grouped GEMM
+    int* grouped_layout;
+    // Only used for masked layout
+    uint32_t curr_group_idx, curr_cumsum;
+
+    __device__ __forceinline__ explicit SchedulerLargeBlockTile(const uint32_t shape_n,
+                                                  int* grouped_layout = nullptr) {
+        num_aligned_n_blocks = ceil_div(shape_n, BLOCK_N);
+        if constexpr (kGemmType == GemmType::Normal) {
+            num_blocks = num_aligned_n_blocks * kNumMBlocks;
+        } else if (kGemmType == GemmType::GroupedContiguous) {
+            num_blocks = num_aligned_n_blocks * kNumMBlocks;
+            this->grouped_layout = grouped_layout;
+        } else if (kGemmType == GemmType::GroupedMasked) {
+            curr_group_idx = curr_cumsum = 0;
+            this->grouped_layout = grouped_layout;
+        }
+    }
+
+    __device__ __forceinline__ void get_swizzled_block_idx(const uint32_t num_n_blocks, int block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
+        DG_STATIC_ASSERT(kNumMBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
+
+        // Swizzle for better L2 usages
+        auto num_blocks_per_group = num_n_blocks * kNumMBlocksPerGroup;
+        auto group_idx = block_idx / num_blocks_per_group;
+        auto first_m_block_idx = group_idx * kNumMBlocksPerGroup;
+        auto num_m_blocks_in_group = min(kNumMBlocksPerGroup, kNumMBlocks - first_m_block_idx);
+        // auto num_m_blocks_in_group = kNumMBlocksPerGroup;
+        auto in_group_idx = block_idx % num_blocks_per_group;
+        n_block_idx = in_group_idx / num_m_blocks_in_group;
+        m_block_idx = first_m_block_idx + in_group_idx % num_m_blocks_in_group;
+        // if (threadIdx.x == 256) {
+        //     printf("blockIdx.x: %d group_idx: %d in_group_idx: %d m_block_idx: %d n_block_idx: %d\n", blockIdx.x, group_idx, in_group_idx, m_block_idx, n_block_idx);
+        // }
+    }
+
+    template <bool kIgnoreGroupedForGroupedContiguous=true>
+    __device__ __forceinline__ uint32_t get_global_idx(const uint32_t shape_dim, const uint32_t block_size,
+                                                       const uint32_t& block_idx, const uint32_t& m_block_idx=0) {
+        if constexpr (kGemmType == GemmType::Normal) {
+            return block_idx * block_size;
+        } else if (kGemmType == GemmType::GroupedContiguous) {
+            auto offset = kIgnoreGroupedForGroupedContiguous ? 0 : __ldg(grouped_layout + m_block_idx * BLOCK_M);
+            return offset * shape_dim + block_idx * block_size;
+        } else if (kGemmType == GemmType::GroupedMasked) {
+            return curr_group_idx * shape_dim + block_idx * block_size;
+        }
+    }
+
+    __device__ __forceinline__ bool get_next_block(uint32_t& m_block_idx, uint32_t& n_block_idx) {
+        const auto next_block_idx = (++ current_iter) * gridDim.x + blockIdx.x;
+
+        if constexpr (kGemmType == GemmType::GroupedMasked) {
+            uint32_t num_m_blocks;
+            while (true) {
+                // End of the task
+                if (curr_group_idx == kNumGroups)
+                    return false;
+
+                // Within current group
+                num_m_blocks = ceil_div(static_cast<uint32_t>(__ldg(grouped_layout + curr_group_idx)), BLOCK_M);
+                auto current_m_block_cumsum = curr_cumsum + num_m_blocks;
+                if (next_block_idx < current_m_block_cumsum * kNumMBlocks)
+                    break;
+
+                // Move to check the next group
+                curr_group_idx ++, curr_cumsum = current_m_block_cumsum;
+            }
+
+            get_swizzled_block_idx(num_m_blocks, next_block_idx - curr_cumsum * kNumMBlocks, m_block_idx, n_block_idx);
+        } else {
+            if (next_block_idx >= num_blocks)
+                return false;
+
+            get_swizzled_block_idx(num_aligned_n_blocks, next_block_idx, m_block_idx, n_block_idx);
+        }
+        return true;
+    }
+};
 #pragma clang diagnostic pop

 } // namespace deep_gemm