ignis-dl 0.0.1

data/lib/nnw/ai/loss.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    # Loss functions for training.
+    # Each returns a scalar Tensor with autograd support.
+    module Loss
+      class << self
+        # Cross-entropy loss (classification, language modeling).
+        # Fused log-softmax + NLL for numerical stability.
+        # @param logits [Tensor] model output [batch_size, vocab_size]
+        # @param targets [Tensor] target indices [batch_size] (int32)
+        # @param label_smoothing [Float] label smoothing factor (0.0 = none)
+        # @return [Tensor] scalar loss
+        def cross_entropy(logits, targets, label_smoothing: 0.0)
+          batch_size = logits.shape[0]
+          vocab_size = logits.shape[-1]
+
+          # Allocate outputs
+          losses_nv = Ignis::Shared::NvArray.new(shape: [batch_size], dtype: :float32,
+                                                device_id: logits.device_id)
+          losses_nv.from_host(Array.new(batch_size, 0.0))
+
+          log_softmax_nv = Ignis::Shared::NvArray.new(shape: logits.shape, dtype: :float32,
+                                                      device_id: logits.device_id)
+          log_softmax_nv.from_host(Array.new(logits.numel, 0.0))
+
+          # Forward kernel
+          kernel = Ignis::JIT::Kernels::Loss.cross_entropy_forward
+          kernel.launch(grid: [(batch_size + 255) / 256], block: [256],
+                        args: [logits.data, targets.data, losses_nv, log_softmax_nv,
+                               batch_size, vocab_size, label_smoothing.to_f])
+
+          # Mean reduction
+          mean_nv = Ignis::Shared::NvArray.new(shape: [1], dtype: :float32, device_id: logits.device_id)
+          mean_nv.from_host([0.0])
+          mean_k = Ignis::JIT::Kernels::Loss.mean_reduce
+          mean_k.launch(grid: [1], block: [1], args: [losses_nv, mean_nv, batch_size])
+
+          result = Tensor.new(data: mean_nv, requires_grad: logits.requires_grad, is_leaf: false)
+
+          if logits.requires_grad
+            saved_lsm = log_softmax_nv
+            saved_targets = targets.data
+            Tape.record(result, inputs: [logits]) do |grad|
+              # Scale by the MEAN reduction (1/batch_size) AND the chained upstream
+              # gradient. The upstream grad was previously ignored (hardcoded to
+              # 1/batch_size), so any downstream loss scaling — e.g. the Trainer's
+              # gradient-accumulation division, or loss*k — was silently dropped,
+              # making effective gradients wrong. grad is the scalar [1] cotangent
+              # of this scalar loss.
+              upstream = grad.to_host[0].to_f
+              grad_scale = Ignis::Shared::NvArray.new(shape: [batch_size], dtype: :float32,
+                                                     device_id: logits.device_id)
+              scale_val = upstream / batch_size
+              grad_scale.from_host(Array.new(batch_size, scale_val))
+
+              grad_logits = Ignis::Shared::NvArray.new(shape: logits.shape, dtype: :float32,
+                                                      device_id: logits.device_id)
+              grad_logits.from_host(Array.new(logits.numel, 0.0))
+
+              bk = Ignis::JIT::Kernels::Loss.cross_entropy_backward
+              total = batch_size * vocab_size
+              bk.launch(grid: [(total + 255) / 256], block: [256],
+                        args: [saved_lsm, saved_targets, grad_scale,
+                               grad_logits, batch_size, vocab_size, label_smoothing.to_f])
+              [grad_logits]
+            end
+          end
+
+          result
+        end
+
+        # Mean squared error loss.
+        # @param predictions [Tensor]
+        # @param targets [Tensor]
+        # @return [Tensor] scalar loss
+        def mse(predictions, targets)
+          n = predictions.numel
+
+          losses_nv = Ignis::Shared::NvArray.new(shape: predictions.shape, dtype: :float32,
+                                                device_id: predictions.device_id)
+          losses_nv.from_host(Array.new(n, 0.0))
+
+          kernel = Ignis::JIT::Kernels::Loss.mse_forward
+          kernel.launch(grid: [(n + 255) / 256], block: [256],
+                        args: [predictions.data, targets.data, losses_nv, n])
+
+          # Mean
+          mean_nv = Ignis::Shared::NvArray.new(shape: [1], dtype: :float32,
+                                              device_id: predictions.device_id)
+          mean_nv.from_host([0.0])
+          mean_k = Ignis::JIT::Kernels::Loss.mean_reduce
+          mean_k.launch(grid: [1], block: [1], args: [losses_nv, mean_nv, n])
+
+          result = Tensor.new(data: mean_nv, requires_grad: predictions.requires_grad, is_leaf: false)
+
+          if predictions.requires_grad
+            saved_pred = predictions.data
+            saved_tgt = targets.data
+            Tape.record(result, inputs: [predictions]) do |grad|
+              grad_input = Ignis::Shared::NvArray.new(shape: predictions.shape, dtype: :float32,
+                                                     device_id: predictions.device_id)
+              grad_input.from_host(Array.new(n, 0.0))
+              bk = Ignis::JIT::Kernels::Loss.mse_backward
+              scale = 1.0 / n
+              bk.launch(grid: [(n + 255) / 256], block: [256],
+                        args: [saved_pred, saved_tgt, grad, grad_input, n, scale.to_f])
+              [grad_input]
+            end
+          end
+
+          result
+        end
+
+        # Binary cross-entropy with logits (sigmoid applied inside).
+        # @param logits [Tensor]
+        # @param targets [Tensor] (0.0 or 1.0)
+        # @return [Tensor] scalar loss
+        def binary_cross_entropy(logits, targets)
+          n = logits.numel
+
+          losses_nv = Ignis::Shared::NvArray.new(shape: logits.shape, dtype: :float32,
+                                                device_id: logits.device_id)
+          losses_nv.from_host(Array.new(n, 0.0))
+
+          kernel = Ignis::JIT::Kernels::Loss.bce_forward
+          kernel.launch(grid: [(n + 255) / 256], block: [256],
+                        args: [logits.data, targets.data, losses_nv, n])
+
+          mean_nv = Ignis::Shared::NvArray.new(shape: [1], dtype: :float32,
+                                              device_id: logits.device_id)
+          mean_nv.from_host([0.0])
+          mean_k = Ignis::JIT::Kernels::Loss.mean_reduce
+          mean_k.launch(grid: [1], block: [1], args: [losses_nv, mean_nv, n])
+
+          result = Tensor.new(data: mean_nv, requires_grad: logits.requires_grad, is_leaf: false)
+
+          if logits.requires_grad
+            saved_logits = logits.data
+            saved_targets = targets.data
+            Tape.record(result, inputs: [logits]) do |grad|
+              grad_input = Ignis::Shared::NvArray.new(shape: logits.shape, dtype: :float32,
+                                                     device_id: logits.device_id)
+              grad_input.from_host(Array.new(n, 0.0))
+              bk = Ignis::JIT::Kernels::Loss.bce_backward
+              bk.launch(grid: [(n + 255) / 256], block: [256],
+                        args: [saved_logits, saved_targets, grad, grad_input, n])
+              [grad_input]
+            end
+          end
+
+          result
+        end
+
+        # KL divergence: KL(p || q) = sum(p * log(p/q))
+        # @param log_q [Tensor] log probabilities of model
+        # @param p [Tensor] target distribution
+        # @return [Tensor] scalar loss
+        def kl_divergence(log_q, p)
+          # KL = sum(p * (log(p) - log_q))
+          # Implement via existing tensor ops
+          diff = p * (p.relu + Tensor.from_host([1e-8], shape: [1], device_id: p.device_id)) - log_q
+          loss = (p * diff).sum
+          loss
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/dropout.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Dropout: randomly zeros elements during training.
+      # In eval mode, passes input through unchanged.
+      # Uses inverted dropout: scale by 1/(1-p) during training.
+      class Dropout < Module
+        # @param p [Float] probability of zeroing an element (0.0 to 1.0)
+        def initialize(p: 0.1)
+          super()
+          @p = p
+        end
+
+        # Forward pass.
+        # In training mode: apply random mask and scale by 1/(1-p).
+        # In eval mode: pass through.
+        # @param x [Tensor]
+        # @return [Tensor]
+        def forward(x)
+          return x unless @training
+          return x if @p == 0.0
+
+          # Generate bernoulli mask on host, transfer to GPU
+          n = x.numel
+          scale = 1.0 / (1.0 - @p)
+          mask_data = Array.new(n) { Kernel.rand >= @p ? scale : 0.0 }
+
+          mask_nv = Ignis::Shared::NvArray.new(shape: x.shape, dtype: x.dtype, device_id: x.device_id)
+          mask_nv.from_host(mask_data)
+
+          # Elementwise multiply
+          result_nv = Ignis::Shared::NvArray.new(shape: x.shape, dtype: x.dtype, device_id: x.device_id)
+          result_nv.from_host(Array.new(n, 0.0))
+
+          kernel = Ignis::JIT::Kernels::Elementwise.mul_forward
+          kernel.launch(grid: [(n + 255) / 256], block: [256],
+                        args: [x.data, mask_nv, result_nv, n])
+
+          result = Tensor.new(data: result_nv, requires_grad: x.requires_grad, is_leaf: false)
+
+          if x.requires_grad
+            Tape.record(result, inputs: [x]) do |grad|
+              # Backward: multiply grad by same mask
+              grad_in = Ignis::Shared::NvArray.new(shape: grad.shape, dtype: grad.dtype,
+                                                   device_id: grad.device_id)
+              grad_in.from_host(Array.new(grad.numel, 0.0))
+              mk = Ignis::JIT::Kernels::Elementwise.mul_forward
+              mk.launch(grid: [(grad.numel + 255) / 256], block: [256],
+                        args: [grad, mask_nv, grad_in, grad.numel])
+              [grad_in]
+            end
+          end
+
+          result
+        end
+
+        # @return [String]
+        def to_s
+          "Dropout(p=#{@p})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/embedding.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Embedding layer: maps integer indices to dense vectors.
+      # Forward uses gather_rows JIT kernel.
+      # Backward uses scatter_add with atomicAdd.
+      class Embedding < Module
+        # @return [Tensor] weight matrix [num_embeddings, embedding_dim]
+        attr_reader :weight
+
+        # @param num_embeddings [Integer] vocabulary size
+        # @param embedding_dim [Integer] dimension of embedding vectors
+        # @param device_id [Integer]
+        def initialize(num_embeddings, embedding_dim, device_id: 0)
+          super()
+          @num_embeddings = num_embeddings
+          @embedding_dim = embedding_dim
+
+          # Initialize uniform[-scale, scale] (scale = 1/sqrt(embedding_dim)) via a
+          # device kernel. The old host Array.new(num_embeddings*embedding_dim) was a
+          # 262M-element / ~10GB array for a 128k-vocab model — infeasible. The kaiming
+          # uniform kernel with bound=scale produces the same [-scale, scale] range.
+          scale = 1.0 / Math.sqrt(embedding_dim)
+          weight_nv = Ignis::Shared::NvArray.new(shape: [num_embeddings, embedding_dim],
+                                                dtype: :float32, device_id: device_id)
+          weight_nv.to_device
+          n = num_embeddings * embedding_dim
+          init_kernel = Ignis::JIT::Kernels::Elementwise.kaiming_uniform_init
+          init_kernel.launch(grid: [(n + 255) / 256], block: [256],
+                             args: [weight_nv, scale.to_f, Ignis::JIT::Kernel::U64.new(::Random.new.rand(2**64)), n])
+
+          @weight = register_parameter("weight",
+                     Tensor.new(data: weight_nv, requires_grad: true))
+        end
+
+        # Forward pass: gather rows from weight table.
+        # @param indices [Tensor] integer indices [batch_size, seq_len] (int32 on GPU)
+        # @return [Tensor] embeddings [batch_size, seq_len, embedding_dim]
+        def forward(indices)
+          num_indices = indices.numel
+          output_shape = indices.shape + [@embedding_dim]
+          output_nv = Ignis::Shared::NvArray.new(shape: output_shape, dtype: :float32,
+                                                device_id: @weight.device_id)
+          output_nv.from_host(Array.new(num_indices * @embedding_dim, 0.0))
+
+          kernel = Ignis::JIT::Kernels::Elementwise.gather_rows
+          total = num_indices * @embedding_dim
+          kernel.launch(grid: [(total + 255) / 256], block: [256],
+                        args: [@weight.data, indices.data, output_nv, num_indices, @embedding_dim])
+
+          result = Tensor.new(data: output_nv,
+                              requires_grad: @weight.requires_grad,
+                              is_leaf: false)
+
+          if @weight.requires_grad
+            saved_indices = indices.data
+            saved_weight = @weight
+            Tape.record(result, inputs: [@weight]) do |grad|
+              # scatter_add: accumulate gradients for each embedding index
+              grad_weight = Ignis::Shared::NvArray.new(
+                shape: [@num_embeddings, @embedding_dim],
+                dtype: :float32, device_id: @weight.device_id)
+              grad_weight.from_host(Array.new(@num_embeddings * @embedding_dim, 0.0))
+
+              scatter_k = Ignis::JIT::Kernels::Elementwise.scatter_add
+              scatter_k.launch(grid: [(total + 255) / 256], block: [256],
+                               args: [grad, saved_indices, grad_weight, num_indices, @embedding_dim])
+              [grad_weight]
+            end
+          end
+
+          result
+        end
+
+        # @return [String]
+        def to_s
+          "Embedding(num=#{@num_embeddings}, dim=#{@embedding_dim})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/layer_norm.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Layer normalization: y = gamma * (x - mean) / sqrt(var + eps) + beta
+      # Normalizes along the last dimension(s).
+      class LayerNorm < Module
+        # @return [Tensor] gamma (scale)
+        attr_reader :weight
+
+        # @return [Tensor] beta (shift)
+        attr_reader :bias
+
+        # @param normalized_shape [Integer] size of the last dimension
+        # @param eps [Float] epsilon for numerical stability
+        # @param device_id [Integer]
+        def initialize(normalized_shape, eps: 1e-5, device_id: 0)
+          super()
+          @normalized_shape = normalized_shape
+          @eps = eps
+
+          # Initialize weight (gamma) to ones
+          weight_nv = Ignis::Shared::NvArray.new(shape: [normalized_shape],
+                                                dtype: :float32, device_id: device_id)
+          weight_nv.from_host(Array.new(normalized_shape, 1.0))
+          @weight = register_parameter("weight",
+                     Tensor.new(data: weight_nv, requires_grad: true))
+
+          # Initialize bias (beta) to zeros
+          bias_nv = Ignis::Shared::NvArray.new(shape: [normalized_shape],
+                                              dtype: :float32, device_id: device_id)
+          bias_nv.from_host(Array.new(normalized_shape, 0.0))
+          @bias = register_parameter("bias",
+                   Tensor.new(data: bias_nv, requires_grad: true))
+        end
+
+        # Forward pass: applies layer normalization.
+        # @param x [Tensor] input tensor [*, normalized_shape]
+        # @return [Tensor] normalized tensor
+        def forward(x)
+          x.layer_norm(@weight, @bias, eps: @eps)
+        end
+
+        # @return [String]
+        def to_s
+          "LayerNorm(#{@normalized_shape}, eps=#{@eps})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/linear.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Linear layer: y = x @ W^T + b
+      # Uses cuBLAS via Ignis::LinAlg::Matmul for the hot path.
+      class Linear < Module
+        # @return [Tensor] weight matrix [out_features, in_features]
+        attr_reader :weight
+
+        # @return [Tensor, nil] bias vector [out_features]
+        attr_reader :bias
+
+        # @param in_features [Integer]
+        # @param out_features [Integer]
+        # @param bias [Boolean] whether to include bias
+        # @param device_id [Integer]
+        def initialize(in_features, out_features, bias: true, device_id: 0)
+          super()
+          @in_features = in_features
+          @out_features = out_features
+
+          # Kaiming uniform initialization: bound = sqrt(6 / in_features)
+          bound = Math.sqrt(6.0 / in_features)
+
+          weight_nv = Ignis::Shared::NvArray.new(shape: [out_features, in_features],
+                                                dtype: :float32, device_id: device_id)
+          # Allocate on device (no host Array — for a 128k-vocab tied head that array
+          # would be 262M Ruby Floats / ~10GB). The kaiming kernel writes every element.
+          weight_nv.to_device
+
+          # Initialize with Kaiming uniform via JIT kernel
+          init_kernel = Ignis::JIT::Kernels::Elementwise.kaiming_uniform_init
+          n = out_features * in_features
+          seed = ::Random.new.rand(2**64)  # stdlib RNG (Ignis::Random is the cuRAND module)
+          init_kernel.launch(grid: [(n + 255) / 256], block: [256],
+                             args: [weight_nv, bound.to_f, Ignis::JIT::Kernel::U64.new(seed), n])
+
+          @weight = register_parameter("weight",
+                     Tensor.new(data: weight_nv, requires_grad: true))
+
+          if bias
+            bias_nv = Ignis::Shared::NvArray.new(shape: [out_features],
+                                                dtype: :float32, device_id: device_id)
+            bias_bound = 1.0 / Math.sqrt(in_features)
+            bias_nv.to_device
+            init_kernel.launch(grid: [(out_features + 255) / 256], block: [256],
+                               args: [bias_nv, bias_bound.to_f, Ignis::JIT::Kernel::U64.new(seed + 1), out_features])
+
+            @bias = register_parameter("bias",
+                     Tensor.new(data: bias_nv, requires_grad: true))
+          else
+            @bias = nil
+          end
+        end
+
+        # Forward pass: x @ W^T + b
+        # @param x [Tensor] input [*, in_features]
+        # @return [Tensor] output [*, out_features]
+        def forward(x)
+          # x @ W^T, with cuBLAS doing the transpose in the GEMM (transpose_b) —
+          # avoids materializing W^T every forward (the LM head's was a 765ms
+          # transpose of a 38M-element weight).
+          out = x.matmul(@weight, transpose_b: true)
+          # Bias is [out_features]; broadcast-add it across rows of [*, out_features].
+          out = out.add_bias(@bias) if @bias
+          out
+        end
+
+        # @return [String]
+        def to_s
+          "Linear(in=#{@in_features}, out=#{@out_features}, bias=#{!@bias.nil?})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/module.rb

