From a1a2de395314fe4479df635d30dc41a18a0fd79c Mon Sep 17 00:00:00 2001 From: Anki77134 <2577484662@qq.com> Date: Wed, 4 Feb 2026 23:41:33 +0800 Subject: [PATCH 1/2] Add CPU implementations and Qwen2 model support - Add CPU implementations for all ops (argmax, embedding, linear, rms_norm, rope, self_attention, swiglu) - Implement Qwen2 model in both Python and C++ - Enhance tensor operations and type support - Add test runner script - Update build configuration --- python/llaisys/libllaisys/llaisys_types.py | 4 + python/llaisys/libllaisys/qwen2.py | 63 +++++ python/llaisys/models/qwen2.py | 176 +++++++++++- python/llaisys/tensor.py | 13 +- run_all_tests.sh | 44 +++ src/llaisys/models/qwen2.cc | 102 +++++++ src/models/qwen2/qwen2.cpp | 254 ++++++++++++++++++ src/models/qwen2/qwen2.hpp | 98 +++++++ src/ops/argmax/cpu/argmax_cpu.cpp | 65 +++++ src/ops/argmax/cpu/argmax_cpu.hpp | 7 + src/ops/argmax/op.cpp | 29 +- src/ops/embedding/cpu/embedding_cpu.cpp | 43 +++ src/ops/embedding/cpu/embedding_cpu.hpp | 7 + src/ops/embedding/op.cpp | 35 ++- src/ops/linear/cpu/linear_cpu.cpp | 77 ++++++ src/ops/linear/cpu/linear_cpu.hpp | 7 + src/ops/linear/op.cpp | 48 +++- src/ops/ops.hpp | 17 ++ src/ops/rms_norm/cpu/rms_norm_cpu.cpp | 74 +++++ src/ops/rms_norm/cpu/rms_norm_cpu.hpp | 7 + src/ops/rms_norm/op.cpp | 34 ++- src/ops/rope/cpu/rope_cpu.cpp | 88 ++++++ src/ops/rope/cpu/rope_cpu.hpp | 7 + src/ops/rope/op.cpp | 38 ++- .../self_attention/cpu/self_attention_cpu.cpp | 129 +++++++++ .../self_attention/cpu/self_attention_cpu.hpp | 8 + src/ops/self_attention/op.cpp | 52 +++- src/ops/swiglu/cpu/swiglu_cpu.cpp | 50 ++++ src/ops/swiglu/cpu/swiglu_cpu.hpp | 7 + src/ops/swiglu/op.cpp | 25 +- src/tensor/tensor.cpp | 94 ++++++- test/test_infer.py | 4 +- xmake.lua | 20 +- 33 files changed, 1700 insertions(+), 26 deletions(-) create mode 100644 python/llaisys/libllaisys/qwen2.py create mode 100644 run_all_tests.sh create mode 100644 src/llaisys/models/qwen2.cc create mode 100644 src/models/qwen2/qwen2.cpp create mode 100644 src/models/qwen2/qwen2.hpp create mode 100644 src/ops/argmax/cpu/argmax_cpu.cpp create mode 100644 src/ops/argmax/cpu/argmax_cpu.hpp create mode 100644 src/ops/embedding/cpu/embedding_cpu.cpp create mode 100644 src/ops/embedding/cpu/embedding_cpu.hpp create mode 100644 src/ops/linear/cpu/linear_cpu.cpp create mode 100644 src/ops/linear/cpu/linear_cpu.hpp create mode 100644 src/ops/ops.hpp create mode 100644 src/ops/rms_norm/cpu/rms_norm_cpu.cpp create mode 100644 src/ops/rms_norm/cpu/rms_norm_cpu.hpp create mode 100644 src/ops/rope/cpu/rope_cpu.cpp create mode 100644 src/ops/rope/cpu/rope_cpu.hpp create mode 100644 src/ops/self_attention/cpu/self_attention_cpu.cpp create mode 100644 src/ops/self_attention/cpu/self_attention_cpu.hpp create mode 100644 src/ops/swiglu/cpu/swiglu_cpu.cpp create mode 100644 src/ops/swiglu/cpu/swiglu_cpu.hpp diff --git a/python/llaisys/libllaisys/llaisys_types.py b/python/llaisys/libllaisys/llaisys_types.py index c5a0b467..08bd88d2 100644 --- a/python/llaisys/libllaisys/llaisys_types.py +++ b/python/llaisys/libllaisys/llaisys_types.py @@ -52,6 +52,9 @@ class MemcpyKind(IntEnum): # Stream type (opaque pointer) llaisysStream_t = ctypes.c_void_p +# Tensor type (opaque pointer) +llaisysTensor_t = ctypes.c_void_p + __all__ = [ "llaisysDeviceType_t", "DeviceType", @@ -60,4 +63,5 @@ class MemcpyKind(IntEnum): "llaisysMemcpyKind_t", "MemcpyKind", "llaisysStream_t", + "llaisysTensor_t", ] diff --git a/python/llaisys/libllaisys/qwen2.py b/python/llaisys/libllaisys/qwen2.py new file mode 100644 index 00000000..996aad5a --- /dev/null +++ b/python/llaisys/libllaisys/qwen2.py @@ -0,0 +1,63 @@ +""" +C API bindings for Qwen2 model +""" +from . import LIB_LLAISYS +import ctypes +from .llaisys_types import llaisysDataType_t, llaisysDeviceType_t, llaisysTensor_t + +class LlaisysQwen2Meta(ctypes.Structure): + _fields_ = [ + ("dtype", llaisysDataType_t), + ("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", llaisysTensor_t), + ("out_embed", llaisysTensor_t), + ("out_norm_w", llaisysTensor_t), + ("attn_norm_w", ctypes.POINTER(llaisysTensor_t)), + ("attn_q_w", ctypes.POINTER(llaisysTensor_t)), + ("attn_q_b", ctypes.POINTER(llaisysTensor_t)), + ("attn_k_w", ctypes.POINTER(llaisysTensor_t)), + ("attn_k_b", ctypes.POINTER(llaisysTensor_t)), + ("attn_v_w", ctypes.POINTER(llaisysTensor_t)), + ("attn_v_b", ctypes.POINTER(llaisysTensor_t)), + ("attn_o_w", ctypes.POINTER(llaisysTensor_t)), + ("mlp_norm_w", ctypes.POINTER(llaisysTensor_t)), + ("mlp_gate_w", ctypes.POINTER(llaisysTensor_t)), + ("mlp_up_w", ctypes.POINTER(llaisysTensor_t)), + ("mlp_down_w", ctypes.POINTER(llaisysTensor_t)), + ] + +# Define function signatures +LIB_LLAISYS.llaisysQwen2ModelCreate.argtypes = [ + ctypes.POINTER(LlaisysQwen2Meta), + llaisysDeviceType_t, + ctypes.POINTER(ctypes.c_int), + ctypes.c_int, +] +LIB_LLAISYS.llaisysQwen2ModelCreate.restype = ctypes.c_void_p + +LIB_LLAISYS.llaisysQwen2ModelDestroy.argtypes = [ctypes.c_void_p] +LIB_LLAISYS.llaisysQwen2ModelDestroy.restype = None + +LIB_LLAISYS.llaisysQwen2ModelWeights.argtypes = [ctypes.c_void_p] +LIB_LLAISYS.llaisysQwen2ModelWeights.restype = ctypes.POINTER(LlaisysQwen2Weights) + +LIB_LLAISYS.llaisysQwen2ModelInfer.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int64), + ctypes.c_size_t, +] +LIB_LLAISYS.llaisysQwen2ModelInfer.restype = ctypes.c_int64 diff --git a/python/llaisys/models/qwen2.py b/python/llaisys/models/qwen2.py index 0d07b0b2..fdfba2a5 100644 --- a/python/llaisys/models/qwen2.py +++ b/python/llaisys/models/qwen2.py @@ -1,23 +1,156 @@ from typing import Sequence from ..libllaisys import LIB_LLAISYS from ..libllaisys import DeviceType +from ..libllaisys.qwen2 import LlaisysQwen2Meta, LlaisysQwen2Weights +from ..tensor import Tensor from pathlib import Path import safetensors +import json +import ctypes class Qwen2: def __init__(self, model_path, device: DeviceType = DeviceType.CPU): - # TODO: Implement model constructor - model_path = Path(model_path) + # Load config + config_path = model_path / "config.json" + with open(config_path, 'r') as f: + config = json.load(f) + + # Create model meta + meta = LlaisysQwen2Meta() + meta.dtype = 19 # LLAISYS_DTYPE_BF16 + meta.nlayer = config["num_hidden_layers"] + meta.hs = config["hidden_size"] + meta.nh = config["num_attention_heads"] + meta.nkvh = config["num_key_value_heads"] + meta.dh = config["hidden_size"] // config["num_attention_heads"] + meta.di = config["intermediate_size"] + meta.maxseq = config.get("max_position_embeddings", 32768) + meta.voc = config["vocab_size"] + meta.epsilon = config["rms_norm_eps"] + meta.theta = config.get("rope_theta", 10000.0) + meta.end_token = config.get("eos_token_id", 151643) + + # Create model + device_id = 0 + device_ids = (ctypes.c_int * 1)(device_id) + self._model = LIB_LLAISYS.llaisysQwen2ModelCreate( + ctypes.byref(meta), device.value, device_ids, 1 + ) + + if not self._model: + raise RuntimeError("Failed to create Qwen2 model") + + # Get weights structure + weights_ptr = LIB_LLAISYS.llaisysQwen2ModelWeights(self._model) + if not weights_ptr: + raise RuntimeError("Failed to get model weights") + + self._weights = weights_ptr.contents + self._meta = meta + self._device = device + + # Load weights from safetensors + self._load_weights(model_path) + + def _load_weights(self, model_path): + """Load weights from safetensors files""" + weight_map = {} + 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 + with safetensors.safe_open(file, framework="numpy", device="cpu") as f: + for name in f.keys(): + weight_map[name] = f.get_tensor(name) + + # Load embedding + if "model.embed_tokens.weight" in weight_map: + embed_data = weight_map["model.embed_tokens.weight"] + tensor = Tensor.from_ptr(self._weights.in_embed) + tensor.load(embed_data.ctypes.data) + + # Load output norm and lm_head + if "model.norm.weight" in weight_map: + norm_data = weight_map["model.norm.weight"] + tensor = Tensor.from_ptr(self._weights.out_norm_w) + tensor.load(norm_data.ctypes.data) + + if "lm_head.weight" in weight_map: + lm_head_data = weight_map["lm_head.weight"] + tensor = Tensor.from_ptr(self._weights.out_embed) + tensor.load(lm_head_data.ctypes.data) + + # Load per-layer weights + for layer_idx in range(self._meta.nlayer): + prefix = f"model.layers.{layer_idx}" + + # Attention norm + if f"{prefix}.input_layernorm.weight" in weight_map: + data = weight_map[f"{prefix}.input_layernorm.weight"] + tensor = Tensor.from_ptr(self._weights.attn_norm_w[layer_idx]) + tensor.load(data.ctypes.data) + + # Q, K, V projections + if f"{prefix}.self_attn.q_proj.weight" in weight_map: + data = weight_map[f"{prefix}.self_attn.q_proj.weight"] + tensor = Tensor.from_ptr(self._weights.attn_q_w[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.self_attn.q_proj.bias" in weight_map: + data = weight_map[f"{prefix}.self_attn.q_proj.bias"] + tensor = Tensor.from_ptr(self._weights.attn_q_b[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.self_attn.k_proj.weight" in weight_map: + data = weight_map[f"{prefix}.self_attn.k_proj.weight"] + tensor = Tensor.from_ptr(self._weights.attn_k_w[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.self_attn.k_proj.bias" in weight_map: + data = weight_map[f"{prefix}.self_attn.k_proj.bias"] + tensor = Tensor.from_ptr(self._weights.attn_k_b[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.self_attn.v_proj.weight" in weight_map: + data = weight_map[f"{prefix}.self_attn.v_proj.weight"] + tensor = Tensor.from_ptr(self._weights.attn_v_w[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.self_attn.v_proj.bias" in weight_map: + data = weight_map[f"{prefix}.self_attn.v_proj.bias"] + tensor = Tensor.from_ptr(self._weights.attn_v_b[layer_idx]) + tensor.load(data.ctypes.data) + + # O projection + if f"{prefix}.self_attn.o_proj.weight" in weight_map: + data = weight_map[f"{prefix}.self_attn.o_proj.weight"] + tensor = Tensor.from_ptr(self._weights.attn_o_w[layer_idx]) + tensor.load(data.ctypes.data) + + # MLP norm + if f"{prefix}.post_attention_layernorm.weight" in weight_map: + data = weight_map[f"{prefix}.post_attention_layernorm.weight"] + tensor = Tensor.from_ptr(self._weights.mlp_norm_w[layer_idx]) + tensor.load(data.ctypes.data) + + # MLP projections + if f"{prefix}.mlp.gate_proj.weight" in weight_map: + data = weight_map[f"{prefix}.mlp.gate_proj.weight"] + tensor = Tensor.from_ptr(self._weights.mlp_gate_w[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.mlp.up_proj.weight" in weight_map: + data = weight_map[f"{prefix}.mlp.up_proj.weight"] + tensor = Tensor.from_ptr(self._weights.mlp_up_w[layer_idx]) + tensor.load(data.ctypes.data) + + if f"{prefix}.mlp.down_proj.weight" in weight_map: + data = weight_map[f"{prefix}.mlp.down_proj.weight"] + tensor = Tensor.from_ptr(self._weights.mlp_down_w[layer_idx]) + tensor.load(data.ctypes.data) def generate( self, @@ -27,7 +160,34 @@ def generate( top_p: float = 0.8, temperature: float = 0.8, ): + """Generate tokens using the model""" + # For now, only support greedy decoding (top_k=1) + if top_k != 1: + raise NotImplementedError("Only greedy decoding (top_k=1) is supported") + + generated = list(inputs) + max_gen = max_new_tokens if max_new_tokens else 100 + + for _ in range(max_gen): + # Convert to ctypes array + input_array = (ctypes.c_int64 * len(generated))(*generated) + + # Call inference + next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( + self._model, input_array, len(generated) + ) + + if next_token < 0: + break + + generated.append(next_token) + + # Check for EOS + if next_token == self._meta.end_token: + break - # TODO: Implement generate function + return generated - return [] + def __del__(self): + if hasattr(self, '_model') and self._model: + LIB_LLAISYS.llaisysQwen2ModelDestroy(self._model) diff --git a/python/llaisys/tensor.py b/python/llaisys/tensor.py index 1466d851..f0da4a90 100644 --- a/python/llaisys/tensor.py +++ b/python/llaisys/tensor.py @@ -33,9 +33,20 @@ def __init__( c_int(device_id), ) + @staticmethod + def from_ptr(tensor_ptr: llaisysTensor_t): + """Create a Tensor wrapper from an existing C pointer without taking ownership""" + tensor = Tensor.__new__(Tensor) + tensor._tensor = tensor_ptr + # Mark as non-owning by setting a flag + tensor._owns_ptr = False + return tensor + def __del__(self): if hasattr(self, "_tensor") and self._tensor is not None: - LIB_LLAISYS.tensorDestroy(self._tensor) + # Only destroy if we own the pointer + if not hasattr(self, "_owns_ptr") or self._owns_ptr: + LIB_LLAISYS.tensorDestroy(self._tensor) self._tensor = None def shape(self) -> Tuple[int]: diff --git a/run_all_tests.sh b/run_all_tests.sh new file mode 100644 index 00000000..a10ec413 --- /dev/null +++ b/run_all_tests.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# llaisys 作业测试脚本 + +cd /root/autodl-tmp/llaisys + +echo "==========================================" +echo "作业 #1: 张量操作测试" +echo "==========================================" +python test/test_tensor.py +echo "" + +echo "==========================================" +echo "作业 #2: CPU 算子测试" +echo "==========================================" + +echo "--- argmax ---" +python test/ops/argmax.py + +echo "--- embedding ---" +python test/ops/embedding.py + +echo "--- swiglu ---" +python test/ops/swiglu.py + +echo "--- rms_norm ---" +python test/ops/rms_norm.py + +echo "--- linear ---" +python test/ops/linear.py + +echo "--- self_attention ---" +python test/ops/self_attention.py + +echo "--- rope (可能有微小精度差异) ---" +python test/ops/rope.py || echo "Note: rope has minor floating point differences" + +echo "" +echo "==========================================" +echo "作业 #3: 模型推理测试" +echo "==========================================" +echo "请运行: python test/test_infer.py --test" +echo "注意: 首次运行会自动下载约 3GB 的模型" +echo "" +echo "所有基础测试完成！" diff --git a/src/llaisys/models/qwen2.cc b/src/llaisys/models/qwen2.cc new file mode 100644 index 00000000..4c1b24e0 --- /dev/null +++ b/src/llaisys/models/qwen2.cc @@ -0,0 +1,102 @@ +#include "llaisys/models/qwen2.h" +#include "../../models/qwen2/qwen2.hpp" +#include "../../core/llaisys_core.hpp" + +using namespace llaisys; + +struct LlaisysQwen2Model { + std::unique_ptr model; +}; + +__export struct LlaisysQwen2Model *llaisysQwen2ModelCreate( + const LlaisysQwen2Meta *meta, + llaisysDeviceType_t device, + int *device_ids, + int ndevice) { + + // Convert meta to config + models::Qwen2Config config; + config.dtype = meta->dtype; + config.n_layers = meta->nlayer; + config.hidden_size = meta->hs; + config.n_heads = meta->nh; + config.n_kv_heads = meta->nkvh; + config.head_dim = meta->dh; + config.intermediate_size = meta->di; + config.max_seq_len = meta->maxseq; + config.vocab_size = meta->voc; + config.rms_norm_eps = meta->epsilon; + config.rope_theta = meta->theta; + config.eos_token_id = meta->end_token; + + // For now, only support single device + int device_id = (ndevice > 0 && device_ids != nullptr) ? device_ids[0] : 0; + + auto model_wrapper = new LlaisysQwen2Model(); + model_wrapper->model = std::make_unique(config, device, device_id); + + return model_wrapper; +} + +__export void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model *model) { + delete model; +} + +__export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model *model) { + if (!model || !model->model) { + return nullptr; + } + + auto &cpp_weights = model->model->weights(); + const auto &config = model->model->config(); + + // Allocate C struct + auto weights = new LlaisysQwen2Weights(); + + // Copy pointers (shallow copy) + weights->in_embed = reinterpret_cast(cpp_weights.embed_tokens.get()); + weights->out_embed = reinterpret_cast(cpp_weights.lm_head.get()); + weights->out_norm_w = reinterpret_cast(cpp_weights.norm_weight.get()); + + // Allocate arrays for per-layer weights + size_t n_layers = config.n_layers; + + weights->attn_norm_w = new llaisysTensor_t[n_layers]; + weights->attn_q_w = new llaisysTensor_t[n_layers]; + weights->attn_q_b = new llaisysTensor_t[n_layers]; + weights->attn_k_w = new llaisysTensor_t[n_layers]; + weights->attn_k_b = new llaisysTensor_t[n_layers]; + weights->attn_v_w = new llaisysTensor_t[n_layers]; + weights->attn_v_b = new llaisysTensor_t[n_layers]; + weights->attn_o_w = new llaisysTensor_t[n_layers]; + weights->mlp_norm_w = new llaisysTensor_t[n_layers]; + weights->mlp_gate_w = new llaisysTensor_t[n_layers]; + weights->mlp_up_w = new llaisysTensor_t[n_layers]; + weights->mlp_down_w = new llaisysTensor_t[n_layers]; + + for (size_t i = 0; i < n_layers; i++) { + weights->attn_norm_w[i] = reinterpret_cast(cpp_weights.attn_norm_weight[i].get()); + weights->attn_q_w[i] = reinterpret_cast(cpp_weights.attn_q_weight[i].get()); + weights->attn_q_b[i] = reinterpret_cast(cpp_weights.attn_q_bias[i].get()); + weights->attn_k_w[i] = reinterpret_cast(cpp_weights.attn_k_weight[i].get()); + weights->attn_k_b[i] = reinterpret_cast(cpp_weights.attn_k_bias[i].get()); + weights->attn_v_w[i] = reinterpret_cast(cpp_weights.attn_v_weight[i].get()); + weights->attn_v_b[i] = reinterpret_cast(cpp_weights.attn_v_bias[i].get()); + weights->attn_o_w[i] = reinterpret_cast(cpp_weights.attn_o_weight[i].get()); + weights->mlp_norm_w[i] = reinterpret_cast(cpp_weights.mlp_norm_weight[i].get()); + weights->mlp_gate_w[i] = reinterpret_cast(cpp_weights.mlp_gate_weight[i].get()); + weights->mlp_up_w[i] = reinterpret_cast(cpp_weights.mlp_up_weight[i].get()); + weights->mlp_down_w[i] = reinterpret_cast(cpp_weights.mlp_down_weight[i].get()); + } + + return weights; +} + +__export int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model *model, int64_t *token_ids, size_t ntoken) { + if (!model || !model->model || !token_ids || ntoken == 0) { + return -1; + } + + std::vector input_ids(token_ids, token_ids + ntoken); + return model->model->generate_next_token(input_ids); +} diff --git a/src/models/qwen2/qwen2.cpp b/src/models/qwen2/qwen2.cpp new file mode 100644 index 00000000..96844345 --- /dev/null +++ b/src/models/qwen2/qwen2.cpp @@ -0,0 +1,254 @@ +#include "qwen2.hpp" +#include +#include +#include + +namespace llaisys::models { + +Qwen2Model::Qwen2Model(const Qwen2Config &config, llaisysDeviceType_t device_type, int device_id) + : config_(config), kv_cache_(config.n_layers), device_type_(device_type), device_id_(device_id) { + + core::context().setDevice(device_type_, device_id_); + + // Initialize embedding weights + weights_.embed_tokens = Tensor::create({config.vocab_size, config.hidden_size}, config.dtype, device_type, device_id); + weights_.lm_head = Tensor::create({config.vocab_size, config.hidden_size}, config.dtype, device_type, device_id); + weights_.norm_weight = Tensor::create({config.hidden_size}, config.dtype, device_type, device_id); + + // Initialize per-layer weights + weights_.attn_norm_weight.resize(config.n_layers); + weights_.attn_q_weight.resize(config.n_layers); + weights_.attn_q_bias.resize(config.n_layers); + weights_.attn_k_weight.resize(config.n_layers); + weights_.attn_k_bias.resize(config.n_layers); + weights_.attn_v_weight.resize(config.n_layers); + weights_.attn_v_bias.resize(config.n_layers); + weights_.attn_o_weight.resize(config.n_layers); + + weights_.mlp_norm_weight.resize(config.n_layers); + weights_.mlp_gate_weight.resize(config.n_layers); + weights_.mlp_up_weight.resize(config.n_layers); + weights_.mlp_down_weight.resize(config.n_layers); + + for (size_t i = 0; i < config.n_layers; i++) { + // Attention weights + weights_.attn_norm_weight[i] = Tensor::create({config.hidden_size}, config.dtype, device_type, device_id); + weights_.attn_q_weight[i] = Tensor::create({config.hidden_size, config.hidden_size}, config.dtype, device_type, device_id); + weights_.attn_q_bias[i] = Tensor::create({config.hidden_size}, config.dtype, device_type, device_id); + weights_.attn_k_weight[i] = Tensor::create({config.n_kv_heads * config.head_dim, config.hidden_size}, config.dtype, device_type, device_id); + weights_.attn_k_bias[i] = Tensor::create({config.n_kv_heads * config.head_dim}, config.dtype, device_type, device_id); + weights_.attn_v_weight[i] = Tensor::create({config.n_kv_heads * config.head_dim, config.hidden_size}, config.dtype, device_type, device_id); + weights_.attn_v_bias[i] = Tensor::create({config.n_kv_heads * config.head_dim}, config.dtype, device_type, device_id); + weights_.attn_o_weight[i] = Tensor::create({config.hidden_size, config.hidden_size}, config.dtype, device_type, device_id); + + // MLP weights + weights_.mlp_norm_weight[i] = Tensor::create({config.hidden_size}, config.dtype, device_type, device_id); + weights_.mlp_gate_weight[i] = Tensor::create({config.intermediate_size, config.hidden_size}, config.dtype, device_type, device_id); + weights_.mlp_up_weight[i] = Tensor::create({config.intermediate_size, config.hidden_size}, config.dtype, device_type, device_id); + weights_.mlp_down_weight[i] = Tensor::create({config.hidden_size, config.intermediate_size}, config.dtype, device_type, device_id); + } +} + +void Qwen2Model::reset_cache() { + kv_cache_.current_seq_len = 0; + for (auto &k : kv_cache_.k_cache) { + k = nullptr; + } + for (auto &v : kv_cache_.v_cache) { + v = nullptr; + } +} + +tensor_t Qwen2Model::forward_attention(int layer_idx, tensor_t hidden_states, tensor_t position_ids) { + // hidden_states: [seq_len, hidden_size] + size_t seq_len = hidden_states->shape()[0]; + + // Input norm + auto normed = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::rms_norm(normed, hidden_states, weights_.attn_norm_weight[layer_idx], config_.rms_norm_eps); + + // Q, K, V projections + auto q = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + auto k = Tensor::create({seq_len, config_.n_kv_heads * config_.head_dim}, config_.dtype, device_type_, device_id_); + auto v = Tensor::create({seq_len, config_.n_kv_heads * config_.head_dim}, config_.dtype, device_type_, device_id_); + + ops::linear(q, normed, weights_.attn_q_weight[layer_idx], weights_.attn_q_bias[layer_idx]); + ops::linear(k, normed, weights_.attn_k_weight[layer_idx], weights_.attn_k_bias[layer_idx]); + ops::linear(v, normed, weights_.attn_v_weight[layer_idx], weights_.attn_v_bias[layer_idx]); + + // Reshape to [seq_len, n_heads, head_dim] + auto q_reshaped = q->view({seq_len, config_.n_heads, config_.head_dim}); + auto k_reshaped = k->view({seq_len, config_.n_kv_heads, config_.head_dim}); + auto v_reshaped = v->view({seq_len, config_.n_kv_heads, config_.head_dim}); + + // Apply RoPE + auto q_rope = Tensor::create({seq_len, config_.n_heads, config_.head_dim}, config_.dtype, device_type_, device_id_); + auto k_rope = Tensor::create({seq_len, config_.n_kv_heads, config_.head_dim}, config_.dtype, device_type_, device_id_); + + ops::rope(q_rope, q_reshaped, position_ids, config_.rope_theta); + ops::rope(k_rope, k_reshaped, position_ids, config_.rope_theta); + + // Update KV cache + tensor_t k_full, v_full; + if (kv_cache_.k_cache[layer_idx] == nullptr) { + // First iteration + k_full = k_rope; + v_full = v_reshaped; + kv_cache_.k_cache[layer_idx] = k_rope; + kv_cache_.v_cache[layer_idx] = v_reshaped; + } else { + // Concat with previous cache + size_t prev_len = kv_cache_.k_cache[layer_idx]->shape()[0]; + size_t total_len = prev_len + seq_len; + + k_full = Tensor::create({total_len, config_.n_kv_heads, config_.head_dim}, config_.dtype, device_type_, device_id_); + v_full = Tensor::create({total_len, config_.n_kv_heads, config_.head_dim}, config_.dtype, device_type_, device_id_); + + // Copy previous cache + auto k_prev_slice = k_full->slice(0, 0, prev_len); + auto k_new_slice = k_full->slice(0, prev_len, total_len); + auto v_prev_slice = v_full->slice(0, 0, prev_len); + auto v_new_slice = v_full->slice(0, prev_len, total_len); + + // Manual copy (since we don't have a copy operator) + size_t elem_size = config_.n_kv_heads * config_.head_dim; + memcpy(k_prev_slice->data(), kv_cache_.k_cache[layer_idx]->data(), + prev_len * elem_size * k_prev_slice->elementSize()); + memcpy(k_new_slice->data(), k_rope->data(), + seq_len * elem_size * k_new_slice->elementSize()); + memcpy(v_prev_slice->data(), kv_cache_.v_cache[layer_idx]->data(), + prev_len * elem_size * v_prev_slice->elementSize()); + memcpy(v_new_slice->data(), v_reshaped->data(), + seq_len * elem_size * v_new_slice->elementSize()); + + kv_cache_.k_cache[layer_idx] = k_full; + kv_cache_.v_cache[layer_idx] = v_full; + } + + // Self-attention + auto attn_output = Tensor::create({seq_len, config_.n_heads, config_.head_dim}, config_.dtype, device_type_, device_id_); + float scale = 1.0f / std::sqrt(static_cast(config_.head_dim)); + + ops::self_attention(attn_output, q_rope, k_full, v_full, scale); + + // Reshape back: [seq_len, n_heads, head_dim] -> [seq_len, hidden_size] + auto attn_flat = attn_output->view({seq_len, config_.hidden_size}); + + // Output projection + auto output = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::linear(output, attn_flat, weights_.attn_o_weight[layer_idx], nullptr); + + // Residual connection + auto result = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::add(result, hidden_states, output); + + return result; +} + +tensor_t Qwen2Model::forward_mlp(int layer_idx, tensor_t hidden_states) { + // hidden_states: [seq_len, hidden_size] + size_t seq_len = hidden_states->shape()[0]; + + // Post-attention norm + auto normed = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::rms_norm(normed, hidden_states, weights_.mlp_norm_weight[layer_idx], config_.rms_norm_eps); + + // Gate and Up projections + auto gate = Tensor::create({seq_len, config_.intermediate_size}, config_.dtype, device_type_, device_id_); + auto up = Tensor::create({seq_len, config_.intermediate_size}, config_.dtype, device_type_, device_id_); + + ops::linear(gate, normed, weights_.mlp_gate_weight[layer_idx], nullptr); + ops::linear(up, normed, weights_.mlp_up_weight[layer_idx], nullptr); + + // SwiGLU activation + auto activated = Tensor::create({seq_len, config_.intermediate_size}, config_.dtype, device_type_, device_id_); + ops::swiglu(activated, gate, up); + + // Down projection + auto mlp_output = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::linear(mlp_output, activated, weights_.mlp_down_weight[layer_idx], nullptr); + + // Residual connection + auto result = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::add(result, hidden_states, mlp_output); + + return result; +} + +tensor_t Qwen2Model::forward_layer(int layer_idx, tensor_t hidden_states, tensor_t position_ids) { + // Attention block + auto attn_output = forward_attention(layer_idx, hidden_states, position_ids); + + // MLP block + auto mlp_output = forward_mlp(layer_idx, attn_output); + + return mlp_output; +} + +tensor_t Qwen2Model::forward(const std::vector &input_ids) { + size_t seq_len = input_ids.size(); + + core::context().setDevice(device_type_, device_id_); + + // Create input tensor + auto input_tensor = Tensor::create({seq_len}, LLAISYS_DTYPE_I64, device_type_, device_id_); + input_tensor->load(input_ids.data()); + + // Embedding + auto hidden_states = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::embedding(hidden_states, input_tensor, weights_.embed_tokens); + + // Create position IDs + std::vector pos_ids(seq_len); + for (size_t i = 0; i < seq_len; i++) { + pos_ids[i] = kv_cache_.current_seq_len + i; + } + auto position_ids = Tensor::create({seq_len}, LLAISYS_DTYPE_I64, device_type_, device_id_); + position_ids->load(pos_ids.data()); + + // Forward through all layers + for (size_t layer = 0; layer < config_.n_layers; layer++) { + hidden_states = forward_layer(layer, hidden_states, position_ids); + } + + // Final norm + auto normed = Tensor::create({seq_len, config_.hidden_size}, config_.dtype, device_type_, device_id_); + ops::rms_norm(normed, hidden_states, weights_.norm_weight, config_.rms_norm_eps); + + // LM head (only need last token for generation) + auto last_hidden = normed->slice(0, seq_len - 1, seq_len); + auto last_hidden_2d = last_hidden->view({1, config_.hidden_size}); + + auto logits = Tensor::create({1, config_.vocab_size}, config_.dtype, device_type_, device_id_); + ops::linear(logits, last_hidden_2d, weights_.lm_head, nullptr); + + // Update cache length + kv_cache_.current_seq_len += seq_len; + + return logits; +} + +int64_t Qwen2Model::generate_next_token(const std::vector &input_ids) { + auto logits = forward(input_ids); + + // Argmax to get next token + auto logits_1d = logits->view({config_.vocab_size}); + auto max_idx = Tensor::create({1}, LLAISYS_DTYPE_I64, device_type_, device_id_); + auto max_val = Tensor::create({1}, config_.dtype, device_type_, device_id_); + + ops::argmax(max_idx, max_val, logits_1d); + + // Read result + int64_t next_token; + if (device_type_ == LLAISYS_DEVICE_CPU) { + next_token = *reinterpret_cast(max_idx->data()); + } else { + // Copy from device to host + core::context().runtime().api()->memcpy_sync( + &next_token, max_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_D2H); + } + + return next_token; +} + +} // namespace llaisys::models diff --git a/src/models/qwen2/qwen2.hpp b/src/models/qwen2/qwen2.hpp new file mode 100644 index 00000000..7bdb6b31 --- /dev/null +++ b/src/models/qwen2/qwen2.hpp @@ -0,0 +1,98 @@ +#pragma once + +#include "../../core/llaisys_core.hpp" +#include "../../tensor/tensor.hpp" +#include "../../ops/ops.hpp" + +#include +#include + +namespace llaisys::models { + +struct Qwen2Config { + llaisysDataType_t dtype; + size_t n_layers; // nlayer + size_t hidden_size; // hs + size_t n_heads; // nh + size_t n_kv_heads; // nkvh + size_t head_dim; // dh + size_t intermediate_size; // di + size_t max_seq_len; // maxseq + size_t vocab_size; // voc + float rms_norm_eps; // epsilon + float rope_theta; // theta + int64_t eos_token_id; // end_token +}; + +struct Qwen2Weights { + // Embedding + tensor_t embed_tokens; + + // Output + tensor_t lm_head; + tensor_t norm_weight; + + // Per-layer weights + std::vector attn_norm_weight; + std::vector attn_q_weight; + std::vector attn_q_bias; + std::vector attn_k_weight; + std::vector attn_k_bias; + std::vector attn_v_weight; + std::vector attn_v_bias; + std::vector attn_o_weight; + + std::vector mlp_norm_weight; + std::vector mlp_gate_weight; + std::vector mlp_up_weight; + std::vector mlp_down_weight; +}; + +struct KVCache { + std::vector k_cache; // [n_layers] + std::vector v_cache; // [n_layers] + size_t current_seq_len; + + KVCache(size_t n_layers) : current_seq_len(0) { + k_cache.resize(n_layers); + v_cache.resize(n_layers); + } +}; + +class Qwen2Model { +private: + Qwen2Config config_; + Qwen2Weights weights_; + KVCache kv_cache_; + llaisysDeviceType_t device_type_; + int device_id_; + + // Forward pass for one layer + tensor_t forward_layer(int layer_idx, tensor_t hidden_states, tensor_t position_ids); + + // Attention + tensor_t forward_attention(int layer_idx, tensor_t hidden_states, tensor_t position_ids); + + // MLP + tensor_t forward_mlp(int layer_idx, tensor_t hidden_states); + +public: + Qwen2Model(const Qwen2Config &config, llaisysDeviceType_t device_type, int device_id); + + ~Qwen2Model() = default; + + Qwen2Weights& weights() { return weights_; } + + const Qwen2Config& config() const { return config_; } + + // Forward pass: input_ids -> logits + tensor_t forward(const std::vector &input_ids); + + // Generate next token (argmax) + int64_t generate_next_token(const std::vector &input_ids); + + // Reset KV cache + void reset_cache(); +}; + +} // namespace llaisys::models diff --git a/src/ops/argmax/cpu/argmax_cpu.cpp b/src/ops/argmax/cpu/argmax_cpu.cpp new file mode 100644 index 00000000..00695ec5 --- /dev/null +++ b/src/ops/argmax/cpu/argmax_cpu.cpp @@ -0,0 +1,65 @@ +#include "argmax_cpu.hpp" +#include "../../../utils.hpp" + +template +void argmax_(IdxT *max_idx, ValT *max_val, const ValT *vals, size_t numel) { + // Find max value and its index + ValT max_v = vals[0]; + IdxT max_i = 0; + + for (size_t i = 1; i < numel; i++) { + if constexpr (std::is_same_v || std::is_same_v) { + float curr = llaisys::utils::cast(vals[i]); + float max_f = llaisys::utils::cast(max_v); + if (curr > max_f) { + max_v = vals[i]; + max_i = static_cast(i); + } + } else { + if (vals[i] > max_v) { + max_v = vals[i]; + max_i = static_cast(i); + } + } + } + + max_idx[0] = max_i; + max_val[0] = max_v; +} + +template +void argmax_dispatch_idx(std::byte *max_idx, std::byte *max_val, const std::byte *vals, + llaisysDataType_t idx_type, size_t numel) { + switch (idx_type) { + case LLAISYS_DTYPE_I32: + return argmax_( + reinterpret_cast(max_idx), + reinterpret_cast(max_val), + reinterpret_cast(vals), + numel); + case LLAISYS_DTYPE_I64: + return argmax_( + reinterpret_cast(max_idx), + reinterpret_cast(max_val), + reinterpret_cast(vals), + numel); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(idx_type); + } +} + +namespace llaisys::ops::cpu { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, + llaisysDataType_t idx_type, llaisysDataType_t val_type, size_t numel) { + switch (val_type) { + case LLAISYS_DTYPE_F32: + return argmax_dispatch_idx(max_idx, max_val, vals, idx_type, numel); + case LLAISYS_DTYPE_BF16: + return argmax_dispatch_idx(max_idx, max_val, vals, idx_type, numel); + case LLAISYS_DTYPE_F16: + return argmax_dispatch_idx(max_idx, max_val, vals, idx_type, numel); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(val_type); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/argmax/cpu/argmax_cpu.hpp b/src/ops/argmax/cpu/argmax_cpu.hpp new file mode 100644 index 00000000..299b7b20 --- /dev/null +++ b/src/ops/argmax/cpu/argmax_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, + llaisysDataType_t idx_type, llaisysDataType_t val_type, size_t numel); +} diff --git a/src/ops/argmax/op.cpp b/src/ops/argmax/op.cpp index 6dc37d42..3e1f8b14 100644 --- a/src/ops/argmax/op.cpp +++ b/src/ops/argmax/op.cpp @@ -1,7 +1,34 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/argmax_cpu.hpp" + namespace llaisys::ops { void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals) { - TO_BE_IMPLEMENTED(); + // Check all tensors on same device + CHECK_SAME_DEVICE(max_idx, max_val, vals); + + // vals should be 1D, max_idx and max_val should have 1 element + ASSERT(vals->ndim() == 1, "vals must be 1D tensor"); + ASSERT(max_idx->numel() == 1, "max_idx must have 1 element"); + ASSERT(max_val->numel() == 1, "max_val must have 1 element"); + ASSERT(max_val->dtype() == vals->dtype(), "max_val and vals must have same dtype"); + + llaisys::core::context().setDevice(vals->deviceType(), vals->deviceId()); + + switch (vals->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), + max_idx->dtype(), vals->dtype(), vals->numel()); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/embedding/cpu/embedding_cpu.cpp b/src/ops/embedding/cpu/embedding_cpu.cpp new file mode 100644 index 00000000..8d172e96 --- /dev/null +++ b/src/ops/embedding/cpu/embedding_cpu.cpp @@ -0,0 +1,43 @@ +#include "embedding_cpu.hpp" +#include "../../../utils.hpp" +#include + +template +void embedding_(T *out, const int64_t *index, const T *weight, + size_t seq_len, size_t hidden_size) { + // Copy rows from weight matrix according to index + for (size_t i = 0; i < seq_len; i++) { + int64_t idx = index[i]; + const T *src = weight + idx * hidden_size; + T *dst = out + i * hidden_size; + std::memcpy(dst, src, hidden_size * sizeof(T)); + } +} + +namespace llaisys::ops::cpu { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t dtype, size_t seq_len, size_t hidden_size) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return embedding_( + reinterpret_cast(out), + reinterpret_cast(index), + reinterpret_cast(weight), + seq_len, hidden_size); + case LLAISYS_DTYPE_BF16: + return embedding_( + reinterpret_cast(out), + reinterpret_cast(index), + reinterpret_cast(weight), + seq_len, hidden_size); + case LLAISYS_DTYPE_F16: + return embedding_( + reinterpret_cast(out), + reinterpret_cast(index), + reinterpret_cast(weight), + seq_len, hidden_size); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/embedding/cpu/embedding_cpu.hpp b/src/ops/embedding/cpu/embedding_cpu.hpp new file mode 100644 index 00000000..268ad125 --- /dev/null +++ b/src/ops/embedding/cpu/embedding_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t dtype, size_t seq_len, size_t hidden_size); +} diff --git a/src/ops/embedding/op.cpp b/src/ops/embedding/op.cpp index 84b9a5d0..2d885849 100644 --- a/src/ops/embedding/op.cpp +++ b/src/ops/embedding/op.cpp @@ -1,7 +1,40 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/embedding_cpu.hpp" + namespace llaisys::ops { void embedding(tensor_t out, tensor_t index, tensor_t weight) { - TO_BE_IMPLEMENTED(); + // Check device + CHECK_SAME_DEVICE(out, index, weight); + + // index: (seq_len,), weight: (vocab_size, hidden_size), out: (seq_len, hidden_size) + ASSERT(index->ndim() == 1, "index must be 1D"); + ASSERT(weight->ndim() == 2, "weight must be 2D"); + ASSERT(out->ndim() == 2, "out must be 2D"); + ASSERT(index->dtype() == LLAISYS_DTYPE_I64, "index must be int64"); + ASSERT(out->dtype() == weight->dtype(), "out and weight must have same dtype"); + + size_t seq_len = index->shape()[0]; + size_t hidden_size = weight->shape()[1]; + ASSERT(out->shape()[0] == seq_len && out->shape()[1] == hidden_size, + "out shape mismatch"); + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::embedding(out->data(), index->data(), weight->data(), + out->dtype(), seq_len, hidden_size); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/linear/cpu/linear_cpu.cpp b/src/ops/linear/cpu/linear_cpu.cpp new file mode 100644 index 00000000..761aebb2 --- /dev/null +++ b/src/ops/linear/cpu/linear_cpu.cpp @@ -0,0 +1,77 @@ +#include "linear_cpu.hpp" +#include "../../../utils.hpp" + +template +void linear_(T *out, const T *in, const T *weight, const T *bias, + size_t batch, size_t in_features, size_t out_features, bool has_bias) { + // out = in @ weight^T + bias + // in: (batch, in_features) + // weight: (out_features, in_features) - note: NOT transposed + // out: (batch, out_features) + + for (size_t b = 0; b < batch; b++) { + for (size_t o = 0; o < out_features; o++) { + float sum = 0.0f; + + // Dot product: in[b, :] @ weight[o, :] + for (size_t i = 0; i < in_features; i++) { + float in_val, weight_val; + if constexpr (std::is_same_v || std::is_same_v) { + in_val = llaisys::utils::cast(in[b * in_features + i]); + weight_val = llaisys::utils::cast(weight[o * in_features + i]); + } else { + in_val = static_cast(in[b * in_features + i]); + weight_val = static_cast(weight[o * in_features + i]); + } + sum += in_val * weight_val; + } + + // Add bias if present + if (has_bias) { + if constexpr (std::is_same_v || std::is_same_v) { + sum += llaisys::utils::cast(bias[o]); + } else { + sum += static_cast(bias[o]); + } + } + + // Store result + if constexpr (std::is_same_v || std::is_same_v) { + out[b * out_features + o] = llaisys::utils::cast(sum); + } else { + out[b * out_features + o] = static_cast(sum); + } + } + } +} + +namespace llaisys::ops::cpu { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t dtype, size_t batch, size_t in_features, size_t out_features, bool has_bias) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return linear_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + reinterpret_cast(bias), + batch, in_features, out_features, has_bias); + case LLAISYS_DTYPE_BF16: + return linear_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + reinterpret_cast(bias), + batch, in_features, out_features, has_bias); + case LLAISYS_DTYPE_F16: + return linear_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + reinterpret_cast(bias), + batch, in_features, out_features, has_bias); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/linear/cpu/linear_cpu.hpp b/src/ops/linear/cpu/linear_cpu.hpp new file mode 100644 index 00000000..f7c2c202 --- /dev/null +++ b/src/ops/linear/cpu/linear_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t dtype, size_t batch, size_t in_features, size_t out_features, bool has_bias); +} diff --git a/src/ops/linear/op.cpp b/src/ops/linear/op.cpp index 97d1f865..1195c488 100644 --- a/src/ops/linear/op.cpp +++ b/src/ops/linear/op.cpp @@ -1,7 +1,53 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/linear_cpu.hpp" + namespace llaisys::ops { void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias) { - TO_BE_IMPLEMENTED(); + // out = in @ weight^T + bias + // in: (batch, in_features) + // weight: (out_features, in_features) + // out: (batch, out_features) + // bias: (out_features,) or nullptr + + CHECK_SAME_DEVICE(out, in, weight); + CHECK_SAME_DTYPE(out->dtype(), in->dtype(), weight->dtype()); + + ASSERT(in->ndim() == 2, "input must be 2D"); + ASSERT(weight->ndim() == 2, "weight must be 2D"); + ASSERT(out->ndim() == 2, "output must be 2D"); + + size_t batch = in->shape()[0]; + size_t in_features = in->shape()[1]; + size_t out_features = weight->shape()[0]; + + ASSERT(weight->shape()[1] == in_features, "weight shape mismatch"); + ASSERT(out->shape()[0] == batch && out->shape()[1] == out_features, "output shape mismatch"); + + bool has_bias = (bias != nullptr); + if (has_bias) { + CHECK_SAME_DEVICE(out, bias); + CHECK_SAME_DTYPE(out->dtype(), bias->dtype()); + ASSERT(bias->ndim() == 1 && bias->shape()[0] == out_features, "bias shape mismatch"); + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::linear(out->data(), in->data(), weight->data(), + has_bias ? bias->data() : nullptr, + out->dtype(), batch, in_features, out_features, has_bias); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/ops.hpp b/src/ops/ops.hpp new file mode 100644 index 00000000..ae17a281 --- /dev/null +++ b/src/ops/ops.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../tensor/tensor.hpp" + +namespace llaisys::ops { + +// All operator declarations +void add(tensor_t c, tensor_t a, tensor_t b); +void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals); +void embedding(tensor_t out, tensor_t index, tensor_t weight); +void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias); +void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps); +void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta); +void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale); +void swiglu(tensor_t out, tensor_t gate, tensor_t up); + +} // namespace llaisys::ops diff --git a/src/ops/rms_norm/cpu/rms_norm_cpu.cpp b/src/ops/rms_norm/cpu/rms_norm_cpu.cpp new file mode 100644 index 00000000..4e11e2d8 --- /dev/null +++ b/src/ops/rms_norm/cpu/rms_norm_cpu.cpp @@ -0,0 +1,74 @@ +#include "rms_norm_cpu.hpp" +#include "../../../utils.hpp" +#include + +template +void rms_norm_(T *out, const T *in, const T *weight, size_t rows, size_t cols, float eps) { + // For each row: Y_i = (W_i * X_i) / sqrt(mean(X^2) + eps) + for (size_t r = 0; r < rows; r++) { + const T *in_row = in + r * cols; + T *out_row = out + r * cols; + + // Calculate RMS: sqrt(mean(x^2) + eps) + float sum_sq = 0.0f; + for (size_t c = 0; c < cols; c++) { + float val; + if constexpr (std::is_same_v || std::is_same_v) { + val = llaisys::utils::cast(in_row[c]); + } else { + val = static_cast(in_row[c]); + } + sum_sq += val * val; + } + + float rms = std::sqrt(sum_sq / cols + eps); + + // Normalize and apply weight + for (size_t c = 0; c < cols; c++) { + float x, w; + if constexpr (std::is_same_v || std::is_same_v) { + x = llaisys::utils::cast(in_row[c]); + w = llaisys::utils::cast(weight[c]); + } else { + x = static_cast(in_row[c]); + w = static_cast(weight[c]); + } + + float result = (w * x) / rms; + + if constexpr (std::is_same_v || std::is_same_v) { + out_row[c] = llaisys::utils::cast(result); + } else { + out_row[c] = static_cast(result); + } + } + } +} + +namespace llaisys::ops::cpu { +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + llaisysDataType_t dtype, size_t rows, size_t cols, float eps) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return rms_norm_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + rows, cols, eps); + case LLAISYS_DTYPE_BF16: + return rms_norm_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + rows, cols, eps); + case LLAISYS_DTYPE_F16: + return rms_norm_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + rows, cols, eps); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/rms_norm/cpu/rms_norm_cpu.hpp b/src/ops/rms_norm/cpu/rms_norm_cpu.hpp new file mode 100644 index 00000000..d9df0582 --- /dev/null +++ b/src/ops/rms_norm/cpu/rms_norm_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + llaisysDataType_t dtype, size_t rows, size_t cols, float eps); +} diff --git a/src/ops/rms_norm/op.cpp b/src/ops/rms_norm/op.cpp index 529553d9..b8e2c393 100644 --- a/src/ops/rms_norm/op.cpp +++ b/src/ops/rms_norm/op.cpp @@ -1,7 +1,39 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/rms_norm_cpu.hpp" + namespace llaisys::ops { void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps) { - TO_BE_IMPLEMENTED(); + // Check device + CHECK_SAME_DEVICE(out, in, weight); + CHECK_SAME_DTYPE(out->dtype(), in->dtype(), weight->dtype()); + + // in and out: 2D tensors, weight: 1D tensor + ASSERT(in->ndim() == 2, "input must be 2D"); + ASSERT(out->ndim() == 2, "output must be 2D"); + ASSERT(weight->ndim() == 1, "weight must be 1D"); + CHECK_SAME_SHAPE(out->shape(), in->shape()); + + size_t rows = in->shape()[0]; + size_t cols = in->shape()[1]; + ASSERT(weight->shape()[0] == cols, "weight shape mismatch"); + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::rms_norm(out->data(), in->data(), weight->data(), + out->dtype(), rows, cols, eps); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/rope/cpu/rope_cpu.cpp b/src/ops/rope/cpu/rope_cpu.cpp new file mode 100644 index 00000000..e0695903 --- /dev/null +++ b/src/ops/rope/cpu/rope_cpu.cpp @@ -0,0 +1,88 @@ +#include "rope_cpu.hpp" +#include "../../../utils.hpp" +#include + +template +void rope_(T *out, const T *in, const int64_t *pos_ids, + size_t seq_len, size_t n_heads, size_t head_dim, float theta) { + // RoPE: Rotary Position Embedding + // Split input into two halves: [a0, a1, ..., a_{d/2-1}, b0, b1, ..., b_{d/2-1}] + // For each position and each dimension pair: + // a'[j] = a[j] * cos(angle) - b[j] * sin(angle) + // b'[j] = b[j] * cos(angle) + a[j] * sin(angle) + // where angle = pos / (theta^(2j/d)) + + size_t half_dim = head_dim / 2; + + // Precompute frequency divisors to improve numerical stability + std::vector inv_freq(half_dim); + for (size_t j = 0; j < half_dim; j++) { + inv_freq[j] = 1.0 / std::pow(theta, (2.0 * j) / static_cast(head_dim)); + } + + for (size_t s = 0; s < seq_len; s++) { + float pos = static_cast(pos_ids[s]); + + for (size_t h = 0; h < n_heads; h++) { + for (size_t j = 0; j < half_dim; j++) { + // Calculate angle + float angle = pos * inv_freq[j]; + float cos_angle = std::cos(angle); + float sin_angle = std::sin(angle); + + // Get indices - first half is 'a', second half is 'b' + size_t idx_a = s * n_heads * head_dim + h * head_dim + j; + size_t idx_b = s * n_heads * head_dim + h * head_dim + half_dim + j; + + float a, b; + if constexpr (std::is_same_v || std::is_same_v) { + a = llaisys::utils::cast(in[idx_a]); + b = llaisys::utils::cast(in[idx_b]); + } else { + a = static_cast(in[idx_a]); + b = static_cast(in[idx_b]); + } + + // Apply rotation + float a_new = a * cos_angle - b * sin_angle; + float b_new = b * cos_angle + a * sin_angle; + + if constexpr (std::is_same_v || std::is_same_v) { + out[idx_a] = llaisys::utils::cast(a_new); + out[idx_b] = llaisys::utils::cast(b_new); + } else { + out[idx_a] = static_cast(a_new); + out[idx_b] = static_cast(b_new); + } + } + } + } +} + +namespace llaisys::ops::cpu { +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t dtype, size_t seq_len, size_t n_heads, size_t head_dim, float theta) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return rope_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(pos_ids), + seq_len, n_heads, head_dim, theta); + case LLAISYS_DTYPE_BF16: + return rope_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(pos_ids), + seq_len, n_heads, head_dim, theta); + case LLAISYS_DTYPE_F16: + return rope_( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(pos_ids), + seq_len, n_heads, head_dim, theta); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/rope/cpu/rope_cpu.hpp b/src/ops/rope/cpu/rope_cpu.hpp new file mode 100644 index 00000000..e9ac68ea --- /dev/null +++ b/src/ops/rope/cpu/rope_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t dtype, size_t seq_len, size_t n_heads, size_t head_dim, float theta); +} diff --git a/src/ops/rope/op.cpp b/src/ops/rope/op.cpp index d60dbe64..e726eb3a 100644 --- a/src/ops/rope/op.cpp +++ b/src/ops/rope/op.cpp @@ -1,7 +1,43 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/rope_cpu.hpp" + namespace llaisys::ops { void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta) { - TO_BE_IMPLEMENTED(); + // in/out: (seq_len, n_heads, head_dim) + // pos_ids: (seq_len,) + + CHECK_SAME_DEVICE(out, in, pos_ids); + CHECK_SAME_DTYPE(out->dtype(), in->dtype()); + ASSERT(pos_ids->dtype() == LLAISYS_DTYPE_I64, "pos_ids must be int64"); + + ASSERT(in->ndim() == 3, "input must be 3D"); + ASSERT(out->ndim() == 3, "output must be 3D"); + ASSERT(pos_ids->ndim() == 1, "pos_ids must be 1D"); + CHECK_SAME_SHAPE(out->shape(), in->shape()); + + size_t seq_len = in->shape()[0]; + size_t n_heads = in->shape()[1]; + size_t head_dim = in->shape()[2]; + + ASSERT(pos_ids->shape()[0] == seq_len, "pos_ids length mismatch"); + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::rope(out->data(), in->data(), pos_ids->data(), + out->dtype(), seq_len, n_heads, head_dim, theta); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/self_attention/cpu/self_attention_cpu.cpp b/src/ops/self_attention/cpu/self_attention_cpu.cpp new file mode 100644 index 00000000..0331ff42 --- /dev/null +++ b/src/ops/self_attention/cpu/self_attention_cpu.cpp @@ -0,0 +1,129 @@ +#include "self_attention_cpu.hpp" +#include "../../../utils.hpp" +#include +#include +#include + +template +void self_attention_(T *attn_val, const T *q, const T *k, const T *v, + size_t seq_len, size_t total_len, + size_t n_heads, size_t n_kv_heads, size_t head_dim, size_t v_dim, float scale) { + // q: (seq_len, n_heads, head_dim) + // k: (total_len, n_kv_heads, head_dim) + // v: (total_len, n_kv_heads, v_dim) + // attn_val: (seq_len, n_heads, v_dim) + + // Group query attention: each kv head serves multiple q heads + size_t heads_per_kv = n_heads / n_kv_heads; + + for (size_t s = 0; s < seq_len; s++) { + for (size_t h = 0; h < n_heads; h++) { + // Which KV head to use + size_t kv_head = h / heads_per_kv; + + // Calculate attention scores: Q @ K^T * scale + std::vector scores(total_len); + for (size_t t = 0; t < total_len; t++) { + float score = 0.0f; + for (size_t d = 0; d < head_dim; d++) { + float q_val, k_val; + size_t q_idx = s * n_heads * head_dim + h * head_dim + d; + size_t k_idx = t * n_kv_heads * head_dim + kv_head * head_dim + d; + + if constexpr (std::is_same_v || std::is_same_v) { + q_val = llaisys::utils::cast(q[q_idx]); + k_val = llaisys::utils::cast(k[k_idx]); + } else { + q_val = static_cast(q[q_idx]); + k_val = static_cast(k[k_idx]); + } + score += q_val * k_val; + } + scores[t] = score * scale; + } + + // Apply causal mask and softmax + // Causal mask: only attend to positions <= current_position + // current_position in full context is: total_len - seq_len + s + size_t current_pos = total_len - seq_len + s; + + // Find max for numerical stability + float max_score = -INFINITY; + for (size_t t = 0; t <= current_pos; t++) { + max_score = std::max(max_score, scores[t]); + } + + // Compute exp and sum + float exp_sum = 0.0f; + for (size_t t = 0; t <= current_pos; t++) { + scores[t] = std::exp(scores[t] - max_score); + exp_sum += scores[t]; + } + + // Normalize + for (size_t t = 0; t <= current_pos; t++) { + scores[t] /= exp_sum; + } + + // Set masked positions to 0 + for (size_t t = current_pos + 1; t < total_len; t++) { + scores[t] = 0.0f; + } + + // Multiply with V: scores @ V + for (size_t d = 0; d < v_dim; d++) { + float sum = 0.0f; + for (size_t t = 0; t < total_len; t++) { + float v_val; + size_t v_idx = t * n_kv_heads * v_dim + kv_head * v_dim + d; + + if constexpr (std::is_same_v || std::is_same_v) { + v_val = llaisys::utils::cast(v[v_idx]); + } else { + v_val = static_cast(v[v_idx]); + } + sum += scores[t] * v_val; + } + + size_t out_idx = s * n_heads * v_dim + h * v_dim + d; + if constexpr (std::is_same_v || std::is_same_v) { + attn_val[out_idx] = llaisys::utils::cast(sum); + } else { + attn_val[out_idx] = static_cast(sum); + } + } + } + } +} + +namespace llaisys::ops::cpu { +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t dtype, size_t seq_len, size_t total_len, + size_t n_heads, size_t n_kv_heads, size_t head_dim, size_t v_dim, float scale) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return self_attention_( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seq_len, total_len, n_heads, n_kv_heads, head_dim, v_dim, scale); + case LLAISYS_DTYPE_BF16: + return self_attention_( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seq_len, total_len, n_heads, n_kv_heads, head_dim, v_dim, scale); + case LLAISYS_DTYPE_F16: + return self_attention_( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seq_len, total_len, n_heads, n_kv_heads, head_dim, v_dim, scale); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/self_attention/cpu/self_attention_cpu.hpp b/src/ops/self_attention/cpu/self_attention_cpu.hpp new file mode 100644 index 00000000..b1eede9f --- /dev/null +++ b/src/ops/self_attention/cpu/self_attention_cpu.hpp @@ -0,0 +1,8 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t dtype, size_t seq_len, size_t total_len, + size_t n_heads, size_t n_kv_heads, size_t head_dim, size_t v_dim, float scale); +} diff --git a/src/ops/self_attention/op.cpp b/src/ops/self_attention/op.cpp index 43d62014..910fbe9e 100644 --- a/src/ops/self_attention/op.cpp +++ b/src/ops/self_attention/op.cpp @@ -1,7 +1,57 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/self_attention_cpu.hpp" + namespace llaisys::ops { void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale) { - TO_BE_IMPLEMENTED(); + // q: (seq_len, n_heads, head_dim) + // k: (total_len, n_kv_heads, head_dim) + // v: (total_len, n_kv_heads, v_dim) + // attn_val: (seq_len, n_heads, v_dim) + + CHECK_SAME_DEVICE(attn_val, q, k, v); + CHECK_SAME_DTYPE(attn_val->dtype(), q->dtype(), k->dtype(), v->dtype()); + + ASSERT(q->ndim() == 3, "q must be 3D"); + ASSERT(k->ndim() == 3, "k must be 3D"); + ASSERT(v->ndim() == 3, "v must be 3D"); + ASSERT(attn_val->ndim() == 3, "attn_val must be 3D"); + + size_t seq_len = q->shape()[0]; + size_t n_heads = q->shape()[1]; + size_t head_dim = q->shape()[2]; + + size_t total_len = k->shape()[0]; + size_t n_kv_heads = k->shape()[1]; + ASSERT(k->shape()[2] == head_dim, "k head_dim mismatch"); + + ASSERT(v->shape()[0] == total_len, "v total_len mismatch"); + ASSERT(v->shape()[1] == n_kv_heads, "v n_kv_heads mismatch"); + size_t v_dim = v->shape()[2]; + + ASSERT(attn_val->shape()[0] == seq_len, "attn_val seq_len mismatch"); + ASSERT(attn_val->shape()[1] == n_heads, "attn_val n_heads mismatch"); + ASSERT(attn_val->shape()[2] == v_dim, "attn_val v_dim mismatch"); + + ASSERT(n_heads % n_kv_heads == 0, "n_heads must be multiple of n_kv_heads"); + + llaisys::core::context().setDevice(attn_val->deviceType(), attn_val->deviceId()); + + switch (attn_val->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::self_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seq_len, total_len, + n_heads, n_kv_heads, head_dim, v_dim, scale); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/swiglu/cpu/swiglu_cpu.cpp b/src/ops/swiglu/cpu/swiglu_cpu.cpp new file mode 100644 index 00000000..f38c0c36 --- /dev/null +++ b/src/ops/swiglu/cpu/swiglu_cpu.cpp @@ -0,0 +1,50 @@ +#include "swiglu_cpu.hpp" +#include "../../../utils.hpp" +#include + +template +void swiglu_(T *out, const T *gate, const T *up, size_t numel) { + // out[i] = up[i] * (gate[i] / (1 + exp(-gate[i]))) + // This is: up[i] * gate[i] * sigmoid(gate[i]) + for (size_t i = 0; i < numel; i++) { + if constexpr (std::is_same_v || std::is_same_v) { + float g = llaisys::utils::cast(gate[i]); + float u = llaisys::utils::cast(up[i]); + float sigmoid_g = 1.0f / (1.0f + std::exp(-g)); + out[i] = llaisys::utils::cast(u * g * sigmoid_g); + } else { + float g = static_cast(gate[i]); + float u = static_cast(up[i]); + float sigmoid_g = 1.0f / (1.0f + std::exp(-g)); + out[i] = static_cast(u * g * sigmoid_g); + } + } +} + +namespace llaisys::ops::cpu { +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, + llaisysDataType_t dtype, size_t numel) { + switch (dtype) { + case LLAISYS_DTYPE_F32: + return swiglu_( + reinterpret_cast(out), + reinterpret_cast(gate), + reinterpret_cast(up), + numel); + case LLAISYS_DTYPE_BF16: + return swiglu_( + reinterpret_cast(out), + reinterpret_cast(gate), + reinterpret_cast(up), + numel); + case LLAISYS_DTYPE_F16: + return swiglu_( + reinterpret_cast(out), + reinterpret_cast(gate), + reinterpret_cast(up), + numel); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(dtype); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/swiglu/cpu/swiglu_cpu.hpp b/src/ops/swiglu/cpu/swiglu_cpu.hpp new file mode 100644 index 00000000..d364990e --- /dev/null +++ b/src/ops/swiglu/cpu/swiglu_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../../../core/llaisys_core.hpp" + +namespace llaisys::ops::cpu { +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, + llaisysDataType_t dtype, size_t numel); +} diff --git a/src/ops/swiglu/op.cpp b/src/ops/swiglu/op.cpp index 47edbcc9..1c0c4208 100644 --- a/src/ops/swiglu/op.cpp +++ b/src/ops/swiglu/op.cpp @@ -1,7 +1,30 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/swiglu_cpu.hpp" + namespace llaisys::ops { void swiglu(tensor_t out, tensor_t gate, tensor_t up) { - TO_BE_IMPLEMENTED(); + // Check device and shapes + CHECK_SAME_DEVICE(out, gate, up); + CHECK_SAME_SHAPE(out->shape(), gate->shape(), up->shape()); + CHECK_SAME_DTYPE(out->dtype(), gate->dtype(), up->dtype()); + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::swiglu(out->data(), gate->data(), up->data(), + out->dtype(), out->numel()); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/tensor/tensor.cpp b/src/tensor/tensor.cpp index 2f594bb6..f105b3d7 100644 --- a/src/tensor/tensor.cpp +++ b/src/tensor/tensor.cpp @@ -164,27 +164,105 @@ void Tensor::debug() const { } bool Tensor::isContiguous() const { - TO_BE_IMPLEMENTED(); + // Check if tensor is contiguous in memory + // A tensor is contiguous if stride[i] = shape[i+1] * shape[i+2] * ... * shape[n-1] + if (_meta.shape.empty()) { + return true; + } + + ptrdiff_t expected_stride = 1; + for (int i = _meta.shape.size() - 1; i >= 0; i--) { + if (_meta.strides[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(_meta.shape[i]); + } return true; } tensor_t Tensor::permute(const std::vector &order) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + // Permute dimensions according to order + // Example: permute(2, 0, 1) on shape (3, 4, 5) -> shape (5, 3, 4) + + CHECK_ARGUMENT(order.size() == _meta.shape.size(), + "Permute order size must match tensor ndim"); + + // Create new shape and strides + TensorMeta new_meta = _meta; + for (size_t i = 0; i < order.size(); i++) { + CHECK_ARGUMENT(order[i] < _meta.shape.size(), + "Permute order index out of range"); + new_meta.shape[i] = _meta.shape[order[i]]; + new_meta.strides[i] = _meta.strides[order[i]]; + } + + return std::shared_ptr(new Tensor(new_meta, _storage, _offset)); } tensor_t Tensor::view(const std::vector &shape) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + // View reshapes tensor without moving data + // Only works if tensor is contiguous or view is compatible with current strides + + // Check if total number of elements matches + size_t new_numel = 1; + for (size_t s : shape) { + new_numel *= s; + } + CHECK_ARGUMENT(new_numel == this->numel(), + "View shape must have same total number of elements"); + + // If tensor is contiguous, we can create any view + if (this->isContiguous()) { + // Calculate new strides + std::vector new_strides(shape.size()); + ptrdiff_t stride = 1; + for (int i = shape.size() - 1; i >= 0; i--) { + new_strides[i] = stride; + stride *= static_cast(shape[i]); + } + + TensorMeta new_meta{_meta.dtype, shape, new_strides}; + return std::shared_ptr(new Tensor(new_meta, _storage, _offset)); + } + + // For non-contiguous tensors, we need to check if the view is compatible + // This is a simplified check - view is allowed only if we're combining or splitting + // contiguous dimensions + + // For now, we'll just throw an error for non-contiguous tensors + // A full implementation would check dimension compatibility + throw std::runtime_error("View is not supported for non-contiguous tensors"); } tensor_t Tensor::slice(size_t dim, size_t start, size_t end) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + // Slice along dimension dim from start (inclusive) to end (exclusive) + CHECK_ARGUMENT(dim < _meta.shape.size(), "Slice dimension out of range"); + CHECK_ARGUMENT(start < end && end <= _meta.shape[dim], "Invalid slice range"); + + // Create new meta with updated shape + TensorMeta new_meta = _meta; + new_meta.shape[dim] = end - start; + + // Calculate new offset: offset += start * stride[dim] * elementSize + size_t new_offset = _offset + start * _meta.strides[dim] * this->elementSize(); + + return std::shared_ptr(new Tensor(new_meta, _storage, new_offset)); } void Tensor::load(const void *src_) { - TO_BE_IMPLEMENTED(); + // Load data from host memory to this tensor + core::context().setDevice(this->deviceType(), this->deviceId()); + const std::byte *src = reinterpret_cast(src_); + size_t bytes = this->numel() * this->elementSize(); + + if (this->deviceType() == LLAISYS_DEVICE_CPU) { + // CPU to CPU: direct memcpy + std::memcpy(this->data(), src, bytes); + } else { + // Host to Device: use H2D memcpy + core::context().runtime().api()->memcpy_sync( + this->data(), src, bytes, LLAISYS_MEMCPY_H2D); + } } tensor_t Tensor::contiguous() const { diff --git a/test/test_infer.py b/test/test_infer.py index 59d06b87..b1110c1a 100644 --- a/test/test_infer.py +++ b/test/test_infer.py @@ -20,7 +20,9 @@ def load_hf_model(model_path=None, device_name="cpu"): if model_path and os.path.isdir(model_path): print(f"Loading model from local path: {model_path}") else: - print(f"Loading model from Hugging Face: {model_id}") + # 设置镜像源环境变量 + os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' + print(f"Loading model from Hugging Face mirror: {model_id}") model_path = snapshot_download(model_id) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( diff --git a/xmake.lua b/xmake.lua index 1f65f7a9..5d5b45e7 100644 --- a/xmake.lua +++ b/xmake.lua @@ -89,12 +89,28 @@ target("llaisys-ops") if not is_plat("windows") then add_cxflags("-fPIC", "-Wno-unknown-pragmas") end - + add_files("src/ops/*/*.cpp") on_install(function (target) end) target_end() +target("llaisys-models") + set_kind("static") + add_deps("llaisys-tensor") + add_deps("llaisys-ops") + + set_languages("cxx17") + set_warnings("all", "error") + if not is_plat("windows") then + add_cxflags("-fPIC", "-Wno-unknown-pragmas") + end + + add_files("src/models/**/*.cpp") + + on_install(function (target) end) +target_end() + target("llaisys") set_kind("shared") add_deps("llaisys-utils") @@ -102,10 +118,12 @@ target("llaisys") add_deps("llaisys-core") add_deps("llaisys-tensor") add_deps("llaisys-ops") + add_deps("llaisys-models") set_languages("cxx17") set_warnings("all", "error") add_files("src/llaisys/*.cc") + add_files("src/llaisys/**/*.cc") set_installdir(".") From 8336eb002df97b159a96f24c813485424b36313e Mon Sep 17 00:00:00 2001 From: Anki77134 <2577484662@qq.com> Date: Tue, 10 Feb 2026 04:43:36 +0800 Subject: [PATCH 2/2] feat: implement Qwen2 inference logic and fixed RoPE CPU op --- python/llaisys/models/qwen2.py | 63 ++++++++++++++++++++++------------ src/llaisys/models/qwen2.cc | 33 +++++++++--------- src/ops/rope/cpu/rope_cpu.cpp | 13 +++---- test/test_infer.py | 10 +++--- 4 files changed, 67 insertions(+), 52 deletions(-) diff --git a/python/llaisys/models/qwen2.py b/python/llaisys/models/qwen2.py index fdfba2a5..4d7749f0 100644 --- a/python/llaisys/models/qwen2.py +++ b/python/llaisys/models/qwen2.py @@ -59,29 +59,36 @@ def __init__(self, model_path, device: DeviceType = DeviceType.CPU): def _load_weights(self, model_path): """Load weights from safetensors files""" + import torch + import numpy as np + import ctypes + self.weight_tensors = [] # Keep references to prevent garbage collection weight_map = {} for file in sorted(model_path.glob("*.safetensors")): - with safetensors.safe_open(file, framework="numpy", device="cpu") as f: + with safetensors.safe_open(file, framework="pt", device="cpu") as f: for name in f.keys(): - weight_map[name] = f.get_tensor(name) + # Load as contiguous PyTorch tensor + pt_tensor = f.get_tensor(name).contiguous() + self.weight_tensors.append(pt_tensor) # Keep alive + weight_map[name] = pt_tensor # Load embedding if "model.embed_tokens.weight" in weight_map: embed_data = weight_map["model.embed_tokens.weight"] tensor = Tensor.from_ptr(self._weights.in_embed) - tensor.load(embed_data.ctypes.data) + tensor.load(ctypes.c_void_p(embed_data.data_ptr())) # Load output norm and lm_head if "model.norm.weight" in weight_map: norm_data = weight_map["model.norm.weight"] tensor = Tensor.from_ptr(self._weights.out_norm_w) - tensor.load(norm_data.ctypes.data) + tensor.load(ctypes.c_void_p(norm_data.data_ptr())) if "lm_head.weight" in weight_map: lm_head_data = weight_map["lm_head.weight"] tensor = Tensor.from_ptr(self._weights.out_embed) - tensor.load(lm_head_data.ctypes.data) + tensor.load(ctypes.c_void_p(lm_head_data.data_ptr())) # Load per-layer weights for layer_idx in range(self._meta.nlayer): @@ -91,66 +98,66 @@ def _load_weights(self, model_path): if f"{prefix}.input_layernorm.weight" in weight_map: data = weight_map[f"{prefix}.input_layernorm.weight"] tensor = Tensor.from_ptr(self._weights.attn_norm_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) # Q, K, V projections if f"{prefix}.self_attn.q_proj.weight" in weight_map: data = weight_map[f"{prefix}.self_attn.q_proj.weight"] tensor = Tensor.from_ptr(self._weights.attn_q_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.self_attn.q_proj.bias" in weight_map: data = weight_map[f"{prefix}.self_attn.q_proj.bias"] tensor = Tensor.from_ptr(self._weights.attn_q_b[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.self_attn.k_proj.weight" in weight_map: data = weight_map[f"{prefix}.self_attn.k_proj.weight"] tensor = Tensor.from_ptr(self._weights.attn_k_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.self_attn.k_proj.bias" in weight_map: data = weight_map[f"{prefix}.self_attn.k_proj.bias"] tensor = Tensor.from_ptr(self._weights.attn_k_b[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.self_attn.v_proj.weight" in weight_map: data = weight_map[f"{prefix}.self_attn.v_proj.weight"] tensor = Tensor.from_ptr(self._weights.attn_v_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.self_attn.v_proj.bias" in weight_map: data = weight_map[f"{prefix}.self_attn.v_proj.bias"] tensor = Tensor.from_ptr(self._weights.attn_v_b[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) # O projection if f"{prefix}.self_attn.o_proj.weight" in weight_map: data = weight_map[f"{prefix}.self_attn.o_proj.weight"] tensor = Tensor.from_ptr(self._weights.attn_o_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) # MLP norm if f"{prefix}.post_attention_layernorm.weight" in weight_map: data = weight_map[f"{prefix}.post_attention_layernorm.weight"] tensor = Tensor.from_ptr(self._weights.mlp_norm_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) # MLP projections if f"{prefix}.mlp.gate_proj.weight" in weight_map: data = weight_map[f"{prefix}.mlp.gate_proj.weight"] tensor = Tensor.from_ptr(self._weights.mlp_gate_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.mlp.up_proj.weight" in weight_map: data = weight_map[f"{prefix}.mlp.up_proj.weight"] tensor = Tensor.from_ptr(self._weights.mlp_up_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) if f"{prefix}.mlp.down_proj.weight" in weight_map: data = weight_map[f"{prefix}.mlp.down_proj.weight"] tensor = Tensor.from_ptr(self._weights.mlp_down_w[layer_idx]) - tensor.load(data.ctypes.data) + tensor.load(ctypes.c_void_p(data.data_ptr())) def generate( self, @@ -168,13 +175,25 @@ def generate( generated = list(inputs) max_gen = max_new_tokens if max_new_tokens else 100 - for _ in range(max_gen): - # Convert to ctypes array - input_array = (ctypes.c_int64 * len(generated))(*generated) + # First forward pass with full prompt + input_array = (ctypes.c_int64 * len(generated))(*generated) + next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( + self._model, input_array, len(generated) + ) + + if next_token < 0 or next_token == self._meta.end_token: + generated.append(next_token) + return generated + + generated.append(next_token) + + # Subsequent passes: only send new token (using KV cache) + for _ in range(max_gen - 1): + # Only send the last generated token + input_array = (ctypes.c_int64 * 1)(generated[-1]) - # Call inference next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( - self._model, input_array, len(generated) + self._model, input_array, 1 ) if next_token < 0: diff --git a/src/llaisys/models/qwen2.cc b/src/llaisys/models/qwen2.cc index 4c1b24e0..a245e842 100644 --- a/src/llaisys/models/qwen2.cc +++ b/src/llaisys/models/qwen2.cc @@ -1,6 +1,7 @@ #include "llaisys/models/qwen2.h" #include "../../models/qwen2/qwen2.hpp" #include "../../core/llaisys_core.hpp" +#include "../llaisys_tensor.hpp" using namespace llaisys; @@ -53,10 +54,10 @@ __export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen // Allocate C struct auto weights = new LlaisysQwen2Weights(); - // Copy pointers (shallow copy) - weights->in_embed = reinterpret_cast(cpp_weights.embed_tokens.get()); - weights->out_embed = reinterpret_cast(cpp_weights.lm_head.get()); - weights->out_norm_w = reinterpret_cast(cpp_weights.norm_weight.get()); + // Wrap tensor_t (shared_ptr) in LlaisysTensor struct + weights->in_embed = new LlaisysTensor{cpp_weights.embed_tokens}; + weights->out_embed = new LlaisysTensor{cpp_weights.lm_head}; + weights->out_norm_w = new LlaisysTensor{cpp_weights.norm_weight}; // Allocate arrays for per-layer weights size_t n_layers = config.n_layers; @@ -75,18 +76,18 @@ __export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen weights->mlp_down_w = new llaisysTensor_t[n_layers]; for (size_t i = 0; i < n_layers; i++) { - weights->attn_norm_w[i] = reinterpret_cast(cpp_weights.attn_norm_weight[i].get()); - weights->attn_q_w[i] = reinterpret_cast(cpp_weights.attn_q_weight[i].get()); - weights->attn_q_b[i] = reinterpret_cast(cpp_weights.attn_q_bias[i].get()); - weights->attn_k_w[i] = reinterpret_cast(cpp_weights.attn_k_weight[i].get()); - weights->attn_k_b[i] = reinterpret_cast(cpp_weights.attn_k_bias[i].get()); - weights->attn_v_w[i] = reinterpret_cast(cpp_weights.attn_v_weight[i].get()); - weights->attn_v_b[i] = reinterpret_cast(cpp_weights.attn_v_bias[i].get()); - weights->attn_o_w[i] = reinterpret_cast(cpp_weights.attn_o_weight[i].get()); - weights->mlp_norm_w[i] = reinterpret_cast(cpp_weights.mlp_norm_weight[i].get()); - weights->mlp_gate_w[i] = reinterpret_cast(cpp_weights.mlp_gate_weight[i].get()); - weights->mlp_up_w[i] = reinterpret_cast(cpp_weights.mlp_up_weight[i].get()); - weights->mlp_down_w[i] = reinterpret_cast(cpp_weights.mlp_down_weight[i].get()); + weights->attn_norm_w[i] = new LlaisysTensor{cpp_weights.attn_norm_weight[i]}; + weights->attn_q_w[i] = new LlaisysTensor{cpp_weights.attn_q_weight[i]}; + weights->attn_q_b[i] = new LlaisysTensor{cpp_weights.attn_q_bias[i]}; + weights->attn_k_w[i] = new LlaisysTensor{cpp_weights.attn_k_weight[i]}; + weights->attn_k_b[i] = new LlaisysTensor{cpp_weights.attn_k_bias[i]}; + weights->attn_v_w[i] = new LlaisysTensor{cpp_weights.attn_v_weight[i]}; + weights->attn_v_b[i] = new LlaisysTensor{cpp_weights.attn_v_bias[i]}; + weights->attn_o_w[i] = new LlaisysTensor{cpp_weights.attn_o_weight[i]}; + weights->mlp_norm_w[i] = new LlaisysTensor{cpp_weights.mlp_norm_weight[i]}; + weights->mlp_gate_w[i] = new LlaisysTensor{cpp_weights.mlp_gate_weight[i]}; + weights->mlp_up_w[i] = new LlaisysTensor{cpp_weights.mlp_up_weight[i]}; + weights->mlp_down_w[i] = new LlaisysTensor{cpp_weights.mlp_down_weight[i]}; } return weights; diff --git a/src/ops/rope/cpu/rope_cpu.cpp b/src/ops/rope/cpu/rope_cpu.cpp index e0695903..26768cd7 100644 --- a/src/ops/rope/cpu/rope_cpu.cpp +++ b/src/ops/rope/cpu/rope_cpu.cpp @@ -14,19 +14,16 @@ void rope_(T *out, const T *in, const int64_t *pos_ids, size_t half_dim = head_dim / 2; - // Precompute frequency divisors to improve numerical stability - std::vector inv_freq(half_dim); - for (size_t j = 0; j < half_dim; j++) { - inv_freq[j] = 1.0 / std::pow(theta, (2.0 * j) / static_cast(head_dim)); - } - for (size_t s = 0; s < seq_len; s++) { float pos = static_cast(pos_ids[s]); for (size_t h = 0; h < n_heads; h++) { for (size_t j = 0; j < half_dim; j++) { - // Calculate angle - float angle = pos * inv_freq[j]; + // Calculate angle: pos / (theta^(2j/d)) to match PyTorch exactly + // Using the same calculation order as PyTorch + float exponent = (2.0f * static_cast(j)) / static_cast(head_dim); + float divisor = std::pow(theta, exponent); + float angle = pos / divisor; float cos_angle = std::cos(angle); float sin_angle = std::sin(angle); diff --git a/test/test_infer.py b/test/test_infer.py index b1110c1a..a89f2258 100644 --- a/test/test_infer.py +++ b/test/test_infer.py @@ -14,15 +14,13 @@ sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8") -def load_hf_model(model_path=None, device_name="cpu"): +def load_hf_model(model_path="/root/autodl-tmp/llaisys/model", device_name="cpu"): model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" if model_path and os.path.isdir(model_path): print(f"Loading model from local path: {model_path}") else: - # 设置镜像源环境变量 - os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' - print(f"Loading model from Hugging Face mirror: {model_id}") + print(f"Loading model from Hugging Face: {model_id}") model_path = snapshot_download(model_id) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( @@ -84,7 +82,7 @@ def llaisys_infer( if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) - parser.add_argument("--model", default=None, type=str) + parser.add_argument("--model", default="/root/autodl-tmp/llaisys/model", type=str) parser.add_argument("--prompt", default="Who are you?", type=str) parser.add_argument("--max_steps", default=128, type=int) parser.add_argument("--top_p", default=0.8, type=float) @@ -148,4 +146,4 @@ def llaisys_infer( if args.test: assert llaisys_tokens == tokens - print("\033[92mTest passed!\033[0m\n") + print("\033[92mTest passed!\033[0m\n") \ No newline at end of file