mlx 0.30.7.2 → 0.30.7.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 215a912d2353fd5edaa60e320a5b857aa13009e55f3a190b21b2ffe5735f37af
-  data.tar.gz: 1c9b4279f8077e3cd067354ea692e1b24248afe7262d88d4569c566c52f5a158
+  metadata.gz: 25d582e4816d69b27713a4027534b75cd00ca72557e69681daf07146d3e79ef2
+  data.tar.gz: c010252aa355370a531fa4f3b9bf8cc729876d2f7fb9ae8b8e0d6a1eb6cb57c4
 SHA512:
-  metadata.gz: 66abcbd58ccfc04186df11b0d2b6445c7d1e0ab4a36451742755f6fcf41022363403536c3d27640935364c58231c4c4a03e39ceba97617a59f2ad69acf23dc16
-  data.tar.gz: ba7ad07ccd31e94bdf3fdee73117f069c6ee22c11cdfc3f2470eedee9ee0bc9e976970bab21f156c003b658723a0e3fb2f63f62d0e11d33de3a61fc1d9121711
+  metadata.gz: 53e629e845342f173c04c7c6d9d976a29dd5492ae945239897d3168a586288ec958ba58753345317f301220ac5f4b91a22f97731ab799fcea5d59f3d19e48214
+  data.tar.gz: 5b04f2e63e3dcdb6a0282184a310600f4fb72e606b45e8e7a27a7b9461abef3a598afe2eb5525e66457da45b8a84a6c2f07c871c955e9cddd4308971496c7fd1

data/ext/mlx/native.cpp

