From ea507739f50089b44b2125da0b21b41df46abe06 Mon Sep 17 00:00:00 2001 From: Andrew Fish Date: Wed, 30 Sep 2026 11:48:32 -0700 Subject: [PATCH 1/2] [Fix][Relax] Dequantize fp8 as a float, not an integer dequantize_compute chose its intermediate dtype with matches_code(FLOAT, BFLOAT). Those are DLPack codes 2 and 4, while every narrow float carries its own code instead -- Float8E3M4 at 7 through Float4E2M1FN at 17. An fp8 input therefore took the int32 branch and was truncated to an integer before the zero point was subtracted and the scale applied, so a stored -0.9375 became 0. This is reachable through the public API: relax.quantize already accepts float8_e4m3fn and float8_e5m2 as output dtypes. It stays hidden wherever the stored value is already integral, which at E4M3's spacing covers most of the top of the range. Ask the dtype whether it is a float instead of enumerating codes, via a PrimType.is_float that forwards to the one tvm_ffi already provides. The float code list keeps a single home, so a future narrow float needs no edit here. quantize is unaffected -- it decides the same question with "int" in out_dtype, which excludes the fp8 dtypes correctly. --- python/tvm/relax/transform/legalize_ops/qdq.py | 7 +------ tirx/python/tvm_tirx/ir/type.py | 9 +++++++++ 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/python/tvm/relax/transform/legalize_ops/qdq.py b/python/tvm/relax/transform/legalize_ops/qdq.py index 6596d0c7843a..8dc6bb59111f 100644 --- a/python/tvm/relax/transform/legalize_ops/qdq.py +++ b/python/tvm/relax/transform/legalize_ops/qdq.py @@ -21,7 +21,6 @@ from tvm import te from tvm_tirx import tirx from tvm_tirx.ir import Call -from tvm_tirx.runtime import DataTypeCode from ...block_builder import BlockBuilder from ...expr import Expr @@ -143,11 +142,7 @@ def dequantize_compute(*indices): zp_value = zp[(0,) * len(zp.shape)] else: zp_value = zp[indices[axis]] - dtype = ( - "float32" - if data.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) - else "int32" - ) + dtype = "float32" if data.dtype.is_float else "int32" sub = data[indices].astype(dtype) - zp_value out = sub * scale_value.astype("float32") if out_dtype == "float32": diff --git a/tirx/python/tvm_tirx/ir/type.py b/tirx/python/tvm_tirx/ir/type.py index fcaf4cbf36ba..ee3363a84cbb 100644 --- a/tirx/python/tvm_tirx/ir/type.py +++ b/tirx/python/tvm_tirx/ir/type.py @@ -147,6 +147,15 @@ def __hash__(self): def __str__(self): return str(self.dtype) + @property + def is_float(self) -> bool: + """Return whether this type stores floating-point values. + + True for every float width, including the narrow ones, which each carry + their own DLPack dtype code. + """ + return self.dtype.is_float + def matches_code(self, *codes) -> bool: """Return whether this type has any of the given DLPack dtype codes.""" type_code = self.dtype.type_code From bdecbf20e51dd3ff8ee36defc62e68db88d26add Mon Sep 17 00:00:00 2001 From: Andrew Fish Date: Wed, 30 Sep 2026 11:48:37 -0700 Subject: [PATCH 2/2] [Relax] Test fp8 dequantize legalization Every dequantize case in this file starts from int8, and the fp8 coverage in test_op_qdq.py stops at type inference, so nothing checked the TIR generated for an fp8 input. --- .../relax/test_transform_legalize_ops_qdq.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/python/relax/test_transform_legalize_ops_qdq.py b/tests/python/relax/test_transform_legalize_ops_qdq.py index bdd9be03ed8c..4d7c48e27c4d 100644 --- a/tests/python/relax/test_transform_legalize_ops_qdq.py +++ b/tests/python/relax/test_transform_legalize_ops_qdq.py @@ -580,5 +580,52 @@ def main(data: R.Tensor((2, 4), dtype="int8")) -> R.Tensor((2, 4), dtype="float1 tvm_tirx.ir.assert_structural_equal(mod, Expected) +def test_dequantize_float8_e4m3fn_to_fp32(): + @tvm_tirx.script.ir_module + class Dequantize: + @R.function + def main( + data: R.Tensor((2, 4), "float8_e4m3fn"), + scale: R.Tensor((2,), "float32"), + zp: R.Tensor((2,), "float16"), + ) -> R.Tensor((2, 4), "float32"): + out = R.dequantize(data, scale, zp, axis=0, out_dtype="float32") + return out + + @tvm_tirx.script.ir_module + class Expected: + @Ts.function(private=True) + def dequantize( + A: T.Tensor((T.int64(2), T.int64(4)), "float8_e4m3fn"), + B: T.Tensor((T.int64(2),), "float32"), + C: T.Tensor((T.int64(2),), "float16"), + dequantized: T.Tensor((T.int64(2), T.int64(4)), "float32"), + ): + T.func_attr({"tirx.noalias": True}) + # with Ts.sblock("root"): + for i0, i1 in T.grid(T.int64(2), T.int64(4)): + with Ts.sblock("dequantized"): + v_i0, v_i1 = Ts.axis.remap("SS", [i0, i1]) + Ts.reads(A[v_i0, v_i1], C[v_i0], B[v_i0]) + Ts.writes(dequantized[v_i0, v_i1]) + dequantized[v_i0, v_i1] = ( + T.Cast(A[v_i0, v_i1], "float32") - T.Cast(C[v_i0], "float32") + ) * B[v_i0] + + @R.function + def main( + data: R.Tensor((2, 4), dtype="float8_e4m3fn"), + scale: R.Tensor((2,), dtype="float32"), + zp: R.Tensor((2,), dtype="float16"), + ) -> R.Tensor((2, 4), dtype="float32"): + out = R.call_tir( + Expected.dequantize, (data, scale, zp), ty_args=[R.Tensor((2, 4), dtype="float32")] + ) + return out + + mod = LegalizeOps()(Dequantize) + tvm_tirx.ir.assert_structural_equal(mod, Expected) + + if __name__ == "__main__": tvm_tirx.testing.main()