From 3a297646cd5ad5996291a194f816b5947e7cef5f Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Thu, 13 Nov 2025 22:15:56 +0800 Subject: [PATCH] add max_block_n. --- csrc/jit_kernels/heuristics/common.hpp | 2 +- csrc/jit_kernels/heuristics/sm100.hpp | 2 +- csrc/jit_kernels/heuristics/sm90.hpp | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/csrc/jit_kernels/heuristics/common.hpp b/csrc/jit_kernels/heuristics/common.hpp index 969056e..5427d13 100644 --- a/csrc/jit_kernels/heuristics/common.hpp +++ b/csrc/jit_kernels/heuristics/common.hpp @@ -165,7 +165,7 @@ static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& k block_ms = std::vector{get_mk_alignment_for_contiguous_layout()}; if (gemm_type == GemmType::MGroupedMasked) // Exclude 256 for performance block_ms = std::vector{64, 128}; - const auto block_ns = ArchSpec::get_block_n_candidates(cd_dtype); + const auto block_ns = ArchSpec::get_block_n_candidates(cd_dtype, max_block_n); // K block size is selected in a fixed manner const auto& block_k = 128 / static_cast(c10::elementSize(ab_dtype)); diff --git a/csrc/jit_kernels/heuristics/sm100.hpp b/csrc/jit_kernels/heuristics/sm100.hpp index e62a13c..d0d1698 100644 --- a/csrc/jit_kernels/heuristics/sm100.hpp +++ b/csrc/jit_kernels/heuristics/sm100.hpp @@ -12,7 +12,7 @@ namespace deep_gemm { struct SM100ArchSpec { static constexpr int smem_capacity = 232448; - static std::vector get_block_n_candidates(const at::ScalarType& cd_dtype) { + static std::vector get_block_n_candidates(const at::ScalarType& cd_dtype, const int& max_block_n) { // 16 is for better SM usage // Stride 32 is due to low-performance swizzle-16/32B std::vector candidates = {16}; diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index 133e2da..d411206 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -11,11 +11,11 @@ namespace deep_gemm { struct SM90ArchSpec { static constexpr int smem_capacity = 232448; - static std::vector get_block_n_candidates(const at::ScalarType& cd_dtype) { + static std::vector get_block_n_candidates(const at::ScalarType& cd_dtype, const int& max_block_n) { // Avoid bank conflicts for FP32 output const auto& start = cd_dtype == torch::kFloat ? 8 : 16; std::vector candidates; - for (int i = start; i <= 256; i += 16) + for (int i = start; i <= max_block_n; i += 16) candidates.push_back(i); return candidates; }