[Common] MXFP8 dequantization CuTeDSL kernels

tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py

@@ -431,3 +431,43 @@ def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset
         swizzled,
         input_offset,
     )
+
+
+OUT_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
+DEQUANT_SHAPES = MATRIX_SIZES + [(100, 48)]
+
+
+# Test cases for dequantization (mirrors C++ test's OperatorTest_DequantizeMXFP8).
+# With both data copies present, dequantize reads the rowwise one.
+@pytest.mark.parametrize("shape", DEQUANT_SHAPES, ids=get_shape_id)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id)
+@pytest.mark.parametrize("out_dtype", OUT_DTYPES, ids=get_dtype_id)
+@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id)
+@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id)
+def test_dequantize(swizzled, fp8_dtype, out_dtype, block_size, shape):
+    M, N = shape
+    rowwise = block_size[1] != 1
+    columnwise = block_size[0] != 1
+    x, _ = generate_inputs(M, N, torch.bfloat16)
+    set_cutedsl_backend(False)
+    q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise)
+    q.optimize_for_gemm = swizzled
+    qx = q(x)
+    out_cuda = qx.dequantize(dtype=out_dtype)
+
+    set_cutedsl_backend(True)
+    try:
+        out_cutedsl = qx.dequantize(dtype=out_dtype)
+    finally:
+        set_cutedsl_backend(False)
+
+    major, minor = device_compute_capability()
+    key = (
+        f"cutedsl_mxfp8_dequant_sm{major * 10 + minor}_{FP8_TO_KEY[fp8_dtype]}_"
+        f"{DTYPE_TO_STR[out_dtype]}_{int(rowwise)}_{int(columnwise)}_{int(swizzled)}"
+    )
+    assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, (
+        f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back "
+        "to CUDA and this case compared CUDA against itself"
+    )
+    assert torch.equal(out_cutedsl, out_cuda)

transformer_engine/common/CuTeDSL/__init__.py

@@ -10,6 +10,9 @@
 
 import tvm_ffi
 
+from transformer_engine.common.CuTeDSL.cast.mxfp8.dequantize_mxfp8 import (
+    get_mxfp8_dequantization_function,
+)
 from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import (
     get_mxfp8_quantization_function,
 )
@@ -25,6 +28,9 @@ def register_cutedsl_backends():
     tvm_ffi.register_global_func(
         "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True
     )
+    tvm_ffi.register_global_func(
+        "get_mxfp8_dequantization_function", get_mxfp8_dequantization_function, override=True
+    )
     tvm_ffi.register_global_func(
         "get_mxfp8_requantization_function", get_mxfp8_requantization_function, override=True
     )

transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py

@@ -2,6 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""CuTeDSL MXFP8 quantization kernels."""
+"""CuTeDSL MXFP8 quantization and dequantization kernels."""
 
+from . import dequantize_mxfp8
 from . import quantize_mxfp8_common