Finish Assignment-Zane

include/llaisys.h

@@ -24,6 +24,7 @@ typedef enum {
     LLAISYS_DEVICE_CPU = 0,
     //// TODO: Add more device types here. Numbers need to be consecutive.
     LLAISYS_DEVICE_NVIDIA = 1,
+    LLAISYS_DEVICE_ILUVATAR = 2,
     LLAISYS_DEVICE_TYPE_COUNT
 } llaisysDeviceType_t;
 

include/llaisys/models/qwen2.h

@@ -37,6 +37,18 @@ __C {
 
     __export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model * model);
 
+    // Load one safetensors entry into the model.  The bytes are copied before
+    // this function returns, so the caller owns the input buffer.
+    __export void llaisysQwen2ModelLoadWeight(
+        struct LlaisysQwen2Model *model,
+        const char *name,
+        const void *data,
+        const size_t *shape,
+        size_t ndim,
+        llaisysDataType_t dtype);
+
     __export int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken);
+
+    __export void llaisysQwen2ModelReset(struct LlaisysQwen2Model * model);
 }
 #endif // LLAISYS_MODELS_QWEN2_H

include/llaisys/runtime.h

@@ -42,6 +42,7 @@ __C {
 
     // Llaisys API for switching device context
     __export void llaisysSetContextRuntime(llaisysDeviceType_t, int);
+
 }
 
 #endif // LLAISYS_RUNTIME_H

python/llaisys/libllaisys/__init__.py

@@ -12,6 +12,9 @@
 from .tensor import llaisysTensor_t
 from .tensor import load_tensor
 from .ops import load_ops
+from .qwen2 import load_qwen2
+
+_DLL_DIRECTORY_HANDLES = []
 
 
 def load_shared_library():
@@ -21,6 +24,11 @@ def load_shared_library():
         libname = "libllaisys.so"
     elif sys.platform == "win32":
         libname = "llaisys.dll"
+        cuda_path = os.environ.get("CUDA_PATH")
+        if cuda_path and hasattr(os, "add_dll_directory"):
+            cuda_bin = Path(cuda_path) / "bin"
+            if cuda_bin.is_dir():
+                _DLL_DIRECTORY_HANDLES.append(os.add_dll_directory(str(cuda_bin)))
     elif sys.platform == "darwin":
         libname = "llaisys.dylib"
     else:
@@ -38,6 +46,7 @@ def load_shared_library():
 load_runtime(LIB_LLAISYS)
 load_tensor(LIB_LLAISYS)
 load_ops(LIB_LLAISYS)