@@ -6625,7 +6625,8 @@ static VALUE core_clear_cache(VALUE) {
 
 static VALUE core_metal_is_available(VALUE) {
   try {
-    return mxmetal::is_available() ? Qtrue : Qfalse;
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    return mx::is_available(gpu_device) ? Qtrue : Qfalse;
   } catch (const std::exception& error) {
     raise_std_exception(error);
     return Qnil;
@@ -6654,7 +6655,12 @@ static VALUE core_metal_stop_capture(VALUE) {
 
 static VALUE core_metal_device_info(VALUE) {
   try {
-    const auto& info = mxmetal::device_info();
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    if (!mx::is_available(gpu_device)) {
+      rb_raise(rb_eRuntimeError, "[metal_device_info] Metal GPU device is not available");
+    }
+
+    const auto& info = mx::device_info(gpu_device);
     VALUE hash = rb_hash_new();
     for (const auto& [key, value] : info) {
       VALUE ruby_key = rb_utf8_str_new(key.c_str(), static_cast<long>(key.size()));

data/lib/mlx/distributed_utils/common.rb

@@ -5,7 +5,7 @@ require "json"
 
 module MLX
   module DistributedUtils
-    Host = Struct.new(:rank, :ssh_hostname, :ips, :rdma, keyword_init: true)
+    Host = Data.define(:rank, :ssh_hostname, :ips, :rdma)
 
     class Hostfile
       attr_accessor :hosts, :backend, :envs

data/lib/mlx/distributed_utils/config.rb

@@ -8,13 +8,14 @@ require "shellwords"
 
 module MLX
   module DistributedUtils
-    SSHInfo = Struct.new(:can_ssh, :has_sudo, keyword_init: true) do
+    SSHInfo = Data.define(:can_ssh, :has_sudo) do
       def to_bool
         can_ssh
       end
     end
-    ThunderboltPort = Struct.new(:iface, :uuid, :connected_to, keyword_init: true)
-    ThunderboltHost = Struct.new(:name, :ports, keyword_init: true)
+    ThunderboltPort = Data.define(:iface, :uuid, :connected_to)
+    ThunderboltHost = Data.define(:name, :ports)
+    CommandResult = Data.define(:stdout, :stderr, :status)
 
     class IPConfigurator
       attr_reader :ips, :hosts, :tb_hosts
@@ -509,6 +510,8 @@ module MLX
       end
 
       def config_main(argv = ARGV, runner: nil)
+        Process.warmup if Process.respond_to?(:warmup)
+
         opts = {
           verbose: false,
           hosts: "127.0.0.1",
@@ -577,7 +580,7 @@ module MLX
         return runner.call(cmd) unless runner.nil?
 
         stdout, stderr, status = Open3.capture3(*cmd)
-        Struct.new(:stdout, :stderr, :status, keyword_init: true).new(stdout: stdout, stderr: stderr, status: status)
+        CommandResult.new(stdout: stdout, stderr: stderr, status: status)
       end
 
       def stdout_for(result)

data/lib/mlx/distributed_utils/launch.rb

@@ -314,6 +314,8 @@ module MLX
       end
 
       def main(argv = ARGV)
+        Process.warmup if Process.respond_to?(:warmup)
+
         opts = {
           print_python: false,
           verbose: false,

data/lib/mlx/dsl/attention.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Attention < MLX::NN::Module
+      def initialize(
+        dims:,
+        num_heads:,
+        kv_heads: nil,
+        qkv_bias: false,
+        backend: :sdpa,
+        rope: nil,
+        cache: false
+      )
+        super()
+
+        @dims = Integer(dims)
+        @num_heads = Integer(num_heads)
+        @kv_heads = kv_heads.nil? ? @num_heads : Integer(kv_heads)
+        if (@dims % @num_heads) != 0
+          raise ArgumentError, "dims must be divisible by num_heads"
+        end
+        if (@num_heads % @kv_heads) != 0
+          raise ArgumentError, "num_heads must be divisible by kv_heads"
+        end
+
+        @head_dim = @dims / @num_heads
+        @kv_repeats = @num_heads / @kv_heads
+        @backend = backend.to_sym
+        @cache_enabled = !!cache
+        @scale = Math.sqrt(1.0 / @head_dim)
+
+        self.query_proj = MLX::NN::Linear.new(@dims, @num_heads * @head_dim, bias: qkv_bias)
+        self.key_proj = MLX::NN::Linear.new(@dims, @kv_heads * @head_dim, bias: qkv_bias)
+        self.value_proj = MLX::NN::Linear.new(@dims, @kv_heads * @head_dim, bias: qkv_bias)
+        self.out_proj = MLX::NN::Linear.new(@num_heads * @head_dim, @dims, bias: qkv_bias)
+        self.rope = __dsl_build_rope(rope)
+      end
+
+      def call(queries, keys = nil, values = nil, mask: nil, cache: nil)
+        keys ||= queries
+        values ||= keys
+        q_was_2d = queries.ndim == 2
+
+        queries = MLX::Core.expand_dims(queries, 0) if q_was_2d
+        keys = MLX::Core.expand_dims(keys, 0) if keys.ndim == 2
+        values = MLX::Core.expand_dims(values, 0) if values.ndim == 2
+
+        batch_size, q_len, = queries.shape
+
+        q = __dsl_pack_heads(query_proj.call(queries), @num_heads)
+        k = __dsl_pack_heads(key_proj.call(keys), @kv_heads)
+        v = __dsl_pack_heads(value_proj.call(values), @kv_heads)
+
+        offset = cache.nil? ? 0 : cache[0].shape[2]
+        if !rope.nil?
+          if offset.zero?
+            q = rope.call(q)
+            k = rope.call(k)
+          else
+            q = rope.call(q, offset: offset)
+            k = rope.call(k, offset: offset)
+          end
+        end
+
+        unless cache.nil?
+          key_cache, value_cache = cache
+          k = MLX::Core.concatenate([key_cache, k], 2)
+          v = MLX::Core.concatenate([value_cache, v], 2)
+        end
+        next_cache = [k, v]
+
+        k_for_attn = __dsl_repeat_kv(k)
+        v_for_attn = __dsl_repeat_kv(v)
+        out = __dsl_attention(q, k_for_attn, v_for_attn, mask)
+        out = MLX::Core.transpose(out, [0, 2, 1, 3])
+        out = MLX::Core.reshape(out, [batch_size, q_len, @num_heads * @head_dim])
+        out = out_proj.call(out)
+        out = MLX::Core.squeeze(out, 0) if q_was_2d
+
+        if @cache_enabled || !cache.nil?
+          [out, next_cache]
+        else
+          out
+        end
+      end
+
+      private
+
+      def __dsl_build_rope(config)
+        return nil if config.nil?
+
+        opts = config.transform_keys(&:to_sym)
+        rope_kwargs = {
+          traditional: opts.fetch(:traditional, false),
+          base: opts.fetch(:base, 10_000.0)
+        }
+        rope_kwargs[:scale] = opts[:scale] if opts.key?(:scale)
+        MLX::NN::RoPE.new(@head_dim, **rope_kwargs)
+      end
+
+      def __dsl_pack_heads(x, heads)
+        batch, length, = x.shape
+        x = MLX::Core.reshape(x, [batch, length, heads, @head_dim])
+        MLX::Core.transpose(x, [0, 2, 1, 3])
+      end
+
+      def __dsl_repeat_kv(x)
+        return x if @kv_repeats == 1
+
+        batch, _heads, length, dim = x.shape
+        expanded = MLX::Core.expand_dims(x, 2)
+        repeated = MLX::Core.concatenate(Array.new(@kv_repeats, expanded), 2)
+        MLX::Core.reshape(repeated, [batch, @num_heads, length, dim])
+      end
+
+      def __dsl_attention(q, k, v, mask)
+        if @backend == :sdpa && MLX::Core.respond_to?(:scaled_dot_product_attention)
+          return MLX::Core.scaled_dot_product_attention(q, k, v, @scale, mask)
+        end
+
+        scores = MLX::Core.matmul(
+          MLX::Core.multiply(q, @scale),
+          MLX::Core.transpose(k, [0, 1, 3, 2])
+        )
+        scores = MLX::Core.add(scores, mask.astype(scores.dtype)) unless mask.nil?
+        probs = MLX::Core.softmax(scores.astype(MLX::Core.float32), -1).astype(scores.dtype)
+        MLX::Core.matmul(probs, v)
+      end
+    end
+  end
+end

data/lib/mlx/dsl/builder.rb

@@ -266,6 +266,14 @@ module MLX
         push(MLX::NN::Transformer.new(*args, **kwargs))
       end
 
+      def attention(*args, **kwargs)
+        push(MLX::DSL::Attention.new(*args, **kwargs))
+      end
+
+      def transformer_block(*args, **kwargs)
+        push(MLX::DSL::TransformerBlock.new(*args, **kwargs))
+      end
+
       def rope(*args, **kwargs)
         push(MLX::NN::RoPE.new(*args, **kwargs))
       end

data/lib/mlx/dsl/config_schema.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module ConfigSchema
+      UNSET = Object.new.freeze
+      class DefaultContext
+        def initialize(values)
+          @values = values
+        end
+
+        def method_missing(name, *args, &block)
+          if args.empty? && block.nil? && @values.key?(name.to_sym)
+            return @values[name.to_sym]
+          end
+
+          super
+        end
+
+        def respond_to_missing?(name, include_private = false)
+          @values.key?(name.to_sym) || super
+        end
+      end
+
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      module ClassMethods
+        def field(name, type = nil, required: false, default: UNSET, &validator)
+          key = name.to_sym
+          config_schema_fields[key] = {
+            type: type,
+            required: !!required,
+            default: default,
+            validator: validator
+          }
+          attr_accessor key unless method_defined?(key) && method_defined?(:"#{key}=")
+        end
+
+        def config_schema_fields
+          @config_schema_fields ||= {}
+        end
+
+        def inherited(subclass)
+          super
+          copied = config_schema_fields.each_with_object({}) do |(key, value), out|
+            out[key] = value.dup
+          end
+          subclass.instance_variable_set(:@config_schema_fields, copied)
+        end
+
+        def from_hash(raw)
+          source = (raw || {}).each_with_object({}) do |(key, value), out|
+            out[key.to_sym] = value
+          end
+          new(**source)
+        end
+
+        private
+
+        def __dsl_call_default(default, resolved)
+          context = DefaultContext.new(resolved)
+          return default unless default.respond_to?(:call)
+          return default.call(context) if default.is_a?(Proc) && default.arity == 1
+          return default.call if !default.is_a?(Proc)
+          return default.call if default.arity.zero?
+
+          default.call(context)
+        end
+
+        def __dsl_validate_field(name, value, spec)
+          type = spec.fetch(:type)
+          if !type.nil? && !value.nil?
+            allowed_types = type.is_a?(Array) ? type : [type]
+            unless allowed_types.any? { |klass| value.is_a?(klass) }
+              raise TypeError,
+                    "config field #{name} must be #{allowed_types.map(&:to_s).join(' or ')}, got #{value.class}"
+            end
+          end
+
+          validator = spec.fetch(:validator)
+          unless validator.nil?
+            if validator.arity == 2
+              validator.call(value, name)
+            else
+              validator.call(value)
+            end
+          end
+
+          value
+        end
+      end
+
+      def initialize(**kwargs)
+        source = kwargs.each_with_object({}) do |(key, value), out|
+          out[key.to_sym] = value
+        end
+        resolved = {}
+        unknown = source.keys - self.class.config_schema_fields.keys
+        unless unknown.empty?
+          names = unknown.map(&:to_s).sort.join(", ")
+          raise ArgumentError, "unknown config field(s): #{names}"
+        end
+
+        self.class.config_schema_fields.each do |name, spec|
+          if source.key?(name)
+            value = source.fetch(name)
+          else
+            default = spec.fetch(:default)
+            if default.equal?(UNSET)
+              if spec.fetch(:required)
+                raise ArgumentError, "missing required config field: #{name}"
+              end
+              next
+            end
+            value = self.class.send(:__dsl_call_default, default, resolved)
+          end
+
+          value = self.class.send(:__dsl_validate_field, name, value, spec)
+          resolved[name] = value
+          public_send(:"#{name}=", value)
+        end
+      end
+
+      def to_h
+        self.class.config_schema_fields.keys.each_with_object({}) do |name, out|
+          out[name.to_s] = public_send(name)
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/generate.rb

@@ -0,0 +1,193 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Generate
+      def initialize(
+        model:,
+        tokenizer: nil,
+        eos_id: nil,
+        sampler: nil,
+        mode: :decoder_only,
+        decoder_start_id: nil
+      )
+        @model = model
+        @tokenizer = tokenizer
+        @eos_id = eos_id
+        @sampler = { strategy: :argmax }.merge((sampler || {}).transform_keys(&:to_sym))
+        @mode = mode.to_sym
+        @decoder_start_id = decoder_start_id
+      end
+
+      def each_token(prompt: nil, input_ids: nil, max_tokens: 128, **kwargs)
+        return enum_for(__method__, prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) unless block_given?
+
+        case @mode
+        when :decoder_only
+          __dsl_each_decoder_only(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        when :encoder_decoder
+          __dsl_each_encoder_decoder(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        else
+          raise ArgumentError, "unsupported generation mode: #{@mode.inspect}"
+        end
+
+        self
+      end
+
+      private
+
+      def __dsl_each_decoder_only(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        model_input = __dsl_input_array(tokens)
+        logits, cache = __dsl_decode_step(model_input, cache: nil, **kwargs)
+
+        max_tokens.to_i.times do
+          token = __dsl_sample(__dsl_last_logits(logits))
+          token_id = __dsl_token_id(token)
+          chunk = __dsl_decode_token(token_id)
+          yield token_id, chunk
+          break if !@eos_id.nil? && token_id == @eos_id
+
+          next_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          logits, cache = __dsl_decode_step(next_input, cache: cache, **kwargs)
+        end
+      end
+
+      def __dsl_each_encoder_decoder(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        source = __dsl_input_array(tokens)
+
+        if @model.respond_to?(:encode) && @model.respond_to?(:decode)
+          memory = @model.encode(source)
+          start_id = __dsl_decoder_start_id
+          decoder_input = MLX::Core.array([[start_id]], MLX::Core.int32)
+          cache = nil
+
+          max_tokens.to_i.times do
+            decoded = @model.decode(decoder_input, memory, cache: cache, **kwargs)
+            logits, cache = __dsl_split_logits_and_cache(decoded, cache)
+            token = __dsl_sample(__dsl_last_logits(logits))
+            token_id = __dsl_token_id(token)
+            chunk = __dsl_decode_token(token_id)
+            yield token_id, chunk
+            break if !@eos_id.nil? && token_id == @eos_id
+
+            decoder_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          end
+          return
+        end
+
+        # Fallback path for model.call-style APIs.
+        __dsl_each_decoder_only(prompt: prompt, input_ids: tokens, max_tokens: max_tokens, **kwargs) do |id, chunk|
+          yield id, chunk
+        end
+      end
+
+      def __dsl_decode_step(input_ids, cache:, **kwargs)
+        output = @model.call(input_ids, cache: cache, **kwargs)
+        __dsl_split_logits_and_cache(output, cache)
+      end
+
+      def __dsl_split_logits_and_cache(output, fallback_cache)
+        if output.is_a?(Array) && output.length == 2
+          [output[0], output[1]]
+        else
+          [output, fallback_cache]
+        end
+      end
+
+      def __dsl_last_logits(logits)
+        return logits if logits.ndim == 2
+        return logits if logits.ndim == 1
+
+        index = MLX::Core.array([logits.shape[1] - 1], MLX::Core.int32)
+        MLX::Core.squeeze(MLX::Core.take(logits, index, 1), 1)
+      end
+
+      def __dsl_sample(logits)
+        strategy = @sampler.fetch(:strategy, :argmax).to_sym
+        temperature = @sampler.fetch(:temperature, 1.0).to_f
+
+        return MLX::Core.argmax(logits, -1) if strategy == :argmax || temperature.zero?
+
+        case strategy
+        when :top_k
+          __dsl_top_k_sample(logits, k: Integer(@sampler.fetch(:k, 40)), temperature: temperature)
+        when :temperature, :categorical
+          __dsl_temperature_sample(logits, temperature: temperature)
+        else
+          raise ArgumentError, "unsupported sampler strategy: #{strategy.inspect}"
+        end
+      end
+
+      def __dsl_temperature_sample(logits, temperature:)
+        scaled = if temperature == 1.0
+          logits
+        else
+          MLX::Core.multiply(logits, 1.0 / temperature)
+        end
+        MLX::Core.categorical(scaled)
+      end
+
+      def __dsl_top_k_sample(logits, k:, temperature:)
+        rows = logits.ndim == 1 ? [logits.to_a] : logits.to_a
+        masked = rows.map do |row|
+          pairs = row.each_with_index.sort_by { |(value, _index)| -value }
+          keep = pairs.first([k, row.length].min).map(&:last)
+          filtered = Array.new(row.length, -Float::INFINITY)
+          keep.each { |idx| filtered[idx] = row[idx] }
+          filtered
+        end
+        masked_logits = MLX::Core.array(masked, logits.dtype)
+        __dsl_temperature_sample(masked_logits, temperature: temperature)
+      end
+
+      def __dsl_encode(prompt)
+        raise ArgumentError, "prompt/input_ids required when tokenizer is unavailable" if @tokenizer.nil?
+
+        @tokenizer.encode(prompt.to_s)
+      end
+
+      def __dsl_input_array(tokens)
+        if tokens.is_a?(MLX::Core::Array)
+          return tokens if tokens.ndim > 1
+
+          return MLX::Core.expand_dims(tokens.astype(MLX::Core.int32), 0)
+        end
+
+        arr = tokens.to_a
+        nested = arr.empty? ? [[]] : (arr.first.is_a?(Array) ? arr : [arr])
+        MLX::Core.array(nested, MLX::Core.int32)
+      end
+
+      def __dsl_token_id(token)
+        return token.item.to_i if token.respond_to?(:item)
+
+        value = token.to_a
+        if value.is_a?(Array)
+          first = value.first
+          return first.first.to_i if first.is_a?(Array)
+          return first.to_i
+        end
+        value.to_i
+      end
+
+      def __dsl_decode_token(token_id)
+        return nil if @tokenizer.nil? || !@tokenizer.respond_to?(:decode)
+
+        @tokenizer.decode([token_id])
+      end
+
+      def __dsl_decoder_start_id
+        return @decoder_start_id unless @decoder_start_id.nil?
+        return @tokenizer.decoder_start_id if !@tokenizer.nil? && @tokenizer.respond_to?(:decoder_start_id)
+
+        raise ArgumentError, "decoder_start_id is required for encoder-decoder mode"
+      end
+    end
+  end
+end

data/lib/mlx/dsl/kv_cache.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class KVCache
+      attr_reader :num_layers
+
+      def initialize(num_layers:)
+        @num_layers = Integer(num_layers)
+        raise ArgumentError, "num_layers must be non-negative" if @num_layers.negative?
+
+        @layers = Array.new(@num_layers)
+      end
+
+      def layer(index)
+        @layers.fetch(__dsl_index(index))
+      end
+
+      def []=(index, value)
+        @layers[__dsl_index(index)] = value
+      end
+
+      def offset(layer:)
+        state = self.layer(layer)
+        return 0 if state.nil?
+
+        keys, = state
+        keys.shape[2]
+      end
+
+      def append(layer:, keys:, values:)
+        idx = __dsl_index(layer)
+        current = @layers[idx]
+        if current.nil?
+          @layers[idx] = [keys, values]
+          return @layers[idx]
+        end
+
+        key_cache, value_cache = current
+        next_keys = MLX::Core.concatenate([key_cache, keys], 2)
+        next_values = MLX::Core.concatenate([value_cache, values], 2)
+        @layers[idx] = [next_keys, next_values]
+      end
+
+      def truncate!(tokens:, layer: nil)
+        keep = Integer(tokens)
+        if layer.nil?
+          @layers.each_index { |idx| __dsl_truncate_layer!(idx, keep) }
+        else
+          __dsl_truncate_layer!(__dsl_index(layer), keep)
+        end
+        self
+      end
+
+      def reset!(layer: nil)
+        if layer.nil?
+          @layers.map! { nil }
+        else
+          @layers[__dsl_index(layer)] = nil
+        end
+        self
+      end
+
+      private
+
+      def __dsl_index(index)
+        idx = Integer(index)
+        if idx.negative? || idx >= @num_layers
+          raise IndexError, "layer index #{idx} out of range (0...#{@num_layers})"
+        end
+
+        idx
+      end
+
+      def __dsl_truncate_layer!(idx, keep)
+        state = @layers[idx]
+        return if state.nil?
+
+        if keep <= 0
+          @layers[idx] = nil
+          return
+        end
+
+        keys, values = state
+        total = keys.shape[2]
+        return if keep >= total
+
+        start = total - keep
+        indices = MLX::Core.arange(start, total, 1, MLX::Core.int32)
+        trimmed_keys = MLX::Core.take(keys, indices, 2)
+        trimmed_values = MLX::Core.take(values, indices, 2)
+        @layers[idx] = [trimmed_keys, trimmed_values]
+      end
+    end
+  end
+end

data/lib/mlx/dsl/masks.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Masks
+      module_function
+
+      def causal(length:, offset: 0, dtype: MLX::Core.float32)
+        length = Integer(length)
+        offset = Integer(offset)
+        raise ArgumentError, "length must be non-negative" if length.negative?
+        raise ArgumentError, "offset must be non-negative" if offset.negative?
+
+        rinds = MLX::Core.arange(0, offset + length, 1)
+        linds = if offset.zero?
+          rinds
+        else
+          MLX::Core.arange(offset, offset + length, 1)
+        end
+        lhs = MLX::Core.expand_dims(linds, 1)
+        rhs = MLX::Core.expand_dims(rinds, 0)
+        mask = MLX::Core.less(lhs, rhs).astype(dtype)
+        min_value = if MLX::Core.respond_to?(:finfo)
+          MLX::Core.finfo(dtype).min
+        else
+          -1e9
+        end
+        MLX::Core.multiply(mask, min_value)
+      end
+    end
+  end
+end

data/lib/mlx/dsl/positions.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Positions
+      module_function
+
+      def ids_like(input_ids, offset: 0, dtype: nil)
+        shape = input_ids.shape
+        seq_len = shape[-1]
+        dtype ||= input_ids.respond_to?(:dtype) ? input_ids.dtype : MLX::Core.int32
+
+        base = MLX::Core.arange(offset.to_i, offset.to_i + seq_len, 1, dtype)
+        return base if shape.length == 1
+
+        reshape_dims = Array.new(shape.length, 1)
+        reshape_dims[-1] = seq_len
+        expanded = MLX::Core.reshape(base, reshape_dims)
+        MLX::Core.broadcast_to(expanded, shape)
+      end
+
+      def offset_from_cache(cache, layer: 0)
+        return 0 if cache.nil?
+        return cache.offset(layer: layer) if cache.respond_to?(:offset)
+
+        if cache.respond_to?(:[]) && !cache[layer].nil?
+          keys, = cache[layer]
+          return keys.shape[2]
+        end
+
+        0
+      end
+    end
+  end
+end

data/lib/mlx/dsl/run_stack.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    def self.run_stack(layers, input, cache: nil, **kwargs)
+      modules = layers.to_a
+      if cache.is_a?(MLX::DSL::KVCache)
+        hidden = input
+        modules.each_with_index do |layer, index|
+          hidden, next_cache = __dsl_run_stack_layer(
+            layer,
+            hidden,
+            kwargs,
+            cache: cache.layer(index),
+            use_cache: true
+          )
+          cache[index] = next_cache
+        end
+        return [hidden, cache]
+      end
+
+      use_cache = !cache.nil?
+      cache_state = if use_cache
+        entries = cache.to_a
+        entries.length < modules.length ? entries + Array.new(modules.length - entries.length) : entries.dup
+      else
+        nil
+      end
+
+      hidden = input
+      modules.each_with_index do |layer, index|
+        layer_cache = use_cache ? cache_state[index] : nil
+        hidden, next_cache = __dsl_run_stack_layer(
+          layer,
+          hidden,
+          kwargs,
+          cache: layer_cache,
+          use_cache: use_cache
+        )
+        cache_state[index] = next_cache if use_cache
+      end
+
+      use_cache ? [hidden, cache_state] : hidden
+    end
+
+    def self.__dsl_run_stack_layer(layer, hidden, kwargs, cache:, use_cache:)
+      call_kwargs = kwargs.dup
+      call_kwargs[:cache] = cache if use_cache
+      result = layer.call(hidden, **call_kwargs)
+
+      if use_cache && result.is_a?(Array) && result.length == 2
+        [result[0], result[1]]
+      else
+        [result, cache]
+      end
+    rescue ArgumentError => e
+      if use_cache && e.message.include?("unknown keyword: :cache")
+        result = layer.call(hidden, **kwargs)
+        if result.is_a?(Array) && result.length == 2
+          return [result[0], result[1]]
+        end
+        return [result, cache]
+      end
+      raise
+    end
+    private_class_method :__dsl_run_stack_layer
+  end
+end

data/lib/mlx/dsl/tensor.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Tensor
+      module_function
+
+      def scatter_rows(base:, row_indices:, values:, axis: nil)
+        axis = __dsl_default_scatter_axis(base, axis)
+        row_indices = __dsl_to_index_array(row_indices)
+        values = __dsl_to_array(values, dtype: base.dtype)
+
+        if base.ndim != values.ndim
+          raise ArgumentError, "base and values must have the same rank"
+        end
+
+        unless base.shape.each_with_index.all? { |dim, idx| idx == axis || dim == values.shape[idx] }
+          raise ArgumentError, "values shape must match base shape except along axis #{axis}"
+        end
+
+        indices = __dsl_expand_indices_for_axis(
+          row_indices: row_indices,
+          values_shape: values.shape,
+          axis: axis
+        )
+        MLX::Core.put_along_axis(base, indices, values, axis)
+      end
+
+      def where_labels(base:, labels:, mapping:, mode: :add_or_replace)
+        mode = mode.to_sym
+        unless [:add_or_replace, :replace].include?(mode)
+          raise ArgumentError, "mode must be :add_or_replace or :replace"
+        end
+
+        out = base
+        mapping.each do |label_value, mapped_value|
+          mask = MLX::Core.equal(labels, label_value)
+          mask = __dsl_expand_trailing_dims(mask, out.ndim)
+          mapped = __dsl_broadcast_mapping(mapped_value, base: out, labels_ndim: labels.ndim)
+          replacement = if mode == :replace
+            mapped
+          else
+            MLX::Core.add(out, mapped)
+          end
+          out = MLX::Core.where(mask, replacement, out)
+        end
+        out
+      end
+
+      def __dsl_default_scatter_axis(base, axis)
+        return axis.to_i unless axis.nil?
+        return 1 if base.ndim == 3
+
+        0
+      end
+      private_class_method :__dsl_default_scatter_axis
+
+      def __dsl_to_index_array(indices)
+        if indices.is_a?(MLX::Core::Array)
+          return indices.astype(MLX::Core.int32)
+        end
+
+        MLX::Core.array(indices, MLX::Core.int32)
+      end
+      private_class_method :__dsl_to_index_array
+
+      def __dsl_to_array(value, dtype:)
+        if value.is_a?(MLX::Core::Array)
+          return value.astype(dtype)
+        end
+
+        MLX::Core.array(value, dtype)
+      end
+      private_class_method :__dsl_to_array
+
+      def __dsl_expand_indices_for_axis(row_indices:, values_shape:, axis:)
+        if row_indices.ndim == 1
+          expected = values_shape[axis]
+          if row_indices.shape[0] != expected
+            raise ArgumentError, "row_indices length must match values shape at axis #{axis}"
+          end
+
+          reshape = Array.new(values_shape.length, 1)
+          reshape[axis] = expected
+          base = MLX::Core.reshape(row_indices, reshape)
+          return MLX::Core.broadcast_to(base, values_shape)
+        end
+
+        if row_indices.shape == values_shape
+          return row_indices
+        end
+
+        raise ArgumentError, "row_indices must be rank-1 or match values shape"
+      end
+      private_class_method :__dsl_expand_indices_for_axis
+
+      def __dsl_expand_trailing_dims(array, target_ndim)
+        out = array
+        while out.ndim < target_ndim
+          out = MLX::Core.expand_dims(out, out.ndim)
+        end
+        out
+      end
+      private_class_method :__dsl_expand_trailing_dims
+
+      def __dsl_broadcast_mapping(value, base:, labels_ndim:)
+        mapped = __dsl_to_array(value, dtype: base.dtype)
+        return mapped if mapped.shape == base.shape
+
+        if mapped.ndim == 1 && mapped.shape[0] == base.shape[-1]
+          reshape = Array.new(labels_ndim, 1) + [mapped.shape[0]]
+          mapped = MLX::Core.reshape(mapped, reshape)
+          return MLX::Core.broadcast_to(mapped, base.shape)
+        end
+
+        if mapped.ndim == labels_ndim && mapped.shape == base.shape[0...labels_ndim]
+          mapped = __dsl_expand_trailing_dims(mapped, base.ndim)
+          return MLX::Core.broadcast_to(mapped, base.shape)
+        end
+
+        MLX::Core.broadcast_to(mapped, base.shape)
+      end
+      private_class_method :__dsl_broadcast_mapping
+    end
+  end
+end

data/lib/mlx/dsl/transformer_block.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class FeedForward < MLX::NN::Module
+      def initialize(dims:, hidden_dims:, kind: :gelu, bias: false)
+        super()
+        @kind = kind.to_sym
+        @dims = Integer(dims)
+        @hidden_dims = Integer(hidden_dims)
+
+        case @kind
+        when :swiglu
+          self.gate_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.down_proj = MLX::NN::Linear.new(@hidden_dims, @dims, bias: bias)
+        else
+          self.in_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.out_proj = MLX::NN::Linear.new(@hidden_dims, @dims, bias: bias)
+        end
+      end
+
+      def call(x)
+        case @kind
+        when :swiglu
+          gated = MLX::NN.silu(gate_proj.call(x))
+          down_proj.call(MLX::Core.multiply(gated, up_proj.call(x)))
+        when :relu
+          out_proj.call(MLX::NN.relu(in_proj.call(x)))
+        else
+          out_proj.call(MLX::NN.gelu(in_proj.call(x)))
+        end
+      end
+    end
+
+    class TransformerBlock < MLX::NN::Module
+      def initialize(
+        dims:,
+        num_heads:,
+        kv_heads: nil,
+        ffn_dims: nil,
+        norm: :rms,
+        norm_eps: 1e-5,
+        ffn: nil,
+        rope: nil,
+        qkv_bias: false,
+        backend: :sdpa,
+        cache: false
+      )
+        super()
+
+        ffn_config = (ffn || {}).transform_keys(&:to_sym)
+        ffn_kind = ffn_config.fetch(:kind, :gelu)
+        hidden_dims = ffn_dims || ffn_config.fetch(:hidden_dims, Integer(dims) * 4)
+        ffn_bias = ffn_config.fetch(:bias, false)
+
+        self.attention_norm = __dsl_build_norm(norm, Integer(dims), norm_eps)
+        self.attention = MLX::DSL::Attention.new(
+          dims: dims,
+          num_heads: num_heads,
+          kv_heads: kv_heads,
+          qkv_bias: qkv_bias,
+          backend: backend,
+          rope: rope,
+          cache: cache
+        )
+        self.ffn_norm = __dsl_build_norm(norm, Integer(dims), norm_eps)
+        self.feed_forward = MLX::DSL::FeedForward.new(
+          dims: dims,
+          hidden_dims: hidden_dims,
+          kind: ffn_kind,
+          bias: ffn_bias
+        )
+
+        @cache_enabled = !!cache
+      end
+
+      def call(x, mask: nil, cache: nil, **_kwargs)
+        attn_input = attention_norm.call(x)
+        attn_result = attention.call(attn_input, attn_input, attn_input, mask: mask, cache: cache)
+        if attn_result.is_a?(Array) && attn_result.length == 2
+          attn_out, next_cache = attn_result
+        else
+          attn_out = attn_result
+          next_cache = cache
+        end
+
+        hidden = MLX::Core.add(x, attn_out)
+        ffn_out = feed_forward.call(ffn_norm.call(hidden))
+        output = MLX::Core.add(hidden, ffn_out)
+
+        if @cache_enabled || !cache.nil?
+          [output, next_cache]
+        else
+          output
+        end
+      end
+
+      private
+
+      def __dsl_build_norm(kind, dims, eps)
+        case kind.to_sym
+        when :layer, :layer_norm
+          MLX::NN::LayerNorm.new(dims, eps: eps)
+        when :rms, :rms_norm
+          MLX::NN::RMSNorm.new(dims, eps: eps)
+        else
+          raise ArgumentError, "unsupported norm kind: #{kind.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/weight_map.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    def self.weight_map(&block)
+      mapper = WeightMap.new
+      mapper.instance_eval(&block) if block_given?
+      mapper
+    end
+
+    class WeightMap
+      def initialize
+        @rules = []
+      end
+
+      def strip_prefix(prefix)
+        @rules << [:strip_prefix, prefix.to_s]
+        self
+      end
+
+      def rename(from = nil, to = nil)
+        if from.is_a?(Hash)
+          from.each do |src, dst|
+            @rules << [:rename, src.to_s, dst.to_s]
+          end
+          return self
+        end
+
+        if from.nil? || to.nil?
+          raise ArgumentError, "rename requires either a mapping hash or from/to arguments"
+        end
+
+        @rules << [:rename, from.to_s, to.to_s]
+        self
+      end
+
+      def regex(pattern, replacement = nil, &block)
+        if pattern.nil?
+          raise ArgumentError, "regex requires a Regexp pattern"
+        end
+        if replacement.nil? && !block_given?
+          raise ArgumentError, "regex requires a replacement argument or block"
+        end
+
+        @rules << [:regex, pattern, replacement, block]
+        self
+      end
+
+      def split_qkv(source, into:, axis: 0)
+        names = into.to_a.map(&:to_s)
+        if names.empty? || names.length < 2
+          raise ArgumentError, "split_qkv :into must include at least two output names"
+        end
+
+        @rules << [:split, source.to_s, names, axis.to_i]
+        self
+      end
+
+      def transpose_if(rank:, order:)
+        @rules << [:transpose_if, rank.to_i, order.to_a]
+        self
+      end
+
+      def apply(weights)
+        entries = if weights.is_a?(Hash)
+          weights.to_a
+        elsif weights.respond_to?(:to_a)
+          weights.to_a
+        else
+          raise ArgumentError, "weights must be a Hash or array-like key/value collection"
+        end
+
+        out = {}
+        entries.each do |entry|
+          key, value = entry
+          __dsl_apply_rules_for_entry(key.to_s, value).each do |mapped_key, mapped_value|
+            out[mapped_key] = mapped_value
+          end
+        end
+        out
+      end
+
+      private
+
+      def __dsl_apply_rules_for_entry(key, value)
+        current_entries = [[key, value]]
+        @rules.each do |rule|
+          current_entries = current_entries.flat_map do |curr_key, curr_value|
+            __dsl_apply_rule(rule, curr_key, curr_value)
+          end
+        end
+        current_entries
+      end
+
+      def __dsl_apply_rule(rule, key, value)
+        kind = rule[0]
+        case kind
+        when :strip_prefix
+          prefix = rule[1]
+          if key.start_with?(prefix)
+            [[key[prefix.length..], value]]
+          else
+            [[key, value]]
+          end
+        when :rename
+          from = rule[1]
+          to = rule[2]
+          [[key.gsub(from, to), value]]
+        when :regex
+          pattern = rule[1]
+          replacement = rule[2]
+          block = rule[3]
+          if block.nil?
+            [[key.gsub(pattern, replacement.to_s), value]]
+          else
+            [[key.gsub(pattern, &block), value]]
+          end
+        when :split
+          source = rule[1]
+          targets = rule[2]
+          axis = rule[3]
+          return [[key, value]] unless key == source
+
+          parts = MLX::Core.split(value, targets.length, axis)
+          targets.each_with_index.map { |target, index| [target, parts[index]] }
+        when :transpose_if
+          target_rank = rule[1]
+          order = rule[2]
+          if value.respond_to?(:shape) && value.shape.length == target_rank
+            [[key, MLX::Core.transpose(value, order)]]
+          else
+            [[key, value]]
+          end
+        else
+          [[key, value]]
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl.rb

@@ -10,6 +10,16 @@ require_relative "dsl/data_pipeline"
 require_relative "dsl/experiment"
 require_relative "dsl/split_plan"
 require_relative "dsl/builder"
+require_relative "dsl/config_schema"
+require_relative "dsl/weight_map"
+require_relative "dsl/kv_cache"
+require_relative "dsl/masks"
+require_relative "dsl/positions"
+require_relative "dsl/tensor"
+require_relative "dsl/run_stack"
+require_relative "dsl/attention"
+require_relative "dsl/transformer_block"
+require_relative "dsl/generate"
 require_relative "dsl/train_step"
 require_relative "dsl/model_mixin"
 require_relative "dsl/model"

data/lib/mlx/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MLX
-  VERSION = "0.30.7.2"
+  VERSION = "0.30.7.3"
 end

metadata

@@ -1,15 +1,57 @@
 --- !ruby/object:Gem::Specification
 name: mlx
 version: !ruby/object:Gem::Version
-  version: 0.30.7.2
+  version: 0.30.7.3
 platform: ruby
 authors:
 - MLX Contributors
 - Aleksey Skryl
 bindir: bin
 cert_chain: []
-date: 2026-02-14 00:00:00.000000000 Z
-dependencies: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: benchmark
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 description: A Ruby wrapper for the native MLX machine learning runtime.
 email:
 - mlx@group.apple.com
@@ -27,15 +69,25 @@ files:
 - lib/mlx/distributed_utils/config.rb
 - lib/mlx/distributed_utils/launch.rb
 - lib/mlx/dsl.rb
+- lib/mlx/dsl/attention.rb
 - lib/mlx/dsl/builder.rb
+- lib/mlx/dsl/config_schema.rb
 - lib/mlx/dsl/data_pipeline.rb
 - lib/mlx/dsl/experiment.rb
+- lib/mlx/dsl/generate.rb
 - lib/mlx/dsl/graph_modules.rb
+- lib/mlx/dsl/kv_cache.rb
+- lib/mlx/dsl/masks.rb
 - lib/mlx/dsl/model.rb
 - lib/mlx/dsl/model_mixin.rb
+- lib/mlx/dsl/positions.rb
+- lib/mlx/dsl/run_stack.rb
 - lib/mlx/dsl/split_plan.rb
+- lib/mlx/dsl/tensor.rb
 - lib/mlx/dsl/train_step.rb
 - lib/mlx/dsl/trainer.rb
+- lib/mlx/dsl/transformer_block.rb
+- lib/mlx/dsl/weight_map.rb
 - lib/mlx/extension.rb
 - lib/mlx/nn.rb
 - lib/mlx/nn/base.rb
@@ -640,14 +692,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3.1'
+      version: '3.3'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.6.2
+rubygems_version: 4.0.3
 specification_version: 4
 summary: Ruby bindings for the native MLX library
 test_files: []