From 1f21a290f9713ab5e64b410db2430ff0bc42a6cc Mon Sep 17 00:00:00 2001
From: Krzysztof Rymski <krymski@google.com>
Date: Tue, 28 Apr 2026 07:04:57 -0700
Subject: [PATCH] Fix rounding in kv transcoding

PiperOrigin-RevId: 906957361
---
 gemma/kv_transcoding.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/gemma/kv_transcoding.cc b/gemma/kv_transcoding.cc
index b81ca817..ca06f0df 100644
--- a/gemma/kv_transcoding.cc
+++ b/gemma/kv_transcoding.cc
@@ -4,6 +4,7 @@
 #include <cstddef>
 #include <cstdlib>
 #include <optional>
+#include <cmath>
 
 #include "compression/types.h"
 #include "gemma/activations.h"
@@ -52,7 +53,7 @@ inline size_t VOffset(bool transposed, size_t qkv_dim, size_t dim,
 }
 
 int8_t Quantize(float v, float inv_scale) {
-  float scaled = v * inv_scale;
+  float scaled = std::nearbyint(v * inv_scale);
   if (scaled > 127.0f) return 127;
   if (scaled < -127.0f) return -127;
   return hwy::ConvertScalarTo<int8_t>(scaled);