From 03e3acb24c724d23139dbce3e8c20a443acd8565 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 14 Nov 2026 14:38:26 -0800 Subject: [PATCH] pybindings: pass float8 tensors through DLPack The pybindings' DLPack tables had no float8 types. A float8 value handed to a method through DLPack was refused. A float8 weight uploaded with from_host(..., dtype="float8_e4m3fn") failed with "ExecuTorch result dtype cannot be represented by DLPack". A raw DLPack capsule, or a wrapper with only __dlpack__, of a CPU torch float8 tensor failed with "Unsupported DLPack dtype". A float8 output could not be read through DLPack either, so torch.from_dlpack() refused a float8 ExecuTorchResult. DLPack 1.1 added type codes for the float8 formats. The pybindings now map Float8_e4m3fn, Float8_e4m3fnuz, Float8_e5m2 and Float8_e5m2fnuz to and from those codes, with 8 bits and 1 lane, the same codes PyTorch uses. Because the codes are from DLPack 1.1, the capsules the pybindings export now say version 1.1, and the pybindings ask producers for up to version 1.1. Both sides still accept any 1.x version. The README now says float8 results keep their dtype through DLPack, as bfloat16 results do, and how to pass a float8 input. Tests: - test_float8_dlpack_input runs float8_e4m3fn and float8_e5m2 tensors through an identity method on the CPU, given as a wrapper with only __dlpack__, a legacy capsule and a versioned capsule. It also checks that the fnuz dtypes are read as themselves from a versioned capsule: the method refuses them with their own dtype names, since exir cannot export them. - test_float8_result_round_trips_through_dlpack gets ExecuTorchResult outputs on the CPU (EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch), reads them with torch.from_dlpack, and passes them back into the method. - test_float8_host_bits_run_through_a_method uploads bits of all four float8 dtypes to CUDA with from_host and checks the dtype torch reads. It runs float8_e4m3fn and float8_e5m2 through a method, since exir cannot export the fnuz dtypes. - In test_pybindings_no_aten.py, the importer test now expects a request for up to version 1.1, and the DLPack export test checks that an exported capsule says version 1.1. --- extension/pybindings/README.md | 9 +- extension/pybindings/pybindings.cpp | 21 +++- extension/pybindings/pybindings_dlpack.h | 9 ++ .../pybindings/pybindings_executorch_result.h | 10 +- extension/pybindings/test/test_pybindings.py | 107 ++++++++++++++++++ .../test/test_pybindings_cuda_dlpack.py | 33 ++++++ .../test/test_pybindings_no_aten.py | 16 ++- 7 files changed, 195 insertions(+), 10 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 7e1ba466e88..71fc589f0f7 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -206,10 +206,11 @@ Other dtypes NumPy lacks, such as float8, are shown the same way, as integers of the same width. NumPy converts that view as integers, so `np.asarray(out, dtype=np.float32)` gives the bit patterns as numbers, not the bfloat16 values. DLPack consumers, such as `torch.from_dlpack(out)`, receive the -original bfloat16 dtype. A bfloat16 result passed back in is read through DLPack -too, so it keeps its dtype, shape, and dim order. A bfloat16 input does not -accept 16-bit integers, so that integer data is never read as bfloat16 by -mistake; pass the result itself or raw bytes. +original bfloat16 or float8 dtype. A bfloat16 or float8 result passed back in is +read through DLPack too, so it keeps its dtype, shape, and dim order. A bfloat16 +input does not accept 16-bit integers, so that integer data is never read as +bfloat16 by mistake; pass the result itself or raw bytes. A float8 input does not +accept 8-bit integers either; pass the result itself or a float8 tensor. ```python import cupy as cp diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 071c8c10657..4af3d9ef921 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1709,8 +1709,8 @@ class BorrowedInput final { }; #ifndef USE_ATEN_LIB -// Asks for a DLPack 1.0 capsule, which can mark memory read-only, and falls -// back to a legacy capsule for producers that predate it. +// Asks for a versioned DLPack capsule, which can mark memory read-only, and +// falls back to a legacy capsule for producers that predate versioned capsules. py::object dlpack_capsule( const py::object& producer, const py::object& stream) { @@ -1720,7 +1720,8 @@ py::object dlpack_capsule( arguments["stream"] = stream; } if (versioned) { - arguments["max_version"] = py::make_tuple(1, 0); + arguments["max_version"] = + py::make_tuple(dlpack::kVersion.major, dlpack::kVersion.minor); } return producer.attr("__dlpack__")(**arguments); }; @@ -1990,6 +1991,20 @@ class DLPackTensorView final { if (dtype.code == dlpack::DataTypeCode::BFloat && dtype.bits == 16) { return ScalarType::BFloat16; } + if (dtype.code == dlpack::DataTypeCode::Float8_e4m3fn && dtype.bits == 8) { + return ScalarType::Float8_e4m3fn; + } + if (dtype.code == dlpack::DataTypeCode::Float8_e4m3fnuz && + dtype.bits == 8) { + return ScalarType::Float8_e4m3fnuz; + } + if (dtype.code == dlpack::DataTypeCode::Float8_e5m2 && dtype.bits == 8) { + return ScalarType::Float8_e5m2; + } + if (dtype.code == dlpack::DataTypeCode::Float8_e5m2fnuz && + dtype.bits == 8) { + return ScalarType::Float8_e5m2fnuz; + } if (dtype.code == dlpack::DataTypeCode::Complex && dtype.bits == 32) { return ScalarType::ComplexHalf; } diff --git a/extension/pybindings/pybindings_dlpack.h b/extension/pybindings/pybindings_dlpack.h index 3e73fe78812..638745358a6 100644 --- a/extension/pybindings/pybindings_dlpack.h +++ b/extension/pybindings/pybindings_dlpack.h @@ -31,6 +31,11 @@ enum DataTypeCode : uint8_t { Complex = 5, BFloat = 4, Bool = 6, + // Since DLPack 1.1. + Float8_e4m3fn = 10, + Float8_e4m3fnuz = 11, + Float8_e5m2 = 12, + Float8_e5m2fnuz = 13, }; struct Device { @@ -65,6 +70,10 @@ struct Version { uint32_t minor; }; +// The version these bindings export and ask producers for: 1.1 added the +// float8 codes. +constexpr Version kVersion{1, 1}; + // DLPack 1.0 capsules carry this struct, named "dltensor_versioned". struct ManagedTensorVersioned { Version version; diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 6a3db374d15..6c5ab0253a8 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -549,7 +549,7 @@ class PyExecuTorchResult final 0}; managed.manager_ctx = this; managed.deleter = release_dlpack_holder; - versioned.version = {1, 0}; + versioned.version = dlpack::kVersion; versioned.manager_ctx = this; versioned.deleter = release_versioned_dlpack_holder; versioned.dl_tensor = managed.dl_tensor; @@ -654,6 +654,14 @@ class PyExecuTorchResult final return {dlpack::DataTypeCode::UInt, 32, 1}; case ScalarType::UInt64: return {dlpack::DataTypeCode::UInt, 64, 1}; + case ScalarType::Float8_e4m3fn: + return {dlpack::DataTypeCode::Float8_e4m3fn, 8, 1}; + case ScalarType::Float8_e4m3fnuz: + return {dlpack::DataTypeCode::Float8_e4m3fnuz, 8, 1}; + case ScalarType::Float8_e5m2: + return {dlpack::DataTypeCode::Float8_e5m2, 8, 1}; + case ScalarType::Float8_e5m2fnuz: + return {dlpack::DataTypeCode::Float8_e5m2fnuz, 8, 1}; default: throw py::value_error( "ExecuTorch result dtype cannot be represented by DLPack"); diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 2a63c4edf53..fab69ed7379 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -2291,6 +2291,113 @@ def forward(self, x): torch.equal(output.view(torch.float16), value.view(torch.float16)) ) + def test_float8_dlpack_input(self): + if self.kernel_mode != "portable": + self.skipTest("only portable bindings read DLPack inputs themselves") + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + for dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype) + program = to_edge(export(Identity(), (value,))).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + ) + method = self.load_prog_fn(program.buffer).load_method("forward") + # The wrapper takes no max_version, so it gets a legacy capsule; raw + # capsules also cover the versioned kind. + for kind, given in ( + ("wrapper", DLPackOnly(value)), + ("legacy capsule", value.__dlpack__()), + ("versioned capsule", value.__dlpack__(max_version=(1, 1))), + ): + with self.subTest(dtype=dtype, input=kind): + output = torch.from_dlpack(method(given)[0]) + self.assertEqual(output.dtype, dtype) + self.assertTrue( + torch.equal(output.view(torch.uint8), value.view(torch.uint8)) + ) + if dtype == torch.float8_e4m3fn: + # exir cannot export fnuz dtypes, so their codes are checked by + # the dtype a float8_e4m3fn method reports when it refuses them, + # read from a versioned capsule. + for fnuz in (torch.float8_e4m3fnuz, torch.float8_e5m2fnuz): + name = str(fnuz).removeprefix("torch.").capitalize() + with self.subTest(dtype=fnuz): + with self.assertRaisesRegex( + ValueError, + f"has dtype {name} but the method expects Float8_e4m3fn", + ): + method(value.to(fnuz).__dlpack__(max_version=(1, 1))) + + def test_float8_result_round_trips_through_dlpack(self): + if self.kernel_mode != "portable": + self.skipTest("only portable bindings return ExecuTorchResult outputs") + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + script = """ +import importlib +import sys +import torch + +runtime = importlib.import_module(sys.argv[1]) +native = sys.modules[runtime.ExecuTorchResult.__module__] +assert native._tensor_output == "executorch" +for path, name in zip(sys.argv[2::2], sys.argv[3::2]): + dtype = getattr(torch, name) + value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype) + with open(path, "rb") as source: + data = source.read() + method = runtime._load_program_from_buffer(data).load_method("forward") + output = method(value)[0] + assert isinstance(output, runtime.ExecuTorchResult), type(output) + tensor = torch.from_dlpack(output) + assert tensor.dtype == dtype, (name, tensor.dtype) + assert torch.equal(tensor.view(torch.uint8), value.view(torch.uint8)), name + roundtrip = torch.from_dlpack(method(output)[0]) + assert roundtrip.dtype == dtype, (name, roundtrip.dtype) + assert torch.equal(roundtrip.view(torch.uint8), value.view(torch.uint8)), name +print("float8 results round trip through DLPack") +""" + arguments = [self.runtime.__name__] + with tempfile.TemporaryDirectory() as directory: + for dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + name = str(dtype).removeprefix("torch.") + value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype) + program = to_edge(export(Identity(), (value,))).to_executorch(config) + path = os.path.join(directory, f"{name}.pte") + with open(path, "wb") as source: + source.write(program.buffer) + arguments += [path, name] + environment = os.environ.copy() + environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = "executorch" + completed = subprocess.run( + [sys.executable, "-c", script, *arguments], + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + self.assertEqual( + completed.returncode, + 0, + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}", + ) + self.assertIn("float8 results round trip through DLPack", completed.stdout) + def test_unsupported_buffer_format_is_rejected(self): exported_program, _ = create_program(ModuleAddSingleInput()) executorch_module = self.load_fn(exported_program.buffer) diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index bea6d7ccf84..ef9b2f13c3c 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -1188,6 +1188,39 @@ def test_bfloat16_host_result_moves_to_the_device(self): with self.assertRaisesRegex(ValueError, "already on a device"): self.bindings.ExecuTorchResult.from_host(on_device, self.device.index) + def test_float8_host_bits_run_through_a_method(self): + # Python buffers have no float8 format, so the values come as bytes. exir + # cannot export fnuz dtypes, so only the others run through a method. + for dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, + torch.float8_e5m2fnuz, + ): + name = str(dtype).removeprefix("torch.") + with self.subTest(dtype=name): + value = torch.tensor([[1.0, -2.0, 0.5], [3.0, 0.0, -0.25]]).to(dtype) + host = memoryview( + bytes(value.view(torch.uint8).flatten().tolist()) + ).cast("B", value.shape) + uploaded = self.bindings.ExecuTorchResult.from_host( + host, self.device.index, dtype=name + ) + self.assertEqual(uploaded.dtype, name) + expected = value.to(self.device).view(torch.uint8) + as_torch = torch.from_dlpack(uploaded) + self.assertEqual(as_torch.dtype, dtype) + torch.testing.assert_close(as_torch.view(torch.uint8), expected) + if name.endswith("fnuz"): + continue + data = self._device_identity(self.device.index, value) + method = self.bindings._load_program_from_buffer(data).load_method( + "forward" + ) + output = self._tensor(method(uploaded)[0]) + self.assertEqual(output.dtype, dtype) + torch.testing.assert_close(output.view(torch.uint8), expected) + def test_bfloat16_host_bits_move_to_the_device(self): # Host data without torch has no bfloat16 buffer format, so it comes as # 16-bit unsigned integers and from_host is told the dtype. diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index a8ee126b4ff..4c51ba41950 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import ctypes import gc import importlib import importlib.util @@ -340,6 +341,16 @@ def test_dlpack_export_keywords(self) -> None: self.assertIn( '"dltensor_versioned"', repr(result.__dlpack__(max_version=(1, 0))) ) + # Versioned capsules say 1.1, the DLPack version that added the float8 + # codes they can carry. + get_pointer = ctypes.PYFUNCTYPE( + ctypes.c_void_p, ctypes.py_object, ctypes.c_char_p + )(("PyCapsule_GetPointer", ctypes.pythonapi)) + capsule = result.__dlpack__(max_version=(1, 1)) + version = (ctypes.c_uint32 * 2).from_address( + get_pointer(capsule, b"dltensor_versioned") + ) + self.assertEqual(tuple(version), (1, 1)) self.assertIn('"dltensor"', repr(result.__dlpack__())) with self.assertRaises(BufferError): result.__dlpack__(dl_device=(2, 0), copy=False) @@ -367,7 +378,8 @@ def test_raw_dlpack_capsule_inputs(self) -> None: output = executor(capsules)[0] np.testing.assert_array_equal(np.from_dlpack(output), [3.0]) - # A producer that knows DLPack 1.0 is asked for a versioned capsule. + # A producer that takes max_version is asked for a versioned capsule, up + # to DLPack 1.1. class VersionedProducer: def __init__(self, array: np.ndarray) -> None: self.array = array @@ -384,7 +396,7 @@ def __dlpack__(self, stream=None, max_version=None): np.testing.assert_array_equal( np.from_dlpack(method((first, producer))[0]), [3.0] ) - self.assertEqual(producer.requests, [(1, 0)]) + self.assertEqual(producer.requests, [(1, 1)]) # A capsule gives its tensor away once. used = first.__dlpack__() method((used, second))