From e92fd4df33b857a610aab2cf693688929cd24c27 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sun, 9 Aug 2026 22:35:46 +0800 Subject: [PATCH] fix(rms_norm): support optional weight --- scripts/generate_wrappers.py | 44 ++++++- src/base/rms_norm.h | 12 +- src/linked/torch/ops/rms_norm.h | 34 ++++-- src/native/cambricon/ops/rms_norm/kernel.mlu | 60 ++++++---- src/native/cambricon/ops/rms_norm/rms_norm.h | 15 ++- src/native/cpu/ops/rms_norm/rms_norm.h | 18 +-- src/native/cuda/ops/rms_norm/kernel.cuh | 8 +- src/native/cuda/ops/rms_norm/kernel.h | 18 ++- src/ninetoothed/ops/rms_norm/rms_norm.h | 22 +++- tests/test_generate_wrappers.py | 120 +++++++++++++++++++ tests/test_rms_norm.py | 43 ++++++- 11 files changed, 321 insertions(+), 73 deletions(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index c4c0ee1a8..001723d3a 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -931,6 +931,8 @@ def _overload_order_key(node): def _generate_legacy_c(operator, paths): op_type = _op_cpp_type(operator.name) symbol_name = _op_symbol_name(operator.name) + optional_tensor_params = _find_optional_tensor_params(operator.name) + optional_non_tensor_params = _find_optional_non_tensor_params(operator.name) def _generate_source(operator): impl_includes = "\n".join( @@ -1060,6 +1062,25 @@ def _generate_call_func_decl(operator): def _generate_destroy_func_decl(operator): return f"infiniStatus_t infiniopDestroy{symbol_name}Descriptor(infiniop{symbol_name}Descriptor_t desc)" + def _is_optional_tensor(arg): + spelling = _strip_top_level_const(arg.type.spelling) + + if spelling.startswith("std::optional<"): + return ( + "Tensor" in spelling or "TensorView" in spelling + ) and "std::vector" not in spelling + + if "Tensor" in spelling or "TensorView" in spelling: + return False + + if _is_known_non_tensor_type(spelling): + return False + + if arg.spelling in optional_non_tensor_params: + return False + + return arg.spelling in optional_tensor_params + def _generate_params(node, call=False): arguments = tuple(node.get_arguments()) @@ -1087,6 +1108,11 @@ def _handle_tensor(spelling): return spelling.replace("Tensor", "infiniopTensorDescriptor_t") + def _handle_optional_tensor(arg): + prefix = "const " if arg.type.spelling.strip().startswith("const ") else "" + tensor_type = "void *" if call else "infiniopTensorDescriptor_t" + return f"{prefix}{tensor_type}" + def _handle_std_optional(spelling): return _unwrap_std_optional(spelling) @@ -1099,7 +1125,11 @@ def _handle_data_type(spelling): return f"{prefix}infiniDtype_t" return ", ".join( - f"{_handle_data_type(_handle_std_optional(_handle_tensor(arg.type.spelling)))} {arg.spelling}" + ( + f"{_handle_optional_tensor(arg)} {arg.spelling}" + if _is_optional_tensor(arg) + else f"{_handle_data_type(_handle_std_optional(_handle_tensor(arg.type.spelling)))} {arg.spelling}" + ) for arg in arguments ) @@ -1108,7 +1138,9 @@ def _generate_arguments(node, is_data=False): f"DataTypeFromInfiniDType({arg.spelling})" if _is_data_type_spelling(arg.type.spelling) else ( - _generate_tensor_caster(arg.spelling, is_data=is_data) + _generate_optional_tensor_caster(arg.spelling, is_data=is_data) + if _is_optional_tensor(arg) + else _generate_tensor_caster(arg.spelling, is_data=is_data) if "Tensor" in arg.type.spelling else arg.spelling ) @@ -1122,6 +1154,14 @@ def _generate_tensor_caster(name, is_data=False): return f"infini::ops::Tensor{{nullptr, {name}->shape(), DataTypeFromInfiniDType({name}->dtype()), infini::ops::Device{{DeviceTypeFromInfiniDevice(handle->device), handle->device_id}}, {name}->strides()}}" + def _generate_optional_tensor_caster(name, is_data=False): + tensor = _generate_tensor_caster(name, is_data=is_data) + + return ( + f"{name} == nullptr ? std::optional{{}} : " + f"std::optional{{{tensor}}}" + ) + return _generate_source(operator), _generate_header(operator) diff --git a/src/base/rms_norm.h b/src/base/rms_norm.h index dc28f0aa1..7bfbbb5b7 100644 --- a/src/base/rms_norm.h +++ b/src/base/rms_norm.h @@ -2,6 +2,7 @@ #define INFINI_OPS_BASE_RMS_NORM_H_ #include +#include #include #include "operator.h" @@ -11,7 +12,8 @@ namespace infini::ops { class RmsNorm : public Operator { public: - RmsNorm(const Tensor input, const Tensor weight, float eps, Tensor out) + RmsNorm(const Tensor input, const std::optional weight, float eps, + Tensor out) : input_shape_{input.shape()}, out_shape_{out.shape()}, input_strides_{input.strides()}, @@ -24,14 +26,16 @@ class RmsNorm : public Operator { assert(input.dtype() == out.dtype()); } - RmsNorm(const Tensor input, const Tensor weight, Tensor out) + RmsNorm(const Tensor input, const std::optional weight, Tensor out) : RmsNorm{input, weight, 1e-6f, out} {} // TODO: Type of `eps` should be `std::optional` instead of `float`. - virtual void operator()(const Tensor input, const Tensor weight, float eps, + virtual void operator()(const Tensor input, + const std::optional weight, float eps, Tensor out) const = 0; - virtual void operator()(const Tensor input, const Tensor weight, + virtual void operator()(const Tensor input, + const std::optional weight, Tensor out) const { return operator()(input, weight, eps_, out); } diff --git a/src/linked/torch/ops/rms_norm.h b/src/linked/torch/ops/rms_norm.h index 2ae538c83..b15e4c260 100644 --- a/src/linked/torch/ops/rms_norm.h +++ b/src/linked/torch/ops/rms_norm.h @@ -1,6 +1,10 @@ #ifndef INFINI_OPS_LINKED_TORCH_OPS_RMS_NORM_H_ #define INFINI_OPS_LINKED_TORCH_OPS_RMS_NORM_H_ +#include + +#include + #include "base/rms_norm.h" #include "torch/tensor_.h" @@ -9,32 +13,42 @@ namespace infini::ops::linked::torch { template class TorchRmsNorm : public ::infini::ops::RmsNorm { public: - TorchRmsNorm(const Tensor input, const Tensor weight, float eps, Tensor out) + TorchRmsNorm(const Tensor input, const std::optional weight, + float eps, Tensor out) : ::infini::ops::RmsNorm{input, weight, eps, out}, - weight_shape_{weight.shape()}, - weight_strides_{weight.strides()}, input_type_{input.dtype()}, - weight_type_{weight.dtype()}, + weight_type_{input.dtype()}, out_type_{out.dtype()}, is_input_contiguous_{input.IsContiguous()}, - is_weight_contiguous_{weight.IsContiguous()}, is_out_contiguous_{out.IsContiguous()}, - device_index_{out.device().index()} {} + device_index_{out.device().index()} { + TORCH_CHECK(weight.has_value(), + "Linked `RmsNorm` does not support `weight=None`"); + weight_shape_ = weight->shape(); + weight_strides_ = weight->strides(); + weight_type_ = weight->dtype(); + is_weight_contiguous_ = weight->IsContiguous(); + } - TorchRmsNorm(const Tensor input, const Tensor weight, Tensor out) + TorchRmsNorm(const Tensor input, const std::optional weight, + Tensor out) : TorchRmsNorm{input, weight, 1e-6f, out} {} using ::infini::ops::RmsNorm::operator(); - void operator()(const Tensor input, const Tensor weight, float eps, - Tensor out) const override { + void operator()(const Tensor input, const std::optional weight, + float eps, Tensor out) const override { + TORCH_CHECK(weight.has_value(), + "Linked `RmsNorm` does not support `weight=None`"); + const Tensor& affine_weight = *weight; + const typename Backend::StreamGuard stream_guard{ Backend::GetStreamFromExternal(stream_, device_index_)}; auto at_input = ToAtenTensor( const_cast(input.data()), input_shape_, input_strides_, input_type_, device_index_); auto at_weight = ToAtenTensor( - const_cast(weight.data()), weight_shape_, weight_strides_, + const_cast(affine_weight.data()), weight_shape_, weight_strides_, weight_type_, device_index_); auto at_out = ToAtenTensor( out.data(), out_shape_, out_strides_, out_type_, device_index_); diff --git a/src/native/cambricon/ops/rms_norm/kernel.mlu b/src/native/cambricon/ops/rms_norm/kernel.mlu index b4d7e8d8a..9c5309b20 100644 --- a/src/native/cambricon/ops/rms_norm/kernel.mlu +++ b/src/native/cambricon/ops/rms_norm/kernel.mlu @@ -134,7 +134,9 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, if (vector_size <= max_batch_size) { __memcpy(input_cache, input + input_offset, vector_size * sizeof(T), GDRAM2NRAM); - __memcpy(weight_cache, weight, vector_size * sizeof(TW), GDRAM2NRAM); + if (weight != nullptr) { + __memcpy(weight_cache, weight, vector_size * sizeof(TW), GDRAM2NRAM); + } if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), @@ -146,18 +148,20 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, NRAM2NRAM); } - if constexpr (std::is_same::value) { - __bang_half2float(weight_float_buffer, - reinterpret_cast(weight_cache), vector_size); - } else if constexpr (std::is_same::value) { - __bang_bfloat162float(weight_float_buffer, weight_cache, vector_size); - } else { - __memcpy(weight_float_buffer, weight_cache, vector_size * sizeof(float), - NRAM2NRAM); + if (weight != nullptr) { + if constexpr (std::is_same::value) { + __bang_half2float(weight_float_buffer, + reinterpret_cast(weight_cache), vector_size); + } else if constexpr (std::is_same::value) { + __bang_bfloat162float(weight_float_buffer, weight_cache, vector_size); + } else { + __memcpy(weight_float_buffer, weight_cache, + vector_size * sizeof(float), NRAM2NRAM); + } + __bang_mul(float_buffer, float_buffer, weight_float_buffer, + vector_size); } - // Multiply by weight and apply normalization. - __bang_mul(float_buffer, float_buffer, weight_float_buffer, vector_size); __bang_mul_scalar(float_buffer, float_buffer, inv_rms, vector_size); if constexpr (std::is_same::value) { @@ -179,13 +183,15 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, size_t current_batch = std::min((size_t)max_batch_size, vector_size - processed_elements); - // Load input and weight data. + // Load input and optional weight data. __memcpy(input_cache, input + input_offset + processed_elements * input_strides[num_dims - 1], current_batch * sizeof(T), GDRAM2NRAM); - __memcpy(weight_cache, weight + processed_elements, - current_batch * sizeof(TW), GDRAM2NRAM); + if (weight != nullptr) { + __memcpy(weight_cache, weight + processed_elements, + current_batch * sizeof(TW), GDRAM2NRAM); + } if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), @@ -197,20 +203,22 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, NRAM2NRAM); } - if constexpr (std::is_same::value) { - __bang_half2float(weight_float_buffer, - reinterpret_cast(weight_cache), - current_batch); - } else if constexpr (std::is_same::value) { - __bang_bfloat162float(weight_float_buffer, weight_cache, - current_batch); - } else { - __memcpy(weight_float_buffer, weight_cache, - current_batch * sizeof(float), NRAM2NRAM); + if (weight != nullptr) { + if constexpr (std::is_same::value) { + __bang_half2float(weight_float_buffer, + reinterpret_cast(weight_cache), + current_batch); + } else if constexpr (std::is_same::value) { + __bang_bfloat162float(weight_float_buffer, weight_cache, + current_batch); + } else { + __memcpy(weight_float_buffer, weight_cache, + current_batch * sizeof(float), NRAM2NRAM); + } + __bang_mul(float_buffer, float_buffer, weight_float_buffer, + current_batch); } - __bang_mul(float_buffer, float_buffer, weight_float_buffer, - current_batch); __bang_mul_scalar(float_buffer, float_buffer, inv_rms, current_batch); if constexpr (std::is_same::value) { diff --git a/src/native/cambricon/ops/rms_norm/rms_norm.h b/src/native/cambricon/ops/rms_norm/rms_norm.h index 6a9aed098..815ef497d 100644 --- a/src/native/cambricon/ops/rms_norm/rms_norm.h +++ b/src/native/cambricon/ops/rms_norm/rms_norm.h @@ -3,6 +3,7 @@ #include #include +#include #include #include "base/rms_norm.h" @@ -21,15 +22,16 @@ void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, template <> class Operator : public RmsNorm { public: - Operator(const Tensor input, const Tensor weight, float eps, Tensor out) + Operator(const Tensor input, const std::optional weight, float eps, + Tensor out) : RmsNorm{input, weight, eps, out} { cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster, &cluster_count); cnrtMalloc(&default_workspace_, workspace_size_in_bytes()); } - void operator()(const Tensor input, const Tensor weight, float eps, - Tensor out) const override { + void operator()(const Tensor input, const std::optional weight, + float eps, Tensor out) const override { auto queue = static_cast(stream_ ? stream_ : 0); auto workspace{workspace_ ? workspace_ : default_workspace_}; @@ -37,15 +39,16 @@ class Operator : public RmsNorm { Device::Type::kCambricon, List, List>( - {input.dtype(), weight.dtype()}, + {input.dtype(), weight.has_value() ? weight->dtype() : input.dtype()}, [&](auto input_tag, auto weight_tag) { using InputT = typename decltype(input_tag)::type; using WeightT = typename decltype(weight_tag)::type; RmsNormUnion( workspace, core_per_cluster, cluster_count, queue, out.data(), - input.data(), weight.data(), out_shape_.data(), - out_strides_.data(), input_strides_.data(), eps, ndim_); + input.data(), weight.has_value() ? weight->data() : nullptr, + out_shape_.data(), out_strides_.data(), input_strides_.data(), + eps, ndim_); }, "CambriconRmsNorm::operator() - output dispatch"); } diff --git a/src/native/cpu/ops/rms_norm/rms_norm.h b/src/native/cpu/ops/rms_norm/rms_norm.h index c3f091cb4..bb1506697 100644 --- a/src/native/cpu/ops/rms_norm/rms_norm.h +++ b/src/native/cpu/ops/rms_norm/rms_norm.h @@ -17,8 +17,8 @@ class Operator : public RmsNorm, public: using RmsNorm::RmsNorm; - void operator()(const Tensor input, const Tensor weight, float eps, - Tensor out) const override { + void operator()(const Tensor input, const std::optional weight, + float eps, Tensor out) const override { DispatchFunc( out.dtype(), [&](auto tag) { @@ -30,11 +30,12 @@ class Operator : public RmsNorm, private: template - void Compute(const Tensor input, const Tensor weight, float eps, - Tensor out) const { + void Compute(const Tensor input, const std::optional weight, + float eps, Tensor out) const { auto* out_ptr = static_cast(out.data()); const auto* input_ptr = static_cast(input.data()); - const auto* weight_ptr = static_cast(weight.data()); + const auto* weight_ptr = + weight.has_value() ? static_cast(weight->data()) : nullptr; auto stride_input_batch = input_strides_.size() > 1 ? input_strides_[0] : 0; auto stride_input_nhead = @@ -57,8 +58,11 @@ class Operator : public RmsNorm, float rms = 1.f / std::sqrt(ss / static_cast(dim_) + eps); for (Tensor::Size k = 0; k < dim_; ++k) { - out_row[k] = Cast(Cast(input_row[k]) * - Cast(weight_ptr[k]) * rms); + float value = Cast(input_row[k]) * rms; + if (weight_ptr != nullptr) { + value *= Cast(weight_ptr[k]); + } + out_row[k] = Cast(value); } } } diff --git a/src/native/cuda/ops/rms_norm/kernel.cuh b/src/native/cuda/ops/rms_norm/kernel.cuh index 91036a37f..f76ec25da 100644 --- a/src/native/cuda/ops/rms_norm/kernel.cuh +++ b/src/native/cuda/ops/rms_norm/kernel.cuh @@ -53,9 +53,11 @@ __global__ void RmsNormKernel(TData* __restrict__ y, int64_t stride_y_batch, __syncthreads(); for (size_t i = threadIdx.x; i < dim; i += block_size) { - y_ptr[i] = Caster::template Cast( - Caster::template Cast(x_ptr[i]) * - Caster::template Cast(w_ptr[i]) * rms); + TCompute value = Caster::template Cast(x_ptr[i]) * rms; + if (w_ptr != nullptr) { + value *= Caster::template Cast(w_ptr[i]); + } + y_ptr[i] = Caster::template Cast(value); } } diff --git a/src/native/cuda/ops/rms_norm/kernel.h b/src/native/cuda/ops/rms_norm/kernel.h index 0cd3be915..bb95476c4 100644 --- a/src/native/cuda/ops/rms_norm/kernel.h +++ b/src/native/cuda/ops/rms_norm/kernel.h @@ -3,6 +3,7 @@ #include #include +#include #include "base/rms_norm.h" #include "data_type.h" @@ -18,8 +19,8 @@ class CudaRmsNorm : public RmsNorm { public: using RmsNorm::RmsNorm; - void operator()(const Tensor input, const Tensor weight, float eps, - Tensor out) const override { + void operator()(const Tensor input, const std::optional weight, + float eps, Tensor out) const override { auto cuda_stream = static_cast(stream_ ? stream_ : 0); @@ -40,20 +41,25 @@ class CudaRmsNorm : public RmsNorm { ConcatType, ReducedFloatTypes>, AllCudaBlockSizes>( {static_cast(out.dtype()), - static_cast(weight.dtype()), block_size}, + static_cast(weight.has_value() ? weight->dtype() + : input.dtype()), + block_size}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; using TWeight = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<2>(list_tag); + auto weight_data = + weight.has_value() + ? reinterpret_cast(weight->data()) + : nullptr; RmsNormKernel <<>>( reinterpret_cast(out.data()), stride_out_batch, stride_out_nhead, reinterpret_cast(input.data()), - stride_input_batch, stride_input_nhead, - reinterpret_cast(weight.data()), nhead_, dim_, - eps); + stride_input_batch, stride_input_nhead, weight_data, nhead_, + dim_, eps); }, "CudaRmsNorm::operator()"); } diff --git a/src/ninetoothed/ops/rms_norm/rms_norm.h b/src/ninetoothed/ops/rms_norm/rms_norm.h index 9cd884b45..df9610867 100644 --- a/src/ninetoothed/ops/rms_norm/rms_norm.h +++ b/src/ninetoothed/ops/rms_norm/rms_norm.h @@ -3,6 +3,8 @@ #include #include +#include +#include #include #include "base/rms_norm.h" @@ -18,15 +20,23 @@ class Operator : public RmsNorm { using RmsNorm::RmsNorm; using RmsNorm::operator(); - void operator()(const Tensor input, const Tensor weight, float eps, - Tensor out) const override { - assert(input.dtype() == out.dtype() && out.dtype() == weight.dtype() && + void operator()(const Tensor input, const std::optional weight, + float eps, Tensor out) const override { + if (!weight.has_value()) { + assert(false && "NineToothed `RmsNorm` does not support `weight=None`"); + std::abort(); + } + const Tensor& affine_weight = *weight; + + assert(input.dtype() == out.dtype() && + out.dtype() == affine_weight.dtype() && "operator `RmsNorm` requires all input and output tensors to have " "the same dtype"); assert(input.shape() == out.shape() && "NineToothed `RmsNorm` requires input and output tensors with the " "same shape"); - assert(weight.ndim() == 1 && weight.size(-1) == out.size(-1) && + assert(affine_weight.ndim() == 1 && + affine_weight.size(-1) == out.size(-1) && "NineToothed `RmsNorm` requires a 1D weight matching the last " "dimension"); assert( @@ -45,7 +55,7 @@ class Operator : public RmsNorm { weight_sizes.assign(out.shape().begin(), out.shape().end()); weight_strides.assign(out.ndim(), 0); weight_strides.back() = - weight.strides().empty() ? 1 : weight.strides().back(); + affine_weight.strides().empty() ? 1 : affine_weight.strides().back(); const int dtype_index = ninetoothed::DataTypeIndex(out.dtype()); assert( @@ -53,7 +63,7 @@ class Operator : public RmsNorm { "NineToothed `RmsNorm` supports only float16, bfloat16, and float32"); ninetoothed::Tensor input_tensor(input); - ninetoothed::Tensor weight_tensor(const_cast(weight.data()), + ninetoothed::Tensor weight_tensor(const_cast(affine_weight.data()), weight_sizes.data(), weight_strides.data()); ninetoothed::Tensor eps_tensor(eps_value, empty_shape, empty_strides); diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 998109de6..8663ad28f 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -156,6 +156,126 @@ class Clamp { ) in text +def test_rms_norm_optional_weight_is_preserved_across_generated_wrappers(): + module = _load_generator_module() + operator = module._parse_operator_header("rms_norm") + + binding = module._generate_pybind11(operator) + optional_conversion = "OptionalTensorFromPybind11Handle(weight)" + + assert binding.count("std::optional weight") == 6 + assert binding.count(optional_conversion) == 6 + assert "TensorFromPybind11Handle(weight)" not in binding.replace( + optional_conversion, "" + ) + + dispatch_declarations, _ = module._generate_generated_dispatch_entries(operator) + dispatch_text = "\n".join(dispatch_declarations) + + assert ( + "void CallRmsNorm(const Handle& handle, const Config& config, " + "Tensor input, std::optional weight, float eps, Tensor out);" + ) in dispatch_text + assert ( + "void CallRmsNorm(const Handle& handle, const Config& config, " + "Tensor input, std::optional weight, Tensor out);" + ) in dispatch_text + + instantiation_declarations, _ = ( + module._generate_operator_call_instantiation_entries(operator) + ) + instantiation_text = "\n".join(instantiation_declarations) + + assert ( + "Operator<::infini::ops::RmsNorm>::Call<" + "Tensor, std::optional, float, Tensor>" + ) in instantiation_text + assert ( + "Operator<::infini::ops::RmsNorm>::Call, Tensor>" + ) in instantiation_text + + legacy_source, legacy_header = module._generate_legacy_c(operator, ()) + + assert "std::optional" not in legacy_header + assert "const infiniopTensorDescriptor_t weight" in legacy_header + assert "const void * weight" in legacy_header + + nullable_weight_prefix = ( + "weight == nullptr ? std::optional{} : " + "std::optional{" + ) + assert ( + nullable_weight_prefix + "infini::ops::Tensor{nullptr, weight->shape(), " + ) in legacy_source + assert ( + nullable_weight_prefix + "infini::ops::Tensor(const_cast(weight), " + ) in legacy_source + + for line in legacy_source.splitlines(): + if "weight->shape()" in line: + assert nullable_weight_prefix in line + + +def test_legacy_c_uses_selected_overload_type_for_reused_optional_name( + monkeypatch, tmp_path +): + module = _load_generator_module() + base_header = tmp_path / "legacy_optional.h" + base_header.write_text( + """ +class LegacyOptional { + public: + virtual void operator()(const Tensor input, + const std::optional value, + Tensor out) const = 0; + virtual void operator()(const Tensor input, const Tensor value, + Tensor out) const = 0; + virtual void operator()(const Tensor input, const int64_t value, + Tensor out) const = 0; +}; +""" + ) + monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header) + + required_tensor = module._ParsedFunction( + [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("const Tensor", "value"), + module._ParsedArgument("Tensor", "out"), + ] + ) + required_scalar = module._ParsedFunction( + [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("const int64_t", "value"), + module._ParsedArgument("Tensor", "out"), + ] + ) + + tensor_operator = module._Operator( + "legacy_optional", + constructors=[required_tensor], + calls=[required_tensor], + ) + tensor_source, tensor_header = module._generate_legacy_c(tensor_operator, ()) + + assert "const infiniopTensorDescriptor_t value" in tensor_header + assert "const void * value" in tensor_header + assert "value == nullptr" not in tensor_source + assert "infini::ops::Tensor{nullptr, value->shape()" in tensor_source + + scalar_operator = module._Operator( + "legacy_optional", + constructors=[required_scalar], + calls=[required_scalar], + ) + scalar_source, scalar_header = module._generate_legacy_c(scalar_operator, ()) + + assert "const int64_t value" in scalar_header + assert "value == nullptr" not in scalar_source + assert "value->shape()" not in scalar_source + + def test_extractor_prefers_header_types_for_reused_parameter_names( monkeypatch, tmp_path ): diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 707f9a85f..bdf7bef51 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -55,7 +55,38 @@ def test_rms_norm( ) +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-4, 1e-4), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 2e-2, 1e-2), + ), +) +def test_rms_norm_without_weight(dtype, device, rtol, atol): + input = torch.randn((7, 769), dtype=dtype, device=device) + out = torch.empty_like(input) + + return Payload( + _rms_norm, + _torch_rms_norm, + (input, None), + {"eps": 1e-6, "out": out}, + rtol=rtol, + atol=atol, + ) + + def test_rms_norm_non_default_stream(device, implementation_index): + _run_rms_norm_non_default_stream(device, implementation_index, has_weight=True) + + +def test_rms_norm_without_weight_non_default_stream(device): + _run_rms_norm_non_default_stream(device, 0, has_weight=False) + + +def _run_rms_norm_non_default_stream(device, implementation_index, *, has_weight): if device == "cuda": accelerator = torch.cuda stream_attribute = "cuda_stream" @@ -69,7 +100,9 @@ def test_rms_norm_non_default_stream(device, implementation_index): pytest.skip("non-default streams require an accelerator backend") input = torch.randn((32, 128), dtype=torch.float16, device=device) - weight = torch.randn((128,), dtype=torch.float16, device=device) + weight = ( + torch.randn((128,), dtype=torch.float16, device=device) if has_weight else None + ) out = torch.zeros_like(input) expected = _torch_rms_norm(input, weight, out=torch.empty_like(out)).cpu() accelerator.synchronize() @@ -110,11 +143,15 @@ def _rms_norm(input, weight, *, eps=1e-6, out=None, implementation_index=0): def _torch_rms_norm(input, weight, *, eps=1e-6, out=None): - # Fallback for `torch<2.3`: `rms_norm = (x / sqrt(mean(x^2) + eps)) * weight`. + # Fallback for `torch<2.3`: normalize first, then apply weight when present. def _fallback(input, _normalized_shape, weight, *, eps=1e-6): rms = torch.sqrt(torch.mean(input * input, dim=-1, keepdim=True) + eps) + result = input / rms + + if weight is not None: + result = result * weight - return (input / rms) * weight + return result rms_norm_fn = getattr(torch.nn.functional, "rms_norm", _fallback)