Cleanup

Author: bkryu

flashinfer/cute_dsl/__init__.py

@@ -36,13 +36,6 @@
         AddRMSNormFP4QuantKernel,
     )
 
-    # MXFP8 Quantization - re-exported from quantization module for backward compatibility
-    from ..quantization.mxfp8_quantize_cute_dsl import (
-        mxfp8_quantize_cute_dsl,
-        MXFP8QuantizeLinearKernel,
-        MXFP8QuantizeSwizzledKernel,
-    )
-
 __all__ = [
     # Utils (always available)
     "is_cute_dsl_available",
@@ -63,8 +56,4 @@
         # Add + RMSNorm + FP4 Quantization
         "add_rmsnorm_fp4quant",
         "AddRMSNormFP4QuantKernel",
-        # MXFP8 Quantization
-        "mxfp8_quantize_cute_dsl",
-        "MXFP8QuantizeLinearKernel",
-        "MXFP8QuantizeSwizzledKernel",
     ]

flashinfer/quantization/__init__.py

@@ -36,16 +36,12 @@
 )
 
 # CuTe-DSL kernels (conditionally exported)
+# Note: is_cute_dsl_available is used internally but not re-exported;
+# users should import from flashinfer.cute_dsl
 from ..cute_dsl import is_cute_dsl_available
 
 if is_cute_dsl_available():
-    from .mxfp8_quantize_cute_dsl import (
-        mxfp8_quantize_cute_dsl,
-        MXFP8QuantizeLinearKernel,
-        MXFP8QuantizeSwizzledKernel,
-        _get_compiled_kernel_linear,
-        _get_compiled_kernel_swizzled,
-    )
+    from .mxfp8_quantize_cute_dsl import mxfp8_quantize_cute_dsl
 
 __all__ = [
     # Packbits
@@ -69,15 +65,9 @@
     "shuffle_matrix_sf_a",
     "scaled_fp4_grouped_quantize",
     "get_fp4_quantization_module",
-    # CuTe-DSL availability check
-    "is_cute_dsl_available",
 ]
 
 if is_cute_dsl_available():
     __all__ += [
         "mxfp8_quantize_cute_dsl",
-        "MXFP8QuantizeLinearKernel",
-        "MXFP8QuantizeSwizzledKernel",
-        "_get_compiled_kernel_linear",
-        "_get_compiled_kernel_swizzled",
     ]