assignment1: implement load/isContiguous/view/permute/slice

.gitignore

@@ -31,6 +31,10 @@ cache/
 #GGUF
 *.gguf
 
+# Downloaded model weights
+model/
+*.safetensors
+
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

include/llaisys/models/qwen2.h

@@ -37,6 +37,10 @@ __C {
 
     __export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model * model);
 
+    // Appends ntoken tokens after whatever is already in the KV cache and returns the
+    // argmax-sampled next token. Call llaisysQwen2ModelResetCache before a new sequence.
     __export int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken);
+
+    __export void llaisysQwen2ModelResetCache(struct LlaisysQwen2Model * model);
 }
 #endif // LLAISYS_MODELS_QWEN2_H

python/llaisys/libllaisys/__init__.py

@@ -12,6 +12,8 @@
 from .tensor import llaisysTensor_t
 from .tensor import load_tensor
 from .ops import load_ops
+from .qwen2 import load_qwen2
+from .qwen2 import LlaisysQwen2Meta, LlaisysQwen2Weights
 
 
 def load_shared_library():
@@ -38,6 +40,7 @@ def load_shared_library():
 load_runtime(LIB_LLAISYS)
 load_tensor(LIB_LLAISYS)
 load_ops(LIB_LLAISYS)
+load_qwen2(LIB_LLAISYS)
 
 
 __all__ = [
@@ -52,4 +55,6 @@ def load_shared_library():
     "llaisysMemcpyKind_t",
     "MemcpyKind",
     "llaisysStream_t",
+    "LlaisysQwen2Meta",
+    "LlaisysQwen2Weights",
 ]

python/llaisys/libllaisys/qwen2.py

@@ -0,0 +1,78 @@
+from ctypes import (
+    POINTER,
+    Structure,
+    c_float,
+    c_int,
+    c_int64,
+    c_size_t,
+    c_void_p,
+)
+
+from .llaisys_types import llaisysDataType_t, llaisysDeviceType_t
+from .tensor import llaisysTensor_t
+
+
+class LlaisysQwen2Meta(Structure):
+    _fields_ = [
+        ("dtype", llaisysDataType_t),
+        ("nlayer", c_size_t),
+        ("hs", c_size_t),
+        ("nh", c_size_t),
+        ("nkvh", c_size_t),
+        ("dh", c_size_t),
+        ("di", c_size_t),
+        ("maxseq", c_size_t),
+        ("voc", c_size_t),
+        ("epsilon", c_float),
+        ("theta", c_float),
+        ("end_token", c_int64),
+    ]
+
+
+class LlaisysQwen2Weights(Structure):
+    _fields_ = [
+        ("in_embed", llaisysTensor_t),
+        ("out_embed", llaisysTensor_t),
+        ("out_norm_w", llaisysTensor_t),
+        ("attn_norm_w", POINTER(llaisysTensor_t)),
+        ("attn_q_w", POINTER(llaisysTensor_t)),
+        ("attn_q_b", POINTER(llaisysTensor_t)),
+        ("attn_k_w", POINTER(llaisysTensor_t)),
+        ("attn_k_b", POINTER(llaisysTensor_t)),
+        ("attn_v_w", POINTER(llaisysTensor_t)),
+        ("attn_v_b", POINTER(llaisysTensor_t)),
+        ("attn_o_w", POINTER(llaisysTensor_t)),
+        ("mlp_norm_w", POINTER(llaisysTensor_t)),
+        ("mlp_gate_w", POINTER(llaisysTensor_t)),
+        ("mlp_up_w", POINTER(llaisysTensor_t)),
+        ("mlp_down_w", POINTER(llaisysTensor_t)),
+    ]
+
+
+llaisysQwen2Model_p = c_void_p
+
+
+def load_qwen2(lib):
+    lib.llaisysQwen2ModelCreate.argtypes = [
+        POINTER(LlaisysQwen2Meta),
+        llaisysDeviceType_t,
+        POINTER(c_int),
+        c_int,
+    ]
+    lib.llaisysQwen2ModelCreate.restype = llaisysQwen2Model_p
+
+    lib.llaisysQwen2ModelDestroy.argtypes = [llaisysQwen2Model_p]
+    lib.llaisysQwen2ModelDestroy.restype = None
+
+    lib.llaisysQwen2ModelWeights.argtypes = [llaisysQwen2Model_p]
+    lib.llaisysQwen2ModelWeights.restype = POINTER(LlaisysQwen2Weights)
+
+    lib.llaisysQwen2ModelInfer.argtypes = [
+        llaisysQwen2Model_p,
+        POINTER(c_int64),
+        c_size_t,
+    ]
+    lib.llaisysQwen2ModelInfer.restype = c_int64
+
+    lib.llaisysQwen2ModelResetCache.argtypes = [llaisysQwen2Model_p]
+    lib.llaisysQwen2ModelResetCache.restype = None

