feat: 完成作业1-4 - LLAISYS 全链路实现及 CUDA 集成

build_nvidia.cmd

@@ -0,0 +1,8 @@
+@echo off
+set CUDA_PATH=D:\ANACONDA\envs\llaisys\Library
+set CUDA_HOME=%CUDA_PATH%
+set PATH=%CUDA_PATH%\bin;%CUDA_PATH%\lib;%PATH%
+cd /d E:\githubwork\llaisys
+rmdir /S /Q build 2>nul
+xmake f --nv-gpu=y -c
+xmake

python/llaisys/__init__.py

@@ -5,9 +5,8 @@
 from .libllaisys import llaisysStream_t as Stream
 from .tensor import Tensor
 from .ops import Ops
-from . import models
 from .models import *
-
+from .models.qwen2 import Qwen2
 __all__ = [
     "RuntimeAPI",
     "DeviceType",

python/llaisys/libllaisys/__init__.py

@@ -2,7 +2,6 @@
 import sys
 import ctypes
 from pathlib import Path
-
 from .runtime import load_runtime
 from .runtime import LlaisysRuntimeAPI
 from .llaisys_types import llaisysDeviceType_t, DeviceType

python/llaisys/libllaisys/model.py

@@ -0,0 +1,76 @@
+import ctypes
+from . import LIB_LLAISYS as _lib
+
+class LlaisysQwen2Meta(ctypes.Structure):
+    _fields_ = [
+        ("dtype", ctypes.c_int),
+        ("nlayer", ctypes.c_size_t),
+        ("hs", ctypes.c_size_t),
+        ("nh", ctypes.c_size_t),
+        ("nkvh", ctypes.c_size_t),
+        ("dh", ctypes.c_size_t),
+        ("di", ctypes.c_size_t),
+        ("maxseq", ctypes.c_size_t),
+        ("voc", ctypes.c_size_t),
+        ("epsilon", ctypes.c_float),
+        ("theta", ctypes.c_float),
+        ("end_token", ctypes.c_int64),
+    ]
+
+class LlaisysQwen2Weights(ctypes.Structure):
+    _fields_ = [
+        ("in_embed", ctypes.c_void_p),
+        ("out_embed", ctypes.c_void_p),
+        ("out_norm_w", ctypes.c_void_p),
+        ("attn_norm_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_q_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_q_b", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_k_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_k_b", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_v_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_v_b", ctypes.POINTER(ctypes.c_void_p)),
+        ("attn_o_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("mlp_norm_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("mlp_gate_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("mlp_up_w", ctypes.POINTER(ctypes.c_void_p)),
+        ("mlp_down_w", ctypes.POINTER(ctypes.c_void_p)),
+    ]
+
+_lib.llaisysQwen2ModelCreate.argtypes = [
+    ctypes.POINTER(LlaisysQwen2Meta),
+    ctypes.c_int,
+    ctypes.POINTER(ctypes.c_int),
+    ctypes.c_int,
+]
+_lib.llaisysQwen2ModelCreate.restype = ctypes.c_void_p
+
+_lib.llaisysQwen2ModelDestroy.argtypes = [ctypes.c_void_p]
+_lib.llaisysQwen2ModelDestroy.restype = None
+
+_lib.llaisysQwen2ModelWeights.argtypes = [ctypes.c_void_p]
+_lib.llaisysQwen2ModelWeights.restype = ctypes.POINTER(LlaisysQwen2Weights)
+
+_lib.llaisysQwen2ModelInfer.argtypes = [
+    ctypes.c_void_p,
+    ctypes.POINTER(ctypes.c_int64),
+    ctypes.c_size_t,
+]
+_lib.llaisysQwen2ModelInfer.restype = ctypes.c_int64
+
+def create_model(meta, device, device_ids):
+    return _lib.llaisysQwen2ModelCreate(
+        ctypes.byref(meta),
+        device,
+        (ctypes.c_int * len(device_ids))(*device_ids),
+        len(device_ids)
+    )
+
+def destroy_model(handle):
+    _lib.llaisysQwen2ModelDestroy(handle)
+
+def get_weights(handle):
+    return _lib.llaisysQwen2ModelWeights(handle)
+
+def infer(handle, token_ids):
+    arr = (ctypes.c_int64 * len(token_ids))(*token_ids)
+    return _lib.llaisysQwen2ModelInfer(handle, arr, len(token_ids))

python/llaisys/models/qwen2.py

@@ -1,33 +1,128 @@
 from typing import Sequence
-from ..libllaisys import LIB_LLAISYS
-from ..libllaisys import DeviceType
-
 from pathlib import Path
+import json
+import ctypes
 import safetensors
+import safetensors.torch
+import torch
+import numpy as np
 
+from llaisys.libllaisys import DeviceType, DataType
+from llaisys.libllaisys.model import (
+    create_model,
+    get_weights,
+    infer,
+    LlaisysQwen2Meta,
+    LlaisysQwen2Weights,
+)
+from llaisys.tensor import Tensor
 
 class Qwen2:
-
     def __init__(self, model_path, device: DeviceType = DeviceType.CPU):
-        # TODO: Implement model constructor
-
         model_path = Path(model_path)
+        self.device = DeviceType.CPU  # Always use CPU for weights (GPU compute via CUDA kernels on-the-fly)
+        self.device_id = 0
+
+        with open(model_path / "config.json", "r") as f:
+            config = json.load(f)
+        self.config = config
+
+        meta = LlaisysQwen2Meta()
+        meta.dtype = 0  # F32
+        meta.nlayer = config["num_hidden_layers"]
+        meta.hs = config["hidden_size"]
+        meta.nh = config["num_attention_heads"]
+        meta.nkvh = config.get("num_key_value_heads", meta.nh)
+        meta.dh = meta.hs // meta.nh
+        meta.di = config["intermediate_size"]
+        meta.maxseq = config.get("max_position_embeddings", 4096)
+        meta.voc = config["vocab_size"]
+        meta.epsilon = config.get("rms_norm_eps", 1e-6)
+        meta.theta = config.get("rope_theta", 10000.0)
+        meta.end_token = config.get("eos_token_id", 151645)
+
+        self.handle = create_model(meta, device.value, [self.device_id])
+        self.weights_ptr = get_weights(self.handle)
+        self._load_weights(model_path)
+
+    def _load_weights(self, model_path):
+        w = self.weights_ptr.contents
+
+        def to_tensor(np_arr):
+            shape = list(np_arr.shape)
+            dtype_map = {
+                np.float32: DataType.F32,
+                np.float16: DataType.F16,
+                np.int64: DataType.I64,
+            }
+            dtype = dtype_map.get(np_arr.dtype.type, DataType.F32)
+            tensor = Tensor(shape, dtype, self.device, self.device_id)
+            tensor.load(np_arr.ctypes.data_as(ctypes.c_void_p))
+            if not hasattr(self, '_tensor_refs'):
+                self._tensor_refs = []
+            self._tensor_refs.append(tensor)
+            return tensor.lib_tensor()
 
         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
-
-    def generate(
-        self,
-        inputs: Sequence[int],
-        max_new_tokens: int = None,
-        top_k: int = 1,
-        top_p: float = 0.8,
-        temperature: float = 0.8,
-    ):
-
-        # TODO: Implement generate function
-
-        return []
+            data = safetensors.torch.load_file(str(file))
+            for name, arr in data.items():
+                if arr.dtype == torch.bfloat16:
+                    arr = arr.float()
+                arr = arr.cpu().numpy()
+                if name == "model.embed_tokens.weight":
+                    w.in_embed = to_tensor(arr)
+                elif name == "model.norm.weight":
+                    w.out_norm_w = to_tensor(arr)
+                elif name == "lm_head.weight":
+                    w.out_embed = to_tensor(arr)
+                elif name.startswith("model.layers."):
+                    parts = name.split('.')
+                    layer_idx = int(parts[2])
+                    suffix = parts[3]
+                    if suffix == "input_layernorm":
+                        w.attn_norm_w[layer_idx] = to_tensor(arr)
+                    elif suffix == "self_attn":
+                        proj = parts[4]
+                        if proj == "q_proj":
+                            if name.endswith(".weight"):
+                                w.attn_q_w[layer_idx] = to_tensor(arr)
+                            elif name.endswith(".bias"):
+                                w.attn_q_b[layer_idx] = to_tensor(arr)
+                        elif proj == "k_proj":
+                            if name.endswith(".weight"):
+                                w.attn_k_w[layer_idx] = to_tensor(arr)
+                            elif name.endswith(".bias"):
+                                w.attn_k_b[layer_idx] = to_tensor(arr)
+                        elif proj == "v_proj":
+                            if name.endswith(".weight"):
+                                w.attn_v_w[layer_idx] = to_tensor(arr)
+                            elif name.endswith(".bias"):
+                                w.attn_v_b[layer_idx] = to_tensor(arr)
+                        elif proj == "o_proj":
+                            if name.endswith(".weight"):
+                                w.attn_o_w[layer_idx] = to_tensor(arr)
+                    elif suffix == "post_attention_layernorm":
+                        w.mlp_norm_w[layer_idx] = to_tensor(arr)
+                    elif suffix == "mlp":
+                        proj = parts[4]
+                        if proj == "gate_proj":
+                            w.mlp_gate_w[layer_idx] = to_tensor(arr)
+                        elif proj == "up_proj":
+                            w.mlp_up_w[layer_idx] = to_tensor(arr)
+                        elif proj == "down_proj":
+                            w.mlp_down_w[layer_idx] = to_tensor(arr)
+
+    def generate(self, inputs: Sequence[int], max_new_tokens: int = 20, **kwargs):
+        print(f"generate called with {inputs}")
+        if not inputs:
+            return []
+        input_ids = list(inputs)
+        generated = []
+        for step in range(max_new_tokens):
+            next_token = infer(self.handle, input_ids)
+            if next_token == self.config.get("eos_token_id", -1):
+                break
+            generated.append(next_token)
+            input_ids.append(next_token)
+        print(f"generate returning {generated}")
+        return generated

python/llaisys/tensor.py

@@ -70,7 +70,12 @@ def debug(self):
         LIB_LLAISYS.tensorDebug(self._tensor)
 
     def __repr__(self):
-        return f"<Tensor shape={self.shape}, dtype={self.dtype}, device={self.device_type}:{self.device_id}>"
+        return (
+            f"<Tensor "
+            f"shape={self.shape()}, "
+            f"dtype={self.dtype()}, "
+            f"device={self.device_type()}:{self.device_id()}>"
+            )
 
     def load(self, data: c_void_p):
         LIB_LLAISYS.tensorLoad(self._tensor, data)

src/device/nvidia/nvidia_runtime_api.cu

@@ -1,56 +1,80 @@
 #include "../runtime_api.hpp"
+#include "nvidia_resource.cuh"
 
+#include <cuda_runtime.h>
 #include <cstdlib>
-#include <cstring>
 
 namespace llaisys::device::nvidia {
 
 namespace runtime_api {
 int getDeviceCount() {
-    TO_BE_IMPLEMENTED();
+    int count = 0;
+    cudaGetDeviceCount(&count);
+    return count;
 }
 
-void setDevice(int) {
-    TO_BE_IMPLEMENTED();
+void setDevice(int device) {
+    cudaSetDevice(device);
 }
 
 void deviceSynchronize() {
-    TO_BE_IMPLEMENTED();
+    cudaDeviceSynchronize();
 }
 
 llaisysStream_t createStream() {
-    TO_BE_IMPLEMENTED();
+    cudaStream_t stream = nullptr;
+    cudaStreamCreate(&stream);
+    return reinterpret_cast<llaisysStream_t>(stream);
 }
 
 void destroyStream(llaisysStream_t stream) {
-    TO_BE_IMPLEMENTED();
+    cudaStreamDestroy(reinterpret_cast<cudaStream_t>(stream));
 }
+
 void streamSynchronize(llaisysStream_t stream) {
-    TO_BE_IMPLEMENTED();
+    cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream));
 }
 
 void *mallocDevice(size_t size) {
-    TO_BE_IMPLEMENTED();
+    void *ptr = nullptr;
+    cudaMalloc(&ptr, size);
+    return ptr;
 }
 
 void freeDevice(void *ptr) {
-    TO_BE_IMPLEMENTED();
+    cudaFree(ptr);
 }
 
 void *mallocHost(size_t size) {
-    TO_BE_IMPLEMENTED();
+    void *ptr = nullptr;
+    cudaMallocHost(&ptr, size);
+    return ptr;
 }
 
 void freeHost(void *ptr) {
-    TO_BE_IMPLEMENTED();
+    cudaFreeHost(ptr);
 }
 
 void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
-    TO_BE_IMPLEMENTED();
+    cudaMemcpyKind cuda_kind = cudaMemcpyDefault;
+    switch (kind) {
+        case LLAISYS_MEMCPY_H2H: cuda_kind = cudaMemcpyHostToHost; break;
+        case LLAISYS_MEMCPY_H2D: cuda_kind = cudaMemcpyHostToDevice; break;
+        case LLAISYS_MEMCPY_D2H: cuda_kind = cudaMemcpyDeviceToHost; break;
+        case LLAISYS_MEMCPY_D2D: cuda_kind = cudaMemcpyDeviceToDevice; break;
+    }
+    cudaMemcpy(dst, src, size, cuda_kind);
 }
 
-void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
-    TO_BE_IMPLEMENTED();
+void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
+    cudaMemcpyKind cuda_kind = cudaMemcpyDefault;
+    switch (kind) {
+        case LLAISYS_MEMCPY_H2H: cuda_kind = cudaMemcpyHostToHost; break;
+        case LLAISYS_MEMCPY_H2D: cuda_kind = cudaMemcpyHostToDevice; break;
+        case LLAISYS_MEMCPY_D2H: cuda_kind = cudaMemcpyDeviceToHost; break;
+        case LLAISYS_MEMCPY_D2D: cuda_kind = cudaMemcpyDeviceToDevice; break;
+    }
+    cudaMemcpyAsync(dst, src, size, cuda_kind, reinterpret_cast<cudaStream_t>(stream));
 }
 
 static const LlaisysRuntimeAPI RUNTIME_API = {
@@ -72,4 +96,4 @@ static const LlaisysRuntimeAPI RUNTIME_API = {
 const LlaisysRuntimeAPI *getRuntimeAPI() {
     return &runtime_api::RUNTIME_API;
 }
-} // namespace llaisys::device::nvidia
+} // namespace llaisys::device::nvidia