# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: F401, F821, F841 from itertools import product import numpy as np import pytest import tvm import tvm.testing from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx, topi from tvm.s_tir import dlight as dl from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T from tvm.testing import env try: import ml_dtypes except ImportError: ml_dtypes = None @pytest.mark.parametrize( "input", [ ("float8_e4m3fn", "__nv_fp8_e4m3"), ("float8_e5m2", "__nv_fp8_e5m2"), ], ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_fp8_conversions(input): dtype, nv_dtype = input def _create_mod(dtype): @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main( A: T.Buffer((64,), dtype), B: T.Buffer((64,), dtype), C: T.Buffer((64,), dtype), ): T.func_attr({"tirx.noalias": True}) for i_0 in T.thread_binding(2, thread="blockIdx.x"): for i_1 in T.thread_binding(32, thread="threadIdx.x"): with T.sblock("C"): v_i = T.axis.spatial(64, i_0 * 32 + i_1) T.reads(A[v_i], B[v_i]) T.writes(C[v_i]) C[v_i] = T.Cast( dtype, T.Cast("float16", A[v_i]) + T.Cast("float16", B[v_i]) ) return Module mod = _create_mod(dtype) target = "cuda" fadd = tvm.tirx.build(mod, target=target) cuda_src = fadd.imports[0].inspect_source() assert nv_dtype in cuda_src, f"{nv_dtype} datatype not found in generated CUDA" dev = tvm.device(target, 0) a = tvm.runtime.tensor(np.random.uniform(low=0, high=5, size=64).astype(dtype), dev) b = tvm.runtime.tensor(np.random.uniform(low=0, high=5, size=64).astype(dtype), dev) c = tvm.runtime.tensor(np.zeros(64, dtype=dtype), dev) fadd(a, b, c) tvm.testing.assert_allclose( c.numpy().astype("float16"), (a.numpy() + b.numpy()).astype("float16") ) @pytest.mark.parametrize( "dtype", ["float8_e4m3fn", "float8_e5m2", "float8_e8m0fnu"], ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_fp8_packing(dtype): length = 64 vector_length = 4 native_dtype, packed_dtype = (f"{dtype}x{vector_length}", "uint32") def _create_mod(native_dtype, packed_dtype, length): @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main( A: T.Buffer((length,), native_dtype), R: T.Buffer((length,), packed_dtype), B: T.Buffer((length,), native_dtype), ): T.func_attr({"tirx.noalias": True}) for i_0 in T.thread_binding(2, thread="blockIdx.x"): for i_1 in T.thread_binding(32, thread="threadIdx.x"): with T.sblock("R"): v_i = T.axis.spatial(length, i_0 * 32 + i_1) T.reads(A[v_i]) T.writes(R[v_i]) R[v_i] = T.reinterpret(packed_dtype, A[v_i]) for i_0 in T.thread_binding(2, thread="blockIdx.x"): for i_1 in T.thread_binding(32, thread="threadIdx.x"): with T.sblock("B"): v_i = T.axis.spatial(length, i_0 * 32 + i_1) T.reads(R[v_i]) T.writes(B[v_i]) B[v_i] = T.reinterpret(native_dtype, R[v_i]) return Module mod = _create_mod(native_dtype, packed_dtype, length) target = "cuda" f = tvm.compile(mod, target=target) dev = tvm.device(target, 0) np_shape = (length, vector_length) a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(dtype) a = tvm.runtime.empty(shape=(length,), dtype=native_dtype, device=dev) r = tvm.runtime.empty(shape=(length,), dtype=packed_dtype, device=dev) b = tvm.runtime.empty(shape=(length,), dtype=native_dtype, device=dev) a.copyfrom(a_np) f(a, r, b) tvm.testing.assert_allclose(a.numpy().astype("float16"), b.numpy().astype("float16")) @pytest.mark.parametrize( "native_dtype,promoted_dtype,numpytype", [ ("float8_e4m3fn", "float32", "float8_e4m3fn"), ("float8_e4m3fn", "float16", "float8_e4m3fn"), ("float8_e4m3fnx2", "float32x2", "float8_e4m3fn"), ("float8_e4m3fnx2", "float16x2", "float8_e4m3fn"), ("float8_e4m3fnx4", "float32x4", "float8_e4m3fn"), # Supported via half4 vector type extension in codegen ("float8_e4m3fnx4", "float16x4", "float8_e4m3fn"), ("float8_e5m2", "float32", "float8_e5m2"), ("float8_e5m2", "float16", "float8_e5m2"), ("float8_e5m2x2", "float32x2", "float8_e5m2"), ("float8_e5m2x2", "float16x2", "float8_e5m2"), ("float8_e5m2x4", "float32x4", "float8_e5m2"), ("float8_e5m2x4", "float16x4", "float8_e5m2"), ], ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_fp8_vector_conversions(native_dtype, promoted_dtype, numpytype): vector_length = 64 def _create_mod(native_dtype, promoted_dtype): @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main( A: T.Buffer((64,), native_dtype), B: T.Buffer((64,), native_dtype), C: T.Buffer((64,), native_dtype), ): T.func_attr({"tirx.noalias": True}) for i_0 in T.thread_binding(2, thread="blockIdx.x"): for i_1 in T.thread_binding(32, thread="threadIdx.x"): with T.sblock("C"): v_i = T.axis.spatial(64, i_0 * 32 + i_1) T.reads(A[v_i], B[v_i]) T.writes(C[v_i]) C[v_i] = T.Cast( native_dtype, T.Cast(promoted_dtype, A[v_i]) + T.Cast(promoted_dtype, B[v_i]), ) return Module mod = _create_mod(native_dtype, promoted_dtype) target = "cuda" fadd = tvm.tirx.build(mod, target=target) cuda_src = fadd.imports[0].inspect_source() dev = tvm.device(target, 0) if "x" in native_dtype: lanes = int(native_dtype.split("x")[-1]) else: lanes = 1 if "x" in promoted_dtype: promoted_base_dtype = promoted_dtype.split("x")[0] else: promoted_base_dtype = promoted_dtype np_shape = (vector_length, lanes) if lanes > 1 else (vector_length,) a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(numpytype) a = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev) a.copyfrom(a_np) b_np = np.random.uniform(low=0, high=5, size=np_shape).astype(numpytype) b = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev) b.copyfrom(b_np) c = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev) fadd(a, b, c) tvm.testing.assert_allclose( c.numpy().astype(promoted_base_dtype), (a_np + b_np).astype(promoted_base_dtype) ) bcast_length = tvm.testing.parameter(2, 4, 6, 8) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0") def test_half_broadcast(bcast_length): dtype = "float16" def _create_mod(bcast_length, dtype): @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main(a: T.Buffer((), dtype), vec: T.Buffer((bcast_length,), dtype)): for i_0 in T.thread_binding(1, thread="blockIdx.x"): for i_1 in T.thread_binding(1, thread="threadIdx.x"): with T.sblock("broadcast"): vec[0:bcast_length] = T.broadcast(a[()], bcast_length) return Module mod = _create_mod(bcast_length, dtype) target = "cuda" func = tvm.compile(mod, target=target) dev = tvm.device(target, 0) a_np = np.random.uniform(low=0, high=4, size=()).astype(dtype) a = tvm.runtime.tensor(a_np, device=dev) b = tvm.runtime.empty((bcast_length,), dtype=dtype, device=dev) func(a, b) b_np = np.full((bcast_length,), a_np) tvm.testing.assert_allclose(b.numpy(), b_np) vector_length = tvm.testing.parameter(2, 4) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0") def test_half_misaligned_vector_load(vector_length): dtype = "float16" vec_dtype = dtype + "x" + str(vector_length) length = 256 @T.prim_func(s_tir=True) def vector_load( A: T.Buffer((length,), dtype), B: T.Buffer((length // vector_length,), vec_dtype) ): for b in T.thread_binding(1, thread="blockIdx.x"): for i in T.thread_binding(length // vector_length, thread="threadIdx.x"): vec_index = T.ramp((i + 1) * vector_length - 1, -1, vector_length) B[i] = A[vec_index] target = "cuda" f = tvm.compile(vector_load, target=target) dev = tvm.device(target, 0) a_np = np.random.uniform(low=0, high=1, size=(length,)).astype(dtype) a = tvm.runtime.tensor(a_np, device=dev) b = tvm.runtime.empty((length // vector_length,), dtype=vec_dtype, device=dev) f(a, b) b_np = np.empty((length // vector_length, vector_length), dtype=dtype) for i in range(length // vector_length): start_index = (i + 1) * vector_length - 1 b_np[i, :] = a_np[start_index - vector_length + 1 : start_index + 1][::-1] tvm.testing.assert_allclose(b.numpy(), b_np) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0") def test_half4_vector_add(): dtype = "float16" length = 64 vector_length = 4 vec_dtype = dtype + "x" + str(vector_length) @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main( A: T.Buffer((64,), "float16x4"), B: T.Buffer((64,), "float16x4"), C: T.Buffer((64,), "float16x4"), ): T.func_attr({"tirx.noalias": True}) for i_0 in T.thread_binding(2, thread="blockIdx.x"): for i_1 in T.thread_binding(32, thread="threadIdx.x"): with T.sblock("C"): v_i = T.axis.spatial(64, i_0 * 32 + i_1) T.reads(A[v_i], B[v_i]) T.writes(C[v_i]) C[v_i] = A[v_i] + B[v_i] target = "cuda" fadd = tvm.compile(Module, target=target) dev = tvm.device(target, 0) a_np = np.random.uniform(-1, 1, (length, vector_length)).astype(dtype) a = tvm.runtime.empty(shape=(length,), dtype=vec_dtype, device=dev) a.copyfrom(a_np) b_np = np.random.uniform(-1, 1, (length, vector_length)).astype(dtype) b = tvm.runtime.empty(shape=(length,), dtype=vec_dtype, device=dev) b.copyfrom(b_np) c = tvm.runtime.empty(shape=(length,), dtype=vec_dtype, device=dev) fadd(a, b, c) c_expected = a_np + b_np tvm.testing.assert_allclose(c.numpy(), c_expected, atol=1e-5, rtol=1e-5) class BaseFP8E4M3QuantScaleOnly: @classmethod def create_quantize_func( cls, weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, max_int_value, axis, output_transpose, ) -> IRModule: if DataType(quantize_dtype).type_code == DataTypeCode.Float8E4M3FN: quantize_func = cls.quantize_fp8x4_e4m3 else: assert NotImplementedError() bb = relax.BlockBuilder() # pylint: disable=invalid-name weight_var = relax.Var("weight", relax.TensorType(weight_shape, model_dtype)) compute_scale, compute_quantize, compute_transpose = quantize_func( weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, max_int_value, axis, output_transpose, ) with bb.function(name="main", params=[weight_var]): with bb.dataflow(): lv_scale = bb.emit_te(compute_scale, weight_var) lv_quantized_weight = compute_quantize(bb, (weight_var, lv_scale)) if compute_transpose: lv_output = bb.emit_te(compute_transpose, lv_quantized_weight, lv_scale) lv_quantized_weight = lv_output[0] lv_scale = lv_output[1] tuple_output = bb.emit((lv_quantized_weight, lv_scale)) gv = bb.emit_output(tuple_output) bb.emit_func_output(gv) return bb.finalize() @classmethod def create_dequantize_func( cls, packed_weight_shape, scale_shape, dequantized_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, axis, ) -> IRModule: if DataType(quantize_dtype).type_code == DataTypeCode.Float8E4M3FN: dequantize_func = cls.dequantize_fp8x4_e4m3 else: assert NotImplementedError() bb = relax.BlockBuilder() # pylint: disable=invalid-name packed_weight_var = relax.Var( "weight", relax.TensorType(packed_weight_shape, storage_dtype) ) scale_var = relax.Var("scale", relax.TensorType(scale_shape, model_dtype)) compute_dequantize = dequantize_func( packed_weight_shape, scale_shape, dequantized_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, axis, ) with bb.function(name="main", params=[packed_weight_var, scale_var]): with bb.dataflow(): lv = compute_dequantize(bb, (packed_weight_var, scale_var)) gv = bb.emit_output(lv) bb.emit_func_output(gv) return bb.finalize() @classmethod def quantize_fp8x4_e4m3( # pylint: disable=too-many-locals cls, weight_shape: list[tirx.Expr], model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, max_int_value, axis: int = -1, output_transpose: bool = False, ) -> tuple[te.Tensor, te.Tensor]: """Group quantization for weight tensor, defined in tensor expression.""" max_int = tirx.const(max_int_value, model_dtype) shape = weight_shape # pylint: disable=invalid-name axis = axis if axis >= 0 else len(shape) + axis k = shape[axis] quantize_dtype = DataType(quantize_dtype) # compute scale per group r = te.reduce_axis((0, group_size), name="r") # pylint: disable=invalid-name num_group = tirx.ceildiv(k, group_size) # (4096, 4096) -> quantize axis = 0, group size = 32 -> (128, 4096) # for channel quant group_size = 4096 -> (1, 4096) scale_shape = (*shape[:axis], num_group, *shape[axis + 1 :]) def compute_scale(weight: te.Tensor): min_scaling_factor = tirx.const(1.0 / (max_int_value * 512.0), model_dtype) max_abs = te.compute( shape=scale_shape, fcompute=lambda *idx: te.max( tirx.if_then_else( idx[axis] * group_size + r < k, te.abs(weight(*idx[:axis], idx[axis] * group_size + r, *idx[axis + 1 :])), te.min_value(model_dtype), ), axis=r, ), name="max_abs_value", ) scale = te.compute( scale_shape, lambda *idx: te.max( max_abs(*idx).astype(model_dtype) / max_int, min_scaling_factor ), name="scale", ) return scale def compute_quantize_weight(bb: relax.BlockBuilder, args: relax.expr.Expr): # compute scaled weight packed_shape = (weight_shape[0], weight_shape[1] // num_elem_per_storage) quant = cls.quant_and_pack_fp8x4_e4m3_sm90( weight_shape, packed_shape, scale_shape, group_size, axis, model_dtype, storage_dtype, quantize_dtype, ) # quant.show() global_var = bb.add_func(quant, "quantized_weight") lv_quantized_weight = bb.emit( relax.call_tir(global_var, args, relax.TensorType(packed_shape, storage_dtype)) ) return lv_quantized_weight compute_transpose = None if output_transpose: def compute_transpose(quantized_weight: te.Tensor, scale: te.Tensor): if len(quantized_weight.shape) != 2 or len(scale.shape) != 2: raise ValueError( "Does not support transpose output quantized weight with ndim != 2" ) quantized_weight = topi.transpose(quantized_weight) scale = topi.transpose(scale) return quantized_weight, scale return compute_scale, compute_quantize_weight, compute_transpose @classmethod def dequantize_fp8x4_e4m3( # pylint: disable=too-many-locals cls, packed_weight_shape: list[tirx.Expr], scale_shape, dequant_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_elem_per_storage, axis: int = -1, ) -> tuple[te.Tensor, te.Tensor]: """Group quantization for weight tensor, defined in tensor expression.""" axis = axis if axis >= 0 else len(shape) + axis def compute_dequantize_weight(bb: relax.BlockBuilder, args: relax.expr.Expr): dequant = cls.dequant_fp8x4_e4m3_sm90( packed_weight_shape, scale_shape, dequant_shape, group_size, axis, model_dtype, storage_dtype, quantize_dtype, ) global_var = bb.add_func(dequant, "dequantize_weight") lv_dequantized_weight = bb.emit( relax.call_tir(global_var, args, relax.TensorType(dequant_shape, model_dtype)) ) return lv_dequantized_weight return compute_dequantize_weight @classmethod def quant_and_pack_fp8x4_e4m3_sm90( cls, weight_shape, packed_shape, scale_shape, group_size, axis, model_dtype, storage_dtype, quantized_dtype, ): vector_length = 4 vec_quantized_dtype = f"{quantized_dtype}x{vector_length}" vec_model_dtype = f"{model_dtype}x{vector_length}" num_elem_per_storage = vector_length # TODO(csullivan) assert on storage dtype / quantize type bytes == vector length assert group_size % vector_length == 0, ( f"Number of elements in a group must be divisible by fp8 vector length {vector_length}" ) @T.prim_func(private=True, s_tir=True) def quant_pack( A: T.Buffer(weight_shape, model_dtype), scale: T.Buffer(scale_shape, model_dtype), compute: T.Buffer( packed_shape, storage_dtype, ), ): # with T.sblock("root"): # test = T.sblock_alloc_buffer(1, dtype=vec_model_dtype, scope="local") for i0, i1 in T.grid( T.int64(weight_shape[0]), T.int64(weight_shape[1] // vector_length) ): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads( A[v_i0, v_i1 : v_i1 + vector_length], scale[v_i0, v_i1 * T.int64(vector_length) // T.int64(group_size)], ) T.writes(compute[v_i0, v_i1 * vector_length]) compute[v_i0, v_i1] = T.reinterpret( storage_dtype, T.Cast( vec_quantized_dtype, A[v_i0, T.ramp(v_i1 * vector_length, 1, vector_length)] / scale[v_i0, v_i1 * T.int64(vector_length) // T.int64(group_size)], ), ) return quant_pack @classmethod def dequant_fp8x4_e4m3_sm90( cls, packed_weight_shape, scale_shape, out_shape, group_size, axis, model_dtype, storage_dtype, quantized_dtype, ): vector_length = 4 vec_quantized_dtype = f"{quantized_dtype}x{vector_length}" vec_model_dtype = f"{model_dtype}x{vector_length}" num_elem_per_storage = vector_length @T.prim_func(s_tir=True) def dequant( packed_weight: T.Buffer(packed_weight_shape, storage_dtype), scale: T.Buffer(scale_shape, model_dtype), dequantize: T.Buffer(out_shape, model_dtype), ): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for i0, i1 in T.grid(T.int64(packed_weight_shape[0]), T.int64(packed_weight_shape[1])): with T.sblock("dequantize"): v_i0 = T.axis.spatial(T.int64(packed_weight_shape[0]), i0) v_i1 = T.axis.spatial(T.int64(packed_weight_shape[1]), i1) T.reads( packed_weight[v_i0, v_i1], scale[v_i0, v_i1 * T.int64(vector_length) // T.int64(group_size)], ) dequantize[v_i0, T.ramp(v_i1 * vector_length, 1, vector_length)] = T.Cast( vec_model_dtype, T.reinterpret(vec_quantized_dtype, packed_weight[v_i0, v_i1]), ) * T.Broadcast( scale[v_i0, v_i1 * T.int64(vector_length) // T.int64(group_size)], vector_length, ) return dequant @classmethod def compile_quant_and_dequant_by_scale( cls, weight_shape, scales_shape, quant_weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_el_per_storage, max_int_value, axis, target_str, dev, ): quant_mod = cls.create_quantize_func( weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_el_per_storage, max_int_value, axis, output_transpose=False, ) # quant_mod.show() target = tvm.target.Target(target_str) with target: quant_mod = dl.ApplyDefaultSchedule( dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), )(quant_mod) ex_1 = tvm.compile(quant_mod, target=target) vm_1 = relax.VirtualMachine(ex_1, dev) dequant_mod = cls.create_dequantize_func( quant_weight_shape, scales_shape, weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_el_per_storage, axis, ) # dequant_mod.show() with target: dequant_mod = dl.ApplyDefaultSchedule( dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), )(dequant_mod) dequant_mod.show() ex_2 = tvm.compile(dequant_mod, target=target) vm_2 = relax.VirtualMachine(ex_2, dev) def print_cuda(target, mod, name=None): if name: mod = mod[name] f = tvm.tirx.build(mod, target=target) cuda_src = f.imports[0].inspect_source() print(cuda_src) print_cuda(target, dequant_mod, name="dequant") return vm_1["main"], vm_2["main"] class TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly): # weight_shape = tvm.testing.parameter((32000, 4096), (4096, 14336)) weight_shape = tvm.testing.parameter((128, 256), (128, 64)) @tvm.testing.fixture def group_size(self): return 64 @tvm.testing.fixture def axis(self): return 1 @tvm.testing.fixture def model_dtype(self): return "float16" @tvm.testing.fixture def storage_dtype(self): return "uint32" @tvm.testing.fixture def quantize_dtype(self): return "float8_e4m3fn" @tvm.testing.fixture def num_el_per_storage(self): return 4 @tvm.testing.fixture def max_int_value(self): return 448 @tvm.testing.fixture def target_str(self): return "cuda" @tvm.testing.fixture def scale_shape(self, weight_shape, group_size, axis): return [ (d + group_size - 1) // group_size if axis == i else d for i, d in enumerate(weight_shape) ] @tvm.testing.fixture def quant_weight_shape(self, weight_shape, num_el_per_storage, axis): return [ (d + num_el_per_storage - 1) // num_el_per_storage if axis == i else d for i, d in enumerate(weight_shape) ] @tvm.testing.fixture def compiled_functions( self, weight_shape, scale_shape, quant_weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_el_per_storage, max_int_value, axis, target_str, ): dev = tvm.device(target_str, 0) return self.compile_quant_and_dequant_by_scale( weight_shape, scale_shape, quant_weight_shape, model_dtype, quantize_dtype, storage_dtype, group_size, num_el_per_storage, max_int_value, axis, target_str, dev, ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9") def test_main(self, weight_shape, model_dtype, target_str, compiled_functions): quant, dequant = compiled_functions dev = tvm.device(target_str, 0) weight_np = np.random.uniform(-100, 100, weight_shape).astype(model_dtype) weight = tvm.runtime.tensor(weight_np, device=dev) quant_weight, scales = quant(weight) quant_weight_np, scales_np = quant_weight.numpy(), scales.numpy() dequant_weight = dequant(quant_weight, scales) dequant_weight_np = dequant_weight.numpy() tvm.testing.assert_allclose(weight_np, dequant_weight_np, atol=10, rtol=5e-2) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") @pytest.mark.parametrize("dtype", ["float8_e5m2", "float8_e4m3fn", "float8_e8m0fnu"]) def test_const(dtype): @T.prim_func(s_tir=True) def func(A: T.Buffer((4,), dtype)) -> None: A_local = T.sblock_alloc_buffer((4,), dtype=dtype, scope="local") for tx in T.thread_binding(0, 4, "threadIdx.x"): for i in T.vectorized(4): A_local[i] = T.float32(1.0).astype(dtype) A[tx] = A_local[tx] mod = tvm.IRModule({"main": func}) tvm.compile(mod, target="cuda") @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9") @pytest.mark.parametrize("dtype", ["float8_e5m2", "float8_e4m3fn"]) @pytest.mark.parametrize("vec_len", [2, 4, 8, 16]) def test_copy(dtype, vec_len): @T.prim_func(s_tir=True) def func( A: T.Buffer( ( 4, vec_len, ), dtype, ), B: T.Buffer( ( 4, vec_len, ), dtype, ), ) -> None: for tx in T.thread_binding(0, 4, "threadIdx.x"): for i in T.vectorized(vec_len): B[tx, i] = A[tx, i] mod = tvm.IRModule({"main": func}) rtmod = tvm.compile(mod, target="cuda") num_experts = 8 reduce_size = 1792 spatial_size = 4096 @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") @pytest.mark.skipif(ml_dtypes is None, reason="Requires ml_dtypes to be installed") def test_moe_gemv_shfl_down_illegal_instr(): global num_experts global reduce_size global spatial_size @I.ir_module(s_tir=True) class SingleBatchMoE_float8_e4m3: @T.prim_func(private=True, s_tir=True) def moe_dequantize_gemv( x_handle: T.handle, w: T.Buffer((num_experts, spatial_size, reduce_size), "float8_e4m3fn"), scale: T.Buffer((1,), "float16"), indptr: T.Buffer((1, 2), "int32"), o: T.Buffer((2, spatial_size), "float16"), ): T.func_attr({"op_pattern": 4, "tirx.noalias": True}) num_seq = T.int64() x = T.match_buffer(x_handle, (num_seq, reduce_size), "float16") for expert_id in T.thread_binding(2, thread="blockIdx.y"): with T.sblock("gemv_o"): e = T.axis.spatial(2, expert_id) T.reads( w[indptr[0, e], 0:spatial_size, 0:reduce_size], indptr[0, e], scale[0], x[e, 0:reduce_size], ) T.writes(o[e, 0:spatial_size]) y = T.sblock_alloc_buffer((spatial_size, reduce_size), "float16") for i1, i2 in T.grid(spatial_size, reduce_size): with T.sblock("dequantize"): i, j = T.axis.remap("SS", [i1, i2]) T.reads(w[indptr[0, e], i, j], indptr[0, e], scale[0]) T.writes(y[i, j]) y[i, j] = T.Cast("float16", w[indptr[0, e], i, j]) * scale[0] for i1, i2 in T.grid(spatial_size, reduce_size): with T.sblock("gemv"): i, j = T.axis.remap("SR", [i1, i2]) T.reads(x[e, j], y[i, j]) T.writes(o[e, i]) with T.init(): o[e, i] = T.float16(0) o[e, i] = o[e, i] + x[e, j] * y[i, j] @R.function def main( x: R.Tensor(("num_seq", reduce_size), dtype="float16"), indptr: R.Tensor((1, 2), dtype="int32"), weight: R.Tensor((num_experts, spatial_size, reduce_size), dtype="float8_e4m3fn"), scale: R.Tensor((1,), dtype="float32"), ) -> R.Tensor((2, spatial_size), dtype="float16"): num_seq = T.int64() R.func_attr({"num_input": 2}) cls = SingleBatchMoE_float8_e4m3 with R.dataflow(): astype: R.Tensor((1,), dtype="float16") = R.astype(scale, dtype="float16") lv = R.call_tir( cls.moe_dequantize_gemv, (x, weight, astype, indptr), out_ty=R.Tensor((2, spatial_size), dtype="float16"), ) gv: R.Tensor((2, spatial_size), dtype="float16") = lv R.output(gv) return gv def _pipeline(mod: tvm.ir.IRModule) -> tvm.ir.IRModule: seq = tvm.transform.Sequential( [ tvm.relax.transform.LegalizeOps(), dl.ApplyDefaultSchedule( dl.gpu.Matmul(), dl.gpu.GEMV(), dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), ), ] ) mod = seq(mod) return mod mod = SingleBatchMoE_float8_e4m3 target = tvm.target.Target("cuda") with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": False}) and target: mod = _pipeline(mod) rt_mod = tvm.compile(mod, target=target) x_data = np.zeros((1, reduce_size), dtype=np.float16) indptr_data = np.zeros((1, 2), dtype=np.int32) weight_data = np.zeros((num_experts, spatial_size, reduce_size), dtype="float8_e4m3fn") scale_data = np.zeros((1,), dtype=np.float32) def run_and_check(): dev = tvm.cuda(0) x = tvm.runtime.tensor(x_data, device=dev) indptr = tvm.runtime.tensor(indptr_data, device=dev) weight = tvm.runtime.tensor(weight_data, device=dev) scale = tvm.runtime.tensor(scale_data, device=dev) vm = relax.VirtualMachine(rt_mod, dev) # Ensure this runs without failure. Utilizing dlight thread extents TS, TR = 4, 64 # in GEMV scheduling will yield: CUDA: an illegal instruction was encountered. vm["main"](x, indptr, weight, scale) dev.sync() tvm.testing.run_with_gpu_lock(run_and_check) @pytest.mark.parametrize("vec_length", [2, 4]) @pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9") def test_fp8_fp16_bf16_vectorize_arith(vec_length, dtype): def _create_mod(vec_length, dtype): num_threads = 128 // vec_length @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) def main( A: T.Buffer((128,), "float8_e4m3fn"), B: T.Buffer((128,), dtype), C: T.Buffer((128,), dtype), ) -> None: for i_0 in T.thread_binding(num_threads, thread="threadIdx.x"): for i_1 in T.vectorized(vec_length): with T.sblock("compute"): vi = T.axis.spatial(128, i_0 * vec_length + i_1) C[vi] = (A[vi].astype(dtype) * B[vi]) + T.bfloat16(3.0) return Module mod = _create_mod(vec_length, dtype) target = tvm.target.Target.from_device(tvm.cuda()) f = tvm.tirx.build(mod, target=target) a_np = np.random.rand(128).astype("float8_e4m3fn") b_np = np.random.rand(128).astype(dtype) c_np = (a_np.astype(dtype) * b_np) + 3 def run_and_check(): device = tvm.cuda() a_tvm = tvm.runtime.tensor(a_np, device=device) b_tvm = tvm.runtime.tensor(b_np, device=device) c_tvm = tvm.runtime.empty((128,), dtype=dtype, device=device) f(a_tvm, b_tvm, c_tvm) actual = c_tvm.numpy() tvm.testing.assert_allclose( actual.astype(np.float32), c_np.astype(np.float32), atol=5e-1, rtol=1e-2 ) tvm.testing.run_with_gpu_lock(run_and_check) if __name__ == "__main__": tvm.testing.main()