python/llaisys/models/qwen2.py

@@ -1,23 +1,139 @@
+import ctypes
+import json
+import mmap
+from ctypes import POINTER, byref, c_char, c_int, c_int64
+from pathlib import Path
 from typing import Sequence
-from ..libllaisys import LIB_LLAISYS
-from ..libllaisys import DeviceType
 
-from pathlib import Path
-import safetensors
+from ..libllaisys import LIB_LLAISYS, DataType, DeviceType, LlaisysQwen2Meta
 
 
-class Qwen2:
+class _SafeTensorsFile:
+    """Minimal safetensors reader.
+
+    The safetensors numpy backend cannot represent bfloat16 and torch is not allowed
+    here, so raw bytes are mapped out of the file and handed to tensorLoad, which
+    copies them verbatim.
+    """
+
+    def __init__(self, path: Path):
+        self._file = open(path, "rb")
+        header_len = int.from_bytes(self._file.read(8), "little")
+        self._header = json.loads(self._file.read(header_len))
+        self._data_start = 8 + header_len
+        # ACCESS_COPY keeps the mapping writable so from_buffer can alias it without
+        # copying; nothing here ever writes to it.
+        self._mm = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_COPY)
+
+    def keys(self):
+        return [k for k in self._header if k != "__metadata__"]
 
-    def __init__(self, model_path, device: DeviceType = DeviceType.CPU):
-        # TODO: Implement model constructor
+    def dtype(self, name: str) -> str:
+        return self._header[name]["dtype"]
+
+    def buffer(self, name: str):
+        begin, end = self._header[name]["data_offsets"]
+        return (c_char * (end - begin)).from_buffer(self._mm, self._data_start + begin)
+
+    def close(self):
+        self._mm.close()
+        self._file.close()
+
+
+class Qwen2:
 
+    def __init__(
+        self, model_path, device: DeviceType = DeviceType.CPU, max_seq: int = None
+    ):
         model_path = Path(model_path)
