diff --git a/fms_mo/aiu_addons/__init__.py b/fms_mo/aiu_addons/__init__.py
index 1d9c5cf..3ab3f3b 100644
--- a/fms_mo/aiu_addons/__init__.py
+++ b/fms_mo/aiu_addons/__init__.py
@@ -30,6 +30,7 @@ def _infer_quantization_config(quant_config: dict) -> dict | None:
                 )
             # First, import required FP8 linear classes from fms-mo
             # Local
+            import fms_mo.aiu_addons.fp8.fp8_attn  # pylint: disable=unused-import
             import fms_mo.aiu_addons.fp8.fp8_adapter  # pylint: disable=unused-import
             import fms_mo.aiu_addons.fp8.fp8_linear  # pylint: disable=unused-import