Skip to content

Commit edf1a0b

Browse files
Shl1015CSShl1015CS
committed
homework
1 parent f3466ec commit edf1a0b

27 files changed

Lines changed: 1347 additions & 32 deletions

File tree

‎python/llaisys/libllaisys/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from .tensor import llaisysTensor_t
1313
from .tensor import load_tensor
1414
from .ops import load_ops
15+
from .models import load_qwen2_models
1516

1617

1718
def load_shared_library():
@@ -38,6 +39,7 @@ def load_shared_library():
3839
load_runtime(LIB_LLAISYS)
3940
load_tensor(LIB_LLAISYS)
4041
load_ops(LIB_LLAISYS)
42+
QWEN2_MODELS = load_qwen2_models(LIB_LLAISYS)
4143

4244

4345
__all__ = [

‎python/llaisys/libllaisys/models.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import ctypes
2+
from .llaisys_types import llaisysDeviceType_t, llaisysDataType_t
3+
from .tensor import llaisysTensor_t
4+
5+
# Qwen2 Meta structure
6+
class LlaisysQwen2Meta(ctypes.Structure):
7+
_fields_ = [
8+
("dtype", llaisysDataType_t),
9+
("nlayer", ctypes.c_size_t),
10+
("hs", ctypes.c_size_t), # hidden size
11+
("nh", ctypes.c_size_t), # number of heads
12+
("nkvh", ctypes.c_size_t), # number of kv heads
13+
("dh", ctypes.c_size_t), # head dimension
14+
("di", ctypes.c_size_t), # intermediate size
15+
("maxseq", ctypes.c_size_t), # max sequence length
16+
("voc", ctypes.c_size_t), # vocabulary size
17+
("epsilon", ctypes.c_float),
18+
("theta", ctypes.c_float),
19+
("end_token", ctypes.c_int64),
20+
]
21+
22+
# Qwen2 Weights structure
23+
class LlaisysQwen2Weights(ctypes.Structure):
24+
_fields_ = [
25+
("in_embed", llaisysTensor_t),
26+
("out_embed", llaisysTensor_t),
27+
("out_norm_w", llaisysTensor_t),
28+
("attn_norm_w", ctypes.POINTER(llaisysTensor_t)),
29+
("attn_q_w", ctypes.POINTER(llaisysTensor_t)),
30+
("attn_q_b", ctypes.POINTER(llaisysTensor_t)),
31+
("attn_k_w", ctypes.POINTER(llaisysTensor_t)),
32+
("attn_k_b", ctypes.POINTER(llaisysTensor_t)),
33+
("attn_v_w", ctypes.POINTER(llaisysTensor_t)),
34+
("attn_v_b", ctypes.POINTER(llaisysTensor_t)),
35+
("attn_o_w", ctypes.POINTER(llaisysTensor_t)),
36+
("mlp_norm_w", ctypes.POINTER(llaisysTensor_t)),
37+
("mlp_gate_w", ctypes.POINTER(llaisysTensor_t)),
38+
("mlp_up_w", ctypes.POINTER(llaisysTensor_t)),
39+
("mlp_down_w", ctypes.POINTER(llaisysTensor_t)),
40+
]
41+
42+
# Opaque model pointer
43+
LlaisysQwen2Model = ctypes.c_void_p
44+
45+
def load_qwen2_models(lib):
46+
"""Load Qwen2 model functions from the shared library"""
47+
48+
# llaisysQwen2ModelCreate
49+
lib.llaisysQwen2ModelCreate.argtypes = [
50+
ctypes.POINTER(LlaisysQwen2Meta),
51+
llaisysDeviceType_t,
52+
ctypes.POINTER(ctypes.c_int),
53+
ctypes.c_int
54+
]
55+
lib.llaisysQwen2ModelCreate.restype = LlaisysQwen2Model
56+
57+
# llaisysQwen2ModelDestroy
58+
lib.llaisysQwen2ModelDestroy.argtypes = [LlaisysQwen2Model]
59+
lib.llaisysQwen2ModelDestroy.restype = None
60+
61+
# llaisysQwen2ModelWeights
62+
lib.llaisysQwen2ModelWeights.argtypes = [LlaisysQwen2Model]
63+
lib.llaisysQwen2ModelWeights.restype = ctypes.POINTER(LlaisysQwen2Weights)
64+
65+
# llaisysQwen2ModelInfer
66+
lib.llaisysQwen2ModelInfer.argtypes = [
67+
LlaisysQwen2Model,
68+
ctypes.POINTER(ctypes.c_int64),
69+
ctypes.c_size_t
70+
]
71+
lib.llaisysQwen2ModelInfer.restype = ctypes.c_int64
72+
73+
return {
74+
'LlaisysQwen2Meta': LlaisysQwen2Meta,
75+
'LlaisysQwen2Weights': LlaisysQwen2Weights,
76+
'LlaisysQwen2Model': LlaisysQwen2Model,
77+
'qwen2_model_create': lib.llaisysQwen2ModelCreate,
78+
'qwen2_model_destroy': lib.llaisysQwen2ModelDestroy,
79+
'qwen2_model_weights': lib.llaisysQwen2ModelWeights,
80+
'qwen2_model_infer': lib.llaisysQwen2ModelInfer,
81+
}

‎python/llaisys/models/qwen2.py‎

Lines changed: 67 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,84 @@
11
from typing import Sequence
2-
from ..libllaisys import LIB_LLAISYS
3-
from ..libllaisys import DeviceType
2+
from ..libllaisys import LIB_LLAISYS, QWEN2_MODELS
3+
from ..libllaisys import DeviceType, DataType
4+
from .. import Tensor
45

56
from pathlib import Path
67
import safetensors
8+
import json
9+
import ctypes
10+
import numpy as np
711

812

913
class Qwen2:
1014

1115
def __init__(self, model_path, device: DeviceType = DeviceType.CPU):
12-
# TODO: Implement model constructor
13-
1416
model_path = Path(model_path)
15-
16-
for file in sorted(model_path.glob("*.safetensors")):
17-
data_ = safetensors.safe_open(file, framework="numpy", device="cpu")
18-
for name_ in data_.keys():
19-
## TODO: load the model weights
20-
pass
17+
18+
# Load config
19+
with open(model_path / "config.json", "r") as f:
20+
config = json.load(f)
21+
22+
# Create meta
23+
meta = QWEN2_MODELS['LlaisysQwen2Meta']()
24+
meta.dtype = DataType.BF16 # Use BF16 for inference
25+
meta.nlayer = config["num_hidden_layers"]
26+
meta.hs = config["hidden_size"]
27+
meta.nh = config["num_attention_heads"]
28+
meta.nkvh = config.get("num_key_value_heads", config["num_attention_heads"])
29+
meta.dh = config["hidden_size"] // config["num_attention_heads"]
30+
meta.di = config["intermediate_size"]
31+
meta.maxseq = 2048 # Max context length
32+
meta.voc = config["vocab_size"]
33+
meta.epsilon = config.get("rms_norm_eps", 1e-6)
34+
meta.theta = config.get("rope_theta", 10000.0)
35+
meta.end_token = config.get("eos_token_id", 151645)
36+
37+
# Create model
38+
device_id = 0
39+
self.model = QWEN2_MODELS['qwen2_model_create'](
40+
ctypes.byref(meta), device, ctypes.byref(ctypes.c_int(device_id)), 1
41+
)
42+
43+
if not self.model:
44+
raise RuntimeError("Failed to create Qwen2 model")
45+
46+
self.meta = meta
47+
self.device = device
48+
49+
print(f"✅ Qwen2 model created with {meta.nlayer} layers, vocab_size={meta.voc}")
50+
51+
# Placeholder for weight loading - to be implemented later
52+
# For now, we use the dummy implementation in C++
2153

2254
def generate(
2355
self,
2456
inputs: Sequence[int],
25-
max_new_tokens: int = None,
57+
max_new_tokens: int = 128,
2658
top_k: int = 1,
27-
top_p: float = 0.8,
59+
top_p: float = 0.8,
2860
temperature: float = 0.8,
2961
):
30-
31-
# TODO: Implement generate function
32-
33-
return []
62+
# Convert inputs to list if not already
63+
token_list = list(inputs)
64+
65+
# Generate tokens one by one
66+
for _ in range(max_new_tokens):
67+
# Convert to ctypes array
68+
token_array = (ctypes.c_int64 * len(token_list))(*token_list)
69+
70+
# Get next token
71+
next_token = QWEN2_MODELS['qwen2_model_infer'](
72+
self.model, token_array, len(token_list)
73+
)
74+
75+
if next_token == self.meta.end_token or next_token < 0:
76+
break
77+
78+
token_list.append(next_token)
79+
80+
return token_list
81+
82+
def __del__(self):
83+
if hasattr(self, 'model') and self.model:
84+
QWEN2_MODELS['qwen2_model_destroy'](self.model)

‎src/llaisys/models/qwen2.cc‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
#include "../../../include/llaisys/models/qwen2.h"
2+
3+
#include <cstdlib>
4+
#include <random>
5+
6+
struct LlaisysQwen2Model {
7+
LlaisysQwen2Meta meta;
8+
LlaisysQwen2Weights weights;
9+
llaisysDeviceType_t device;
10+
11+
LlaisysQwen2Model(const LlaisysQwen2Meta &m, llaisysDeviceType_t dev)
12+
: meta(m), device(dev) {
13+
// Initialize weights structure to null
14+
weights.in_embed = nullptr;
15+
weights.out_embed = nullptr;
16+
weights.out_norm_w = nullptr;
17+
weights.attn_norm_w = nullptr;
18+
weights.attn_q_w = nullptr;
19+
weights.attn_q_b = nullptr;
20+
weights.attn_k_w = nullptr;
21+
weights.attn_k_b = nullptr;
22+
weights.attn_v_w = nullptr;
23+
weights.attn_v_b = nullptr;
24+
weights.attn_o_w = nullptr;
25+
weights.mlp_norm_w = nullptr;
26+
weights.mlp_gate_w = nullptr;
27+
weights.mlp_up_w = nullptr;
28+
weights.mlp_down_w = nullptr;
29+
}
30+
};
31+
32+
__export struct LlaisysQwen2Model *llaisysQwen2ModelCreate(
33+
const LlaisysQwen2Meta *meta,
34+
llaisysDeviceType_t device,
35+
int *device_ids,
36+
int ndevice
37+
) {
38+
try {
39+
return new LlaisysQwen2Model(*meta, device);
40+
} catch (...) {
41+
return nullptr;
42+
}
43+
}
44+
45+
__export void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model *model) {
46+
delete model;
47+
}
48+
49+
__export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model *model) {
50+
return &model->weights;
51+
}
52+
53+
__export int64_t llaisysQwen2ModelInfer(
54+
struct LlaisysQwen2Model *model,
55+
int64_t *token_ids,
56+
size_t ntoken
57+
) {
58+
// For now, return a dummy token to test basic functionality
59+
// This is a placeholder implementation
60+
61+
if (ntoken == 0) {
62+
return -1; // Error case
63+
}
64+
65+
// Simple dummy response - return the last token + 1 (for testing)
66+
// In a real implementation, this would do full model inference
67+
int64_t last_token = token_ids[ntoken - 1];
68+
69+
// Return some pattern based on input for testing
70+
return (last_token + 1) % model->meta.voc;
71+
}

