InfiniTensor · voltjia · Aug 9, 2026
diff --git 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<infini::ops::Tensor>{{}} : "
+            f"std::optional<infini::ops::Tensor>{{{tensor}}}"
+        )
+
     return _generate_source(operator), _generate_header(operator)
 
 

diff --git a/src/base/rms_norm.h b/src/base/rms_norm.h
@@ -2,6 +2,7 @@
 #define INFINI_OPS_BASE_RMS_NORM_H_
 
 #include <cstddef>
+#include <optional>
 #include <vector>
 
 #include "operator.h"
@@ -11,7 +12,8 @@ namespace infini::ops {
 
 class RmsNorm : public Operator<RmsNorm> {
  public:
-  RmsNorm(const Tensor input, const Tensor weight, float eps, Tensor out)
+  RmsNorm(const Tensor input, const std::optional<Tensor> 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<RmsNorm> {
     assert(input.dtype() == out.dtype());
   }
 
-  RmsNorm(const Tensor input, const Tensor weight, Tensor out)
+  RmsNorm(const Tensor input, const std::optional<Tensor> weight, Tensor out)
       : RmsNorm{input, weight, 1e-6f, out} {}
 
   // TODO: Type of `eps` should be `std::optional<float>` instead of `float`.
-  virtual void operator()(const Tensor input, const Tensor weight, float eps,
+  virtual void operator()(const Tensor input,
+                          const std::optional<Tensor> 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<Tensor> 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
@@ -1,6 +1,10 @@
 #ifndef INFINI_OPS_LINKED_TORCH_OPS_RMS_NORM_H_
 #define INFINI_OPS_LINKED_TORCH_OPS_RMS_NORM_H_
 
+#include <c10/util/Exception.h>
+
+#include <optional>
+
 #include "base/rms_norm.h"
 #include "torch/tensor_.h"
 
@@ -9,32 +13,42 @@ namespace infini::ops::linked::torch {
 template <typename Backend>
 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<Tensor> 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<Tensor> 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<Tensor> 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<Backend::kDeviceType>(
         const_cast<void*>(input.data()), input_shape_, input_strides_,
         input_type_, device_index_);
     auto at_weight = ToAtenTensor<Backend::kDeviceType>(
-        const_cast<void*>(weight.data()), weight_shape_, weight_strides_,
+        const_cast<void*>(affine_weight.data()), weight_shape_, weight_strides_,
         weight_type_, device_index_);
     auto at_out = ToAtenTensor<Backend::kDeviceType>(
         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
@@ -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<T, __half>::value) {
         __bang_half2float(float_buffer, reinterpret_cast<half*>(input_cache),
@@ -146,18 +148,20 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output,
                  NRAM2NRAM);
       }
 
-      if constexpr (std::is_same<TW, __half>::value) {
-        __bang_half2float(weight_float_buffer,
-                          reinterpret_cast<half*>(weight_cache), vector_size);
-      } else if constexpr (std::is_same<TW, __bang_bfloat16>::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<TW, __half>::value) {
+          __bang_half2float(weight_float_buffer,
+                            reinterpret_cast<half*>(weight_cache), vector_size);
+        } else if constexpr (std::is_same<TW, __bang_bfloat16>::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<T, __half>::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<T, __half>::value) {
           __bang_half2float(float_buffer, reinterpret_cast<half*>(input_cache),
@@ -197,20 +203,22 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output,
                    NRAM2NRAM);
         }
 
-        if constexpr (std::is_same<TW, __half>::value) {
-          __bang_half2float(weight_float_buffer,
-                            reinterpret_cast<half*>(weight_cache),
-                            current_batch);
-        } else if constexpr (std::is_same<TW, __bang_bfloat16>::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<TW, __half>::value) {
+            __bang_half2float(weight_float_buffer,
+                              reinterpret_cast<half*>(weight_cache),
+                              current_batch);
+          } else if constexpr (std::is_same<TW, __bang_bfloat16>::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<T, __half>::value) {

diff --git 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 <cstddef>
 #include <cstdint>
+#include <optional>
 #include <vector>
 
 #include "base/rms_norm.h"
@@ -21,31 +22,33 @@ void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count,
 template <>
 class Operator<RmsNorm, Device::Type::kCambricon> : public RmsNorm {
  public:
-  Operator(const Tensor input, const Tensor weight, float eps, Tensor out)
+  Operator(const Tensor input, const std::optional<Tensor> 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<Tensor> weight,
+                  float eps, Tensor out) const override {
     auto queue = static_cast<cnrtQueue_t>(stream_ ? stream_ : 0);
     auto workspace{workspace_ ? workspace_ : default_workspace_};
 
     DispatchFunc<
         Device::Type::kCambricon,
         List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>,
         List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>>(
-        {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<InputT, WeightT>(
               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
@@ -17,8 +17,8 @@ class Operator<RmsNorm, Device::Type::kCpu> : 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<Tensor> weight,
+                  float eps, Tensor out) const override {
     DispatchFunc<Device::Type::kCpu, AllFloatTypes>(
         out.dtype(),
         [&](auto tag) {
@@ -30,11 +30,12 @@ class Operator<RmsNorm, Device::Type::kCpu> : public RmsNorm,
 
  private:
   template <typename T>
-  void Compute(const Tensor input, const Tensor weight, float eps,
-               Tensor out) const {
+  void Compute(const Tensor input, const std::optional<Tensor> weight,
+               float eps, Tensor out) const {
     auto* out_ptr = static_cast<T*>(out.data());
     const auto* input_ptr = static_cast<const T*>(input.data());
-    const auto* weight_ptr = static_cast<const T*>(weight.data());
+    const auto* weight_ptr =
+        weight.has_value() ? static_cast<const T*>(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<RmsNorm, Device::Type::kCpu> : public RmsNorm,
         float rms = 1.f / std::sqrt(ss / static_cast<float>(dim_) + eps);
 
         for (Tensor::Size k = 0; k < dim_; ++k) {
-          out_row[k] = Cast<T>(Cast<float>(input_row[k]) *
-                               Cast<float>(weight_ptr[k]) * rms);
+          float value = Cast<float>(input_row[k]) * rms;
+          if (weight_ptr != nullptr) {
+            value *= Cast<float>(weight_ptr[k]);
+          }
+          out_row[k] = Cast<T>(value);
         }
       }
     }

diff --git 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<kDev>::template Cast<TData>(
-        Caster<kDev>::template Cast<TCompute>(x_ptr[i]) *
-        Caster<kDev>::template Cast<TCompute>(w_ptr[i]) * rms);
+    TCompute value = Caster<kDev>::template Cast<TCompute>(x_ptr[i]) * rms;
+    if (w_ptr != nullptr) {
+      value *= Caster<kDev>::template Cast<TCompute>(w_ptr[i]);
+    }
+    y_ptr[i] = Caster<kDev>::template Cast<TData>(value);
   }
 }