ignis-dl 0.0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 2e2b7a2b0882360bf9c07ca60b72aa71c3a43dbb9739aa9dde31182e8e9068ca
+  data.tar.gz: a34343e8f5e2334a79b8a5b6dc935891d8776ffe5b64f47832dccb2f87ce9fcc
+SHA512:
+  metadata.gz: 525582d82f78c61c0fedf0a31273f6010453c5f3e1753ff2908b6ba402f03de909aa6a7eb1dc12c0feb773e7a65c519f6be11a4dc44db6307f5aed4fe1baa65d
+  data.tar.gz: f7c0f62d4cdcb6ddcf8daf3af78a1d8580663f91dab1ef6ac8219c9b2ce445e6858368b9782ce779b0710ca403579c58505ff2079297599e6d454b3cad5c2d0e

data/README.md

@@ -0,0 +1,15 @@
+# ignis-dl
+
+Neural-net layers + transformers for Ruby, on the [Ignis](https://github.com/tigel-agm/Ignis) GPU/autograd stack. Installing this pulls the whole chain (`ignis` + `ignis-autograd`).
+
+- **NN:** Linear, Embedding, LayerNorm, RMSNorm, Dropout; SGD/Adam/AdamW; losses.
+- **Transformers:** multi-head + grouped-query attention, RoPE, SwiGLU, KV cache.
+- **Loaders:** real GPT-2 and Llama (HuggingFace `safetensors`) — matches HF logits.
+
+```ruby
+require "ignis-dl"
+model = Ignis::AI::LlamaLoader.from_pretrained("path/to/Llama-3.2-1B")
+# ... matches HuggingFace; "The capital of France is" => " Paris"
+```
+
+Loads real GPT-2 and Llama-3.2 (cosine 1.0 vs HuggingFace) and trains transformers from scratch. MIT.

data/lib/ignis-dl.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+# ignis-dl — neural-net layers + transformers for Ruby, on the Ignis GPU/autograd
+# stack. Loads real GPT-2 and Llama checkpoints (matching HuggingFace) and trains
+# transformers from scratch. Requiring this pulls the whole chain (ignis +
+# ignis-autograd), so it doubles as the meta-gem.
+
+require "ignis-autograd"
+
+# --- NN modules ---
+require_relative "nnw/ai/nn/module"
+require_relative "nnw/ai/nn/linear"
+require_relative "nnw/ai/nn/embedding"
+require_relative "nnw/ai/nn/layer_norm"
+require_relative "nnw/ai/nn/rms_norm"
+require_relative "nnw/ai/nn/dropout"
+require_relative "nnw/ai/nn/sequential"
+
+# --- Transformer ---
+require_relative "nnw/ai/transformer/attention"
+require_relative "nnw/ai/transformer/feed_forward"
+require_relative "nnw/ai/transformer/swiglu"
+require_relative "nnw/ai/transformer/block"
+require_relative "nnw/ai/transformer/model"
+require_relative "nnw/ai/transformer/modern"
+require_relative "nnw/ai/kv_cache"
+
+# --- Loss ---
+require_relative "nnw/ai/loss"
+
+# --- Optimizers ---
+require_relative "nnw/ai/optim/base"
+require_relative "nnw/ai/optim/sgd"
+require_relative "nnw/ai/optim/adam"
+require_relative "nnw/ai/optim/adamw"
+require_relative "nnw/ai/optim/lr_scheduler"
+
+# --- Weight loading + tokenizer ---
+require_relative "nnw/ai/safetensors"
+require_relative "nnw/ai/weight_map"
+require_relative "nnw/ai/gpt2_loader"
+require_relative "nnw/ai/llama_loader"
+require_relative "nnw/ai/tokenizer"
+
+# --- Inference + training ---
+require_relative "nnw/ai/inference"
+require_relative "nnw/ai/trainer"
+require_relative "nnw/ai/server"

data/lib/nnw/ai/gpt2_loader.rb

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Ignis
+  module AI
+    # Loads HuggingFace GPT-2 safetensors checkpoints into an Ignis Transformer::Model.
+    #
+    # GPT-2 is NOT a drop-in for nn.Linear: its attention/MLP projections are
+    # `Conv1D` layers whose weights are stored as [in, out] (transposed vs.
+    # nn.Linear's [out, in]), the QKV projection is a single fused [embed, 3*embed]
+    # tensor, and the LM head is tied to the token embedding. This loader applies
+    # those transforms while copying weights in.
+    #
+    # @example
+    #   model = Ignis::AI::Transformer::Model.gpt2_small
+    #   Ignis::AI::GPT2Loader.load(model, "gpt2")
+    #   model.eval!
+    module GPT2Loader
+      module_function
+
+      # Load weights from <dir>/model.safetensors into +model+ (in place).
+      # @param model [Ignis::AI::Transformer::Model]
+      # @param dir [String] directory holding model.safetensors
+      # @return [model]
+      def load(model, dir)
+        path = File.join(dir, "model.safetensors")
+        raise ArgumentError, "not found: #{path}" unless File.exist?(path)
+
+        header, data_offset = read_header(path)
+        params = model.named_parameters
+        embed = model.embed_dim
+        n_layer = model.num_layers
+
+        File.open(path, "rb") do |f|
+          get = lambda do |name|
+            meta = header[name] or raise "missing safetensors tensor: #{name}"
+            b, e = meta["data_offsets"]
+            f.seek(data_offset + b)
+            raw = f.read(e - b)
+            { values: raw.unpack("e*"), shape: meta["shape"] } # F32 little-endian
+          end
+
+          set = lambda do |pname, values|
+            p = params[pname] or raise "missing model param: #{pname}"
+            unless p.numel == values.length
+              raise "size mismatch for #{pname}: param has #{p.numel}, data has #{values.length}"
+            end
+            p.data.from_host(values)
+          end
+
+          # Embeddings (direct)
+          set.call("token_embedding.weight", get.call("wte.weight")[:values])
+          set.call("position_embedding.weight", get.call("wpe.weight")[:values])
+
+          n_layer.times do |i|
+            set.call("blocks.#{i}.norm1.weight", get.call("h.#{i}.ln_1.weight")[:values])
+            set.call("blocks.#{i}.norm1.bias",   get.call("h.#{i}.ln_1.bias")[:values])
+            set.call("blocks.#{i}.norm2.weight", get.call("h.#{i}.ln_2.weight")[:values])
+            set.call("blocks.#{i}.norm2.bias",   get.call("h.#{i}.ln_2.bias")[:values])
+
+            # Fused QKV: c_attn.weight is Conv1D [embed, 3*embed]. Split columns into
+            # q|k|v and transpose each into nn.Linear layout [out, in] = [embed, embed].
+            cattn_w = get.call("h.#{i}.attn.c_attn.weight")[:values]
+            cattn_b = get.call("h.#{i}.attn.c_attn.bias")[:values]
+            set.call("blocks.#{i}.attention.q_proj.weight", transpose_slice_cols(cattn_w, embed, 3 * embed, 0 * embed, embed))
+            set.call("blocks.#{i}.attention.k_proj.weight", transpose_slice_cols(cattn_w, embed, 3 * embed, 1 * embed, embed))
+            set.call("blocks.#{i}.attention.v_proj.weight", transpose_slice_cols(cattn_w, embed, 3 * embed, 2 * embed, embed))
+            set.call("blocks.#{i}.attention.q_proj.bias", cattn_b[(0 * embed)...(1 * embed)])
+            set.call("blocks.#{i}.attention.k_proj.bias", cattn_b[(1 * embed)...(2 * embed)])
+            set.call("blocks.#{i}.attention.v_proj.bias", cattn_b[(2 * embed)...(3 * embed)])
+
+            # Attention output projection: Conv1D [embed, embed] -> transpose.
+            aproj = get.call("h.#{i}.attn.c_proj.weight")[:values]
+            set.call("blocks.#{i}.attention.out_proj.weight", transpose(aproj, embed, embed))
+            set.call("blocks.#{i}.attention.out_proj.bias", get.call("h.#{i}.attn.c_proj.bias")[:values])
+
+            # MLP fc1: c_fc Conv1D [embed, ff] -> transpose -> [ff, embed].
+            fc = get.call("h.#{i}.mlp.c_fc.weight")
+            ff = fc[:shape][1]
+            set.call("blocks.#{i}.feed_forward.fc1.weight", transpose(fc[:values], embed, ff))
+            set.call("blocks.#{i}.feed_forward.fc1.bias", get.call("h.#{i}.mlp.c_fc.bias")[:values])
+
+            # MLP fc2: c_proj Conv1D [ff, embed] -> transpose -> [embed, ff].
+            fp = get.call("h.#{i}.mlp.c_proj.weight")[:values]
+            set.call("blocks.#{i}.feed_forward.fc2.weight", transpose(fp, ff, embed))
+            set.call("blocks.#{i}.feed_forward.fc2.bias", get.call("h.#{i}.mlp.c_proj.bias")[:values])
+          end
+
+          set.call("norm.weight", get.call("ln_f.weight")[:values])
+          set.call("norm.bias",   get.call("ln_f.bias")[:values])
+
+          # LM head is tied to the token embedding (logits = x @ wte^T).
+          set.call("head.weight", get.call("wte.weight")[:values])
+        end
+
+        model
+      end
+
+      # Read the safetensors JSON header. @return [[Hash, Integer]] header, data_offset
+      def read_header(path)
+        File.open(path, "rb") do |f|
+          n = f.read(8).unpack1("Q<")
+          h = JSON.parse(f.read(n))
+          h.delete("__metadata__")
+          [h, 8 + n]
+        end
+      end
+
+      # Transpose a row-major [rows, cols] flat array → [cols, rows].
+      def transpose(arr, rows, cols)
+        out = Array.new(rows * cols)
+        r = 0
+        while r < rows
+          base = r * cols
+          c = 0
+          while c < cols
+            out[c * rows + r] = arr[base + c]
+            c += 1
+          end
+          r += 1
+        end
+        out
+      end
+
+      # Slice columns [col_off, col_off+len) from a [rows, total_cols] array, then
+      # transpose → [len, rows]. (Used to split the fused QKV Conv1D weight.)
+      def transpose_slice_cols(arr, rows, total_cols, col_off, len)
+        out = Array.new(rows * len)
+        r = 0
+        while r < rows
+          base = r * total_cols + col_off
+          c = 0
+          while c < len
+            out[c * rows + r] = arr[base + c]
+            c += 1
+          end
+          r += 1
+        end
+        out
+      end
+    end
+  end
+end

data/lib/nnw/ai/inference.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    # TextGenerator — autoregressive inference with KV cache.
+    #
+    # Loads a Transformer model + tokenizer, generates text token-by-token
+    # with top-k/top-p sampling and streaming support via EventBus.
+    class TextGenerator
+      # @return [Transformer::Model]
+      attr_reader :model
+
+      # @return [Tokenizer]
+      attr_reader :tokenizer
+
+      # @param model [Transformer::Model]
+      # @param tokenizer [Tokenizer]
+      def initialize(model, tokenizer)
+        @model = model
+        @tokenizer = tokenizer
+        @model.eval!
+      end
+
+      # Generate text from a prompt.
+      # @param prompt [String] input text
+      # @param max_tokens [Integer] maximum new tokens to generate
+      # @param temperature [Float] sampling temperature (0.0 = greedy)
+      # @param top_k [Integer] top-k sampling (0 = disabled)
+      # @param top_p [Float] nucleus sampling threshold (1.0 = disabled)
+      # @param stop_tokens [Array<String>] strings that trigger early stop
+      # @param stream [Boolean] yield tokens as they are generated
+      # @yield [String] each generated token (when stream: true)
+      # @return [String] full generated text
+      def generate(prompt, max_tokens: 128, temperature: 0.7,
+                   top_k: 50, top_p: 0.9, stop_tokens: [], stream: false, &block)
+        input_ids = @tokenizer.encode(prompt)
+        generated_ids = []
+        cache = @model.make_kv_cache
+
+        Tape.no_grad do
+          # Prefill: stream the prompt through the cache one token at a time. Each
+          # decode_step appends that token's K/V; the last token's logits predict
+          # the first generated token. O(prefix) total, like one full forward.
+          last_logits = nil
+          input_ids.each do |tid|
+            break if cache.full?
+            last_logits = @model.decode_step(tid, cache).to_host
+          end
+
+          # Decode: sample from the running logits, feed the sampled token back to
+          # extend the cache by one (O(prefix) per step, not O(prefix²) re-forward).
+          max_tokens.times do |step|
+            break if last_logits.nil? || cache.full?
+
+            next_id = sample_token(last_logits, temperature: temperature,
+                                   top_k: top_k, top_p: top_p)
+            generated_ids << next_id
+
+            token_text = @tokenizer.decode([next_id], skip_special_tokens: true)
+            block.call(token_text) if stream && block
+
+            if defined?(Ignis::Shared::EventBus)
+              Ignis::Shared::EventBus.publish(:token_generated, {
+                text: token_text, token_id: next_id, step: step
+              })
+            end
+
+            full_text = @tokenizer.decode(generated_ids, skip_special_tokens: true)
+            break if stop_tokens.any? { |s| full_text.include?(s) }
+            break if @tokenizer.special_token_ids.include?(next_id) # EOS
+
+            # Advance: logits conditioned on prompt + everything generated so far.
+            last_logits = @model.decode_step(next_id, cache).to_host
+          end
+        end
+
+        @tokenizer.decode(input_ids + generated_ids, skip_special_tokens: true)
+      end
+
+      # Compute embeddings for text (useful for similarity/search).
+      # @param text [String]
+      # @return [Array<Float>] embedding vector
+      def embed(text)
+        Tape.no_grad do
+          input_ids = @tokenizer.encode(text)
+          input_tensor = Tensor.from_host(input_ids, shape: [1, input_ids.length],
+                                           dtype: :int32, device_id: 0)
+          # Get token embeddings, mean pool
+          tok_emb = @model.instance_variable_get(:@token_embedding).call(input_tensor)
+          tok_emb.mean.to_host
+        end
+      end
+
+      private
+
+      # Sample a token from logits.
+      # @param logits [Array<Float>]
+      # @param temperature [Float]
+      # @param top_k [Integer]
+      # @param top_p [Float]
+      # @return [Integer] sampled token id
+      def sample_token(logits, temperature:, top_k:, top_p:)
+        return logits.each_with_index.max_by { |v, _| v }[1] if temperature <= 0.0
+
+        # Temperature scaling
+        scaled = logits.map { |l| l / temperature }
+
+        # Top-k filtering
+        if top_k > 0 && top_k < scaled.length
+          threshold = scaled.sort.reverse[top_k - 1]
+          scaled = scaled.map { |l| l >= threshold ? l : -Float::INFINITY }
+        end
+
+        # Softmax
+        max_val = scaled.max
+        exps = scaled.map { |l| Math.exp(l - max_val) }
+        sum = exps.sum
+        probs = exps.map { |e| e / sum }
+
+        # Top-p filtering
+        if top_p < 1.0
+          sorted_indices = probs.each_with_index.sort_by { |p, _| -p }.map { |_, i| i }
+          cumsum = 0.0
+          cutoff_idx = sorted_indices.length
+          sorted_indices.each_with_index do |idx, rank|
+            cumsum += probs[idx]
+            if cumsum > top_p
+              cutoff_idx = rank + 1
+              break
+            end
+          end
+          allowed = Set.new(sorted_indices[0...cutoff_idx])
+          probs = probs.each_with_index.map { |p, i| allowed.include?(i) ? p : 0.0 }
+          sum = probs.sum
+          probs = probs.map { |p| p / sum } if sum > 0
+        end
+
+        # Multinomial sampling
+        r = Kernel.rand
+        cumulative = 0.0
+        probs.each_with_index do |p, i|
+          cumulative += p
+          return i if r <= cumulative
+        end
+
+        probs.length - 1  # Fallback
+      end
+    end
+
+    # BatchProcessor — concurrent inference with dynamic batching.
+    class BatchProcessor
+      # @param model [Transformer::Model]
+      # @param tokenizer [Tokenizer]
+      # @param max_batch_size [Integer]
+      # @param max_wait_ms [Integer] max milliseconds to wait for batch fill
+      def initialize(model, tokenizer, max_batch_size: 8, max_wait_ms: 50)
+        @generator = TextGenerator.new(model, tokenizer)
+        @max_batch_size = max_batch_size
+        @max_wait_ms = max_wait_ms
+        @queue = Queue.new
+        @running = false
+      end
+
+      # Submit a request for processing.
+      # @param prompt [String]
+      # @param params [Hash] generation parameters
+      # @return [String] generated text (blocks until complete)
+      def submit(prompt, **params)
+        result_queue = Queue.new
+        @queue << { prompt: prompt, params: params, result: result_queue }
+        result_queue.pop
+      end
+
+      # Start the batch processing loop in a background thread.
+      # @return [Thread]
+      def start!
+        @running = true
+        @thread = Thread.new { batch_loop }
+        @thread
+      end
+
+      # Stop the processor.
+      # @return [void]
+      def stop!
+        @running = false
+        @thread&.join(5)
+      end
+
+      private
+
+      def batch_loop
+        while @running
+          batch = []
+          deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @max_wait_ms / 1000.0
+
+          # Collect up to max_batch_size requests
+          while batch.size < @max_batch_size
+            remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            break if remaining <= 0
+
+            begin
+              request = @queue.pop(true) rescue nil
+              batch << request if request
+            rescue ThreadError
+              sleep(0.001)
+            end
+          end
+
+          next if batch.empty?
+
+          # Process each request (serial for now — batched matmul is future optimization)
+          batch.each do |req|
+            begin
+              result = @generator.generate(req[:prompt], **req[:params])
+              req[:result] << result
+            rescue => e
+              req[:result] << "Error: #{e.message}"
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/kv_cache.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    # KVCache — per-layer key/value cache for O(1)-prefix autoregressive decoding.
+    #
+    # Without a cache, generating token t re-runs attention over the entire prefix
+    # (O(t) work that grows every step → O(n²) generation). The cache stores each
+    # layer's projected K and V once; each new token only projects ITS key/value,
+    # appends them, and attends its single query over all cached keys.
+    #
+    # Buffers are preallocated to [max_seq_len, embed_dim] per layer so appending a
+    # row is an O(row) device→device copy (NvArray#write_rows!) rather than a
+    # realloc + full recopy. A contiguous [length+1, embed] view (slice along dim 0)
+    # exposes the live region — including the row just appended this step.
+    class KVCache
+      # @return [Integer] number of tokens currently cached (positions 0..length-1)
+      attr_reader :length
+      # @return [Integer]
+      attr_reader :max_seq_len, :num_layers, :embed_dim
+
+      # @param num_layers [Integer]
+      # @param max_seq_len [Integer] capacity (generation cannot exceed this)
+      # @param embed_dim [Integer]
+      # @param device_id [Integer]
+      def initialize(num_layers:, max_seq_len:, embed_dim:, device_id: 0)
+        @num_layers = num_layers
+        @max_seq_len = max_seq_len
+        @embed_dim = embed_dim
+        @device_id = device_id
+        @length = 0
+        @k = Array.new(num_layers) { make_buffer }
+        @v = Array.new(num_layers) { make_buffer }
+      end
+
+      # Append a layer's K/V for the current token at row +length+.
+      # @param layer [Integer] layer index
+      # @param k_row [Ignis::Shared::NvArray] [1, embed] key projection of the new token
+      # @param v_row [Ignis::Shared::NvArray] [1, embed] value projection of the new token
+      # @return [void]
+      def append(layer, k_row, v_row)
+        @k[layer].write_rows!(k_row, @length)
+        @v[layer].write_rows!(v_row, @length)
+      end
+
+      # Contiguous [length+1, embed] view of cached keys for +layer+, INCLUDING the
+      # row just appended this step. (Call after #append, before #advance!.)
+      # @param layer [Integer]
+      # @return [Ignis::Shared::NvArray] non-owning view
+      def k_view(layer)
+        @k[layer].slice(0, 0, @length + 1)
+      end
+
+      # @param layer [Integer]
+      # @return [Ignis::Shared::NvArray] non-owning view of cached values
+      def v_view(layer)
+        @v[layer].slice(0, 0, @length + 1)
+      end
+
+      # Advance to the next position after every layer has appended this token.
+      # @return [void]
+      def advance!
+        @length += 1
+      end
+
+      # @return [Boolean] true if the cache is at capacity (no more tokens fit)
+      def full?
+        @length >= @max_seq_len
+      end
+
+      private
+
+      def make_buffer
+        Ignis::Shared::NvArray.new(shape: [@max_seq_len, @embed_dim],
+                                 dtype: :float32, device_id: @device_id).to_device
+      end
+    end
+  end
+end

data/lib/nnw/ai/llama_loader.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Ignis
+  module AI
+    # LlamaLoader — load HuggingFace Llama-family checkpoints (Llama-3.x, and other
+    # LlamaForCausalLM models: SmolLM, TinyLlama, …) into a Transformer::ModernModel.
+    #
+    # Reads config.json to size the model (RoPE base/scaling, GQA head counts, RMSNorm
+    # eps, tied embeddings), then loads model.safetensors, dequantizing bf16 weights to
+    # fp32 on-device. HF's nn.Linear weights are [out, in] and applied as x·Wᵀ — the
+    # SAME convention as Ignis::AI::NN::Linear — so weights map across with no transpose
+    # (unlike GPT-2's Conv1D).
+    module LlamaLoader
+      class << self
+        # Build a ModernModel from config.json and load its weights.
+        # @param dir [String] directory containing config.json + model.safetensors
+        # @param device_id [Integer]
+        # @return [Transformer::ModernModel]
+        def from_pretrained(dir, device_id: 0)
+          cfg = JSON.parse(File.read(File.join(dir, "config.json")))
+          model = Transformer::ModernModel.new(
+            vocab_size:   cfg["vocab_size"],
+            embed_dim:    cfg["hidden_size"],
+            num_heads:    cfg["num_attention_heads"],
+            num_kv_heads: cfg["num_key_value_heads"] || cfg["num_attention_heads"],
+            num_layers:   cfg["num_hidden_layers"],
+            ff_dim:       cfg["intermediate_size"],
+            max_seq_len:  cfg["max_position_embeddings"],
+            rope_base:    (cfg["rope_theta"] || 10000.0).to_f,
+            rope_scaling: cfg["rope_scaling"],
+            head_dim:     cfg["head_dim"],
+            eps:          (cfg["rms_norm_eps"] || 1e-5).to_f,
+            device_id:    device_id
+          )
+          load(model, dir, device_id: device_id)
+          model
+        end
+
+        # Load weights from dir/model.safetensors into an existing ModernModel.
+        # @return [Integer] number of parameters loaded
+        def load(model, dir, device_id: 0)
+          tensors = Safetensors.load(File.join(dir, "model.safetensors"), device_id: device_id)
+          embed_src = tensors["model.embed_tokens.weight"]
+          raise "LlamaLoader: missing model.embed_tokens.weight" unless embed_src
+          lm_head_src = tensors["lm_head.weight"] # nil when embeddings are tied
+
+          count = 0
+          model.named_parameters.each do |name, param|
+            src = if name == "head.weight"
+                    lm_head_src || embed_src   # untied lm_head, else tied to embeddings
+                  else
+                    tensors[hf_name(name)]
+                  end
+            raise "LlamaLoader: no source weight for param #{name} (#{hf_name(name)})" unless src
+            unless param.shape == src.shape
+              raise "LlamaLoader: shape mismatch for #{name}: model #{param.shape} vs file #{src.shape}"
+            end
+            dequant_into!(src.data, param.data)
+            count += 1
+          end
+          Ignis.synchronize
+          count
+        end
+
+        private
+
+        # Translate an Ignis ModernModel parameter name → HuggingFace Llama weight name.
+        # ("head.weight" is resolved in #load, since it may be tied or an explicit lm_head.)
+        def hf_name(name)
+          case name
+          when "token_embedding.weight" then "model.embed_tokens.weight"
+          when "norm.weight"            then "model.norm.weight"
+          when /\Ablocks\.(\d+)\.norm1\.weight\z/ then "model.layers.#{$1}.input_layernorm.weight"
+          when /\Ablocks\.(\d+)\.norm2\.weight\z/ then "model.layers.#{$1}.post_attention_layernorm.weight"
+          when /\Ablocks\.(\d+)\.attn\.([qkvo])_proj\.weight\z/ then "model.layers.#{$1}.self_attn.#{$2}_proj.weight"
+          when /\Ablocks\.(\d+)\.mlp\.(gate|up|down)\.weight\z/ then "model.layers.#{$1}.mlp.#{$2}_proj.weight"
+          else raise "LlamaLoader: unmapped parameter #{name}"
+          end
+        end
+
+        # Dequantize a bf16 source NvArray into an fp32 destination NvArray (on-device).
+        def dequant_into!(src_nv, dst_nv)
+          unless src_nv.numel == dst_nv.numel
+            raise "LlamaLoader: numel mismatch #{src_nv.numel} vs #{dst_nv.numel}"
+          end
+          if src_nv.dtype == :float32
+            # already fp32 (some checkpoints) — straight device→device copy
+            dst_nv.write_rows!(src_nv, 0)
+            return
+          end
+          n = dst_nv.numel
+          k = Ignis::JIT::Kernels::Elementwise.bf16_to_f32
+          k.launch(grid: [(n + 255) / 256], block: [256], args: [src_nv, dst_nv, n])
+        end
+      end
+    end
+  end
+end