update api name

stack-info: PR: #1428, branch: danielvegamyhre/stack/2

Author: danielvegamyhre

torchtitan/components/quantization/mx.py

@@ -63,14 +63,10 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
 
         config = MXLinearConfig.from_recipe_name(NAME_MAP[mx_job_config.recipe_name])
 
-        dim1_cast_kernel_choice_str = (
+        # String to enum
+        config.mxfp8_dim1_cast_kernel_choice = MXFP8Dim1CastKernelChoice[
             mx_job_config.mxfp8_dim1_cast_kernel_choice.upper()
-        )
-        config.mxfp8_dim1_cast_kernel_choice = (
-            MXFP8Dim1CastKernelChoice[dim1_cast_kernel_choice_str]
-            if mx_job_config.mxfp8_dim1_cast_kernel_choice != "NONE"
-            else None
-        )
+        ]
         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:
-    mxfp8_dim1_cast_kernel_choice: Literal["triton", "cuda", "none"] = "triton"
+    mxfp8_dim1_cast_kernel_choice: Literal["triton", "cuda", "torch"] = "triton"
     """Temp work around for inductor performance gap"""
 
     recipe_name: Literal["mxfp8"] = "mxfp8"