From 81953c81d6083a4dd78dba66eba5fa3d57fe8a50 Mon Sep 17 00:00:00 2001
From: Kevin Abraham <kevin.abraham@streamhpc.com>
Date: Fri, 12 Sep 2025 15:35:24 +0000
Subject: [PATCH] add type-conversion to AccDataType and then to CDataType to
 exactly mimic GPU's behavior

---
 .../cpu/reference_batched_gemm.hpp                         | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp
index ca17255b25..1828387e3c 100644
--- a/library/include/ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp
+++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp
@@ -105,12 +105,15 @@ struct ReferenceBatchedGemm : public device::BaseOperator
                 {
                     for(int batchIdx = 0;batchIdx < arg.k_batch_;batchIdx++)
                     {
-                        v_c = ck::type_convert<CDataType>(ck::type_convert<AccDataType>(v_c) + ck::type_convert<AccDataType>(partialSums[batchIdx]));
+                        // mimic the way fp operations would be done on GPU for k-batching
+                        v_c = ck::type_convert<AccDataType>(ck::type_convert<CDataType>(
+                            ck::type_convert<AccDataType>(v_c) +
+                            ck::type_convert<AccDataType>(partialSums[batchIdx])));
                     }
                 }
                 else
                 {
-                    v_c = partialSums[0];
+                    v_c = ck::type_convert<AccDataType>(partialSums[0]);
                 }
 
                 arg.c_g_m_n_(g, m, n) = ck::type_convert<CDataType>(v_c);