make mxfp8 dim1 cast kernel configurable

stack-info: PR: #1427, branch: danielvegamyhre/stack/1

Author: danielvegamyhre

torchtitan/components/quantization/mx.py

@@ -56,12 +56,17 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
         self.filter_fqns = mx_job_config.filter_fqns
 
         # Configure MXFP8
-        from torchao.prototype.mx_formats.config import MXLinearConfig
+        from torchao.prototype.mx_formats.config import (
+            MXFP8Dim1CastKernelChoice,
+            MXLinearConfig,
+        )
 
         config = MXLinearConfig.from_recipe_name(NAME_MAP[mx_job_config.recipe_name])
-        config.use_fp8_dim1_cast_triton_kernel = (
-            mx_job_config.use_fp8_dim1_cast_triton_kernel
-        )
+
+        # String to enum
+        config.mxfp8_dim1_cast_kernel_choice = MXFP8Dim1CastKernelChoice[
+            mx_job_config.mxfp8_dim1_cast_kernel_choice.upper()
+        ]
         self.config = config
 
         logger.info(f"Float8 training active with recipe {mx_job_config.recipe_name}")

torchtitan/config_manager.py

@@ -556,7 +556,7 @@ class Float8:
 
 @dataclass
 class MX:
-    use_fp8_dim1_cast_triton_kernel: bool = True
+    mxfp8_dim1_cast_kernel_choice: Literal["triton", "cuda", "torch"] = "triton"
     """Temp work around for inductor performance gap"""
 
     recipe_name: Literal["mxfp8"] = "mxfp8"