‎src/ops/argmax/cpu/argmax_cpu.cpp‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
#include "argmax_cpu.hpp"
2+
3+
#include "../../../utils.hpp"
4+
5+
#include <limits>
6+
7+
template <typename T>
8+
void argmax_(int64_t *max_idx, T *max_val, const T *vals, size_t numel) {
9+
if (numel == 0) return;
10+
11+
size_t best_idx = 0;
12+
T best_val;
13+
14+
if constexpr (std::is_same_v<T, llaisys::bf16_t> || std::is_same_v<T, llaisys::fp16_t>) {
15+
best_val = vals[0];
16+
float best_float = llaisys::utils::cast<float>(best_val);
17+
18+
for (size_t i = 1; i < numel; i++) {
19+
float current_float = llaisys::utils::cast<float>(vals[i]);
20+
if (current_float > best_float) {
21+
best_float = current_float;
22+
best_val = vals[i];
23+
best_idx = i;
24+
}
25+
}
26+
} else {
27+
best_val = vals[0];
28+
for (size_t i = 1; i < numel; i++) {
29+
if (vals[i] > best_val) {
30+
best_val = vals[i];
31+
best_idx = i;
32+
}
33+
}
34+
}
35+
36+
*max_idx = static_cast<int64_t>(best_idx);
37+
*max_val = best_val;
38+
}
39+
40+
namespace llaisys::ops::cpu {
41+
void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel) {
42+
switch (type) {
43+
case LLAISYS_DTYPE_F32:
44+
return argmax_(
45+
reinterpret_cast<int64_t *>(max_idx),
46+
reinterpret_cast<float *>(max_val),
47+
reinterpret_cast<const float *>(vals),
48+
numel
49+
);
50+
case LLAISYS_DTYPE_BF16:
51+
return argmax_(
52+
reinterpret_cast<int64_t *>(max_idx),
53+
reinterpret_cast<llaisys::bf16_t *>(max_val),
54+
reinterpret_cast<const llaisys::bf16_t *>(vals),
55+
numel
56+
);
57+
case LLAISYS_DTYPE_F16:
58+
return argmax_(
59+
reinterpret_cast<int64_t *>(max_idx),
60+
reinterpret_cast<llaisys::fp16_t *>(max_val),
61+
reinterpret_cast<const llaisys::fp16_t *>(vals),
62+
numel
63+
);
64+
default:
65+
EXCEPTION_UNSUPPORTED_DATATYPE(type);
66+
}
67+
}
68+
} // namespace llaisys::ops::cpu

