From 8da33d6bd976a4d166bf9d48019246d9283b328e Mon Sep 17 00:00:00 2001
From: Chenggang Zhao <chenggangz@deepseek.com>
Date: Wed, 19 Nov 2025 11:00:55 +0800
Subject: [PATCH] Clean up

---
 csrc/apis/gemm.hpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp
index 076179c..bf1f758 100644
--- a/csrc/apis/gemm.hpp
+++ b/csrc/apis/gemm.hpp
@@ -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 = static_cast<uint64_t>(a.first.numel());
-    const auto& sum_nk = static_cast<uint64_t>(b.first.numel());
-    uint64_t sum_k = 0;
+    const auto& sum_mk = a.first.numel();
+    const auto& sum_nk = b.first.numel();
+    int sum_k = 0;
     for (const auto& k: ks)
-        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);
+        sum_k += k;
+    DG_HOST_ASSERT(sum_mk == static_cast<int64_t>(m) * sum_k);
+    DG_HOST_ASSERT(sum_nk == static_cast<int64_t>(n) * sum_k);
 
     // Contiguity checks
     DG_HOST_ASSERT(a.first.is_contiguous());