Add MXFP8 E4M3 flash attention for CDNA4

alto/kernels/dispatch/attention.py

@@ -6,6 +6,7 @@
 from torchtitan.models.common.attention import (ScaledDotProductAttentionWrapper)
 
 from alto.kernels.fp4.mxfp4.triton_flash_attention_mxfp4 import triton_attention_mxfp4
+from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8
 from .config import TrainingOpConfig
 
 __all__ = ["LPScaledDotProductAttentionWrapper"]
@@ -20,6 +21,8 @@ def __init__(self, config: TrainingOpConfig):
 
         if isinstance(config, TrainingOpConfig) and config.precision == "mxfp4":
             self.attn_func = triton_attention_mxfp4
+        elif isinstance(config, TrainingOpConfig) and config.precision == "mxfp8_e4m3":
+            self.attn_func = triton_attention_mxfp8
         else:
             raise ValueError(f"Unsupported SDPA config: {config}")
 

alto/kernels/mxfp8/__init__.py

@@ -1,3 +1,7 @@
 # Copyright (c) 2026 Advanced Micro Devices, Inc.
 #
 # SPDX-License-Identifier: MIT
+
+from .triton_flash_attention_mxfp8 import triton_attention_mxfp8
+
+__all__ = ("triton_attention_mxfp8", )