‎src/ops/argmax/cpu/argmax_cpu.hpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#pragma once
2+
#include "llaisys.h"
3+
4+
#include <cstddef>
5+
6+
namespace llaisys::ops::cpu {
7+
void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel);
8+
}

‎src/ops/argmax/op.cpp‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,45 @@
11
#include "op.hpp"
22

3+
#include "../../core/llaisys_core.hpp"
4+
#include "../../utils.hpp"
5+
6+
#include "cpu/argmax_cpu.hpp"
7+
38
namespace llaisys::ops {
49
void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals) {
5-
TO_BE_IMPLEMENTED();
10+
// Check that tensors are on same device
11+
CHECK_SAME_DEVICE(max_idx, max_val, vals);
12+
13+
// Check data types
14+
ASSERT(max_idx->dtype() == LLAISYS_DTYPE_I64, "Argmax: max_idx must be int64");
15+
CHECK_SAME_DTYPE(max_val->dtype(), vals->dtype());
16+
17+
// Check contiguity
18+
ASSERT(max_idx->isContiguous() && max_val->isContiguous() && vals->isContiguous(),
19+
"Argmax: all tensors must be contiguous");
20+
21+
// For now, assume vals is 1D and results are scalar (single element)
22+
ASSERT(vals->ndim() == 1, "Argmax: vals must be 1D tensor for now");
23+
ASSERT(max_idx->numel() == 1, "Argmax: max_idx must be scalar");
24+
ASSERT(max_val->numel() == 1, "Argmax: max_val must be scalar");
25+
26+
// always support cpu calculation
27+
if (vals->deviceType() == LLAISYS_DEVICE_CPU) {
28+
return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
29+
}
30+
31+
llaisys::core::context().setDevice(vals->deviceType(), vals->deviceId());
32+
33+
switch (vals->deviceType()) {
34+
case LLAISYS_DEVICE_CPU:
35+
return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
36+
#ifdef ENABLE_NVIDIA_API
37+
case LLAISYS_DEVICE_NVIDIA:
38+
TO_BE_IMPLEMENTED();
39+
return;
40+
#endif
41+
default:
42+
EXCEPTION_UNSUPPORTED_DEVICE;
43+
}
644
}
745
} // namespace llaisys::ops

0 commit comments

Comments
 (0)