+load_qwen2(LIB_LLAISYS)
 
 
 __all__ = [

python/llaisys/libllaisys/llaisys_types.py

@@ -6,7 +6,8 @@
 class DeviceType(IntEnum):
     CPU = 0
     NVIDIA = 1
-    COUNT = 2
+    ILUVATAR = 2
+    COUNT = 3
 
 
 llaisysDeviceType_t = ctypes.c_int

python/llaisys/libllaisys/qwen2.py

@@ -0,0 +1,39 @@
+import ctypes
+from ctypes import POINTER, Structure, c_char_p, c_float, c_int, c_int64, c_size_t, c_void_p
+
+
+class LlaisysQwen2Meta(Structure):
+    _fields_ = [
+        ("dtype", c_int),
+        ("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),
+    ]
+
+
+def load_qwen2(lib):
+    lib.llaisysQwen2ModelCreate.argtypes = [
+        POINTER(LlaisysQwen2Meta), c_int, POINTER(c_int), c_int
+    ]
+    lib.llaisysQwen2ModelCreate.restype = c_void_p
+    lib.llaisysQwen2ModelDestroy.argtypes = [c_void_p]
+    lib.llaisysQwen2ModelDestroy.restype = None
+    lib.llaisysQwen2ModelLoadWeight.argtypes = [
+        c_void_p, c_char_p, c_void_p, POINTER(c_size_t), c_size_t, c_int
+    ]
+    lib.llaisysQwen2ModelLoadWeight.restype = None
+    lib.llaisysQwen2ModelReset.argtypes = [c_void_p]
+    lib.llaisysQwen2ModelReset.restype = None
+    lib.llaisysQwen2ModelInfer.argtypes = [c_void_p, POINTER(c_int64), c_size_t]
+    lib.llaisysQwen2ModelInfer.restype = c_int64
+
+
+__all__ = ["LlaisysQwen2Meta", "load_qwen2"]

python/llaisys/models/qwen2.py

@@ -1,23 +1,114 @@
 from typing import Sequence
-from ..libllaisys import LIB_LLAISYS
-from ..libllaisys import DeviceType
-
+import ctypes
+import gc
+import json
+import mmap
+import sys
+import struct
+from ctypes import c_int, c_int64, c_size_t
 from pathlib import Path
-import safetensors
+import numpy as np
+
+from ..libllaisys import LIB_LLAISYS, DeviceType, DataType
+from ..libllaisys.qwen2 import LlaisysQwen2Meta
+
+
+def _trim_windows_working_set():
+    """Release freed PyTorch pages before allocating the native model weights."""
+    if sys.platform != "win32":
+        return
+
+    kernel32 = ctypes.windll.kernel32
+    kernel32.SetProcessWorkingSetSize.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_size_t]
+    kernel32.SetProcessWorkingSetSize.restype = ctypes.c_int
+    current_process = kernel32.GetCurrentProcess()
+    trim = ctypes.c_size_t(-1).value
+    kernel32.SetProcessWorkingSetSize(current_process, trim, trim)
 
 
 class Qwen2:
 
     def __init__(self, model_path, device: DeviceType = DeviceType.CPU):
-        # TODO: Implement model constructor
-
+        gc.collect()
+        _trim_windows_working_set()
         model_path = Path(model_path)
+        with (model_path / "config.json").open("r", encoding="utf-8") as config_file:
+            config = json.load(config_file)
+
+        hidden_size = int(config["hidden_size"])
+        head_count = int(config["num_attention_heads"])
+        head_dim = int(config.get("head_dim", hidden_size // head_count))
+        dtype_name = str(config.get("torch_dtype", "bfloat16")).lower()
+        dtype = {"float32": DataType.F32, "float16": DataType.F16, "bfloat16": DataType.BF16}[dtype_name]
+        eos = config.get("eos_token_id", 2)
+        if isinstance(eos, list):
+            eos = eos[0]
+
+        meta = LlaisysQwen2Meta(
+            int(dtype),
+            int(config["num_hidden_layers"]),
+            hidden_size,
+            head_count,
+            int(config.get("num_key_value_heads", head_count)),
+            head_dim,
+            int(config["intermediate_size"]),
+            int(config.get("max_position_embeddings", 32768)),
+            int(config["vocab_size"]),
+            float(config.get("rms_norm_eps", 1e-6)),
+            float(config.get("rope_theta", 10000.0)),
+            int(eos),
+        )
+        device_ids = (c_int * 1)(0)
+        self._model = LIB_LLAISYS.llaisysQwen2ModelCreate(
+            ctypes.byref(meta), int(device), device_ids, 1
+        )
+        if not self._model:
+            raise RuntimeError("Unable to create Qwen2 model")
+        self._end_token = int(eos)
 
         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
+            self._load_safetensors(file)
+
+    def _load_safetensors(self, path):
+        dtype_info = {
+            "F32": (DataType.F32, np.dtype("<f4")),
+            "F16": (DataType.F16, np.dtype("<f2")),
+            # NumPy has no built-in bfloat16; uint16 preserves its raw bits.
+            "BF16": (DataType.BF16, np.dtype("<u2")),
+        }
+        with path.open("rb") as stream, mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_READ) as mapped:
+            header_size = struct.unpack_from("<Q", mapped, 0)[0]
+            data_start = 8 + header_size
+            header = json.loads(mapped[8:data_start].decode("utf-8"))
+            for name, descriptor in header.items():
+                if name == "__metadata__":
+                    continue
+                dtype_name = descriptor["dtype"]
+                if dtype_name not in dtype_info:
+                    raise TypeError(f"Unsupported Qwen2 weight dtype: {dtype_name}")
+                tensor_dtype, numpy_dtype = dtype_info[dtype_name]
+                shape_values = tuple(int(dim) for dim in descriptor["shape"])
+                begin, end = descriptor["data_offsets"]
+                count = int(np.prod(shape_values, dtype=np.int64))
+                if end - begin != count * numpy_dtype.itemsize:
+                    raise ValueError(f"Invalid safetensors byte range for {name}")
+                values = np.frombuffer(mapped, dtype=numpy_dtype, count=count, offset=data_start + begin)
+                shape = (c_size_t * len(shape_values))(*shape_values)
+                LIB_LLAISYS.llaisysQwen2ModelLoadWeight(
+                    self._model,
+                    name.encode("utf-8"),
+                    values.ctypes.data_as(ctypes.c_void_p),
+                    shape,
+                    c_size_t(len(shape_values)),
+                    int(tensor_dtype),
+                )
+                del values
+
+    def __del__(self):
+        model = getattr(self, "_model", None)
+        if model:
+            LIB_LLAISYS.llaisysQwen2ModelDestroy(model)
+            self._model = None
 
     def generate(
         self,
@@ -28,6 +119,22 @@ def generate(
         temperature: float = 0.8,
     ):
 
-        # TODO: Implement generate function
+        del top_k, top_p, temperature  # Assignment 3 uses deterministic argmax.
+        if max_new_tokens is None:
+            max_new_tokens = 128
+        if max_new_tokens <= 0:
+            return list(inputs)
 
-        return []
+        result = [int(token) for token in inputs]
+        if not result:
+            raise ValueError("inputs must contain at least one token")
+        LIB_LLAISYS.llaisysQwen2ModelReset(self._model)
+        prompt = (c_int64 * len(result))(*result)
+        next_token = int(LIB_LLAISYS.llaisysQwen2ModelInfer(self._model, prompt, c_size_t(len(result))))
+        for step in range(max_new_tokens):
+            result.append(next_token)
+            if next_token == self._end_token or step + 1 == max_new_tokens:
+                break
+            token = c_int64(next_token)
+            next_token = int(LIB_LLAISYS.llaisysQwen2ModelInfer(self._model, ctypes.byref(token), c_size_t(1)))
+        return result

python/llaisys/ops.py

@@ -1,6 +1,6 @@
 from .libllaisys import LIB_LLAISYS
 from .tensor import Tensor
-from ctypes import c_float, c_int
+from ctypes import c_float
 
 
 class Ops:
@@ -21,7 +21,10 @@ def embedding(out: Tensor, index: Tensor, weight: Tensor):
     @staticmethod
     def linear(out: Tensor, inp: Tensor, weight: Tensor, bias: Tensor):
         LIB_LLAISYS.llaisysLinear(
-            out.lib_tensor(), inp.lib_tensor(), weight.lib_tensor(), bias.lib_tensor()
+            out.lib_tensor(),
+            inp.lib_tensor(),
+            weight.lib_tensor(),
+            None if bias is None else bias.lib_tensor(),
         )
 
     @staticmethod

src/core/context/context.cpp

@@ -4,7 +4,7 @@
 
 namespace llaisys::core {
 
-Context::Context() {
+Context::Context() : _current_runtime(nullptr) {
     // All device types, put CPU at the end
     std::vector<llaisysDeviceType_t> device_typs;
     for (int i = 1; i < LLAISYS_DEVICE_TYPE_COUNT; i++) {
@@ -52,7 +52,7 @@ Context::~Context() {
 void Context::setDevice(llaisysDeviceType_t device_type, int device_id) {
     // If doest not match the current runtime.
     if (_current_runtime == nullptr || _current_runtime->deviceType() != device_type || _current_runtime->deviceId() != device_id) {
-        auto runtimes = _runtime_map[device_type];
+        auto &runtimes = _runtime_map[device_type];
         CHECK_ARGUMENT((size_t)device_id < runtimes.size() && device_id >= 0, "invalid device id");
         if (_current_runtime != nullptr) {
             _current_runtime->_deactivate();
@@ -70,10 +70,19 @@ Runtime &Context::runtime() {
     return *_current_runtime;
 }
 
+namespace {
+// Keep the context alive for the process lifetime on Windows. A thread-local
+// smart pointer may run its destructor after the DLL has started unloading,
+// making CUDA cleanup unsafe.
+thread_local Context *thread_context = nullptr;
+}
+
 // Global API to get thread-local context.
 Context &context() {
-    thread_local Context thread_context;
-    return thread_context;
+    if (!thread_context) {
+        thread_context = new Context();
+    }
+    return *thread_context;
 }
 
 } // namespace llaisys::core

src/device/iluvatar/iluvatar_runtime_api.cpp

@@ -0,0 +1 @@
+#include "iluvatar_runtime_api.cu"