mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/ministral3.rb

@@ -0,0 +1,304 @@
+require_relative "activations"
+require_relative "cache"
+require_relative "pipeline"
+require_relative "rope_utils"
+
+module MlxLm
+  module Models
+    module Ministral3
+      def self.llama4_attn_scale(size, offset, beta, max_position_embeddings)
+        mx = MLX::Core
+        positions = mx.arange(size) + offset
+        scale = 1.0 + beta.to_f * mx.log(1.0 + mx.floor(positions / max_position_embeddings.to_f))
+        scale.reshape([size, 1])
+      end
+
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "ministral3"
+        field :hidden_size
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :rms_norm_eps
+        field :vocab_size
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: nil
+        field :num_key_value_heads, default: nil
+        field :rope_parameters, default: nil
+        field :tie_word_embeddings, default: true
+        field :layer_types, default: nil
+        field :sliding_window, default: nil
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+          @rope_parameters = _stringify_keys(@rope_parameters || {})
+          @rope_parameters["rope_theta"] = 10_000.0 unless @rope_parameters.key?("rope_theta")
+          @layer_types ||= Array.new(@num_hidden_layers) { "full_attention" }
+        end
+
+        def rope_parameter(key, default = nil)
+          return default unless @rope_parameters.is_a?(Hash)
+          return @rope_parameters[key.to_s] if @rope_parameters.key?(key.to_s)
+          return @rope_parameters[key.to_sym] if @rope_parameters.key?(key.to_sym)
+
+          default
+        end
+
+        private
+
+        def _stringify_keys(hash)
+          hash.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_parameter("rope_theta", 10_000.0),
+            false,
+            args.rope_parameters,
+            max_position_embeddings: args.max_position_embeddings
+          )
+        end
+
+        def call(x, attn_scale:, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          queries = queries * attn_scale
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: false)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        attr_reader :use_sliding
+
+        def initialize(args, use_sliding: false)
+          super()
+          @use_sliding = use_sliding
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, attn_scale:, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), attn_scale: attn_scale, mask: mask, cache: cache)
+          h = x + r
+          r = mlp.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class LanguageModel < MLX::NN::Module
+        include PipelineMixin
+        attr_reader :sliding_window
+
+        def initialize(args)
+          super()
+          @args = args
+          @sliding_window = args.sliding_window
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = args.layer_types.map do |layer_type|
+            TransformerBlock.new(args, use_sliding: layer_type == "sliding_attention")
+          end
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          h = input_embeddings || embed_tokens.call(inputs)
+          active_layers = pipeline_layers
+          layer_cache = cache || Array.new(active_layers.length)
+
+          first_cache = layer_cache.find { |entry| !entry.nil? }
+          offset = first_cache ? first_cache.offset : 0
+
+          fa_idx = nil
+          swa_idx = nil
+          active_layers.each_with_index do |layer, i|
+            if layer.use_sliding
+              swa_idx ||= i
+            else
+              fa_idx ||= i
+            end
+            break if fa_idx && swa_idx
+          end
+
+          fa_mask = fa_idx.nil? ? nil : _create_attention_mask(h, layer_cache[fa_idx])
+          swa_mask = if swa_idx.nil?
+            nil
+          else
+            _create_attention_mask(h, layer_cache[swa_idx], window_size: @sliding_window)
+          end
+
+          beta = @args.rope_parameter("llama_4_scaling_beta", 0.0).to_f
+          max_pos = @args.rope_parameter(
+            "original_max_position_embeddings",
+            @args.max_position_embeddings || h.shape[1]
+          ).to_i
+          max_pos = 1 if max_pos <= 0
+
+          attn_scale = MlxLm::Models::Ministral3.llama4_attn_scale(
+            inputs.shape[1],
+            offset,
+            beta,
+            max_pos
+          ).astype(h.dtype)
+
+          active_layers.each_with_index do |layer, idx|
+            mask = layer.use_sliding ? swa_mask : fa_mask
+            h = layer.call(h, attn_scale: attn_scale, mask: mask, cache: layer_cache[idx])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache = nil, window_size: nil)
+          n = h.shape[1]
+          if cache && cache.respond_to?(:make_mask)
+            return cache.make_mask(n, window_size: window_size)
+          end
+
+          if window_size
+            offset = cache ? cache.offset : 0
+            if cache && cache.instance_variable_defined?(:@max_size)
+              max_size = cache.instance_variable_get(:@max_size)
+              offset = [max_size - 1, offset].min if max_size && max_size > 0
+            end
+
+            return _create_causal_mask(n, offset: offset, window_size: window_size) if offset + n > window_size
+          end
+
+          return nil if n == 1
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset: 0, window_size: nil)
+          mx = MLX::Core
+          rinds = mx.arange(0, offset + n, 1, mx.int32).reshape([1, offset + n])
+          linds = mx.arange(offset, offset + n, 1, mx.int32).reshape([n, 1])
+
+          mask = mx.greater_equal(linds, rinds)
+          if window_size
+            mask = mx.logical_and(mask, mx.less(linds, mx.add(rinds, window_size)))
+          end
+          mask
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = LanguageModel.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          out = model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          sanitized = weights.reject do |key, _|
+            key_name = key.to_s
+            key_name.include?("self_attn.rotary_emb.inv_freq") || key_name.include?("self_attn.rope.inv_freq")
+          end
+          sanitized.delete("lm_head.weight") if @args.tie_word_embeddings
+
+          new_weights = {}
+          sanitized.each do |key, value|
+            key_name = key.to_s
+            if key_name.include?("weight_scale_inv")
+              wk = key_name.sub("_scale_inv", "")
+              next unless sanitized.key?(wk)
+
+              new_weights[wk] = sanitized[wk] * value
+            elsif key_name.include?("activation_scale")
+              next
+            elsif !new_weights.key?(key)
+              new_weights[key] = value
+            end
+          end
+          new_weights
+        end
+
+        def layers
+          model.pipeline_layers
+        end
+
+        def make_cache
+          max_size = @args.sliding_window || @args.max_position_embeddings || 1
+          layers.map do |layer|
+            if layer.use_sliding
+              MlxLm::RotatingKVCache.new(max_size: max_size)
+            else
+              MlxLm::KVCache.new
+            end
+          end
+        end
+      end
+
+      Models.register("ministral3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/mistral3.rb

@@ -0,0 +1,84 @@
+require_relative "llama"
+
+module MlxLm
+  module Models
+    module Mistral3
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "mistral3"
+        field :text_config, default: nil
+
+        def initialize(**kwargs)
+          super
+          @text_config = (@text_config || {}).dup
+          @text_config["tie_word_embeddings"] = false unless @text_config.key?("tie_word_embeddings")
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+
+          text_config = args.text_config || {}
+          text_model_type = text_config["model_type"]
+
+          if text_model_type == "ministral3" && Models::REGISTRY.key?("ministral3")
+            model_class, args_class = Models.get_classes(text_config)
+            self.language_model = model_class.new(args_class.from_dict(text_config))
+          else
+            self.language_model = Llama::Model.new(Llama::ModelArgs.from_dict(text_config))
+          end
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          supports_input_embeddings = language_model.method(:call).parameters.any? do |_, name|
+            name == :input_embeddings
+          end
+
+          if supports_input_embeddings
+            language_model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          else
+            language_model.call(inputs, cache: cache)
+          end
+        end
+
+        def sanitize(weights)
+          result = {}
+          language_weights = {}
+
+          weights.each do |k, v|
+            next if k == "vision_tower" || k.start_with?("vision_tower.")
+            next if k == "multi_modal_projector" || k.start_with?("multi_modal_projector.")
+
+            if k.start_with?("language_model.")
+              language_weights[k.delete_prefix("language_model.")] = v
+            else
+              result[k] = v
+            end
+          end
+
+          sanitized_language = if language_model.respond_to?(:sanitize)
+            language_model.sanitize(language_weights)
+          else
+            language_weights
+          end
+
+          sanitized_language.each do |k, v|
+            result["language_model.#{k}"] = v
+          end
+
+          result
+        end
+
+        def layers
+          return language_model.model.layers if language_model.respond_to?(:model) && language_model.model.respond_to?(:layers)
+
+          language_model.layers
+        end
+      end
+
+      Models.register("mistral3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/mixtral.rb

@@ -0,0 +1,192 @@
+module MlxLm
+  module Models
+    module Mixtral
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "mixtral"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: 8
+        field :intermediate_size, default: 14336
+        field :vocab_size, default: 32000
+        field :rms_norm_eps, default: 1e-5
+        field :rope_theta, default: 1e6
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :num_local_experts, default: 8
+        field :num_experts_per_tok, default: 2
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = dim / @n_heads
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+
+          self.rope = MLX::NN::RoPE.new(
+            @head_dim,
+            traditional: args.rope_traditional,
+            base: args.rope_theta
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class SparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          @num_experts = args.num_local_experts
+          @num_experts_per_tok = args.num_experts_per_tok
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+
+          self.gate = MLX::NN::Linear.new(dim, @num_experts, bias: false)
+          self.switch_mlp = SwitchLayers::SwitchGLU.new(dim, hidden_dim, @num_experts)
+        end
+
+        def call(x)
+          mx = MLX::Core
+          k = @num_experts_per_tok
+
+          gates = gate.call(x)
+          inds = mx.stop_gradient(mx.argpartition(gates * -1.0, k - 1, -1))
+          take_ids = mx.array((0...k).to_a, dtype: mx.int32)
+          inds = mx.take(inds, take_ids, -1)
+
+          scores = mx.take_along_axis(gates, inds, -1)
+          scores = mx.softmax(scores.astype(mx.float32), -1).astype(gates.dtype)
+
+          y = switch_mlp.call(x, inds)
+          y = mx.sum(y * mx.expand_dims(scores, -1), -2)
+          y
+        end
+      end
+
+      class MixtralDecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.block_sparse_moe = SparseMoeBlock.new(args)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = block_sparse_moe.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class MixtralModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { MixtralDecoderLayer.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model = MixtralModel.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          mx = MLX::Core
+          result = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+
+          # Convert per-expert weights to stacked SwitchGLU format
+          @args.num_hidden_layers.times do |l|
+            prefix = "model.layers.#{l}"
+            [["w1", "gate_proj"], ["w2", "down_proj"], ["w3", "up_proj"]].each do |n, m|
+              ["weight", "scales", "biases"].each do |k|
+                key0 = "#{prefix}.block_sparse_moe.experts.0.#{n}.#{k}"
+                if result.key?(key0)
+                  to_join = (0...@args.num_local_experts).map { |e|
+                    result.delete("#{prefix}.block_sparse_moe.experts.#{e}.#{n}.#{k}")
+                  }
+                  result["#{prefix}.block_sparse_moe.switch_mlp.#{m}.#{k}"] = mx.stack(to_join)
+                end
+              end
+            end
+          end
+
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("mixtral", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/mla.rb

@@ -0,0 +1,75 @@
+module MlxLm
+  module Models
+    module MLA
+      class MultiLinear < MLX::NN::Module
+        def initialize(input_dims, output_dims, num_heads)
+          super()
+          scale = Math.sqrt(1.0 / input_dims)
+          self.weight = MLX::Core.uniform([num_heads, output_dims, input_dims], -scale, scale)
+        end
+
+        def call(x, transpose: true)
+          if transpose
+            MLX::Core.matmul(x, MLX::Core.swapaxes(weight, -1, -2))
+          else
+            MLX::Core.matmul(x, weight)
+          end
+        end
+
+        def to_quantized(group_size: nil, bits: nil, mode: "affine", quantize_input: false)
+          raise ArgumentError, "Quantized input is not supported." if quantize_input
+
+          QuantizedMultiLinear.from_multi_linear(self, group_size, bits, mode: mode)
+        end
+      end
+
+      class QuantizedMultiLinear < MLX::NN::Module
+        attr_reader :group_size, :bits, :mode
+
+        def initialize(input_dims, output_dims, num_heads, group_size = nil, bits = nil, mode: "affine")
+          super()
+
+          @group_size, @bits = MLX::NN.__send__(:defaults_for_mode, mode, group_size, bits)
+          @mode = mode
+
+          scale = Math.sqrt(1.0 / input_dims)
+          weight = MLX::Core.uniform([num_heads, output_dims, input_dims], -scale, scale)
+          q_weight, q_scales, *q_biases = MLX::Core.quantize(weight, @group_size, @bits, @mode)
+          self.weight = q_weight
+          self.scales = q_scales
+          self.biases = q_biases.empty? ? nil : q_biases[0]
+
+          freeze
+        end
+
+        def call(x, transpose: true)
+          MLX::Core.quantized_matmul(
+            x,
+            weight,
+            scales,
+            biases,
+            transpose,
+            @group_size,
+            @bits,
+            @mode
+          )
+        end
+
+        def self.from_multi_linear(multi_linear_layer, group_size = nil, bits = nil, mode: "affine")
+          num_heads, output_dims, input_dims = multi_linear_layer.weight.shape
+          out = new(input_dims, output_dims, num_heads, group_size, bits, mode: mode)
+          q_weight, q_scales, *q_biases = MLX::Core.quantize(
+            multi_linear_layer.weight,
+            out.group_size,
+            out.bits,
+            out.mode
+          )
+          out.weight = q_weight
+          out.scales = q_scales
+          out.biases = q_biases.empty? ? nil : q_biases[0]
+          out
+        end
+      end
+    end
+  end
+end