biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: prevent int32 overflow in k-grouped GEMM size calculations (#226)

Browse Source
This commit is contained in:
Guoteng
2025-11-19 10:52:08 +08:00
committed by GitHub
parent ec5e9ed0b8
commit f63d7f24d6
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
csrc/apis/gemm.hpp
View File

@@ -280,13 +280,13 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
// Shape checks
const auto& [num_groups, m, n] = get_shape<3>(d);
const auto& sum_mk = a.first.numel();
const auto& sum_nk = b.first.numel();
int sum_k = 0;
const auto& sum_mk = static_cast<uint64_t>(a.first.numel());
const auto& sum_nk = static_cast<uint64_t>(b.first.numel());
uint64_t sum_k = 0;
for (const auto& k: ks)
sum_k += k;
DG_HOST_ASSERT(sum_mk == m * sum_k);
DG_HOST_ASSERT(sum_nk == n * sum_k);
sum_k += static_cast<uint64_t>(k);
DG_HOST_ASSERT(sum_mk == static_cast<uint64_t>(m) * sum_k);
DG_HOST_ASSERT(sum_nk == static_cast<uint64_t>(n) * sum_k);
// Contiguity checks
DG_HOST_ASSERT(a.first.is_contiguous());