Update on "Refactor use_triton_kernel to use nvfp4_quantize_kernel_choice"

Summary:
This is to prefer the addition of flashinfer quantize kernel path in next PR

Test Plan:
python test/prototype/mx_formats/test_inference_workflow.py

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

Author: jerryzh168

test/prototype/mx_formats/test_inference_workflow.py

@@ -188,7 +188,7 @@ def test_inference_workflow_nvfp4(
         and quantize_to_nvfp4_kernel_choice == QuantizeToNVFP4KernelChoice.MSLK
     ):
         pytest.skip("unsupported configuration")
-    if use_triton_kernel and not use_dynamic_per_tensor_scale:
+    if quantize_to_nvfp4_kernel_choice == QuantizeToNVFP4KernelChoice.MSLK and not use_dynamic_per_tensor_scale:
         pytest.skip("unsupported configuration")
 
     if use_inference_mode and (

torchao/prototype/mx_formats/__init__.py

@@ -1,8 +1,6 @@
 from torchao.prototype.mx_formats.config import (
-    ScaleCalculationMode,
-    MXLinearConfig,
-    MXLinearRecipeName,
     QuantizeToNVFP4KernelChoice,
+    ScaleCalculationMode,
 )
 
 # Note: Prototype and subject to change