Update the accumulation dtype if FP8 precision is used for AQT.

PiperOrigin-RevId: 686533904
google · Nov 5, 2024 · 05778b8 · 05778b8
1 parent 6134c4b
commit 05778b8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/aqt/jax/v2/aqt_dot_general.py b/aqt/jax/v2/aqt_dot_general.py
@@ -117,6 +117,11 @@ def dot_general_raw_make(
       and 2 <= rhs_bits <= 8
   ):
     dg_accumulator_dtype = jnp.int32
+  elif (
+      lhs_bits in fp8_numerics.fp8_map.keys()
+      or rhs_bits in fp8_numerics.fp8_map.keys()
+  ):
+    dg_accumulator_dtype = jnp.float32
   else:
     dg_accumulator_dtype = None