@@ -0,0 +1,178 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    module NN
+      # Module — base class for all neural network layers.
+      #
+      # Provides parameter management, state dict serialization,
+      # training/eval mode toggling, and automatic parameter collection.
+      #
+      # @example
+      #   class MyModel < Ignis::AI::NN::Module
+      #     def initialize
+      #       super
+      #       @linear = register_module("linear", Linear.new(768, 256))
+      #     end
+      #
+      #     def forward(x)
+      #       @linear.call(x).relu
+      #     end
+      #   end
+      class Module
+        # @return [Boolean] whether in training mode
+        attr_reader :training
+
+        def initialize
+          @training = true
+          @modules = {}    # name → Module
+          @parameters_map = {} # name → Tensor
+        end
+
+        # Override in subclass to define forward pass.
+        # @raise [NotImplementedError]
+        def forward(*args)
+          raise NotImplementedError, "#{self.class}#forward must be implemented"
+        end
+
+        # Make modules callable. Forwards positional AND keyword args (e.g. the
+        # `mask:` passed to Block/MultiHeadAttention) through to #forward.
+        # @return [Tensor]
+        def call(*args, **kwargs, &block)
+          forward(*args, **kwargs, &block)
+        end
+
+        # Collect all leaf parameters (requires_grad: true) recursively.
+        # @return [Array<Tensor>]
+        def parameters
+          params = []
+          @parameters_map.each_value { |p| params << p }
+          @modules.each_value { |m| params.concat(m.parameters) }
+          params
+        end
+
+        # Named parameters as a flat hash.
+        # @param prefix [String] prefix for nested module names
+        # @return [Hash{String => Tensor}]
+        def named_parameters(prefix: "")
+          result = {}
+          @parameters_map.each do |name, param|
+            key = prefix.empty? ? name : "#{prefix}.#{name}"
+            result[key] = param
+          end
+          @modules.each do |name, mod|
+            mod_prefix = prefix.empty? ? name : "#{prefix}.#{name}"
+            result.merge!(mod.named_parameters(prefix: mod_prefix))
+          end
+          result
+        end
+
+        # Switch to training mode (enables dropout, batch norm updates).
+        # @return [self]
+        def train!
+          @training = true
+          @modules.each_value { |m| m.train! }
+          self
+        end
+
+        # Switch to eval mode (disables dropout, uses running stats for BN).
+        # @return [self]
+        def eval!
+          @training = false
+          @modules.each_value { |m| m.eval! }
+          self
+        end
+
+        # Move all parameters to a specific device.
+        # @param device_id [Integer]
+        # @return [self]
+        def to(device_id:)
+          parameters.each do |p|
+            next if p.device_id == device_id
+            host_data = p.to_host
+            new_nv = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: device_id)
+            new_nv.from_host(host_data)
+            p.instance_variable_set(:@data, new_nv)
+          end
+          self
+        end
+
+        # Zero all gradients.
+        # @return [void]
+        def zero_grad!
+          parameters.each(&:zero_grad!)
+        end
+
+        # Export state dict (parameter name → Ruby Array of values).
+        # @return [Hash{String => Array}]
+        def state_dict
+          result = {}
+          named_parameters.each do |name, param|
+            result[name] = param.to_host
+          end
+          result
+        end
+
+        # Load state dict from Hash of {name → Ruby Array}.
+        # @param dict [Hash{String => Array}]
+        # @param strict [Boolean] if true, raises on missing/unexpected keys
+        # @return [self]
+        def load_state_dict(dict, strict: true)
+          params = named_parameters
+
+          if strict
+            missing = params.keys - dict.keys
+            unexpected = dict.keys - params.keys
+            raise KeyError, "Missing keys: #{missing}" unless missing.empty?
+            raise KeyError, "Unexpected keys: #{unexpected}" unless unexpected.empty?
+          end
+
+          dict.each do |name, values|
+            next unless params.key?(name)
+            param = params[name]
+            param.data.from_host(values)
+          end
+
+          self
+        end
+
+        # Total parameter count.
+        # @return [Integer]
+        def num_parameters
+          parameters.sum(&:numel)
+        end
+
+        # String representation.
+        # @return [String]
+        def to_s
+          parts = ["#{self.class.name}("]
+          @modules.each do |name, mod|
+            parts << "  (#{name}): #{mod}"
+          end
+          parts << ")"
+          parts.join("\n")
+        end
+
+        protected
+
+        # Register a parameter (leaf tensor with requires_grad: true).
+        # @param name [String]
+        # @param tensor [Tensor]
+        # @return [Tensor]
+        def register_parameter(name, tensor)
+          @parameters_map[name.to_s] = tensor
+          tensor
+        end
+
+        # Register a submodule.
+        # @param name [String]
+        # @param mod [Module]
+        # @return [Module]
+        def register_module(name, mod)
+          @modules[name.to_s] = mod
+          mod
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/nn/rms_norm.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Root-mean-square layer normalization: y = gamma * x / sqrt(mean(x^2) + eps).
+      # Used by Llama, Qwen, Mistral, SmolLM, Phi. Unlike LayerNorm there is no
+      # mean-subtraction and no bias — only a learned per-feature scale (gamma).
+      class RMSNorm < Module
+        # @return [Tensor] gamma (scale), initialized to ones
+        attr_reader :weight
+
+        # @param normalized_shape [Integer] size of the last dimension
+        # @param eps [Float] epsilon for numerical stability (Llama/Qwen use 1e-6/1e-5)
+        # @param device_id [Integer]
+        def initialize(normalized_shape, eps: 1e-6, device_id: 0)
+          super()
+          @normalized_shape = normalized_shape
+          @eps = eps
+
+          weight_nv = Ignis::Shared::NvArray.new(shape: [normalized_shape],
+                                                dtype: :float32, device_id: device_id)
+          weight_nv.from_host(Array.new(normalized_shape, 1.0))
+          @weight = register_parameter("weight",
+                     Tensor.new(data: weight_nv, requires_grad: true))
+        end
+
+        # @param x [Tensor] input [*, normalized_shape]
+        # @return [Tensor]
+        def forward(x)
+          x.rms_norm(@weight, eps: @eps)
+        end
+
+        # @return [String]
+        def to_s
+          "RMSNorm(#{@normalized_shape}, eps=#{@eps})"
+        end
+      end
+    end
+  end
+end