DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-06-26 23:15:49 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor the loop if/else check

Browse Source 
Keeps DivisibleK and NotDivisibleK logic explicit
Ensures clear separation between the bulk iterations and the last special case
Easier to modify later if conditions change
This commit is contained in:
A-transformer
2025-02-27 22:23:53 +04:00
committed by GitHub
parent a2e0d68eed
commit b0b9e03345
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
deep_gemm/include/deep_gemm/fp8_gemm.cuh
View File

@@ -129,14 +129,17 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
struct DivisibleK {}; struct DivisibleK {};
struct NotDivisibleK {}; struct NotDivisibleK {};
auto launch_k_iterations = [](const auto& func) { auto launch_k_iterations = [](const auto& func) {
if constexpr (SHAPE_K % kFullKOfAllStages == 0) { constexpr bool is_divisible = (SHAPE_K % kFullKOfAllStages == 0);
for (int k_iter = 0; k_iter < kNumIterations; ++ k_iter) constexpr int last_iter = kNumIterations - 1;
func(k_iter, DivisibleK{}); static_assert(kNumIterations > 0, "At least one iteration required");
} else {
for (int k_iter = 0; k_iter < kNumIterations - 1; ++ k_iter) // Handle all but the last iteration with DivisibleK
func(k_iter, DivisibleK{}); for (int k_iter = 0; k_iter < last_iter; ++k_iter) {
func(kNumIterations - 1, NotDivisibleK{}); func(k_iter, DivisibleK{});
} }
// Last iteration: DivisibleK if SHAPE_K is divisible, NotDivisibleK otherwise
func(last_iter, is_divisible ? DivisibleK{} : NotDivisibleK{});
}; };
// Register reconfigurations // Register reconfigurations