+        config = json.loads((model_path / "config.json").read_text(encoding="utf-8"))
+
+        nh = config["num_attention_heads"]
+        hs = config["hidden_size"]
+
+        # The config allows 131072 positions, which would make the KV cache far larger
+        # than any test needs, so cap it.
+        if max_seq is None:
+            max_seq = min(config["max_position_embeddings"], 4096)
+
+        end_token = config.get("eos_token_id", 151643)
+        if isinstance(end_token, list):
+            end_token = end_token[0]
+        self._end_token = end_token
+
+        self._meta = LlaisysQwen2Meta(
+            dtype=DataType.BF16,
+            nlayer=config["num_hidden_layers"],
+            hs=hs,
+            nh=nh,
+            nkvh=config["num_key_value_heads"],
+            dh=config.get("head_dim", hs // nh),
+            di=config["intermediate_size"],
+            maxseq=max_seq,
+            voc=config["vocab_size"],
+            epsilon=config["rms_norm_eps"],
+            theta=config["rope_theta"],
+            end_token=end_token,
+        )
+
+        device_ids = (c_int * 1)(0)
+        self._model = LIB_LLAISYS.llaisysQwen2ModelCreate(
+            byref(self._meta), c_int(device), device_ids, c_int(1)
+        )
+        if not self._model:
+            raise RuntimeError("Failed to create Qwen2 model")
+
+        targets = self._weight_targets(
+            LIB_LLAISYS.llaisysQwen2ModelWeights(self._model).contents
+        )
 
         for file in sorted(model_path.glob("*.safetensors")):
-            data_ = safetensors.safe_open(file, framework="numpy", device="cpu")
-            for name_ in data_.keys():
-                ## TODO: load the model weights
-                pass
+            data_ = _SafeTensorsFile(file)
+            try:
+                for name_ in data_.keys():
+                    target = targets.get(name_)
+                    if target is None:
+                        continue
+                    if data_.dtype(name_) != "BF16":
+                        raise RuntimeError(
+                            f"{name_}: expected BF16, got {data_.dtype(name_)}"
+                        )
+                    buf = data_.buffer(name_)
+                    LIB_LLAISYS.tensorLoad(target, ctypes.addressof(buf))
+                    del buf  # release the mmap export before the file is closed
+            finally:
+                data_.close()
+
+    def _weight_targets(self, weights):
+        """Maps each safetensors name to the tensor handle it should be loaded into."""
+        targets = {
+            "model.embed_tokens.weight": weights.in_embed,
+            "lm_head.weight": weights.out_embed,
+            "model.norm.weight": weights.out_norm_w,
+        }
+        per_layer = {
+            "input_layernorm.weight": weights.attn_norm_w,
+            "self_attn.q_proj.weight": weights.attn_q_w,
+            "self_attn.q_proj.bias": weights.attn_q_b,
+            "self_attn.k_proj.weight": weights.attn_k_w,
+            "self_attn.k_proj.bias": weights.attn_k_b,
+            "self_attn.v_proj.weight": weights.attn_v_w,
+            "self_attn.v_proj.bias": weights.attn_v_b,
+            "self_attn.o_proj.weight": weights.attn_o_w,
+            "post_attention_layernorm.weight": weights.mlp_norm_w,
+            "mlp.gate_proj.weight": weights.mlp_gate_w,
+            "mlp.up_proj.weight": weights.mlp_up_w,
+            "mlp.down_proj.weight": weights.mlp_down_w,
+        }
+        for layer in range(self._meta.nlayer):
+            for suffix, array in per_layer.items():
+                targets[f"model.layers.{layer}.{suffix}"] = array[layer]
+        return targets
+
+    def __del__(self):
+        if getattr(self, "_model", None):
+            LIB_LLAISYS.llaisysQwen2ModelDestroy(self._model)
+            self._model = None
 
     def generate(
         self,
@@ -27,7 +143,28 @@ def generate(
         top_p: float = 0.8,
         temperature: float = 0.8,
     ):
+        # Only argmax sampling is implemented, so top_k/top_p/temperature are unused.
+        tokens = list(inputs)
+        if max_new_tokens is None:
+            max_new_tokens = self._meta.maxseq - len(tokens)
+
+        LIB_LLAISYS.llaisysQwen2ModelResetCache(self._model)
 
-        # TODO: Implement generate function
+        # Prefill the whole prompt, then feed one token per step; the backend KV cache
+        # keeps earlier positions so each step only computes the new token.
+        pending = tokens
+        for _ in range(max_new_tokens):
+            if len(tokens) >= self._meta.maxseq:
+                break
+            buf = (c_int64 * len(pending))(*pending)
+            next_token = int(
+                LIB_LLAISYS.llaisysQwen2ModelInfer(
+                    self._model, ctypes.cast(buf, POINTER(c_int64)), len(pending)
+                )
+            )
+            tokens.append(next_token)
+            if next_token == self._end_token:
+                break
+            pending = [next_token]
 
-        return []
+        return tokens

src/llaisys/models/qwen2.cc

@@ -0,0 +1,85 @@
+#include "llaisys/models/qwen2.h"
+
+#include "../../models/qwen2/qwen2.hpp"
+#include "../llaisys_tensor.hpp"
+
+#include <vector>
+
+// Owns the C++ model plus the LlaisysTensor wrappers handed out through
+// LlaisysQwen2Weights. The wrappers borrow the model's tensors, so callers must not
+// destroy them; they die with the model.
+struct LlaisysQwen2Model {
+    llaisys::models::qwen2::Qwen2 model;
+    LlaisysQwen2Weights weights{};
+    std::vector<LlaisysTensor *> owned;
+    std::vector<std::vector<llaisysTensor_t>> arrays;
+
+    LlaisysQwen2Model(const LlaisysQwen2Meta &meta, llaisysDeviceType_t device, int device_id)
+        : model(meta, device, device_id) {}
+
+    ~LlaisysQwen2Model() {
+        for (auto *w : owned) {
+            delete w;
+        }
+    }
+
+    llaisysTensor_t wrap(const llaisys::tensor_t &t) {
+        auto *w = new LlaisysTensor{t};
+        owned.push_back(w);
+        return w;
+    }
+
+    llaisysTensor_t *wrapAll(const std::vector<llaisys::tensor_t> &v) {
+        arrays.emplace_back();
+        auto &arr = arrays.back();
+        arr.reserve(v.size());
+        for (const auto &t : v) {
+            arr.push_back(wrap(t));
+        }
+        return arr.data();
+    }
+};
+
+__C {
+    struct LlaisysQwen2Model *llaisysQwen2ModelCreate(const LlaisysQwen2Meta *meta, llaisysDeviceType_t device, int *device_ids, int ndevice) {
+        const int device_id = (ndevice > 0 && device_ids != nullptr) ? device_ids[0] : 0;
+        auto *m = new LlaisysQwen2Model(*meta, device, device_id);
+
+        // arrays must not reallocate while wrapAll hands out pointers into it.
+        m->arrays.reserve(12);
+
+        auto &w = m->model.weights();
+        m->weights.in_embed = m->wrap(w.in_embed);
+        m->weights.out_embed = m->wrap(w.out_embed);
+        m->weights.out_norm_w = m->wrap(w.out_norm_w);
+        m->weights.attn_norm_w = m->wrapAll(w.attn_norm_w);
+        m->weights.attn_q_w = m->wrapAll(w.attn_q_w);
+        m->weights.attn_q_b = m->wrapAll(w.attn_q_b);
+        m->weights.attn_k_w = m->wrapAll(w.attn_k_w);
+        m->weights.attn_k_b = m->wrapAll(w.attn_k_b);
+        m->weights.attn_v_w = m->wrapAll(w.attn_v_w);
+        m->weights.attn_v_b = m->wrapAll(w.attn_v_b);
+        m->weights.attn_o_w = m->wrapAll(w.attn_o_w);
+        m->weights.mlp_norm_w = m->wrapAll(w.mlp_norm_w);
+        m->weights.mlp_gate_w = m->wrapAll(w.mlp_gate_w);
+        m->weights.mlp_up_w = m->wrapAll(w.mlp_up_w);
+        m->weights.mlp_down_w = m->wrapAll(w.mlp_down_w);
+        return m;
+    }
+
+    void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model * model) {
+        delete model;
+    }
+
+    struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model * model) {
+        return &model->weights;
+    }
+
+    int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken) {
+        return model->model.infer(token_ids, ntoken);
+    }
+
+    void llaisysQwen2ModelResetCache(struct LlaisysQwen2Model * model) {
+        model->model.resetCache();
+    }
+}