This is an automated email from the ASF dual-hosted git repository. junrushao pushed a commit to branch unity-staging in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 231653cca550cb4f840e00de82642d6fed8e67cd Author: masahi <[email protected]> AuthorDate: Tue Jul 18 04:41:47 2023 +0900 [Unity] fp16 A x int B GEMM update - support int8, more bias shape (#15318) * Support int8 FasterTransformer kernel * test strided bias * test gelu act * update decode check * update FT rev to disable GCC warning --- 3rdparty/cutlass_fpA_intB_gemm | 2 +- python/tvm/contrib/cutlass/build.py | 45 +++-- python/tvm/contrib/cutlass/gemm_operation.py | 26 +-- python/tvm/contrib/cutlass/gen_tensor_op.py | 19 ++- python/tvm/contrib/cutlass/layer_norm_operation.py | 4 +- python/tvm/relax/backend/contrib/cutlass.py | 39 +++-- src/runtime/contrib/cutlass/weight_preprocess.cc | 11 +- tests/python/relax/test_codegen_cutlass.py | 190 ++++++++++++++++++--- 8 files changed, 261 insertions(+), 75 deletions(-) diff --git a/3rdparty/cutlass_fpA_intB_gemm b/3rdparty/cutlass_fpA_intB_gemm index 634184f502..390e821fba 160000 --- a/3rdparty/cutlass_fpA_intB_gemm +++ b/3rdparty/cutlass_fpA_intB_gemm @@ -1 +1 @@ -Subproject commit 634184f50272227f95805fb8ef1eb8c4da373304 +Subproject commit 390e821fbad2356089aab603d7116c6c820eae65 diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index 2d99cca8f9..b6681ae80f 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -700,7 +700,12 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator): arg_idx = _extract_arg_idx(op_type, f) signature = _extract_relax_function_signature(f) lhs_arg = f"arg{arg_idx['lhs']}" + rhs_arg = f"arg{arg_idx['w_encoded']}" lhs_shape = signature[f"{lhs_arg}_shape"] + rhs_shape = signature[f"{rhs_arg}_shape"] + ret_shape = signature["ret_shape"] + N = ret_shape[-1] + attrs = { "op_type": op_type, "lhs_arg_idx": arg_idx["lhs"], @@ -708,8 +713,23 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator): "scales_arg_idx": arg_idx["scales"], "bias_arg_idx": arg_idx.get("bias"), "batch_offset": len(lhs_shape) - 2, + "activation": "identity", } + attrs["bias_stride"] = 0 + + if "bias" in arg_idx: + bias_shape = signature[f"arg{arg_idx['bias']}_shape"] + bias_shape_1d = reduce(operator.mul, bias_shape, 1) + if bias_shape_1d != bias_shape[-1]: + attrs["bias_stride"] = bias_shape[-1] + + if N == rhs_shape[1]: + attrs["weight_nbit"] = 8 + else: + assert N == rhs_shape[1] * 2 + attrs["weight_nbit"] = 4 + if "residual" in op_type: residual_pos = op_type.find("residual_") postfix = op_type[residual_pos + len("residual_") :] @@ -739,6 +759,11 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator): "residual_arg_idx": arg_idx["residual"], } ) + else: + for act in ["relu", "silu", "gelu"]: + if act in op_type: + attrs["activation"] = act + break return f.with_attrs(attrs) @@ -912,18 +937,8 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator): } ) - def handle_layer_norm(self, f, _): - """Annotate a layer norm op.""" - signature = _extract_relax_function_signature(f) - attrs = {} - attrs["M"] = reduce(operator.mul, signature["arg0_shape"][:-1], 1) - attrs["N"] = signature["arg0_shape"][-1] - dtype = signature["arg0_dtype"] - attrs["data_type"] = {"float32": "float", "float16": "cutlass::half_t"}[str(dtype)] - return f.with_attrs(attrs) - - def handle_rms_norm(self, f, _): - """Annotate a rms norm op.""" + def handle_norm(self, f, _): + """Annotate a layer or rms norm op.""" signature = _extract_relax_function_signature(f) attrs = {} attrs["batch_rank"] = len(signature["arg0_shape"][:-1]) @@ -948,10 +963,8 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator): return self.handle_matmul(f, op_type) elif "attention" in op_type: return self.handle_attention(f, op_type) - elif "layer_norm" in op_type: - return self.handle_layer_norm(f, op_type) - elif "rms_norm" in op_type: - return self.handle_rms_norm(f, op_type) + elif "layer_norm" in op_type or "rms_norm" in op_type: + return self.handle_norm(f, op_type) raise ValueError("Unsupported composite {}".format(op_type)) diff --git a/python/tvm/contrib/cutlass/gemm_operation.py b/python/tvm/contrib/cutlass/gemm_operation.py index 3fa6e9be8d..e445316dfd 100644 --- a/python/tvm/contrib/cutlass/gemm_operation.py +++ b/python/tvm/contrib/cutlass/gemm_operation.py @@ -412,14 +412,13 @@ def instantiate_gemm_template(attrs): return substitute_template(template, attrs) -def emit_fp16A_int4B_matmul(attrs): - """Return CUTLASS host code for fp16 A and int4 B GEMM.""" - +def emit_fp16A_intB_matmul(attrs): + """Return CUTLASS host code for fp16 A and int4 or int8 B GEMM.""" attrs["template_common"] = substitute_template( """ using namespace fastertransformer; int m = ${A_arg}->shape[${batch_offset}]; - int n = ${B_arg}->shape[1] * 2; + int n = ${B_arg}->shape[1] * ${float_per_int}; int k = ${B_arg}->shape[0]; auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream"); @@ -451,17 +450,20 @@ def emit_fp16A_int4B_matmul(attrs): m, n, k, nullptr, 0, stream); """ - if "residual_arg" in attrs and "bias_arg" in attrs: - template_residual = substitute_template( - template_residual, {"bias": "static_cast<cutlass::half_t*>(${bias_arg}->data)"} - ) - return substitute_template(template_residual, attrs) - if "residual_arg" in attrs: - template_residual = substitute_template(template_residual, {"bias": "nullptr"}) + if "bias_arg" in attrs: + bias = "static_cast<cutlass::half_t*>(${bias_arg}->data)" + else: + bias = "nullptr" + + template_residual = substitute_template(template_residual, {"bias": bias}) return substitute_template(template_residual, attrs) if "bias_arg" in attrs: - return substitute_template(template_bias, attrs) + template = substitute_template( + template, {"bias": "static_cast<cutlass::half_t*>(${bias_arg}->data)"} + ) + else: + template = substitute_template(template, {"bias": "nullptr"}) return substitute_template(template, attrs) diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py index 8c8bcc20c3..bf02d8f7b8 100644 --- a/python/tvm/contrib/cutlass/gen_tensor_op.py +++ b/python/tvm/contrib/cutlass/gen_tensor_op.py @@ -31,7 +31,7 @@ from tvm.tir import IntImm from . import _ffi_api as ffi from .attention_operation import instantiate_attention_template from .conv2d_operation import instantiate_conv2d_template -from .gemm_operation import instantiate_gemm_template, emit_fp16A_int4B_matmul +from .gemm_operation import instantiate_gemm_template, emit_fp16A_intB_matmul from .layer_norm_operation import instantiate_layer_norm_template from .rms_norm_operation import instantiate_rms_norm_template from .library import ( @@ -526,6 +526,7 @@ def instantiate_template(func_name, annotations, func_args): attrs["scales_arg"] = func_args[scales_arg_idx] attrs["batch_offset"] = _get_optional_int_annotation(annotations, "batch_offset", 0) attrs["activation"] = annotations.get("activation", "identity") + attrs["bias_stride"] = annotations["bias_stride"] if bias_arg_idx is not None: attrs["bias_arg"] = func_args[bias_arg_idx] @@ -535,7 +536,15 @@ def instantiate_template(func_name, annotations, func_args): attrs["binary_op"] = annotations["binary_op"] attrs["unary_op"] = annotations["unary_op"] - code = emit_fp16A_int4B_matmul(attrs) + if annotations["weight_nbit"] == 4: + attrs["weight_dtype"] = "cutlass::uint4b_t" + attrs["float_per_int"] = 2 + else: + assert annotations["weight_nbit"] == 8 + attrs["weight_dtype"] = "uint8_t" + attrs["float_per_int"] = 1 + + code = emit_fp16A_intB_matmul(attrs) return CodegenResult(code, headers) elif "dense" in func_name or "matmul" in func_name: @@ -792,6 +801,12 @@ def instantiate_template(func_name, annotations, func_args): headers.append("cutlass/layout/matrix.h") attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]} attrs.update(dict(annotations)) + + if isinstance(attrs["M"], tvm.tir.Var): + attrs["M"] = " * ".join( + ["{}->shape[{}]".format(func_args[0], i) for i in range(int(attrs["batch_rank"]))] + ) + code = instantiate_layer_norm_template(attrs) return CodegenResult(code, headers) elif "rms_norm" in func_name: diff --git a/python/tvm/contrib/cutlass/layer_norm_operation.py b/python/tvm/contrib/cutlass/layer_norm_operation.py index ad2e730d27..760ddec3bf 100644 --- a/python/tvm/contrib/cutlass/layer_norm_operation.py +++ b/python/tvm/contrib/cutlass/layer_norm_operation.py @@ -28,8 +28,8 @@ def instantiate_layer_norm_template(attrs): using data_type = ${data_type}; using namespace cutlass::layout; - auto M = ${M}; - auto N = ${N}; + int M = ${M}; + int N = ${N}; cutlass::MatrixCoord size(M, N); auto layout_2D = RowMajor::packed(size); auto layout_channels = RowMajor::packed({1, N}); diff --git a/python/tvm/relax/backend/contrib/cutlass.py b/python/tvm/relax/backend/contrib/cutlass.py index 5cb5a6f3d7..cdce3225b5 100644 --- a/python/tvm/relax/backend/contrib/cutlass.py +++ b/python/tvm/relax/backend/contrib/cutlass.py @@ -102,6 +102,15 @@ def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequenc return False +def _is_same_shape(shape1, shape2): + analyzer = tvm.arith.Analyzer() + return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)]) + + +def _is_bias_like(shape, out_channel): + return shape[-1] == out_channel and _shape_1d(shape) == out_channel + + def _check_residual(root_call: Call, context: PatternCheckContext) -> bool: if "residual" in context.annotated_expr: residual = context.annotated_expr["residual"] @@ -116,13 +125,11 @@ def _check_residual(root_call: Call, context: PatternCheckContext) -> bool: # If residual depends on the result of the root call, this cannot be handled by cutlass. return False - shape1 = [int(s) for s in root_var.struct_info.shape] - shape2 = [int(s) for s in residual.struct_info.shape] - + shape1 = root_var.struct_info.shape + shape2 = residual.struct_info.shape out_channel = shape1[-1] - is_bias_like = lambda shape: (shape[-1] == out_channel and _shape_1d(shape) == out_channel) - if shape1 != shape2 and not is_bias_like(shape2): + if not _is_same_shape(shape1, shape2) and not _is_bias_like(shape2, out_channel): return False return True @@ -239,7 +246,6 @@ def _check_decode_matmul(ctx): return False N = root.struct_info.shape[-1] - K = call_tir_decode.struct_info.shape[0] if ctx.annotated_expr["lhs"].struct_info.dtype != "float16": return False @@ -254,14 +260,10 @@ def _check_decode_matmul(ctx): if ( isinstance(packed_weight, Call) and isinstance(packed_weight.args[0], ExternFunc) - and packed_weight.args[0].global_symbol != "cutlass.ft_preprocess_weight_int4" + and packed_weight.args[0].global_symbol != "cutlass.ft_preprocess_weight" ): return False - # packed weight needs to be of shape (K, N // 2) - if packed_weight.struct_info.shape[0] != K or packed_weight.struct_info.shape[1] != N // 2: - return False - scales = ctx.annotated_expr["scales"] if scales.struct_info.dtype != "float16": @@ -271,11 +273,13 @@ def _check_decode_matmul(ctx): if len(scales.struct_info.shape) != 1 or scales.struct_info.shape[0] != N: return False - # bias shape needs to be (N,), possibly with additional axes on the front. if "bias" in ctx.annotated_expr: + out_shape = root.struct_info.shape bias_shape = ctx.annotated_expr["bias"].struct_info.shape - bias_shape_1d = reduce(operator.mul, bias_shape, 1) - if bias_shape_1d != bias_shape[-1]: + + # bias shape needs to be (N,), possibly with additional axes on the front. + # It can also have the same shape as the output. + if not _is_bias_like(bias_shape, N) and not _is_same_shape(out_shape, bias_shape): return False return True @@ -309,16 +313,15 @@ def decode_matmul_patterns(): else: out = matmul - # TODO(masahi): Support more activations - if "silu" in name: - out = is_op("relax.nn.silu")(out) + if "gelu" in name: + out = is_op("relax.nn.gelu")(out) return name, out, annotations, _check_decode_matmul return [ _decode_matmul_pattern("cutlass.decode_matmul"), _decode_matmul_pattern("cutlass.decode_matmul_bias"), - _decode_matmul_pattern("cutlass.decode_matmul_silu"), + _decode_matmul_pattern("cutlass.decode_matmul_bias_gelu"), ] diff --git a/src/runtime/contrib/cutlass/weight_preprocess.cc b/src/runtime/contrib/cutlass/weight_preprocess.cc index 902926ea1b..ef80627cc7 100644 --- a/src/runtime/contrib/cutlass/weight_preprocess.cc +++ b/src/runtime/contrib/cutlass/weight_preprocess.cc @@ -35,8 +35,8 @@ namespace runtime { // black box. // // The preprocessing functions are defined in C++, so we need to copy the input weight to CPU. -TVM_REGISTER_GLOBAL("cutlass.ft_preprocess_weight_int4") - .set_body_typed([](NDArray packed_weight, int sm) { +TVM_REGISTER_GLOBAL("cutlass.ft_preprocess_weight") + .set_body_typed([](NDArray packed_weight, int sm, bool is_int4) { int rows = packed_weight->shape[0]; int cols = packed_weight->shape[1]; std::vector<int8_t> input_cpu(rows * cols); @@ -44,8 +44,11 @@ TVM_REGISTER_GLOBAL("cutlass.ft_preprocess_weight_int4") packed_weight.CopyToBytes(input_cpu.data(), input_cpu.size()); // multiply cols by 2 since the "col" params in preprocess_weights refers to the column of // the unpacked weight. - fastertransformer::preprocess_weights(output_cpu.data(), input_cpu.data(), rows, cols * 2, - true /*is_int4*/, sm); + if (is_int4) { + cols *= 2; + } + fastertransformer::preprocess_weights(output_cpu.data(), input_cpu.data(), rows, cols, + is_int4, sm); auto out = NDArray::Empty(packed_weight.Shape(), packed_weight->dtype, packed_weight->device); out.CopyFromBytes(output_cpu.data(), output_cpu.size()); return out; diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py index 1528141e4a..02f15ad3d7 100644 --- a/tests/python/relax/test_codegen_cutlass.py +++ b/tests/python/relax/test_codegen_cutlass.py @@ -1258,6 +1258,25 @@ def test_attention_rewrite_fp16(): tvm.ir.assert_structural_equal(mod, Expected) +def split_transform_deploy_mod(mod): + mod_transform = tvm.IRModule() + mod_deploy = tvm.IRModule().with_attrs(mod.attrs) + + transform_func_name = None + + for gv, func in mod.functions.items(): + if "transform_params" in gv.name_hint: + transform_func_name = gv.name_hint + mod_transform[gv] = func + elif isinstance(func, tvm.tir.PrimFunc): + mod_transform[gv] = func + else: + mod_deploy[gv] = func + + assert transform_func_name is not None + return mod_transform, mod_deploy, transform_func_name + + def test_fp16A_int4B_gemm(): @I.ir_module class Module: @@ -1378,9 +1397,10 @@ def test_fp16A_int4B_gemm(): ) lv1 = lv[0] lv2 = R.call_pure_packed( - "cutlass.ft_preprocess_weight_int4", + "cutlass.ft_preprocess_weight", lv1, 80, + True, sinfo_args=(R.Tensor((64, 64), dtype="int8"),), ) lv3: R.Tensor((128,), dtype="float16") = lv[1] @@ -1409,9 +1429,10 @@ def test_fp16A_int4B_gemm(): ) lv1 = lv[0] lv2 = R.call_pure_packed( - "cutlass.ft_preprocess_weight_int4", + "cutlass.ft_preprocess_weight", lv1, 80, + True, sinfo_args=(R.Tensor((64, 64), dtype="int8"),), ) lv3: R.Tensor((128,), dtype="float16") = lv[1] @@ -1424,24 +1445,6 @@ def test_fp16A_int4B_gemm(): R.output(lv3_1) return lv3_1 - def split_transform_deploy_mod(mod): - mod_transform = tvm.IRModule() - mod_deploy = tvm.IRModule().with_attrs(mod.attrs) - - transform_func_name = None - - for gv, func in mod.functions.items(): - if "transform_params" in gv.name_hint: - transform_func_name = gv.name_hint - mod_transform[gv] = func - elif isinstance(func, tvm.tir.PrimFunc): - mod_transform[gv] = func - else: - mod_deploy[gv] = func - - assert transform_func_name is not None - return mod_transform, mod_deploy, transform_func_name - x_shape = (64, 64) y_shape = (128, 64) @@ -1496,6 +1499,153 @@ def test_fp16A_int4B_gemm(): tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) +def test_fp16A_int8B_gemm(): + @I.ir_module + class Module: + @T.prim_func + def decode( + A: T.Buffer((T.int64(64), T.int64(64)), "int8"), + B: T.Buffer((T.int64(64),), "float16"), + decode_1: T.Buffer((T.int64(64), T.int64(64)), "float16"), + ): + T.func_attr({"tir.noalias": T.bool(True)}) + # with T.block("root"): + for i, j in T.grid(T.int64(64), T.int64(64)): + with T.block("decode"): + v_i, v_j = T.axis.remap("SS", [i, j]) + T.reads(A[v_i, v_j], B[v_j]) + T.writes(decode_1[v_i, v_j]) + decode_1[v_i, v_j] = T.Cast("float16", A[v_i, v_j]) * B[v_j] + + @T.prim_func + def encode( + A: T.Buffer((T.int64(64), T.int64(64)), "float16"), + w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"), + compute: T.Buffer((T.int64(64),), "float16"), + ): + T.func_attr({"tir.noalias": T.bool(True)}) + # with T.block("root"): + max_abs_value = T.alloc_buffer((T.int64(64),), "float16") + scale = T.alloc_buffer((T.int64(64),)) + for i, k in T.grid(T.int64(64), T.int64(64)): + with T.block("max_abs_value"): + v_i, v_k = T.axis.remap("SR", [i, k]) + T.reads(A[v_i, v_k]) + T.writes(max_abs_value[v_i]) + with T.init(): + max_abs_value[v_i] = T.float16(-65504) + max_abs_value[v_i] = T.max(max_abs_value[v_i], T.fabs(A[v_i, v_k])) + for i in range(T.int64(64)): + with T.block("scale"): + v_i = T.axis.spatial(T.int64(64), i) + T.reads(max_abs_value[v_i]) + T.writes(scale[v_i]) + scale[v_i] = T.max( + T.Cast("float32", max_abs_value[v_i]), T.float32(0.0001) + ) * T.float32(0.0078125) + for j, i in T.grid(T.int64(64), T.int64(64)): + with T.block("w_gathered"): + v_j, v_i = T.axis.remap("SS", [j, i]) + T.reads(A[v_i, v_j], scale[v_i]) + T.writes(w_gathered[v_j, v_i]) + w_gathered[v_j, v_i] = T.Cast( + "int8", + T.min( + T.max( + T.round(T.Cast("float32", A[v_i, v_j]) / scale[v_i]), + T.float32(-128), + ), + T.float32(127), + ), + ) + for i0 in range(T.int64(64)): + with T.block("compute"): + v_i0 = T.axis.spatial(T.int64(64), i0) + T.reads(scale[v_i0]) + T.writes(compute[v_i0]) + compute[v_i0] = T.Cast("float16", scale[v_i0]) + + @R.function + def main( + x: R.Tensor((64, 64), dtype="float16"), + y: R.Tensor((64, 64), dtype="float16"), + bias: R.Tensor((64, 64), dtype="float16"), + ) -> R.Tensor((64, 64), dtype="float16"): + R.func_attr({"num_input": 1}) + cls = Module + with R.dataflow(): + lv = R.call_tir( + cls.encode, + (y,), + out_sinfo=[R.Tensor((64, 64), dtype="int8"), R.Tensor((64,), dtype="float16")], + ) + lv1: R.Tensor((64, 64), dtype="int8") = lv[0] + lv2: R.Tensor((64, 64), dtype="int8") = R.call_pure_packed( + "cutlass.ft_preprocess_weight", + lv1, + R.prim_value(80), + R.prim_value(0), + sinfo_args=(R.Tensor((64, 64), dtype="int8"),), + ) + lv3: R.Tensor((64,), dtype="float16") = lv[1] + lv4: R.Tensor((64, 64), dtype="int8") = R.builtin.stop_lift_params(lv2) + lv5: R.Tensor((64,), dtype="float16") = R.builtin.stop_lift_params(lv3) + lv6 = R.call_tir( + cls.decode, (lv4, lv5), out_sinfo=R.Tensor((64, 64), dtype="float16") + ) + lv1_1: R.Tensor((64, 64), dtype="float16") = R.matmul(x, lv6, out_dtype="float16") + lv2_1: R.Tensor((64, 128), dtype="float16") = R.add(lv1_1, bias) + lv2_2: R.Tensor((64, 128), dtype="float16") = R.nn.gelu(lv2_1) + R.output(lv2_2) + return lv2_2 + + x_shape = (64, 64) + y_shape = (64, 64) + + mod = partition_for_cutlass(Module) + func_names = [name.name_hint for (name, _) in mod.functions.items()] + assert "fused_decode_relax_matmul_relax_add_relax_nn_gelu_cutlass" in func_names + + mod = relax.transform.RunCodegen( + {"cutlass": {"sm": 80, "find_first_valid": False}}, + )(mod) + + x = np.random.randn(*x_shape).astype("float16") + y = np.random.normal(0, 0.002, size=y_shape).astype("float16") + bias = np.random.randn(x_shape[0], y_shape[0]).astype("float16") + + mod = relax.pipeline.get_pipeline()(mod) + mod = relax.transform.LiftTransformParams()(mod) + + mod_transform, mod_deploy, transform_func_name = split_transform_deploy_mod(mod) + + ex = relax.build(mod_transform, target="llvm") + vm = relax.vm.VirtualMachine(ex, tvm.cpu(0)) + + packed_weight, scales, bias_trans = vm[transform_func_name]( + (tvm.nd.array(y), tvm.nd.array(bias)) + ) + + dev = tvm.device("cuda", 0) + ex = relax.build(mod_deploy, target="cuda") + vm = relax.vm.VirtualMachine(ex, dev) + + x_nd = tvm.nd.array(x, dev) + params = (packed_weight.copyto(dev), scales.copyto(dev), bias_trans.copyto(dev)) + inp = [x_nd, params] + out = vm["main"](*inp).numpy() + + def gelu_fp16(x): + erf_inp = x * (0.5**0.5) + from scipy.special import erf + + erf_out = erf(erf_inp.astype("float32")).astype("float16") + return x * 0.5 * (1.0 + erf_out) + + ref = gelu_fp16(np.dot(x, y.transpose()) + bias) + tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) + + def test_rms_norm(): @I.ir_module class Module:
