[change] Reduce calculated bit width of AV

- This assumes that the attention is distributed among different values and not on one token
pulp-platform · Oct 30, 2024 · 93c3986 · 93c3986
1 parent 5942086
commit 93c3986
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/PyITA/ITA.py b/PyITA/ITA.py
@@ -259,7 +259,7 @@ def _initialize_quantization_parameters(self):
             elif i == 3:  # QK
                 max_bit_width = np.log2(self.requant_eps_mult[i, :].astype(np.uint32) * self.P * 2**8).astype(np.uint32)
             elif i == 4:  # AV
-                max_bit_width = np.log2(self.requant_eps_mult[i, :].astype(np.uint32) * self.S * 2**8).astype(np.uint32)
+                max_bit_width = np.log2(self.requant_eps_mult[i, :].astype(np.uint32) * self.S * 2**5).astype(np.uint32)
             elif i == 5:  # OW
                 max_bit_width = np.log2(self.requant_eps_mult[i, :].astype(np.uint32) * self.E * 2**9).astype(np.uint32)
             elif i == 6:  # Sum OW