leavelet/deepgemm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add max_block_n.

Browse Source
This commit is contained in:
Eric Wong
2025-11-13 22:15:56 +08:00
parent 5f8a71a864
commit 3a297646cd
3 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/jit_kernels/heuristics/common.hpp
View File

@@ -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<int>(c10::elementSize(ab_dtype));

2
csrc/jit_kernels/heuristics/sm100.hpp
View File

@@ -12,7 +12,7 @@ namespace deep_gemm {
struct SM100ArchSpec {
static constexpr int smem_capacity = 232448;
static std::vector<int> get_block_n_candidates(const at::ScalarType& cd_dtype) {
static std::vector<int> 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<int> candidates = {16};

4
csrc/jit_kernels/heuristics/sm90.hpp
View File

@@ -11,11 +11,11 @@ namespace deep_gemm {
struct SM90ArchSpec {
static constexpr int smem_capacity = 232448;
static std::vector<int> get_block_n_candidates(const at::ScalarType& cd_dtype) {
static std::vector<int> 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<int> 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;
}