diff --git a/.clang-format b/.clang-format
index a77ae97c3..291c545de 100644
--- a/.clang-format
+++ b/.clang-format
@@ -1,30 +1,30 @@
----
-BasedOnStyle: LLVM
-IndentWidth: 4 # 缩进宽度，LLVM 默认值为 2，改为 4
-AccessModifierOffset: -4 # public/protected/private 访问控制符相对成员的偏移，与 IndentWidth 配合，LLVM 默认值为 -2
-AlignOperands: AlignAfterOperator # 双目运算符的行间对齐，LLVM 默认值为 Align，改为带符号一起换行
-BreakBeforeBinaryOperators: All # 在双目运算符之前换行，LLVM 默认值为 None，改为换行时总是把双目运算符放在行首，包括赋值（=）
-ColumnLimit: 0 # 列宽限制，LLVM 默认值为 80，改为不限制
-AllowShortBlocksOnASingleLine: Always # 是否允许短块（单个语句的块）不换行，LLVM 默认值为 Never，改为允许
-AllowShortLoopsOnASingleLine: true # 是否允许短循环不换行，LLVM 默认值为 false，改为允许
-InsertBraces: true # 是否在 if/for/while/switch 等语句后插入大括号，LLVM 默认值为 false，改为允许
-BreakBeforeBraces: Custom # 大括号换行配置，LLVM 默认值为 LLVM，改为自定义以使 BraceWrapping 生效
-BraceWrapping:
- AfterCaseLabel: false
- AfterClass: false
- AfterControlStatement: Never
- AfterEnum: false
- AfterFunction: false
- AfterNamespace: false
- AfterObjCDeclaration: false
- AfterStruct: false
- AfterUnion: false
- AfterExternBlock: false
- BeforeCatch: false
- BeforeElse: false
- BeforeLambdaBody: false
- BeforeWhile: false
- IndentBraces: false
- SplitEmptyFunction: true
- SplitEmptyRecord: true
- SplitEmptyNamespace: true
+---
+BasedOnStyle: LLVM
+IndentWidth: 4 # 缩进宽度，LLVM 默认值为 2，改为 4
+AccessModifierOffset: -4 # public/protected/private 访问控制符相对成员的偏移，与 IndentWidth 配合，LLVM 默认值为 -2
+AlignOperands: AlignAfterOperator # 双目运算符的行间对齐，LLVM 默认值为 Align，改为带符号一起换行
+BreakBeforeBinaryOperators: All # 在双目运算符之前换行，LLVM 默认值为 None，改为换行时总是把双目运算符放在行首，包括赋值（=）
+ColumnLimit: 0 # 列宽限制，LLVM 默认值为 80，改为不限制
+AllowShortBlocksOnASingleLine: Always # 是否允许短块（单个语句的块）不换行，LLVM 默认值为 Never，改为允许
+AllowShortLoopsOnASingleLine: true # 是否允许短循环不换行，LLVM 默认值为 false，改为允许
+InsertBraces: true # 是否在 if/for/while/switch 等语句后插入大括号，LLVM 默认值为 false，改为允许
+BreakBeforeBraces: Custom # 大括号换行配置，LLVM 默认值为 LLVM，改为自定义以使 BraceWrapping 生效
+BraceWrapping:
+ AfterCaseLabel: false
+ AfterClass: false
+ AfterControlStatement: Never
+ AfterEnum: false
+ AfterFunction: false
+ AfterNamespace: false
+ AfterObjCDeclaration: false
+ AfterStruct: false
+ AfterUnion: false
+ AfterExternBlock: false
+ BeforeCatch: false
+ BeforeElse: false
+ BeforeLambdaBody: false
+ BeforeWhile: false
+ IndentBraces: false
+ SplitEmptyFunction: true
+ SplitEmptyRecord: true
+ SplitEmptyNamespace: true
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 3d31c23bb..bd78cba2b 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -1,60 +1,60 @@
-name: Build and test
-on:
- pull_request:
- push:
- paths-ignore:
- - '**.md'
- - 'LICENSE'
-
-jobs:
- build:
- name: Build
- strategy:
- fail-fast: false
- matrix:
- os: [windows-latest, ubuntu-latest]
- type: [release]
- runs-on: ${{ matrix.os }}
- steps:
-
- - name: checkout code
- uses: actions/checkout@v4
-
- - name: install xmake
- uses: xmake-io/github-action-setup-xmake@v1
- with:
- xmake-version: latest
-
- - name: Xmake Build & Install
- run: |
- xmake
- xmake install
-
- - name: Install Python
- run: |
- cd python
- pip install .
- cd ..
-
- - name: Assignment-0
- run: |
- python test/test_runtime.py --device cpu
-
- - name: Assignment-1
- run: |
- python test/test_tensor.py
-
- - name: Assignment-2
- run: |
- python test/ops/add.py
- python test/ops/argmax.py
- python test/ops/embedding.py
- python test/ops/linear.py
- python test/ops/rms_norm.py
- python test/ops/rope.py
- python test/ops/self_attention.py
- python test/ops/swiglu.py
-
- - name: Assignment-3
- run: |
- python test/test_infer.py --test
+name: Build and test
+on:
+ pull_request:
+ push:
+ paths-ignore:
+ - '**.md'
+ - 'LICENSE'
+
+jobs:
+ build:
+ name: Build
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [windows-latest, ubuntu-latest]
+ type: [release]
+ runs-on: ${{ matrix.os }}
+ steps:
+
+ - name: checkout code
+ uses: actions/checkout@v4
+
+ - name: install xmake
+ uses: xmake-io/github-action-setup-xmake@v1
+ with:
+ xmake-version: latest
+
+ - name: Xmake Build & Install
+ run: |
+ xmake
+ xmake install
+
+ - name: Install Python
+ run: |
+ cd python
+ pip install .
+ cd ..
+
+ - name: Assignment-0
+ run: |
+ python test/test_runtime.py --device cpu
+
+ - name: Assignment-1
+ run: |
+ python test/test_tensor.py
+
+ - name: Assignment-2
+ run: |
+ python test/ops/add.py
+ python test/ops/argmax.py
+ python test/ops/embedding.py
+ python test/ops/linear.py
+ python test/ops/rms_norm.py
+ python test/ops/rope.py
+ python test/ops/self_attention.py
+ python test/ops/swiglu.py
+
+ - name: Assignment-3
+ run: |
+ python test/test_infer.py --test
diff --git a/.gitignore b/.gitignore
index e38cf5747..1c96ab0f8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,90 +1,83 @@
-# Xmake cache
-.xmake/
-build/
-
-# Binaries
-bin/
-lib/
-*.so
-*.dll
-*.dylib
-*.pyd
-
-# MacOS Cache
-.DS_Store
-
-# Vscode
-.vscode/
-
-# Python
-__pycache__/
-
-# Log
-*.log
-
-# Cache
-cache/
-
-# JSON
-*.json
-
-#GGUF
-*.gguf
-
-
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# Distribution / packaging
-build/
-dist/
-*.egg-info/
-.eggs/
-
-# Virtual environments
-.venv/
-env/
-venv/
-ENV/
-*.env
-*.venv
-
-# PyInstaller
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# MyPy and other type checking
-.mypy_cache/
-.dmypy.json
-.pyre/
-
-# Test and coverage
-.coverage
-htmlcov/
-.tox/
-.nox/
-.cache/
-.pytest_cache/
-
-# Jupyter Notebook checkpoints
-.ipynb_checkpoints
-
-# IDE and editor settings
-.vscode/
-.idea/
-*.swp
-*~
-
-# macOS
-.DS_Store
-
-# Windows
-Thumbs.db
-ehthumbs.db
-desktop.ini
\ No newline at end of file
+# Xmake cache
+.xmake/
+build/
+
+# Binaries
+bin/
+lib/
+*.so
+*.dll
+*.dylib
+*.pyd
+
+# MacOS Cache
+.DS_Store
+
+# Vscode
+.vscode/
+
+# Python
+__pycache__/
+
+# Log
+*.log
+
+# Cache
+cache/
+
+# 本地模型目录（体积大，勿提交）
+DeepSeek-R1-Distill-Qwen-1___5B/
+
+# JSON
+*.json
+
+#GGUF
+*.gguf
+
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+.eggs/
+
+# Virtual environments
+.venv/
+env/
+venv/
+ENV/
+*.env
+*.venv
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# MyPy and other type checking
+.mypy_cache/
+.dmypy.json
+.pyre/
+
+# Test and coverage
+.coverage
+htmlcov/
+.tox/
+.nox/
+.cache/
+.pytest_cache/
+
+# Jupyter Notebook checkpoints
+.ipynb_checkpoints
+
+# 本地文档与图片（不随代码库同步；项目完成情况报告.md 需提交则不要写在这里）
+docs/
+
diff --git a/=42 b/=42
new file mode 100644
index 000000000..5bf441788
--- /dev/null
+++ b/=42
@@ -0,0 +1,11 @@
+Looking in indexes: https://mirrors.aliyun.com/pypi/simple/
+Collecting setuptools
+ Downloading https://mirrors.aliyun.com/pypi/packages/e1/c6/76dc613121b793286a3f91621d7b75a2b493e0390ddca50f11993eadf192/setuptools-82.0.0-py3-none-any.whl (1.0 MB)
+ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.0/1.0 MB 24.9 MB/s eta 0:00:00
+Collecting wheel
+ Downloading https://mirrors.aliyun.com/pypi/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl (30 kB)
+Collecting packaging>=24.0 (from wheel)
+ Downloading https://mirrors.aliyun.com/pypi/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl (74 kB)
+ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 74.4/74.4 kB 23.6 MB/s eta 0:00:00
+Installing collected packages: setuptools, packaging, wheel
+Successfully installed packaging-26.0 setuptools-82.0.0 wheel-0.46.3
diff --git a/LICENSE b/LICENSE
index 0e0021080..ebd477723 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,8 +1,8 @@
-The MIT License (MIT)
-Copyright © 2025 InfiniTensor
-
-Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+The MIT License (MIT)
+Copyright © 2025 InfiniTensor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
diff --git a/README.md b/README.md
index 456067c82..44bf70647 100644
--- a/README.md
+++ b/README.md
@@ -1,431 +1,431 @@
-# Welcome to LLAISYS
-
-
-English ｜
-中文
-
-
-## Introduction
-
-LLAISYS (Let's Learn AI SYStem) is an educational project that aims to provide a platform for new and future AI engineers to learn how to build AI systems from scratch. LLAISYS consists of several assignments, which help students learn and build the basic modules, and projects that challenge them to add more fancy features to their systems. LLAISYS uses C++ as primary programming language for system backend, and is compiled into shared libraries exposing C language APIs. Frontend codes are written in Python which calls these APIs to provide more convenient testing and interaction with other architectures such as PyTorch.
-
-### Project Structure Overview
-
-- `\include`: directory that contains of the header files which defines all the C APIs exposed by the shared library. (Functions declarations start with `__export`)
-
-- `\src`: C++ source files.
- - `\src\llaisys` contains all the direct implementation of waht are defined in the header files and follows the same directory structure as the `\include`. This is also as far as C++ codes can go.
- - other directories contain the actual implementaion of different modules.
-
-- `xmake.lua`: build rules for llaisys backend. `\xmake` directory contains the sub-xmake files for different devices. You may add `nvidia.lua` in the directory in the future for instance to support CUDA.
-
-- `\python`: Python source files.
- - `\python\llaisys\libllaisys` contains all the ctypes wrapper functions of llaisys APIs. It basically matches the structure of C header files.
- - `\python\llaisys` contains Python warppers of the ctypes functions to make the package more Python-like.
-
-- `\test`: Python test files that import llaisys python package.
-
-## Assignment #0: Getting Started
-
-### Task-0.1 Install Prerequisites
-
-- Compile Tool: [Xmake](https://xmake.io/)
-- C++ Compiler: MSVC (Windows) or Clang or GCC
-- Python >= 3.9 (PyTorch, Transformers, etc.)
-- Clang-Format-16 (Optional): for formatting C++ codes.
-
-### Task-0.2 Fork and Build LLAISYS
-
-- FORK LLAISYS Repository and Clone it to your local machine. Both Windows and Linux are supported.
-
-- Compile and Install
-
- ```bash
- # compile c++ codes
- xmake
- # install llaisys shared library
- xmake install
- # install llaisys python package
- pip install ./python/
- ```
-
-- Github Auto Tests
-
- LLAISYS uses Github Actions to run automated tests on every push and pull request. You can see testing results on your repo page. All tests should pass once you have finished all assignment tasks.
-
-### Task-0.3 Run LLAISYS for the First Time
-
-- Run cpu runtime tests
-
- ```bash
- python test/test_runtime.py --device cpu
- ```
-
- You should see the test passed.
-
-### Task-0.4 Download test model
-
-- The model we use for assignments is [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B).
-
-- Run an inference test with the model using PyTorch
-
- ```bash
- python test/test_infer.py --model [dir_path/to/model]
- ```
-
- You can see that PyTorch is able to load the model and perform inference with the sample input. You can debug into `transformers` library codes to see how what is going on behind. Right now, your code cannot do anything yet, but you are going to build a system that can achieve the same functionality in the assignments.
-
-## Assignment #1: Tensor
-
-Tensor is a data structure that represents multi-dimensional data. It is the basic building block of LLAISYS, and most AI frameworks such as PyTorch. In this assignment, you will learn how to implement a basic tensor class.
-
-A Tensor object has the following fields:
-
-- `storage`: a shared pointer to a memory block that stores the tensor's data. It can be shared by multiple tensors. Check storage class for more details.
-- `offset`: the starting index (in bytes) of the tensor in the storage.
-- `meta`: metadata that describes the tensor's shape, data type, and strides.
-
-Implement the following functions defined in the `src/tensor/tensor.hpp`:
-
-### Task-1.1
-
-```c++
-void load(const void *src);
-```
-
-Load host (cpu) data to the tensor (can be on device). Check contructor to see how to get runtime apis of the current device context, and do a memcpy from host to device.
-
-### Task-1.2
-
-```c++
-bool isContiguous() const;
-```
-
-Check shape and strides of the tensor, and tell wether it is contiguous in memory.
-
-### Task-1.3
-
-```c++
-tensor_t view(const std::vector &shape) const;
-```
-
-Create a new tensor which reshapes the original tensor to the given shape by splitting or merging the original dimensions. No data transfer is involved. For example change a tensor of shape (2, 3, 5) to (2, 15) by merging the last two dimensions.
-
-This function is not as easy as simply changing the shape of the tensor, although the test will pass. It should raise an error if new view is not compatible with the original tensor. Think about a tensor of shape (2, 3, 5) and strides (30, 10, 1). Can you still reshape it to (2, 15) without data transfer?
-
-### Task-1.4
-
-```c++
-tensor_t permute(const std::vector &order) const;
-```
-
-Create a new tensor which changes the order of the dimensions of original tensor. Transpose can be achieved by this function without moving data around.
-
-### Task-1.5
-
-```c++
-tensor_t slice(size_t dim, size_t start, size_t end) const;
-```
-
-Create a new tensor which slices the original tensor along the given dimension,
-start (inclusive) and end (exclusive) indices.
-
-### Task-1.6
-
-Run tensor tests.
-
-```bash
-python test/test_tensor.py
-```
-
-You should see all tests passed. Commit and push your changes. You should see the auto tests for assignment #1 passed.
-
-## Assignment #2: Operators
-
-In this assignment, you will implement the cpu verision the following operators:
-
-- argmax
-- embedding
-- linear
-- rms_norm
-- rope
-- self_attention
-- swiglu
-
-Read the codes in `src/ops/add/` to see how "add" operator is implemented. Make sure you understand how the operator codes are organized, compiled, linked, and exposed to Python frontend. **Your operators should at least support Float32, Float16 and BFloat16 data types**. A helper function for naive type casting is provided in `src/utils/`. All python tests are in `test/ops`, you implementation should at least pass these tests. Try running the test script for "add" operator for starting.
-
-### Task-2.1 argmax
-
-```c++
-void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
-```
-
-Get the max value and its index of tensor `vals`, and store them in `max_val` and `max_idx` respectively. You can assume that `vals` is a 1D tensor for now, and `max_idx` and `max_val` are both 1D tensors with a single element (, which means the dimension of `vals` is kept).
-
-You should be able to pass the test cases in `test/ops/argmax.py` after you finish the implementation.
-
-### Task-2.2 embedding
-
-```c++
-void embedding(tensor_t out, tensor_t index, tensor_t weight);
-```
-
-Copy the rows in `index` (1-D) from `weight` (2-D) to `output` (2-D). `index` must be of type Int64 (the default data type for int of PyTorch).
-
-You should be able to pass the test cases in `test/ops/embedding.py` after you finish the implementation.
-
-### Task-2.3 linear
-
-```c++
-void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias);
-```
-
-Compute the following:
-
-$$
-Y = xW^T + b
-$$
-
-- `out`: output $Y$ . You can assume output is a 2D contiguous tensor and no broadcasting is involved for now.
-- `input`: input $X$ . You can assume input is a 2D contiguous tensor and no broadcasting is involved for now.
-- `weight`: weight $W$ . 2D contiguous tensor. Note that weight tensor is not transposed. You need to deal with this during your calculation.
-- `bias` (optional): bias $b$ . 1D tensor. You need to support the situation where bias is not provided.
-
-You should be able to pass the test cases in `test/ops/linear.py` after you finish the implementation.
-
-### Task-2.4 rms normalization
-
-```c++
-void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps);
-```
-
-Compute the following for each row:
-
-$$
-Y_i = \frac{W_i \times X_i}{\sqrt{\frac{1}{d}(\sum_{j=1}^d X_j^2) + \epsilon}}
-$$
-
-- `out`: output $Y$ . You can assume output is a 2D contiguous tensor and no broadcasting is involved for now.
-- `input`: input $X$ . You can assume input is a 2D contiguous tensor and no broadcasting is involved for now. The normalization is performed along the last dimension (a.k.a. each row of length $d$ ) of the input tensor.
-- `weight`: weight $W$ . 1D tensor, same length as a row of input tensor.
-- `eps`: small value $\epsilon$ to avoid division by zero.
-
-You should be able to pass the test cases in `test/ops/rms_norm.py` after you finish the implementation.
-
-### Task-2.5 rope
-
-```c++
-void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta);
-```
-
-Compute the following for each vector of input tensor `in`, corresponding to a position id in `pos_ids`:
-
-Let $\mathbf{x}_i = [\mathbf{a}_i, \mathbf{b}_i] \in \mathbb{R}^d$ be the input vector and $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i] \in \mathbb{R}^d$ be the output vector at index $i$, where $\mathbf{a}_i, \mathbf{b}_i,\mathbf{a}'_i, \mathbf{b}'_i \in \mathbb{R}^{d/2}$ .
-
-Let $\theta$ be a fixed base (e.g. $\theta = 10000$) and $j = 0, 1, \ldots, d/2 - 1$.
-
-Let $p_i \in \mathbb{N}$ is the position id for token at input index i.
-
-Then the angle for RoPE is $\phi_{i,j} = \frac{p_i}{\theta^{2j/d}}$
-
-The output vector $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i]$ is computed as follows:
-
-$$a_{i,j}' = a_{i,j} \cos(\phi_{i,j}) - b_{i,j} \sin(\phi_{i,j})$$
-
-$$b_{i,j}' = b_{i,j} \cos(\phi_{i,j}) + a_{i,j} \sin(\phi_{i,j})$$
-
-- `out`: the resulting **q** or **k** tensor. Shape should be [seqlen, nhead, d] or [seqlen, nkvhead, d]. You can assume that the tensor is contiguous for now.
-- `in`: the orignal **q** or **k** tensor. Shape should be [seqlen, nhead, d] or [seqlen, nkvhead, d]. You can assume that the tensor is contiguous for now.
-- `pos_ids`: the position id (index in the whole context) for each token in the input sequence. Shape should be [seqlen,], dtype should be int64.
-- `theta`: the base value for the frequency vector.
-
-You should be able to pass the test cases in `test/ops/rope.py` after you finish the implementation.
-
-### Task-2.6 self-attention
-
-```c++
-void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale);
-```
-
-Compute the self-attention for query tensor `q`, key tensor `k`, and value tensor `v`. You should concat kvcache tensors, if needed, before doing this calculation.
-
-$$
-A = Q K^\top * scale \\
-$$
-
-$$
-Y = \mathrm{causalsoftmax}(A) \cdot V \\
-$$
-
-- `attn_val`: the resulting attention value tensor. Shape should be [seqlen, nhead, dv]. You can assume that the tensor is contiguous for now.
-- `q`: the query tensor. Shape should be [seqlen, nhead, d]. You can assume that the tensor is contiguous for now.
-- `k`: the key tensor. Shape should be [total_len, nkvhead, d]. You can assume that the tensor is contiguous for now.
-- `v`: the value tensor. Shape should be [total_len, nkvhead, dv]. You can assume that the tensor is contiguous for now.
-- `scale`: a scaling factor. It is set to $\frac{1}{\sqrt{d}}$ in most cases.
-
-You should be able to pass the test cases in `test/ops/self_attention.py` after you finish the implementation.
-
-### Task-2.7 swiglu
-
-```c++
-void swiglu(tensor_t out, tensor_t gate, tensor_t up);
-```
-
-This is an element-wise function that computes the following:
-
-$$
-out_{i} = up_{i} \circ \frac { gate_{i}}{1 + e^{-gate_{i}}}
-$$
-
-`out`, `up` and `gate` are 2D contiguous tensors with the same shape [seqlen, intermediate_size].
-
-You should be able to pass the test cases in `test/ops/swiglu.py` after you finish the implementation.
-
-### Task-2.8
-
-Run operator tests.
-
-```bash
-python test/test_ops.py
-```
-
-You should see all tests passed. Commit and push your changes. You should see the auto tests for assignment #2 passed.
-
-### Task-2.9 (Optional) rearrange
-
-This is a bonus task. You may or may not need it for model inference.
-
-```c++
-void rearrange(tensor_t out, tensor_t in);
-```
-
-This operator is used to copy data from a tensor to another tensor with the same shape but different strides. With this, you can easily implement `contiguous` functionality for tensors.
-
-## Assignment #3: Large Language Model Inference
-
-Finally, it is the time for you to achieve text generation with LLAISYS.
-
-- In `test/test_infer.py`, your implementation should be able to generate the same texts as PyTorch, using argmax sampling. The model we use for this assignment is [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B).
-
-- The python wrapper of your implementation is in `python/llaisys/models/qwen2.py`. You are NOT allowed to implement your model infer logic here using any python based frameworks, such as PyTorch. Instead, you need to implement the model with C/C++ in LLAISYS backend. The script loads each tensor in the safetensors file, and you will need to load data from them into your model backend.
-
-- In `include/llaisys/models/qwen2.h`, a prototype is defined for you. Feel free to modify the codes as you want, but you should at least provide basic APIs for model creation, destruction, data loading, and infer. Implement your C APIs in `src/llaisys/` and organize your C++ codes as other modules in `src/`. Remember to define the compiling procedures in `xmake.lua`.
-
-- In `python/llaisys/libllaisys/`, define the ctypes wrapper functions for your C APIs. Implement `python/llaisys/models/qwen2.py` with your wrapper functions.
-
-- You need to implement KV Cache, or your model will be too slow.
-
-- Debug until your model works. Take advantage of tensor's `debug` function which prints the tensor data. It allows you to compare the data of any tensor during the model inference with PyTorch.
-
-After you finish the implementation, you can run the following command to test your model:
-
-```bash
-python test/test_infer.py --model [dir_path/to/model] --test
-```
-
-Commit and push your changes. You should see the auto tests for assignment #3 passed.
-
-
-## You can proceed to the projects only after you finish the assignments.
-
-## Project #1: Optimize LLAISYS for CPU
-You probably have already noticed that your model inference is very slow compared to PyTorch. This is mostly because your operators are not optimized. Run your operater test scripts with "--profile" flag to see how your operators perform. You would probably see that `linear` operation is much slower than PyTorch. This operator is mainly a matrix multiplication, and is the most time consuming operation in transformer-based models.
-
-There are several ways to optimize your operators for CPU:
-
-### SIMD instructions
-
-SIMD (Single Instruction Multiple Data) instructions are instructions that can perform the same operation on multiple data elements in a single instruction. Modern CPUs have support for SIMD instructions. Look for online materials to learn about compiler intrinsics (such as AVX2, AVX-512, NEON, SVE) to vectorize your operations.
-
-### Use OpenMP for parallelism
-
-You can use multi-threading to parallelize your operators. OpenMP is a popular library for multi-threading in C/C++. Add OpenMP support for LLAISYS to parallelize your `linear` and other operators.
-
-### 3rd-party Libraries
-
-There are several libraries that can help you optimize your operators for CPU. Look for libraries like Eigen, OpenBLAS, MKL, etc. to optimize your linear algebra operations. Note that some libraries are supported only for certain hardware platforms. Check their documentations and use them in your codes with care. You can also try to dig out how PyTorch implement these operators and see if you can use them.
-
-Optimize your implementation with any methods you like and report your performance improvement.
-
-## Project #2: Intigrate CUDA into LLAISYS
-
-This project does not depend on **Project #1**. You should choose two CUDA/CUDA-ish hardware platforms from Nvidia, Iluvatar, Metax, and Moore Threads.
-
-This camp session provides computation resources from the four platforms above, access to which is granted based on applications from the official website. You can accelerate your model with CUDA on these GPU platforms. Before doing that, let's dive deeper into LLAISYS framework.
-
-LLAISYS is actually a framework with homogeous hardware support. When using LLAISYS, each thread will create a thread-local `Context` object which manages all the device `Runtime` objects used by this thread. A `Runtime` object is a resource manager for a device, and `Context` will create (with lazy initialization) a single `Runtime` object for each device. You can set and switch between them using `setDevice` function in `Context`. Only one device will be active at a time for each thread. Check `src/core/context.hpp` for more details.
-
-### Implement CUDA Runtime APIs
-Each `Runtime` object is intialized with a set of generic functions called `Runtime APIs`. You will need to implement CUDA version of these APIS. Check `src/device/cpu/cpu_runtime_api.cpp` to see how these functions are implemented for CPU and look for CUDA APIs to use in [`CUDA Runtime documentation`](https://docs.nvidia.com/cuda/cuda-runtime-api/index.html).
-
-You can see in `src/device/runtime_api.hpp` that `nvidia::getRuntimeAPI()` is guarded by `ENABLE_NVIDIA_API` macro.
-
-```c++
-#ifdef ENABLE_NVIDIA_API
-namespace nvidia {
-const LlaisysRuntimeAPI *getRuntimeAPI();
-}
-#endif
-```
-
-This macro is defined in `xmake.lua` as a switch to enable/disable CUDA support. CUDA codes will not be compiled if the switch is off. In `xmake/` directory, create a `nvidia.lua` that configs your compiling process. (Similar to `cpu.lua` for CPU.) Search online to learn how to do it with Xmake.
-
-After you implement the CUDA Runtime APIs, config your xmake with `--nv-gpu=y` to enable CUDA support and recompile your program. Run runtime tests to see if your implementation works.
-
-```bash
-xmake f --nv-gpu=y -cv
-xmake
-xmake install
-python test/test_runtime.py --device nvidia
-```
-
-### Implement CUDA Operators
-Create a `nvdia/` sub-directory in each operator source directory and implement a cuda version. Check `src/ops/add/op.cpp` to see how to include your cuda implementations. Remeber to define the compiling procedures in the xmake files. Run the operator tests with `--device nvidia` flag to test your CUDA implementation.
-
-You can use CUDA libraries like cuBLAS, cuDNN, etc. to accelerate your operators. Check their documentations to see how to use them. You can store extra device resources in `src/device/nvidia/nvidia_resource.cu`.
-
-Modify your model codes to support CUDA inference.
-
-```bash
-python test/test_infer.py --model [dir_path/to/model] --test --device nvidia
-```
-
-## Project #3: Build an AI chatbot
-
-In this project you will build an AI chatbot that can do live conversations with single user with LLAISYS.
-
-### Random Sampling
-
-So far we have been testing our model with argmax sampling. This is good enough for testing, but a chatbot should be able to generate more natural responses. Implement a random sample operator. Try to add supports for **Temperature**, **Top-K** and **Top-P**.
-
-### Build a Chatbot Server
-
-In your Python frontend, implement a server that can receive http requests from user and send responses back. You can use frameworks like FastAPI to build the server. You should follow the OpenAI chat-completion APIs. Try to support streaming responses if you can. You can assume, for now, that the server is only serving one user, and block the endpoint until the previous request is served.
-
-
-### Interactive Chat UI
-
-Build a UI that send requests to and receive responses from the chatbot server. You can build a simple command-line interface or a fancy web interface. You should be able to keep a conversation going with the chatbot by sending messages and receiving responses consecutively.
-
-### (Optional) Chat Session Management
-
-In real-world AI applications, users are allowed to start new conversations and switch between them. Users can also edit a past question and let the AI regenerate an answer. Enhance your UI to support these features. Implement a KV-Cache pool with prefix matching to reuse past results as much as possible.
-
-
-## Project #4: Multi-user Inference Service
-
-You need to finish **Project #2** and achieve streaming response first before proceeding to this project.
-
-### Serving Multiple Users
-
-In real-world scenarios, an inference service will serve multiple users. Requests can come in at any time, and the service should be able to handle them concurrently. Your endpoint should add a new request to a request pool or queue and have a another looping process or thread to serve the requests.
-
-### Continous Batching
-To maximize the throughput of your inference service, you need to batch your requests instead of serving them one by one. Since each request can have different length, you will need a continous and iteration-level batching mechanism. For each interation you extract several requests from pool to form a batch, do one round of batch inference, and then return the unfinished requests back to the pool. Use batched matrix multiplication when possible to speed up your inference. Note that every request in the batch need to bind with a different KV-Cache. You should build a KV-Cache pool with prefix matching to reuse past results as much as possible.
-
-## Project #5: Distributed Inference
-Introduce Tensor Parallelism to LLAISYS. Shard your model across multiple devices and implement distributed model inference. Support NCCL in LLAISYS if your are uing Nvidia GPUs, or MPI if you are using CPUs.
-
-## Project #6: Support New Models
-
-Support another model type than the one we use for homework in LLAISYS.
+# Welcome to LLAISYS
+
+
+English ｜
+中文
+
+
+## Introduction
+
+LLAISYS (Let's Learn AI SYStem) is an educational project that aims to provide a platform for new and future AI engineers to learn how to build AI systems from scratch. LLAISYS consists of several assignments, which help students learn and build the basic modules, and projects that challenge them to add more fancy features to their systems. LLAISYS uses C++ as primary programming language for system backend, and is compiled into shared libraries exposing C language APIs. Frontend codes are written in Python which calls these APIs to provide more convenient testing and interaction with other architectures such as PyTorch.
+
+### Project Structure Overview
+
+- `\include`: directory that contains of the header files which defines all the C APIs exposed by the shared library. (Functions declarations start with `__export`)
+
+- `\src`: C++ source files.
+ - `\src\llaisys` contains all the direct implementation of waht are defined in the header files and follows the same directory structure as the `\include`. This is also as far as C++ codes can go.
+ - other directories contain the actual implementaion of different modules.
+
+- `xmake.lua`: build rules for llaisys backend. `\xmake` directory contains the sub-xmake files for different devices. You may add `nvidia.lua` in the directory in the future for instance to support CUDA.
+
+- `\python`: Python source files.
+ - `\python\llaisys_py\libllaisys` contains all the ctypes wrapper functions of llaisys APIs. It basically matches the structure of C header files.
+ - `\python\llaisys_py` contains Python wrappers of the ctypes functions to make the package more Python-like.
+
+- `\test`: Python test files that import llaisys_py python package.
+
+## Assignment #0: Getting Started
+
+### Task-0.1 Install Prerequisites
+
+- Compile Tool: [Xmake](https://xmake.io/)
+- C++ Compiler: MSVC (Windows) or Clang or GCC
+- Python >= 3.9 (PyTorch, Transformers, etc.)
+- Clang-Format-16 (Optional): for formatting C++ codes.
+
+### Task-0.2 Fork and Build LLAISYS
+
+- FORK LLAISYS Repository and Clone it to your local machine. Both Windows and Linux are supported.
+
+- Compile and Install
+
+ ```bash
+ # compile c++ codes
+ xmake
+ # install llaisys shared library
+ xmake install
+ # install llaisys python package
+ pip install ./python/
+ ```
+
+- Github Auto Tests
+
+ LLAISYS uses Github Actions to run automated tests on every push and pull request. You can see testing results on your repo page. All tests should pass once you have finished all assignment tasks.
+
+### Task-0.3 Run LLAISYS for the First Time
+
+- Run cpu runtime tests
+
+ ```bash
+ python test/test_runtime.py --device cpu
+ ```
+
+ You should see the test passed.
+
+### Task-0.4 Download test model
+
+- The model we use for assignments is [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B).
+
+- Run an inference test with the model using PyTorch
+
+ ```bash
+ python test/test_infer.py --model [dir_path/to/model]
+ ```
+
+ You can see that PyTorch is able to load the model and perform inference with the sample input. You can debug into `transformers` library codes to see how what is going on behind. Right now, your code cannot do anything yet, but you are going to build a system that can achieve the same functionality in the assignments.
+
+## Assignment #1: Tensor
+
+Tensor is a data structure that represents multi-dimensional data. It is the basic building block of LLAISYS, and most AI frameworks such as PyTorch. In this assignment, you will learn how to implement a basic tensor class.
+
+A Tensor object has the following fields:
+
+- `storage`: a shared pointer to a memory block that stores the tensor's data. It can be shared by multiple tensors. Check storage class for more details.
+- `offset`: the starting index (in bytes) of the tensor in the storage.
+- `meta`: metadata that describes the tensor's shape, data type, and strides.
+
+Implement the following functions defined in the `src/tensor/tensor.hpp`:
+
+### Task-1.1
+
+```c++
+void load(const void *src);
+```
+
+Load host (cpu) data to the tensor (can be on device). Check contructor to see how to get runtime apis of the current device context, and do a memcpy from host to device.
+
+### Task-1.2
+
+```c++
+bool isContiguous() const;
+```
+
+Check shape and strides of the tensor, and tell wether it is contiguous in memory.
+
+### Task-1.3
+
+```c++
+tensor_t view(const std::vector &shape) const;
+```
+
+Create a new tensor which reshapes the original tensor to the given shape by splitting or merging the original dimensions. No data transfer is involved. For example change a tensor of shape (2, 3, 5) to (2, 15) by merging the last two dimensions.
+
+This function is not as easy as simply changing the shape of the tensor, although the test will pass. It should raise an error if new view is not compatible with the original tensor. Think about a tensor of shape (2, 3, 5) and strides (30, 10, 1). Can you still reshape it to (2, 15) without data transfer?
+
+### Task-1.4
+
+```c++
+tensor_t permute(const std::vector &order) const;
+```
+
+Create a new tensor which changes the order of the dimensions of original tensor. Transpose can be achieved by this function without moving data around.
+
+### Task-1.5
+
+```c++
+tensor_t slice(size_t dim, size_t start, size_t end) const;
+```
+
+Create a new tensor which slices the original tensor along the given dimension,
+start (inclusive) and end (exclusive) indices.
+
+### Task-1.6
+
+Run tensor tests.
+
+```bash
+python test/test_tensor.py
+```
+
+You should see all tests passed. Commit and push your changes. You should see the auto tests for assignment #1 passed.
+
+## Assignment #2: Operators
+
+In this assignment, you will implement the cpu verision the following operators:
+
+- argmax
+- embedding
+- linear
+- rms_norm
+- rope
+- self_attention
+- swiglu
+
+Read the codes in `src/ops/add/` to see how "add" operator is implemented. Make sure you understand how the operator codes are organized, compiled, linked, and exposed to Python frontend. **Your operators should at least support Float32, Float16 and BFloat16 data types**. A helper function for naive type casting is provided in `src/utils/`. All python tests are in `test/ops`, you implementation should at least pass these tests. Try running the test script for "add" operator for starting.
+
+### Task-2.1 argmax
+
+```c++
+void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
+```
+
+Get the max value and its index of tensor `vals`, and store them in `max_val` and `max_idx` respectively. You can assume that `vals` is a 1D tensor for now, and `max_idx` and `max_val` are both 1D tensors with a single element (, which means the dimension of `vals` is kept).
+
+You should be able to pass the test cases in `test/ops/argmax.py` after you finish the implementation.
+
+### Task-2.2 embedding
+
+```c++
+void embedding(tensor_t out, tensor_t index, tensor_t weight);
+```
+
+Copy the rows in `index` (1-D) from `weight` (2-D) to `output` (2-D). `index` must be of type Int64 (the default data type for int of PyTorch).
+
+You should be able to pass the test cases in `test/ops/embedding.py` after you finish the implementation.
+
+### Task-2.3 linear
+
+```c++
+void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias);
+```
+
+Compute the following:
+
+$$
+Y = xW^T + b
+$$
+
+- `out`: output $Y$ . You can assume output is a 2D contiguous tensor and no broadcasting is involved for now.
+- `input`: input $X$ . You can assume input is a 2D contiguous tensor and no broadcasting is involved for now.
+- `weight`: weight $W$ . 2D contiguous tensor. Note that weight tensor is not transposed. You need to deal with this during your calculation.
+- `bias` (optional): bias $b$ . 1D tensor. You need to support the situation where bias is not provided.
+
+You should be able to pass the test cases in `test/ops/linear.py` after you finish the implementation.
+
+### Task-2.4 rms normalization
+
+```c++
+void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps);
+```
+
+Compute the following for each row:
+
+$$
+Y_i = \frac{W_i \times X_i}{\sqrt{\frac{1}{d}(\sum_{j=1}^d X_j^2) + \epsilon}}
+$$
+
+- `out`: output $Y$ . You can assume output is a 2D contiguous tensor and no broadcasting is involved for now.
+- `input`: input $X$ . You can assume input is a 2D contiguous tensor and no broadcasting is involved for now. The normalization is performed along the last dimension (a.k.a. each row of length $d$ ) of the input tensor.
+- `weight`: weight $W$ . 1D tensor, same length as a row of input tensor.
+- `eps`: small value $\epsilon$ to avoid division by zero.
+
+You should be able to pass the test cases in `test/ops/rms_norm.py` after you finish the implementation.
+
+### Task-2.5 rope
+
+```c++
+void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta);
+```
+
+Compute the following for each vector of input tensor `in`, corresponding to a position id in `pos_ids`:
+
+Let $\mathbf{x}_i = [\mathbf{a}_i, \mathbf{b}_i] \in \mathbb{R}^d$ be the input vector and $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i] \in \mathbb{R}^d$ be the output vector at index $i$, where $\mathbf{a}_i, \mathbf{b}_i,\mathbf{a}'_i, \mathbf{b}'_i \in \mathbb{R}^{d/2}$ .
+
+Let $\theta$ be a fixed base (e.g. $\theta = 10000$) and $j = 0, 1, \ldots, d/2 - 1$.
+
+Let $p_i \in \mathbb{N}$ is the position id for token at input index i.
+
+Then the angle for RoPE is $\phi_{i,j} = \frac{p_i}{\theta^{2j/d}}$
+
+The output vector $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i]$ is computed as follows:
+
+$$a_{i,j}' = a_{i,j} \cos(\phi_{i,j}) - b_{i,j} \sin(\phi_{i,j})$$
+
+$$b_{i,j}' = b_{i,j} \cos(\phi_{i,j}) + a_{i,j} \sin(\phi_{i,j})$$
+
+- `out`: the resulting **q** or **k** tensor. Shape should be [seqlen, nhead, d] or [seqlen, nkvhead, d]. You can assume that the tensor is contiguous for now.
+- `in`: the orignal **q** or **k** tensor. Shape should be [seqlen, nhead, d] or [seqlen, nkvhead, d]. You can assume that the tensor is contiguous for now.
+- `pos_ids`: the position id (index in the whole context) for each token in the input sequence. Shape should be [seqlen,], dtype should be int64.
+- `theta`: the base value for the frequency vector.
+
+You should be able to pass the test cases in `test/ops/rope.py` after you finish the implementation.
+
+### Task-2.6 self-attention
+
+```c++
+void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale);
+```
+
+Compute the self-attention for query tensor `q`, key tensor `k`, and value tensor `v`. You should concat kvcache tensors, if needed, before doing this calculation.
+
+$$
+A = Q K^\top * scale \\
+$$
+
+$$
+Y = \mathrm{causalsoftmax}(A) \cdot V \\
+$$
+
+- `attn_val`: the resulting attention value tensor. Shape should be [seqlen, nhead, dv]. You can assume that the tensor is contiguous for now.
+- `q`: the query tensor. Shape should be [seqlen, nhead, d]. You can assume that the tensor is contiguous for now.
+- `k`: the key tensor. Shape should be [total_len, nkvhead, d]. You can assume that the tensor is contiguous for now.
+- `v`: the value tensor. Shape should be [total_len, nkvhead, dv]. You can assume that the tensor is contiguous for now.
+- `scale`: a scaling factor. It is set to $\frac{1}{\sqrt{d}}$ in most cases.
+
+You should be able to pass the test cases in `test/ops/self_attention.py` after you finish the implementation.
+
+### Task-2.7 swiglu
+
+```c++
+void swiglu(tensor_t out, tensor_t gate, tensor_t up);
+```
+
+This is an element-wise function that computes the following:
+
+$$
+out_{i} = up_{i} \circ \frac { gate_{i}}{1 + e^{-gate_{i}}}
+$$
+
+`out`, `up` and `gate` are 2D contiguous tensors with the same shape [seqlen, intermediate_size].
+
+You should be able to pass the test cases in `test/ops/swiglu.py` after you finish the implementation.
+
+### Task-2.8
+
+Run operator tests.
+
+```bash
+python test/test_ops.py
+```
+
+You should see all tests passed. Commit and push your changes. You should see the auto tests for assignment #2 passed.
+
+### Task-2.9 (Optional) rearrange
+
+This is a bonus task. You may or may not need it for model inference.
+
+```c++
+void rearrange(tensor_t out, tensor_t in);
+```
+
+This operator is used to copy data from a tensor to another tensor with the same shape but different strides. With this, you can easily implement `contiguous` functionality for tensors.
+
+## Assignment #3: Large Language Model Inference
+
+Finally, it is the time for you to achieve text generation with LLAISYS.
+
+- In `test/test_infer.py`, your implementation should be able to generate the same texts as PyTorch, using argmax sampling. The model we use for this assignment is [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B).
+
+- The python wrapper of your implementation is in `python/llaisys_py/models/qwen2.py`. You are NOT allowed to implement your model infer logic here using any python based frameworks, such as PyTorch. Instead, you need to implement the model with C/C++ in LLAISYS backend. The script loads each tensor in the safetensors file, and you will need to load data from them into your model backend.
+
+- In `include/llaisys/models/qwen2.h`, a prototype is defined for you. Feel free to modify the codes as you want, but you should at least provide basic APIs for model creation, destruction, data loading, and infer. Implement your C APIs in `src/llaisys/` and organize your C++ codes as other modules in `src/`. Remember to define the compiling procedures in `xmake.lua`.
+
+- In `python/llaisys_py/libllaisys/`, define the ctypes wrapper functions for your C APIs. Implement `python/llaisys_py/models/qwen2.py` with your wrapper functions.
+
+- You need to implement KV Cache, or your model will be too slow.
+
+- Debug until your model works. Take advantage of tensor's `debug` function which prints the tensor data. It allows you to compare the data of any tensor during the model inference with PyTorch.
+
+After you finish the implementation, you can run the following command to test your model:
+
+```bash
+python test/test_infer.py --model [dir_path/to/model] --test
+```
+
+Commit and push your changes. You should see the auto tests for assignment #3 passed.
+
+
+## You can proceed to the projects only after you finish the assignments.
+
+## Project #1: Optimize LLAISYS for CPU
+You probably have already noticed that your model inference is very slow compared to PyTorch. This is mostly because your operators are not optimized. Run your operater test scripts with "--profile" flag to see how your operators perform. You would probably see that `linear` operation is much slower than PyTorch. This operator is mainly a matrix multiplication, and is the most time consuming operation in transformer-based models.
+
+There are several ways to optimize your operators for CPU:
+
+### SIMD instructions
+
+SIMD (Single Instruction Multiple Data) instructions are instructions that can perform the same operation on multiple data elements in a single instruction. Modern CPUs have support for SIMD instructions. Look for online materials to learn about compiler intrinsics (such as AVX2, AVX-512, NEON, SVE) to vectorize your operations.
+
+### Use OpenMP for parallelism
+
+You can use multi-threading to parallelize your operators. OpenMP is a popular library for multi-threading in C/C++. Add OpenMP support for LLAISYS to parallelize your `linear` and other operators.
+
+### 3rd-party Libraries
+
+There are several libraries that can help you optimize your operators for CPU. Look for libraries like Eigen, OpenBLAS, MKL, etc. to optimize your linear algebra operations. Note that some libraries are supported only for certain hardware platforms. Check their documentations and use them in your codes with care. You can also try to dig out how PyTorch implement these operators and see if you can use them.
+
+Optimize your implementation with any methods you like and report your performance improvement.
+
+## Project #2: Intigrate CUDA into LLAISYS
+
+This project does not depend on **Project #1**. You should choose two CUDA/CUDA-ish hardware platforms from Nvidia, Iluvatar, Metax, and Moore Threads.
+
+This camp session provides computation resources from the four platforms above, access to which is granted based on applications from the official website. You can accelerate your model with CUDA on these GPU platforms. Before doing that, let's dive deeper into LLAISYS framework.
+
+LLAISYS is actually a framework with homogeous hardware support. When using LLAISYS, each thread will create a thread-local `Context` object which manages all the device `Runtime` objects used by this thread. A `Runtime` object is a resource manager for a device, and `Context` will create (with lazy initialization) a single `Runtime` object for each device. You can set and switch between them using `setDevice` function in `Context`. Only one device will be active at a time for each thread. Check `src/core/context.hpp` for more details.
+
+### Implement CUDA Runtime APIs
+Each `Runtime` object is intialized with a set of generic functions called `Runtime APIs`. You will need to implement CUDA version of these APIS. Check `src/device/cpu/cpu_runtime_api.cpp` to see how these functions are implemented for CPU and look for CUDA APIs to use in [`CUDA Runtime documentation`](https://docs.nvidia.com/cuda/cuda-runtime-api/index.html).
+
+You can see in `src/device/runtime_api.hpp` that `nvidia::getRuntimeAPI()` is guarded by `ENABLE_NVIDIA_API` macro.
+
+```c++
+#ifdef ENABLE_NVIDIA_API
+namespace nvidia {
+const LlaisysRuntimeAPI *getRuntimeAPI();
+}
+#endif
+```
+
+This macro is defined in `xmake.lua` as a switch to enable/disable CUDA support. CUDA codes will not be compiled if the switch is off. In `xmake/` directory, create a `nvidia.lua` that configs your compiling process. (Similar to `cpu.lua` for CPU.) Search online to learn how to do it with Xmake.
+
+After you implement the CUDA Runtime APIs, config your xmake with `--nv-gpu=y` to enable CUDA support and recompile your program. Run runtime tests to see if your implementation works.
+
+```bash
+xmake f --nv-gpu=y -cv
+xmake
+xmake install
+python test/test_runtime.py --device nvidia
+```
+
+### Implement CUDA Operators
+Create a `nvdia/` sub-directory in each operator source directory and implement a cuda version. Check `src/ops/add/op.cpp` to see how to include your cuda implementations. Remeber to define the compiling procedures in the xmake files. Run the operator tests with `--device nvidia` flag to test your CUDA implementation.
+
+You can use CUDA libraries like cuBLAS, cuDNN, etc. to accelerate your operators. Check their documentations to see how to use them. You can store extra device resources in `src/device/nvidia/nvidia_resource.cu`.
+
+Modify your model codes to support CUDA inference.
+
+```bash
+python test/test_infer.py --model [dir_path/to/model] --test --device nvidia
+```
+
+## Project #3: Build an AI chatbot
+
+In this project you will build an AI chatbot that can do live conversations with single user with LLAISYS.
+
+### Random Sampling
+
+So far we have been testing our model with argmax sampling. This is good enough for testing, but a chatbot should be able to generate more natural responses. Implement a random sample operator. Try to add supports for **Temperature**, **Top-K** and **Top-P**.
+
+### Build a Chatbot Server
+
+In your Python frontend, implement a server that can receive http requests from user and send responses back. You can use frameworks like FastAPI to build the server. You should follow the OpenAI chat-completion APIs. Try to support streaming responses if you can. You can assume, for now, that the server is only serving one user, and block the endpoint until the previous request is served.
+
+
+### Interactive Chat UI
+
+Build a UI that send requests to and receive responses from the chatbot server. You can build a simple command-line interface or a fancy web interface. You should be able to keep a conversation going with the chatbot by sending messages and receiving responses consecutively.
+
+### (Optional) Chat Session Management
+
+In real-world AI applications, users are allowed to start new conversations and switch between them. Users can also edit a past question and let the AI regenerate an answer. Enhance your UI to support these features. Implement a KV-Cache pool with prefix matching to reuse past results as much as possible.
+
+
+## Project #4: Multi-user Inference Service
+
+You need to finish **Project #2** and achieve streaming response first before proceeding to this project.
+
+### Serving Multiple Users
+
+In real-world scenarios, an inference service will serve multiple users. Requests can come in at any time, and the service should be able to handle them concurrently. Your endpoint should add a new request to a request pool or queue and have a another looping process or thread to serve the requests.
+
+### Continous Batching
+To maximize the throughput of your inference service, you need to batch your requests instead of serving them one by one. Since each request can have different length, you will need a continous and iteration-level batching mechanism. For each interation you extract several requests from pool to form a batch, do one round of batch inference, and then return the unfinished requests back to the pool. Use batched matrix multiplication when possible to speed up your inference. Note that every request in the batch need to bind with a different KV-Cache. You should build a KV-Cache pool with prefix matching to reuse past results as much as possible.
+
+## Project #5: Distributed Inference
+Introduce Tensor Parallelism to LLAISYS. Shard your model across multiple devices and implement distributed model inference. Support NCCL in LLAISYS if your are uing Nvidia GPUs, or MPI if you are using CPUs.
+
+## Project #6: Support New Models
+
+Support another model type than the one we use for homework in LLAISYS.
diff --git a/README_ZN.md b/README_ZN.md
index 7704dbd5b..e3d9e4531 100644
--- a/README_ZN.md
+++ b/README_ZN.md
@@ -1,432 +1,432 @@
-# 欢迎使用 LLAISYS
-
-
-English ｜
-中文
-
-
-## 简介
-
-LLAISYS（Let's Learn AI SYStem）是一个教育项目，旨在为新手和未来的AI工程师提供一个从零开始构建AI系统的学习平台。LLAISYS包含多个作业，帮助学生学习和构建基础模块；以及一些项目挑战，让他们为系统添加更多高级功能。LLAISYS使用C++作为系统后端的主要编程语言，并编译成共享库，提供C语言API。前端代码使用Python编写，调用这些API以提供更便捷的测试和与其他架构（如PyTorch）的交互。
-
-### 项目结构概览
-
-- `\include`：包含所有定义共享库提供的C API的头文件的目录。（函数声明以`__export`开头）
-
-- `\src`：C++源文件。
- - `\src\llaisys`包含头文件中定义的所有直接实现，并遵循与`\include`相同的目录结构。这也是C++代码的边界。
- - 其他目录包含不同模块的实际实现。
-
-- `xmake.lua`：llaisys后端的构建规则。`\xmake`目录包含不同设备的子xmake文件。例如，将来可以在目录中添加`nvidia.lua`来支持CUDA。
-
-- `\python`：Python源文件。
- - `\python\llaisys\libllaisys`包含llaisys API的所有ctypes封装函数。它基本上与C头文件的结构相匹配。
- - `\python\llaisys`包含ctypes函数的Python包装器，使包更符合Python风格。
-
-- `\test`：导入llaisys python包的Python测试文件。
-
-## 作业 #0：入门
-
-### 任务-0.1 安装必备组件
-
-- 编译工具：[Xmake](https://xmake.io/)
-- C++编译器：MSVC（Windows）或Clang或GCC
-- Python >= 3.9（PyTorch、Transformers等）
-- Clang-Format-16（可选）：用于格式化C++代码。
-
-### 任务-0.2 Fork并构建LLAISYS
-
-- Fork LLAISYS仓库并克隆到本地机器。支持Windows和Linux。
-
-- 编译和安装
-
- ```bash
- # 编译c++代码
- xmake
- # 安装llaisys共享库
- xmake install
- # 安装llaisys python包
- pip install ./python/
- ```
-
-- Github自动测试
-
- LLAISYS使用Github Actions在每次推送和拉取请求时运行自动化测试。你可以在仓库页面上看到测试结果。完成所有作业任务后，所有测试都应该通过。
-
-### 任务-0.3 首次运行LLAISYS
-
-- 运行cpu运行时测试
-
- ```bash
- python test/test_runtime.py --device cpu
- ```
-
- 你应该看到测试通过。
-
-### 任务-0.4 下载测试模型
-
-- 我们用于作业的模型是[DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B)。
-
-- 使用PyTorch运行模型推理测试
-
- ```bash
- python test/test_infer.py --model [dir_path/to/model]
- ```
-
- 你可以看到PyTorch能够加载模型并使用示例输入执行推理。你可以调试进入`transformers`库代码来深入查看并了解其内部运作原理。现在，你的代码还无法执行任何操作，但在后续的作业中，你将构建一个能够实现相同功能的系统。
-
-## 作业 #1：张量
-
-张量是表示多维数据的数据结构。它是LLAISYS和大多数AI框架（如PyTorch）的基本构建单元。在这个作业中，你将学习如何实现一个基本的张量类。
-
-张量对象具有以下字段：
-
-- `storage`：指向存储张量数据的内存块的共享指针。它可以被多个张量共享。有关更多详细信息，请查看storage类。
-- `offset`：张量在存储中的起始索引（以字节为单位）。
-- `meta`：描述张量形状、数据类型和步长的元数据。
-
-实现`src/tensor/tensor.hpp`中定义的以下函数：
-
-### 任务-1.1
-
-```c++
-void load(const void *src);
-```
-
-将主机（cpu）数据加载到张量（可以在设备上）。查看构造函数了解如何获取当前设备上下文的运行时API，并执行从主机到设备的内存复制。
-
-### 任务-1.2
-
-```c++
-bool isContiguous() const;
-```
-
-检查张量的形状和步长，判断它在内存中是否连续。
-
-### 任务-1.3
-
-```c++
-tensor_t view(const std::vector &shape) const;
-```
-
-创建一个新张量，通过拆分或合并原始维度将原始张量重塑为给定形状。不涉及数据传输。例如，通过合并最后两个维度，将形状为(2, 3, 5)的张量更改为(2, 15)。
-
-这个函数不是简单地改变张量的形状那么简单，尽管测试会通过。如果新视图与原始张量不兼容，它应该引发错误。想想一个形状为(2, 3, 5)、步长为(30, 10, 1)的张量。你还能在不传输数据的情况下将其重塑为(2, 15)吗？
-
-### 任务-1.4
-
-```c++
-tensor_t permute(const std::vector &order) const;
-```
-
-创建一个新张量，改变原始张量维度的顺序。转置可以通过这个函数实现，而无需移动数据。
-
-### 任务-1.5
-
-```c++
-tensor_t slice(size_t dim, size_t start, size_t end) const;
-```
-
-创建一个新张量，沿给定维度，start（包含）和end（不包含）索引对原始张量进行切片操作。
-
-### 任务-1.6
-
-运行张量测试。
-
-```bash
-python test/test_tensor.py
-```
-
-你应该看到所有测试都通过了。提交并推送你的更改。你应该看到作业#1的自动测试通过了。
-
-## 作业 #2：算子
-
-在这个作业中，你将实现以下算子的cpu版本：
-
-- argmax
-- embedding
-- linear
-- rms_norm
-- rope
-- self_attention
-- swiglu
-
-阅读`src/ops/add/`中的代码，了解"add"算子是如何实现的。确保你理解算子代码是如何组织、编译、链接以及暴露给Python前端的。**你的算子应该至少支持Float32、Float16和BFloat16数据类型**。`src/utils/`中提供了一个用于简单类型转换的辅助函数。所有python测试都在`test/ops`中，你的实现应该至少通过这些测试。首先尝试运行"add"算子的测试脚本。
-
-### 任务-2.1 Argmax
-
-```c++
-void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
-```
-
-获取张量`vals`的最大值及其索引，并分别存储在`max_val`和`max_idx`中。你暂时可以假设`vals`是一个1D张量，`max_idx`和`max_val`都是包含单个元素的1D张量（这意味着保留了`vals`的维度）。
-
-完成实现后，你应该能够通过`test/ops/argmax.py`中的测试用例。
-
-### 任务-2.2 Embedding
-
-```c++
-void embedding(tensor_t out, tensor_t index, tensor_t weight);
-```
-
-从`weight`（2-D）中复制`index`（1-D）中的行到`output`（2-D）。`index`必须是Int64类型（PyTorch中int的默认数据类型）。
-
-完成实现后，你应该能够通过`test/ops/embedding.py`中的测试用例。
-
-### 任务-2.3 Linear
-
-```c++
-void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias);
-```
-
-计算以下内容：
-
-$$
-Y = xW^T + b
-$$
-
-- `out`：输出 $Y$ 。你暂时可以假设输出是一个2D连续张量，不涉及广播。
-- `input`：输入 $X$ 。你暂时可以假设输入是一个2D连续张量，不涉及广播。
-- `weight`：权重 $W$ 。2D连续张量。注意权重张量没有转置。你需要在计算过程中处理这个问题。
-- `bias`（可选）：偏置 $b$ 。1D张量。你需要支持不提供偏置的情况。
-
-完成实现后，你应该能够通过`test/ops/linear.py`中的测试用例。
-
-### 任务-2.4 RMS Normalization
-
-```c++
-void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps);
-```
-
-为每一行计算以下内容：
-
-$$
-Y_i = \frac{W_i \times X_i}{\sqrt{\frac{1}{d}(\sum_{j=1}^d X_j^2) + \epsilon}}
-$$
-
-- `out`：输出 $Y$ 。你暂时可以假设输出是一个2D连续张量，不涉及广播。
-- `input`：输入 $X$ 。你暂时可以假设输入是一个2D连续张量，不涉及广播。标准化沿输入张量的最后一个维度（即每一行，长度为 $d$ ）执行。
-- `weight`：权重 $W$ 。1D张量，与输入张量的一行长度相同。
-- `eps`：小值 $\epsilon$ 以避免除以零。
-
-完成实现后，你应该能够通过`test/ops/rms_norm.py`中的测试用例。
-
-### 任务-2.5 旋转位置编码（RoPE）
-
-```c++
-void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta);
-```
-
-为输入张量`in`的每个向量（这些向量与 pos_ids 中的位置 id 相对应）计算以下内容：
-
-设 $\mathbf{x}_i = [\mathbf{a}_i, \mathbf{b}_i] \in \mathbb{R}^d$ 为输入向量， $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i] \in \mathbb{R}^d$ 为索引 $i$ 处的输出向量，其中 $\mathbf{a}_i, \mathbf{b}_i,\mathbf{a}'_i, \mathbf{b}'_i \in \mathbb{R}^{d/2}$ 。
-
-设 $\theta$ 为固定基数（例如 $\theta = 10000$）， $j = 0, 1, \ldots, d/2 - 1$。
-
-设 $p_i \in \mathbb{N}$ 是输入索引i处token的位置id。
-
-那么RoPE的角度为 $\phi_{i,j} = \frac{p_i}{\theta^{2j/d}}$
-
-输出向量 $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i]$ 计算如下：
-
-$$a_{i,j}' = a_{i,j} \cos(\phi_{i,j}) - b_{i,j} \sin(\phi_{i,j})$$
-
-$$b_{i,j}' = b_{i,j} \cos(\phi_{i,j}) + a_{i,j} \sin(\phi_{i,j})$$
-
-- `out`：结果**q**或**k**张量。形状应该是 [seqlen, nhead, d] 或 [seqlen, nkvhead, d]。你暂时可以假设张量是连续的。
-- `in`：原始**q**或**k**张量。形状应该是 [seqlen, nhead, d] 或 [seqlen, nkvhead, d]。你暂时可以假设张量是连续的。
-- `pos_ids`：输入序列中每个token的位置id（整个上下文中的索引）。形状应该是 [seqlen,]，dtype应该是int64。
-- `theta`：频率向量的基值。
-
-完成实现后，你应该能够通过`test/ops/rope.py`中的测试用例。
-
-### 任务-2.6 自注意力（self-attention）
-
-```c++
-void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale);
-```
-
-为查询张量`q`、键张量`k`和值张量`v`计算自注意力。如果需要，你应该在进行此计算之前连接kvcache张量。
-
-$$
-A = Q K^\top * scale \\
-$$
-
-$$
-Y = \mathrm{causalsoftmax}(A) \cdot V \\
-$$
-
-- `attn_val`：结果注意力值张量。形状应该是[seqlen, nhead, dv]。你暂时可以假设张量是连续的。
-- `q`：查询张量。形状应该是 [seqlen, nhead, d]。你暂时可以假设张量是连续的。
-- `k`：键张量。形状应该是 [total_len, nkvhead, d]。你暂时可以假设张量是连续的。
-- `v`：值张量。形状应该是 [total_len, nkvhead, dv]。你暂时可以假设张量是连续的。
-- `scale`：缩放因子。在大多数情况下取值为 $\frac{1}{\sqrt{d}}$ 。
-
-完成实现后，你应该能够通过`test/ops/self_attention.py`中的测试用例。
-
-### 任务-2.7 SwiGLU
-
-```c++
-void swiglu(tensor_t out, tensor_t gate, tensor_t up);
-```
-
-这是一个逐元素函数，计算以下内容：
-
-$$
-out_{i} = up_{i} \circ \frac { gate_{i}}{1 + e^{-gate_{i}}}
-$$
-
-`out`、`up`和`gate`是具有相同形状 [seqlen, intermediate_size] 的2D连续张量。
-
-完成实现后，你应该能够通过`test/ops/swiglu.py`中的测试用例。
-
-### 任务-2.8
-
-运行算子测试。
-
-```bash
-python test/test_ops.py
-```
-
-你应该看到所有测试都通过了。提交并推送你的更改。你应该看到作业#2的自动测试通过了。
-
-### 任务-2.9（可选）rearrange
-
-这是一个奖励任务。你在模型推理中可能需要也可能不需要它。
-
-```c++
-void rearrange(tensor_t out, tensor_t in);
-```
-
-此算子用于将数据从一个张量复制到另一个具有相同形状但不同步长的张量。有了这个，你可以轻松地为张量实现`contiguous`功能。
-
-## 作业 #3：大语言模型推理
-
-终于，是时候用LLAISYS实现文本生成了。
-
-- 在`test/test_infer.py`中，你的实现应该能够使用argmax采样生成与PyTorch相同的文本。我们用于此作业的模型是[DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B)。
-
-- 你的实现的python包装器在`python/llaisys/models/qwen2.py`中。你不允许在这里使用任何基于python的框架（如PyTorch）实现你的模型推理逻辑。相反，你需要在LLAISYS后端用C/C++实现模型。脚本加载safetensors文件中的每个张量，你需要从它们加载数据到你的模型后端。
-
-- 在`include/llaisys/models/qwen2.h`中，为你定义了一个原型。你可以随意修改代码，但你应该至少提供模型创建、销毁、数据加载和推理的基本API。在`src/llaisys/`中实现你的C API，并像`src/`中的其他模块一样组织你的C++代码。记得在`xmake.lua`中定义编译过程。
-
-- 在`python/llaisys/libllaisys/`中，为你的C API定义ctypes包装函数。使用你的包装函数实现`python/llaisys/models/qwen2.py`。
-
-- 你需要实现 KV-Cache 功能，否则模型推理速度会过慢。
-
-- 调试直到你的模型工作。利用张量的`debug`函数打印张量数据。它允许你在模型推理期间将任何张量的数据与PyTorch进行比较。
-
-完成实现后，你可以运行以下命令来测试你的模型：
-
-```bash
-python test/test_infer.py --model [dir_path/to/model] --test
-```
-
-提交并推送你的更改。你应该看到作业#3的自动测试通过了。
-
-## 只有完成作业后，才能开始做项目。
-
-## 项目#1：优化 LLAISYS 的 CPU 推理
-
-你可能已经注意到，你的模型推理速度相比 PyTorch 非常慢。这主要是因为你的算子没有经过优化。运行算子测试脚本时加上 ``--profile`` 参数，看看算子的性能表现。你可能会发现 ``linear`` 操作比 PyTorch 慢很多。这个算子本质上是矩阵乘法，是 Transformer 模型里最耗时的操作。
-
-以下是几种优化 CPU 算子的方法：
-
-### 使用 SIMD 指令
-
-SIMD（单指令多数据）是一类可以在单条指令中对多个数据元素同时执行相同操作的指令。现代 CPU 都支持 SIMD。你可以查阅相关资料，学习编译器内建函数（如 AVX2、AVX-512、NEON、SVE）来向量化你的算子。
-
-### 使用 OpenMP 实现并行
-
-你可以用多线程来并行化算子。OpenMP 是 C/C++ 中常见的多线程库。为 LLAISYS 增加 OpenMP 支持，使得 ``linear`` 等算子能够并行执行。
-
-### 使用第三方库
-
-有很多库能帮你优化 CPU 上的算子，例如 Eigen、OpenBLAS、MKL 等，它们能高效处理线性代数运算。但要注意，有些库只支持特定硬件平台，需要仔细阅读文档并小心使用。你也可以参考 PyTorch 的算子实现，看是否能复用。
-
-用任何你喜欢的方法优化你的推理实现，并报告性能提升情况。
-
-## 项目#2：在 LLAISYS 中集成 CUDA，适配两款CUDA或类CUDA平台(以下统称CUDA)
-
-这个项目不依赖 ``项目#1``。需要选择 Nvidia、天数、摩尔、沐曦中的至少两款平台。
-
-本次训练营提供了以上四种平台的算力，可以在官方进行申请算力，并用 CUDA 加速模型推理。在动手前，先深入理解 LLAISYS 框架。
-
-事实上，LLAISYS 是一个支持同构硬件的框架。使用时，每个线程会创建一个线程唯一的 **Context** 对象，管理该线程使用的所有设备 **Runtime**。**Runtime** 对象是设备的资源管理器，**Context** 会为每个设备（以延迟初始化的方式）创建唯一的 **Runtime**。你可以用 ``setDevice`` 在不同设备间切换，每个线程同一时间只会激活一个设备。详情见 ``src/core/context.hpp``。
-
-### 实现 CUDA Runtime API
-
-每个 **Runtime** 对象都会初始化一组通用的 **Runtime API**。你需要实现 CUDA 版本的 API。参考 ``src/device/cpu/cpu_runtime_api.cpp`` 看 CPU 的实现方式，查阅 [`CUDA Runtime 文档`](https://docs.nvidia.com/cuda/cuda-runtime-api/index.html) 找到对应 API。
-
-在 ``src/device/runtime_api.hpp`` 中，``nvidia::getRuntimeAPI()`` 被 ``ENABLE_NVIDIA_API`` 宏保护：
-
-```c++
-#ifdef ENABLE_NVIDIA_API
-namespace nvidia {
-const LlaisysRuntimeAPI *getRuntimeAPI();
-}
-#endif
-```
-
-该宏的定义在 ``xmake.lua`` 中，用于开关 CUDA 支持。若关闭，CUDA 代码不会被编译。你需要在 ``xmake/`` 下新建 ``nvidia.lua``，配置编译流程（参考 ``cpu.lua``）。查阅资料学习如何用 Xmake 配置。
-
-完成 CUDA Runtime API 后，用 ``--nv-gpu=y`` 打开 CUDA 支持并重新编译，运行测试：
-
-```bash
-xmake f --nv-gpu=y -cv
-xmake
-xmake install
-python test/test_runtime.py --device nvidia
-```
-
-### 实现 CUDA 算子
-
-在每个算子目录下新建 ``nvidia/`` 子目录，写 CUDA 版本实现。参考 ``src/ops/add/op.cpp`` 看如何包含 CUDA 实现。别忘了在 xmake 文件中定义编译流程。用 ``--device nvidia`` 参数运行测试。
-
-你可以使用 cuBLAS、cuDNN 等 CUDA 库来加速算子，额外的设备资源可以放在 `src/device/nvidia/nvidia_resource.cu`。
-
-最后,修改模型代码，支持 CUDA 推理：
-
-```bash
-python test/test_infer.py --model [dir_path/to/model] --test --device nvidia
-```
-
-## 项目#3：构建 AI 聊天机器人
-
-本项目中，你将用 LLAISYS 构建一个能与单用户实时对话的聊天机器人。
-
-### 随机采样
-
-目前我们只用过 argmax 采样，这在测试时够用，但聊天机器人需要更自然的回复。请实现一个随机采样算子，并尽量支持 **Temperature**、**Top-K**、**Top-P**。
-
-### 搭建聊天服务器
-
-在 Python 前端里，实现一个能接收 HTTP 请求并返回响应的服务器。可以用 FastAPI 等框架。接口最好遵循 OpenAI 的 chat-completion API。如果可以，尽量支持流式输出。你可以先假设只有一个用户在使用，每次请求可以阻塞直到处理完成。
-
-### 交互式聊天 UI
-
-实现一个 UI，能向服务器发送请求并接收回复。可以是命令行界面，也可以是 Web 界面。要能通过连续发送消息与机器人保持对话。
-
-### （可选）会话管理
-
-实际应用中，用户可以开启多个对话并在它们之间切换，还能修改历史问题让 AI 重新生成回答。扩展 UI，支持这些功能。实现一个支持前缀匹配的 KV-Cache 池，尽可能复用已有结果。
-
-## 项目#4：多用户推理服务
-
-在做这个项目之前，你需要完成 ``项目#3`` 并实现流式输出。
-
-### 支持多用户
-
-现实中推理服务要同时为多个用户提供服务，请求可能随时到来。你的服务端需要将请求加入请求池/队列，并用单独的循环线程/进程来处理。
-
-### 连续批处理
-
-为了最大化吞吐量，你需要做批处理，而不是逐一处理。由于每个请求长度不同，需要实现连续的迭代级批处理机制：每轮从池中取出若干请求组成批次（batch），执行一次批量推理，再把未完成的请求放回池中。推理时尽量用批量矩阵乘法加速。注意每个请求需要绑定不同的 KV-Cache，应实现支持前缀匹配的 KV-Cache 池来复用结果。
-
-## 项目#5：分布式推理
-
-在 LLAISYS 中引入张量并行。把模型分片到多个设备上，实现分布式推理。如果用 Nvidia GPU，需要支持 NCCL；如果用 CPU，需要支持 MPI。
-
-## 项目#6：支持新模型
-
-在 LLAISYS 中支持除作业所用模型以外的其他模型。
+# 欢迎使用 LLAISYS
+
+
+English ｜
+中文
+
+
+## 简介
+
+LLAISYS（Let's Learn AI SYStem）是一个教育项目，旨在为新手和未来的AI工程师提供一个从零开始构建AI系统的学习平台。LLAISYS包含多个作业，帮助学生学习和构建基础模块；以及一些项目挑战，让他们为系统添加更多高级功能。LLAISYS使用C++作为系统后端的主要编程语言，并编译成共享库，提供C语言API。前端代码使用Python编写，调用这些API以提供更便捷的测试和与其他架构（如PyTorch）的交互。
+
+### 项目结构概览
+
+- `\include`：包含所有定义共享库提供的C API的头文件的目录。（函数声明以`__export`开头）
+
+- `\src`：C++源文件。
+ - `\src\llaisys`包含头文件中定义的所有直接实现，并遵循与`\include`相同的目录结构。这也是C++代码的边界。
+ - 其他目录包含不同模块的实际实现。
+
+- `xmake.lua`：llaisys后端的构建规则。`\xmake`目录包含不同设备的子xmake文件。例如，将来可以在目录中添加`nvidia.lua`来支持CUDA。
+
+- `\python`：Python源文件。
+ - `\python\llaisys\libllaisys`包含llaisys API的所有ctypes封装函数。它基本上与C头文件的结构相匹配。
+ - `\python\llaisys`包含ctypes函数的Python包装器，使包更符合Python风格。
+
+- `\test`：导入llaisys python包的Python测试文件。
+
+## 作业 #0：入门
+
+### 任务-0.1 安装必备组件
+
+- 编译工具：[Xmake](https://xmake.io/)
+- C++编译器：MSVC（Windows）或Clang或GCC
+- Python >= 3.9（PyTorch、Transformers等）
+- Clang-Format-16（可选）：用于格式化C++代码。
+
+### 任务-0.2 Fork并构建LLAISYS
+
+- Fork LLAISYS仓库并克隆到本地机器。支持Windows和Linux。
+
+- 编译和安装
+
+ ```bash
+ # 编译c++代码
+ xmake
+ # 安装llaisys共享库
+ xmake install
+ # 安装llaisys python包
+ pip install ./python/
+ ```
+
+- Github自动测试
+
+ LLAISYS使用Github Actions在每次推送和拉取请求时运行自动化测试。你可以在仓库页面上看到测试结果。完成所有作业任务后，所有测试都应该通过。
+
+### 任务-0.3 首次运行LLAISYS
+
+- 运行cpu运行时测试
+
+ ```bash
+ python test/test_runtime.py --device cpu
+ ```
+
+ 你应该看到测试通过。
+
+### 任务-0.4 下载测试模型
+
+- 我们用于作业的模型是[DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B)。
+
+- 使用PyTorch运行模型推理测试
+
+ ```bash
+ python test/test_infer.py --model [dir_path/to/model]
+ ```
+
+ 你可以看到PyTorch能够加载模型并使用示例输入执行推理。你可以调试进入`transformers`库代码来深入查看并了解其内部运作原理。现在，你的代码还无法执行任何操作，但在后续的作业中，你将构建一个能够实现相同功能的系统。
+
+## 作业 #1：张量
+
+张量是表示多维数据的数据结构。它是LLAISYS和大多数AI框架（如PyTorch）的基本构建单元。在这个作业中，你将学习如何实现一个基本的张量类。
+
+张量对象具有以下字段：
+
+- `storage`：指向存储张量数据的内存块的共享指针。它可以被多个张量共享。有关更多详细信息，请查看storage类。
+- `offset`：张量在存储中的起始索引（以字节为单位）。
+- `meta`：描述张量形状、数据类型和步长的元数据。
+
+实现`src/tensor/tensor.hpp`中定义的以下函数：
+
+### 任务-1.1
+
+```c++
+void load(const void *src);
+```
+
+将主机（cpu）数据加载到张量（可以在设备上）。查看构造函数了解如何获取当前设备上下文的运行时API，并执行从主机到设备的内存复制。
+
+### 任务-1.2
+
+```c++
+bool isContiguous() const;
+```
+
+检查张量的形状和步长，判断它在内存中是否连续。
+
+### 任务-1.3
+
+```c++
+tensor_t view(const std::vector &shape) const;
+```
+
+创建一个新张量，通过拆分或合并原始维度将原始张量重塑为给定形状。不涉及数据传输。例如，通过合并最后两个维度，将形状为(2, 3, 5)的张量更改为(2, 15)。
+
+这个函数不是简单地改变张量的形状那么简单，尽管测试会通过。如果新视图与原始张量不兼容，它应该引发错误。想想一个形状为(2, 3, 5)、步长为(30, 10, 1)的张量。你还能在不传输数据的情况下将其重塑为(2, 15)吗？
+
+### 任务-1.4
+
+```c++
+tensor_t permute(const std::vector &order) const;
+```
+
+创建一个新张量，改变原始张量维度的顺序。转置可以通过这个函数实现，而无需移动数据。
+
+### 任务-1.5
+
+```c++
+tensor_t slice(size_t dim, size_t start, size_t end) const;
+```
+
+创建一个新张量，沿给定维度，start（包含）和end（不包含）索引对原始张量进行切片操作。
+
+### 任务-1.6
+
+运行张量测试。
+
+```bash
+python test/test_tensor.py
+```
+
+你应该看到所有测试都通过了。提交并推送你的更改。你应该看到作业#1的自动测试通过了。
+
+## 作业 #2：算子
+
+在这个作业中，你将实现以下算子的cpu版本：
+
+- argmax
+- embedding
+- linear
+- rms_norm
+- rope
+- self_attention
+- swiglu
+
+阅读`src/ops/add/`中的代码，了解"add"算子是如何实现的。确保你理解算子代码是如何组织、编译、链接以及暴露给Python前端的。**你的算子应该至少支持Float32、Float16和BFloat16数据类型**。`src/utils/`中提供了一个用于简单类型转换的辅助函数。所有python测试都在`test/ops`中，你的实现应该至少通过这些测试。首先尝试运行"add"算子的测试脚本。
+
+### 任务-2.1 Argmax
+
+```c++
+void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
+```
+
+获取张量`vals`的最大值及其索引，并分别存储在`max_val`和`max_idx`中。你暂时可以假设`vals`是一个1D张量，`max_idx`和`max_val`都是包含单个元素的1D张量（这意味着保留了`vals`的维度）。
+
+完成实现后，你应该能够通过`test/ops/argmax.py`中的测试用例。
+
+### 任务-2.2 Embedding
+
+```c++
+void embedding(tensor_t out, tensor_t index, tensor_t weight);
+```
+
+从`weight`（2-D）中复制`index`（1-D）中的行到`output`（2-D）。`index`必须是Int64类型（PyTorch中int的默认数据类型）。
+
+完成实现后，你应该能够通过`test/ops/embedding.py`中的测试用例。
+
+### 任务-2.3 Linear
+
+```c++
+void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias);
+```
+
+计算以下内容：
+
+$$
+Y = xW^T + b
+$$
+
+- `out`：输出 $Y$ 。你暂时可以假设输出是一个2D连续张量，不涉及广播。
+- `input`：输入 $X$ 。你暂时可以假设输入是一个2D连续张量，不涉及广播。
+- `weight`：权重 $W$ 。2D连续张量。注意权重张量没有转置。你需要在计算过程中处理这个问题。
+- `bias`（可选）：偏置 $b$ 。1D张量。你需要支持不提供偏置的情况。
+
+完成实现后，你应该能够通过`test/ops/linear.py`中的测试用例。
+
+### 任务-2.4 RMS Normalization
+
+```c++
+void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps);
+```
+
+为每一行计算以下内容：
+
+$$
+Y_i = \frac{W_i \times X_i}{\sqrt{\frac{1}{d}(\sum_{j=1}^d X_j^2) + \epsilon}}
+$$
+
+- `out`：输出 $Y$ 。你暂时可以假设输出是一个2D连续张量，不涉及广播。
+- `input`：输入 $X$ 。你暂时可以假设输入是一个2D连续张量，不涉及广播。标准化沿输入张量的最后一个维度（即每一行，长度为 $d$ ）执行。
+- `weight`：权重 $W$ 。1D张量，与输入张量的一行长度相同。
+- `eps`：小值 $\epsilon$ 以避免除以零。
+
+完成实现后，你应该能够通过`test/ops/rms_norm.py`中的测试用例。
+
+### 任务-2.5 旋转位置编码（RoPE）
+
+```c++
+void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta);
+```
+
+为输入张量`in`的每个向量（这些向量与 pos_ids 中的位置 id 相对应）计算以下内容：
+
+设 $\mathbf{x}_i = [\mathbf{a}_i, \mathbf{b}_i] \in \mathbb{R}^d$ 为输入向量， $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i] \in \mathbb{R}^d$ 为索引 $i$ 处的输出向量，其中 $\mathbf{a}_i, \mathbf{b}_i,\mathbf{a}'_i, \mathbf{b}'_i \in \mathbb{R}^{d/2}$ 。
+
+设 $\theta$ 为固定基数（例如 $\theta = 10000$）， $j = 0, 1, \ldots, d/2 - 1$。
+
+设 $p_i \in \mathbb{N}$ 是输入索引i处token的位置id。
+
+那么RoPE的角度为 $\phi_{i,j} = \frac{p_i}{\theta^{2j/d}}$
+
+输出向量 $\mathbf{y}_i = [\mathbf{a}'_i, \mathbf{b}'_i]$ 计算如下：
+
+$$a_{i,j}' = a_{i,j} \cos(\phi_{i,j}) - b_{i,j} \sin(\phi_{i,j})$$
+
+$$b_{i,j}' = b_{i,j} \cos(\phi_{i,j}) + a_{i,j} \sin(\phi_{i,j})$$
+
+- `out`：结果**q**或**k**张量。形状应该是 [seqlen, nhead, d] 或 [seqlen, nkvhead, d]。你暂时可以假设张量是连续的。
+- `in`：原始**q**或**k**张量。形状应该是 [seqlen, nhead, d] 或 [seqlen, nkvhead, d]。你暂时可以假设张量是连续的。
+- `pos_ids`：输入序列中每个token的位置id（整个上下文中的索引）。形状应该是 [seqlen,]，dtype应该是int64。
+- `theta`：频率向量的基值。
+
+完成实现后，你应该能够通过`test/ops/rope.py`中的测试用例。
+
+### 任务-2.6 自注意力（self-attention）
+
+```c++
+void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale);
+```
+
+为查询张量`q`、键张量`k`和值张量`v`计算自注意力。如果需要，你应该在进行此计算之前连接kvcache张量。
+
+$$
+A = Q K^\top * scale \\
+$$
+
+$$
+Y = \mathrm{causalsoftmax}(A) \cdot V \\
+$$
+
+- `attn_val`：结果注意力值张量。形状应该是[seqlen, nhead, dv]。你暂时可以假设张量是连续的。
+- `q`：查询张量。形状应该是 [seqlen, nhead, d]。你暂时可以假设张量是连续的。
+- `k`：键张量。形状应该是 [total_len, nkvhead, d]。你暂时可以假设张量是连续的。
+- `v`：值张量。形状应该是 [total_len, nkvhead, dv]。你暂时可以假设张量是连续的。
+- `scale`：缩放因子。在大多数情况下取值为 $\frac{1}{\sqrt{d}}$ 。
+
+完成实现后，你应该能够通过`test/ops/self_attention.py`中的测试用例。
+
+### 任务-2.7 SwiGLU
+
+```c++
+void swiglu(tensor_t out, tensor_t gate, tensor_t up);
+```
+
+这是一个逐元素函数，计算以下内容：
+
+$$
+out_{i} = up_{i} \circ \frac { gate_{i}}{1 + e^{-gate_{i}}}
+$$
+
+`out`、`up`和`gate`是具有相同形状 [seqlen, intermediate_size] 的2D连续张量。
+
+完成实现后，你应该能够通过`test/ops/swiglu.py`中的测试用例。
+
+### 任务-2.8
+
+运行算子测试。
+
+```bash
+python test/test_ops.py
+```
+
+你应该看到所有测试都通过了。提交并推送你的更改。你应该看到作业#2的自动测试通过了。
+
+### 任务-2.9（可选）rearrange
+
+这是一个奖励任务。你在模型推理中可能需要也可能不需要它。
+
+```c++
+void rearrange(tensor_t out, tensor_t in);
+```
+
+此算子用于将数据从一个张量复制到另一个具有相同形状但不同步长的张量。有了这个，你可以轻松地为张量实现`contiguous`功能。
+
+## 作业 #3：大语言模型推理
+
+终于，是时候用LLAISYS实现文本生成了。
+
+- 在`test/test_infer.py`中，你的实现应该能够使用argmax采样生成与PyTorch相同的文本。我们用于此作业的模型是[DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B)。
+
+- 你的实现的python包装器在`python/llaisys_py/models/qwen2.py`中。你不允许在这里使用任何基于python的框架（如PyTorch）实现你的模型推理逻辑。相反，你需要在LLAISYS后端用C/C++实现模型。脚本加载safetensors文件中的每个张量，你需要从它们加载数据到你的模型后端。
+
+- 在`include/llaisys/models/qwen2.h`中，为你定义了一个原型。你可以随意修改代码，但你应该至少提供模型创建、销毁、数据加载和推理的基本API。在`src/llaisys/`中实现你的C API，并像`src/`中的其他模块一样组织你的C++代码。记得在`xmake.lua`中定义编译过程。
+
+- 在`python/llaisys_py/libllaisys/`中，为你的C API定义ctypes包装函数。使用你的包装函数实现`python/llaisys_py/models/qwen2.py`。
+
+- 你需要实现 KV-Cache 功能，否则模型推理速度会过慢。
+
+- 调试直到你的模型工作。利用张量的`debug`函数打印张量数据。它允许你在模型推理期间将任何张量的数据与PyTorch进行比较。
+
+完成实现后，你可以运行以下命令来测试你的模型：
+
+```bash
+python test/test_infer.py --model [dir_path/to/model] --test
+```
+
+提交并推送你的更改。你应该看到作业#3的自动测试通过了。
+
+## 只有完成作业后，才能开始做项目。
+
+## 项目#1：优化 LLAISYS 的 CPU 推理
+
+你可能已经注意到，你的模型推理速度相比 PyTorch 非常慢。这主要是因为你的算子没有经过优化。运行算子测试脚本时加上 ``--profile`` 参数，看看算子的性能表现。你可能会发现 ``linear`` 操作比 PyTorch 慢很多。这个算子本质上是矩阵乘法，是 Transformer 模型里最耗时的操作。
+
+以下是几种优化 CPU 算子的方法：
+
+### 使用 SIMD 指令
+
+SIMD（单指令多数据）是一类可以在单条指令中对多个数据元素同时执行相同操作的指令。现代 CPU 都支持 SIMD。你可以查阅相关资料，学习编译器内建函数（如 AVX2、AVX-512、NEON、SVE）来向量化你的算子。
+
+### 使用 OpenMP 实现并行
+
+你可以用多线程来并行化算子。OpenMP 是 C/C++ 中常见的多线程库。为 LLAISYS 增加 OpenMP 支持，使得 ``linear`` 等算子能够并行执行。
+
+### 使用第三方库
+
+有很多库能帮你优化 CPU 上的算子，例如 Eigen、OpenBLAS、MKL 等，它们能高效处理线性代数运算。但要注意，有些库只支持特定硬件平台，需要仔细阅读文档并小心使用。你也可以参考 PyTorch 的算子实现，看是否能复用。
+
+用任何你喜欢的方法优化你的推理实现，并报告性能提升情况。
+
+## 项目#2：在 LLAISYS 中集成 CUDA，适配两款CUDA或类CUDA平台(以下统称CUDA)
+
+这个项目不依赖 ``项目#1``。需要选择 Nvidia、天数、摩尔、沐曦中的至少两款平台。
+
+本次训练营提供了以上四种平台的算力，可以在官方进行申请算力，并用 CUDA 加速模型推理。在动手前，先深入理解 LLAISYS 框架。
+
+事实上，LLAISYS 是一个支持同构硬件的框架。使用时，每个线程会创建一个线程唯一的 **Context** 对象，管理该线程使用的所有设备 **Runtime**。**Runtime** 对象是设备的资源管理器，**Context** 会为每个设备（以延迟初始化的方式）创建唯一的 **Runtime**。你可以用 ``setDevice`` 在不同设备间切换，每个线程同一时间只会激活一个设备。详情见 ``src/core/context.hpp``。
+
+### 实现 CUDA Runtime API
+
+每个 **Runtime** 对象都会初始化一组通用的 **Runtime API**。你需要实现 CUDA 版本的 API。参考 ``src/device/cpu/cpu_runtime_api.cpp`` 看 CPU 的实现方式，查阅 [`CUDA Runtime 文档`](https://docs.nvidia.com/cuda/cuda-runtime-api/index.html) 找到对应 API。
+
+在 ``src/device/runtime_api.hpp`` 中，``nvidia::getRuntimeAPI()`` 被 ``ENABLE_NVIDIA_API`` 宏保护：
+
+```c++
+#ifdef ENABLE_NVIDIA_API
+namespace nvidia {
+const LlaisysRuntimeAPI *getRuntimeAPI();
+}
+#endif
+```
+
+该宏的定义在 ``xmake.lua`` 中，用于开关 CUDA 支持。若关闭，CUDA 代码不会被编译。你需要在 ``xmake/`` 下新建 ``nvidia.lua``，配置编译流程（参考 ``cpu.lua``）。查阅资料学习如何用 Xmake 配置。
+
+完成 CUDA Runtime API 后，用 ``--nv-gpu=y`` 打开 CUDA 支持并重新编译，运行测试：
+
+```bash
+xmake f --nv-gpu=y -cv
+xmake
+xmake install
+python test/test_runtime.py --device nvidia
+```
+
+### 实现 CUDA 算子
+
+在每个算子目录下新建 ``nvidia/`` 子目录，写 CUDA 版本实现。参考 ``src/ops/add/op.cpp`` 看如何包含 CUDA 实现。别忘了在 xmake 文件中定义编译流程。用 ``--device nvidia`` 参数运行测试。
+
+你可以使用 cuBLAS、cuDNN 等 CUDA 库来加速算子，额外的设备资源可以放在 `src/device/nvidia/nvidia_resource.cu`。
+
+最后,修改模型代码，支持 CUDA 推理：
+
+```bash
+python test/test_infer.py --model [dir_path/to/model] --test --device nvidia
+```
+
+## 项目#3：构建 AI 聊天机器人
+
+本项目中，你将用 LLAISYS 构建一个能与单用户实时对话的聊天机器人。
+
+### 随机采样
+
+目前我们只用过 argmax 采样，这在测试时够用，但聊天机器人需要更自然的回复。请实现一个随机采样算子，并尽量支持 **Temperature**、**Top-K**、**Top-P**。
+
+### 搭建聊天服务器
+
+在 Python 前端里，实现一个能接收 HTTP 请求并返回响应的服务器。可以用 FastAPI 等框架。接口最好遵循 OpenAI 的 chat-completion API。如果可以，尽量支持流式输出。你可以先假设只有一个用户在使用，每次请求可以阻塞直到处理完成。
+
+### 交互式聊天 UI
+
+实现一个 UI，能向服务器发送请求并接收回复。可以是命令行界面，也可以是 Web 界面。要能通过连续发送消息与机器人保持对话。
+
+### （可选）会话管理
+
+实际应用中，用户可以开启多个对话并在它们之间切换，还能修改历史问题让 AI 重新生成回答。扩展 UI，支持这些功能。实现一个支持前缀匹配的 KV-Cache 池，尽可能复用已有结果。
+
+## 项目#4：多用户推理服务
+
+在做这个项目之前，你需要完成 ``项目#3`` 并实现流式输出。
+
+### 支持多用户
+
+现实中推理服务要同时为多个用户提供服务，请求可能随时到来。你的服务端需要将请求加入请求池/队列，并用单独的循环线程/进程来处理。
+
+### 连续批处理
+
+为了最大化吞吐量，你需要做批处理，而不是逐一处理。由于每个请求长度不同，需要实现连续的迭代级批处理机制：每轮从池中取出若干请求组成批次（batch），执行一次批量推理，再把未完成的请求放回池中。推理时尽量用批量矩阵乘法加速。注意每个请求需要绑定不同的 KV-Cache，应实现支持前缀匹配的 KV-Cache 池来复用结果。
+
+## 项目#5：分布式推理
+
+在 LLAISYS 中引入张量并行。把模型分片到多个设备上，实现分布式推理。如果用 Nvidia GPU，需要支持 NCCL；如果用 CPU，需要支持 MPI。
+
+## 项目#6：支持新模型
+
+在 LLAISYS 中支持除作业所用模型以外的其他模型。
diff --git a/docs/cpu-inference-optimization.md b/docs/cpu-inference-optimization.md
new file mode 100644
index 000000000..77e918b3d
--- /dev/null
+++ b/docs/cpu-inference-optimization.md
@@ -0,0 +1,393 @@
+# LLAISYS CPU 推理优化文档
+
+本文档记录针对项目 #1「优化 LLAISYS 的 CPU 推理」所做的改动，主要包括：**OpenMP 多线程并行**、**AVX2 SIMD 向量化**（FP32 linear），以及为兼容系统头文件所做的 **LLAISYS_EXTERN_C 宏重命名**。
+
+---
+
+## 一、背景与目标
+
+### 1.1 问题
+
+- 未优化前，LLAISYS 的模型推理速度相比 PyTorch 明显更慢。
+- 主要瓶颈在 **linear**（矩阵乘法）算子，该算子在 Transformer 中调用最频繁、耗时占比最高。
+- 原始实现为三重循环的朴素矩阵乘，未利用多核与 SIMD。
+
+### 1.2 优化思路（与 README/项目说明一致）
+
+| 方法 | 说明 |
+|----------------|------|
+| **SIMD** | 使用 AVX2/AVX-512 等指令一次处理多个 float，提高吞吐。 |
+| **OpenMP** | 用多线程并行化算子，使 linear 等能利用多核 CPU。 |
+| **第三方库** | 可选用 OpenBLAS、MKL、Eigen 等（本次未接入，见后续可选）。 |
+
+本次实现采用 **OpenMP + AVX2**，在不引入额外依赖的前提下提升 CPU 推理速度。
+
+---
+
+## 二、优化一：OpenMP 多线程
+
+### 2.1 思路
+
+- linear 计算 `Y = X W^T + b`，其中 `out` 形状为 `(B, M)`，`in` 为 `(B, K)`，`weight` 为 `(M, K)`。
+- 外层按行（B 维）并行：每个线程负责若干行输出，互不写同一位置，无需加锁。
+- 使用 `schedule(static)` 静态划分，便于缓存局部性。
+
+### 2.2 修改内容
+
+**xmake.lua**
+
+- 在 `llaisys-ops` 目标中增加：
+ - `add_cxflags("-fopenmp")`、`add_mxflags("-fopenmp")`、`add_ldflags("-fopenmp")`。
+- 在最终动态库目标 `llaisys` 中增加：
+ - `add_ldflags("-fopenmp")`，以便链接 OpenMP 运行时。
+
+**src/ops/linear/op.cpp**
+
+- 在文件顶部增加（可选）：`#ifdef _OPENMP` 时 `#include `。
+- 在模板函数 `linear_impl` 的外层循环（B 维）前增加：
+
+```cpp
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
+for (size_t i = 0; i < B; i++) {
+ // ...
+}
+```
+
+- 对所有 dtype（F32 / F16 / BF16）均生效；FP32 在启用 AVX2 时走 `linear_f32_avx2`，其内部同样使用上述同一套 OpenMP 并行。
+
+### 2.3 效果
+
+- 多核机器上，B 或 M 较大时能明显利用多核，线性层耗时随核数近似线性下降（受内存带宽限制会有所折扣）。
+
+---
+
+## 三、优化二：AVX2 + FMA（仅 FP32 linear）
+
+### 3.1 思路
+
+- 内层 K 维是连续内存上的点积，适合用 SIMD 一次处理多个 float。
+- AVX2 提供 256 位寄存器，一次处理 **8 个 float**；FMA（Fused Multiply-Add）一条指令完成 `a*b+c`，减少舍入与指令数。
+- 仅对 **FP32** 实现 AVX2 路径；F16/BF16 仍走原有 `linear_impl`（内部用 float 累加），避免重复实现半精度 SIMD。
+
+### 3.2 修改内容
+
+**xmake.lua**
+
+- 在 `llaisys-ops` 中，当 `is_arch("x86_64")` 时增加：
+ - `add_cxflags("-mavx2", "-mfma")`，使编译器生成 AVX2/FMA 指令并定义 `__AVX2__`。
+
+**src/ops/linear/op.cpp**
+
+1. **头文件顺序**
+ - 在包含任何项目头文件之前，先写：
+ - `#ifdef __AVX2__`
+ - `#include `
+ - `#endif`
+ - 避免项目中的宏（见第四节）与 `` 及其间接包含的系统头文件中的符号冲突。
+
+2. **AVX2 版 FP32 实现**
+ - `linear_f32_avx2(out, in, weight, bias, B, M, K)`：
+ - 外层 B 维用 OpenMP 并行（与 `linear_impl` 一致）。
+ - 对每个 `(i, j)`，内层 K 维：
+ - 用 `_mm256_loadu_ps` 每次读 8 个 float；
+ - 用 `_mm256_fmadd_ps(a, b, sum8)` 做乘加；
+ - K 不是 8 的倍数时，剩余标量补齐。
+ - 对 8 路累加结果做水平求和：`hsum_avx(sum8)`，得到标量 `sum`，再加 bias 写入 `out[i*M+j]`。
+ - `hsum_avx(__m256 v)`：将 256 位寄存器中 8 个 float 相加为 1 个 float（用 `_mm256_castps256_ps128`、`_mm256_extractf128_ps`、`_mm_add_ps`、`_mm_movehdup_ps`、`_mm_movehl_ps`、`_mm_add_ss` 等实现）。
+
+3. **分支选择**
+ - 在 `linear_cpu` 的 `LLAISYS_DTYPE_F32` 分支中：
+ - 若定义了 `__AVX2__`，则调用 `linear_f32_avx2(...)`；
+ - 否则调用原有 `linear_impl(...)`。
+ - F16/BF16 仍只走 `linear_impl`。
+
+### 3.3 平台说明
+
+- **x86_64**：默认开启 `-mavx2 -mfma`，FP32 linear 使用 AVX2 路径。
+- **ARM / 其他架构**：不添加上述编译选项，FP32 仍为 OpenMP + 标量三重循环；后续可仿照实现 NEON 等 SIMD 版本。
+
+---
+
+## 四、兼容性修复：LLAISYS_EXTERN_C 宏
+
+### 4.1 问题
+
+- 项目在 `include/llaisys.h` 中用宏 `__C` 表示 `extern "C"`（C++ 时）或空（C 时）。
+- 系统头文件（如 GCC 的 `` 间接包含的 ``）中，`__C` 被用作**参数名**。
+- 开启 AVX2 并包含 `` 后，这些系统头在展开时会把参数名 `__C` 错误替换成 `extern "C"`，导致编译错误。
+
+### 4.2 修改内容
+
+- 在 **include/llaisys.h** 中：
+ - 将 `#define __C extern "C"` 改为 `#define LLAISYS_EXTERN_C extern "C"`；
+ - 将 `#define __C` 改为 `#define LLAISYS_EXTERN_C`。
+- 在所有使用 `__C` 的地方改为 **LLAISYS_EXTERN_C**，涉及文件包括：
+ - **头文件**：`include/llaisys.h`、`include/llaisys/models/qwen2.h`、`include/llaisys/runtime.h`、`include/llaisys/tensor.h`、`include/llaisys/ops.h`；
+ - **实现**：`src/llaisys/llaisys_tensor.hpp`、`src/llaisys/runtime.cc`、`src/llaisys/tensor.cc`、`src/llaisys/ops.cc`、`src/llaisys/qwen2.cc`（两处 `__C {` 均改为 `LLAISYS_EXTERN_C {`）。
+
+这样在任意源文件中先包含 `` 再包含项目头，也不会再与系统头中的 `__C` 冲突。
+
+---
+
+## 五、构建与验证
+
+### 5.1 构建
+
+```bash
+cd /path/to/llaisys
+xmake build llaisys
+xmake install llaisys
+```
+
+- `xmake install` 会将生成的 `libllaisys.so`（或 Windows 下 `llaisys.dll`）复制到 `python/llaisys_py/libllaisys/`，供 Python 调用。
+
+### 5.2 正确性测试
+
+```bash
+export PYTHONPATH="/path/to/llaisys/python:$PYTHONPATH"
+python test/ops/linear.py --device cpu
+```
+
+- 应通过所有 shape 与 dtype 的测试（含 (512, 4096) 等大矩阵）。
+
+### 5.3 性能对比（profile）
+
+```bash
+export PYTHONPATH="/path/to/llaisys/python:$PYTHONPATH"
+python test/ops/linear.py --device cpu --profile
+```
+
+- 脚本会对 PyTorch 与 LLAISYS 的 linear 做 warmup + 多次重复计时，并打印两者耗时（ms）。
+- 大矩阵（如 512×4096 × 4096×4096）下，预期 LLAISYS 相对未优化版本有明显加速（具体倍数与 CPU 核数、是否支持 AVX2 有关）。
+
+### 5.4 实际推理体感
+
+```bash
+.venv/bin/python -m llaisys_py.server --model /path/to/DeepSeek-R1-Distill-Qwen-1___5B --port 8002
+```
+
+- 与优化前对比：首 token 延迟与后续 token 延迟应有所下降，尤其在多核、支持 AVX2 的 x86_64 上。
+
+### 5.5 性能分析：优化前后对比
+
+要量化「OpenMP + AVX2」带来的提升，可采用下面三种方式（由简到繁）。
+
+#### 方法一：用 OMP_NUM_THREADS 看多线程收益（无需改代码、无需两套构建）
+
+在同一台机器、同一套已开启 OpenMP 的构建下，只改变线程数对比耗时：
+
+```bash
+export PYTHONPATH="/path/to/llaisys/python:$PYTHONPATH"
+
+# 单线程（相当于“无多线程优化”的耗时）
+OMP_NUM_THREADS=1 python test/ops/linear_bench.py --device cpu --dtype f32
+
+# 多线程（例如 8 核）
+OMP_NUM_THREADS=8 python test/ops/linear_bench.py --device cpu --dtype f32
+```
+
+保存两次输出的 `lla_ms`，则 **多线程加速比 ≈ 单线程时间 / 多线程时间**。例如单线程 200 ms、8 线程 35 ms，加速比约 5.7x。
+
+#### 方法二：固定 benchmark 脚本，保存“优化前 / 优化后”数据
+
+使用 `test/ops/linear_bench.py` 做**可复现**的计时，输出便于 diff 或写脚本解析：
+
+```bash
+# 优化前：例如先 checkout 到未加 OpenMP/AVX2 的提交，构建并安装后
+xmake build && xmake install
+python test/ops/linear_bench.py --device cpu --repeat 100 --json > baseline.json
+
+# 优化后：切回当前代码，重新构建安装
+xmake build && xmake install
+python test/ops/linear_bench.py --device cpu --repeat 100 --json > optimized.json
+
+# 对比（可用 jq 或手写脚本算 speedup = baseline_ms / optimized_ms）
+```
+
+同一台机器、同一 `--repeat` 下，直接比较各 shape/dtype 的 `lla_ms` 即可得到优化倍数。
+
+#### 方法三：AVX2 开/关对比（需两套构建）
+
+若想单独看 **AVX2 SIMD** 的收益（不含多线程差异），需要两次构建：
+
+1. **无 AVX2 构建**：在 `xmake.lua` 的 `llaisys-ops` 中临时注释掉 `add_cxflags("-mavx2", "-mfma")`，然后 `xmake build && xmake install`，运行 `linear_bench.py` 保存结果（例如 `no_avx2.json`）。
+2. **有 AVX2 构建**：恢复 `-mavx2 -mfma`，重新 `xmake build && xmake install`，再跑一次保存（例如 `with_avx2.json`）。
+
+对比两者在 **FP32**、同一 shape 下的 `lla_ms`，即可得到 AVX2 带来的加速比。F16/BF16 当前无 AVX2 路径，对比意义不大。
+
+#### 建议记录格式
+
+- 每次 benchmark 注明：**机器（CPU 型号、核数）、OMP_NUM_THREADS、repeat、warmup**。
+- 重点关注大矩阵：**out (512, 4096), x (512, 4096), w (4096, 4096)**，dtype **f32**，与 PyTorch 的耗时对比可作为参考（见 5.3）。
+
+#### 如何报告性能提升
+
+1. **生成对比报告**：用两份 JSON 跑报告脚本，直接得到表格和加速比。
+
+```bash
+# 单线程 vs 多线程
+OMP_NUM_THREADS=1 python test/ops/linear_bench.py --device cpu --dtype f32 --json > single.json
+OMP_NUM_THREADS=8 python test/ops/linear_bench.py --device cpu --dtype f32 --json > multi.json
+python test/ops/linear_bench_report.py single.json multi.json
+```
+
+输出示例：
+```
+======================================================================
+Linear 性能对比报告
+======================================================================
+ 基准: single.json (e.g. 优化前 / 单线程)
+ 对比: multi.json (e.g. 优化后 / 多线程)
+
+shape dtype 基准(ms) 对比(ms) 加速比
+----------------------------------------------------------------------
+[512, 4096] @ [4096, 4096] f32 2537.22 82.02 30.94x
+----------------------------------------------------------------------
+说明: 加速比 = 基准耗时 / 对比耗时，>1 表示对比版本更快。
+======================================================================
+```
+
+2. **书面报告建议结构**（可粘贴到 README / 实验报告）：
+ - **环境**：CPU 型号、核数、OMP_NUM_THREADS、repeat/warmup。
+ - **测试内容**：shape（如 512×4096 @ 4096×4096）、dtype（f32）。
+ - **结果**：基准耗时（ms）、优化后耗时（ms）、**加速比**（基准/优化后）。
+ - **结论**：例如「在 8 核机器上，OpenMP 多线程使 linear (f32) 大矩阵耗时从 xxx ms 降至 xxx ms，加速约 x.x 倍。」
+
+---
+
+## 六、涉及文件一览
+
+| 文件 | 修改要点 |
+|------|----------|
+| **xmake.lua** | llaisys-ops：OpenMP 编译/链接选项；x86_64 下 -mavx2 -mfma。llaisys：-fopenmp 链接。 |
+| **src/ops/linear/op.cpp** | 顶部条件包含 immintrin.h；linear_impl 外层 B 维 OpenMP；linear_f32_avx2 + hsum_avx；linear_cpu 中 F32 分支选 AVX2 或标量。 |
+| **src/ops/self_attention/op.cpp** | 条件包含 omp.h；parallel 区内线程私有 scores，对 qlen 做 omp for；typed 写回对 total 做 omp for。 |
+| **src/ops/rms_norm/op.cpp** | 条件包含 omp.h；rms_norm_impl 外层 rows 循环 omp parallel for。 |
+| **src/ops/swiglu/op.cpp** | 条件包含 omp.h；swiglu_impl 外层 n 循环 omp parallel for。 |
+| **src/ops/rope/op.cpp** | 条件包含 omp.h；rope_impl 外层 seq_len 循环 omp parallel for。 |
+| **include/llaisys.h** | __C → LLAISYS_EXTERN_C。 |
+| **include/llaisys/*.h**、**include/llaisys/models/qwen2.h** | 所有 __C 改为 LLAISYS_EXTERN_C。 |
+| **src/llaisys/*.cc**、**src/llaisys/llaisys_tensor.hpp** | 所有 __C 改为 LLAISYS_EXTERN_C。 |
+
+---
+
+## 七、其他算子的 OpenMP 并行（已实现）
+
+在 `self_attention`、`rms_norm`、`swiglu`、`rope` 的外层循环上已增加 OpenMP 并行，与 linear 共用同一套 `-fopenmp` 编译/链接选项，无需额外配置。
+
+| 算子 | 并行维度 | 说明 |
+|------|----------|------|
+| **self_attention** | `qlen`（query 序列长度） | 每个线程私有 `scores` 缓冲区，`#pragma omp parallel` + `#pragma omp for`；F16/BF16 的 cast 写回用 `#pragma omp parallel for`。 |
+| **rms_norm** | `rows`（行数） | 按行独立，`#pragma omp parallel for schedule(static)`。 |
+| **swiglu** | 元素下标 `i`（总元素数 n） | 逐元素独立，`#pragma omp parallel for schedule(static)`。 |
+| **rope** | `seq_len`（序列长度） | 每帧独立，`inv_freq` 只读共享；`#pragma omp parallel for schedule(static)`。 |
+
+验证：`python test/ops/self_attention.py --device cpu`、`test/ops/rms_norm.py`、`test/ops/swiglu.py`、`test/ops/rope.py` 均已通过。
+
+---
+
+## 八、可选后续优化
+
+1. **其他算子**
+ `add`、`embedding`、`rearrange`、`argmax`、`sample` 等若在 profile 中占比高，可同样对外层循环加 `#pragma omp parallel for`。
+
+2. **BLAS 库**
+ 将 FP32 linear 改为调用 OpenBLAS 的 `cblas_sgemm` 或 MKL 的等效接口，在 xmake 中增加对应依赖与链接，通常能获得更好性能，但需处理跨平台与依赖安装。
+
+3. **ARM NEON**
+ 在 ARM 架构下为 FP32 linear 实现 NEON 版本（一次处理 4 个 float），并在对应架构的编译选项中启用。
+
+4. **BF16/F16 SIMD**
+ 若模型以 BF16/F16 为主，可为半精度 linear 增加 AVX2/NEON 的 16 位或 32 位累加路径，以进一步提升半精度推理速度。
+
+---
+
+## 九、参考
+
+- 项目 README / README_ZN 中「项目 #1：优化 LLAISYS 的 CPU 推理」说明。
+- 算子性能分析：`python test/ops/linear.py --device cpu --profile`；优化前后对比：`python test/ops/linear_bench.py --device cpu [--dtype f32] [--json]`，见 5.5 节。
+- OpenMP：/`，被 `src/llaisys/ops.cc` 包含并转成 C API。
+- **C 暴露**：`src/llaisys/ops.cc` 里对每个算子有一个 `llaisysXxx(...)`，内部把 `llaisysTensor_t` 转成 `tensor_t` 再调 `llaisys::ops::xxx(...)`。
+- **Python 封装**：`python/llaisys_py/libllaisys/ops.py` 里用 ctypes 声明 `llaisysAdd` 等，并封装成 `llaisys_py.Ops.add` 等；测试在 `test/ops/.py`。
+
+### 以 add 为例：从 Python 到 C++ 的完整链路
+
+1. **测试脚本** `test/ops/add.py`（约 17–32 行）：
+ - 用 `random_tensor` 造两个同 shape 的 llaisys 张量 `a_`, `b_`，以及同 shape 的 `c_`。
+ - `torch_add(c, a, b)` 得到 PyTorch 结果，`llaisys_py.Ops.add(c_, a_, b_)` 得到 LLAISYS 结果。
+ - `check_equal(c_, c)` 比较两者是否接近（支持 f32/f16/bf16 的 atol/rtol）。
+
+2. **Python 封装** `python/llaisys_py/libllaisys/ops.py`：
+ - 声明 `lib.llaisysAdd(c_, a_, b_)` 的 argtypes/restype（若存在），供上层 `llaisys_py.Ops.add` 调用。
+
+3. **C API** `src/llaisys/ops.cc`（约 16–18 行）：
+ - `llaisysAdd(c, a, b)` 内部调 `llaisys::ops::add(c->tensor, a->tensor, b->tensor)`。
+
+4. **C++ 算子** `src/ops/add/op.cpp`：
+ - 校验设备一致、形状一致、dtype 一致、张量连续（CHECK_SAME_DEVICE / CHECK_SAME_SHAPE 等）。
+ - 若为 CPU，调 `cpu::add(c->data(), a->data(), b->data(), c->dtype(), c->numel())`。
+ - `src/ops/add/cpu/add_cpu.cpp` 里按 dtype 分支，对 f32/f16/bf16 做逐元素加法。
+
+### 每个算子做什么（公式与形状，便于对照 README）
+
+| 算子 | 公式/语义 | 输入输出形状（简要） |
+|------|-----------|----------------------|
+| **argmax** | 在 vals 上取最大值与下标，写入 max_val、max_idx | vals 1D；max_idx、max_val 各 1 元素 |
+| **embedding** | out[i] = weight[index[i]] | index 1D int64；weight [V, D]；out [seq, D] |
+| **linear** | Y = X W^T + b | out/in 2D [B,K]、[B,M]；weight [M,K]；bias [M] 或 null |
+| **rms_norm** | 每行：y = w * x / sqrt(mean(x^2)+eps) | out/in 2D；weight 1D，长度=行宽 |
+| **rope** | 按 pos_ids 和 theta 对 Q/K 做旋转（公式见 README） | in/out [seq, nhead, d]；pos_ids [seq] int64 |
+| **self_attention** | causal softmax(Q K^T * scale) @ V | q [seq, nhead, d]；k/v [total_len, nkvh, d]；attn_val [seq, nhead, d] |
+| **swiglu** | out = up * sigmoid(gate) 逐元素 | out/gate/up 同形 2D [seq, di] |
+
+### 类型与工具（F32/F16/BF16）
+
+- **类型转换**：`src/utils/` 下通常有 `cast()`、`dsize(dtype)` 等；F16/BF16 在算子里常先转成 float 算再转回，避免精度问题（见 `src/ops/linear/op.cpp` 里 `linear_impl` 的 bf16_t/fp16_t 分支）。
+- **头文件**：`include/llaisys/ops.h` 声明所有 `llaisysXxx` 的 C 接口；`src/llaisys/ops.cc` 包含各 `../ops/xxx/op.hpp` 并实现这些 C 函数。
+
+### 建议看哪些文件、看什么（按顺序）
+
+1. **README**：`README.md` “Assignment #2: Operators” 及每个 Task-2.x，看公式、形状、bias 可选等约定。
+2. **add 全链路**：
+ - `src/ops/add/op.hpp`、`op.cpp`（校验 + 调 cpu::add）；
+ - `src/ops/add/cpu/add_cpu.hpp`、`add_cpu.cpp`（按 dtype 分支）；
+ - `src/llaisys/ops.cc` 里 `llaisysAdd`；
+ - `test/ops/add.py` 里 test 与 shape/dtype 组合。
+3. **linear**：`src/ops/linear/op.cpp`（Y=XW^T+b 的实现、B/M/K 含义、bias 为 null 的处理、F16/BF16 用 float 累加）。
+4. **其它算子**：按需看 `src/ops//op.cpp` 和 `test/ops/.py`，对照 README 的公式与形状。
+5. **工具**：`src/utils/types.hpp`、`utils.cpp`（若存在）中的 `cast`、`dsize`、`bf16_t`、`fp16_t` 等。
+
+### 如何验证看懂
+
+- 运行 `python test/ops/add.py`、`python test/ops/linear.py` 等全部通过。
+- 能说出：一个算子从 `test/ops/xxx.py` 到 `ops.cc` 再到 `ops/xxx/op.cpp` 和 `cpu/xxx_cpu.cpp` 的调用链；以及 F16/BF16 为何在 linear 里用 float 累加。
+
+---
+
+## 作业 3：LLM 推理（Qwen2）
+
+### 在干什么（一句话）
+
+用**作业 1 的 Tensor** 和 **作业 2 的算子**，在 C++ 里实现 Qwen2 的**单步前向**（给定当前 token 序列，算下一个 token 的 id），并在 Python 里读 config、从 safetensors 加载权重到 C 侧、循环调用该单步前向直到 EOS 或达到长度。**推理逻辑全部在 C++，Python 不写前向。**
+
+### 核心 C 接口（见 include/llaisys/models/qwen2.h）
+
+| 接口 | 作用 |
+|------|------|
+| **llaisysQwen2ModelCreate(meta, device, ...)** | 根据 meta 分配模型结构、所有权重张量、每层 KV cache；返回模型指针。权重数据由调用方后续灌入。 |
+| **llaisysQwen2ModelDestroy(model)** | 释放所有权重张量和 KV cache，再 delete 模型。 |
+| **llaisysQwen2ModelWeights(model)** | 返回指向 `LlaisysQwen2Weights` 的指针，Python 据此把 safetensors 里每个 key 对应的数据 tensorLoad 到对应句柄。 |
+| **llaisysQwen2ModelInfer(model, token_ids, ntoken)** | 单步推理：输入当前 token 序列与长度，执行一次前向，更新 KV cache，返回**下一个 token 的 id**（int64）。 |
+
+### 单步前向在做什么（数据流）
+
+`llaisysQwen2ModelInfer` 内部（`src/llaisys/qwen2.cc`）大致顺序：
+
+1. **prefill / decode 区分**
+ - `cache_len == 0`：prefill，本步输入整段 `token_ids`，seq_len = ntoken。
+ - 否则：decode，本步只输入最后一个 token，seq_len = 1；KV cache 中已有历史，本步只追加当前步的 K/V。
+
+2. **准备输入**
+ - 把 token id 拷到设备上的 `token_tensor`（长度 seq_len）。
+ - 分配本步用的临时张量：hidden、normed、q_buf、k_buf、v_buf、q_rope、k_rope、attn_val、o_proj_out、res_buf、gate_buf、up_buf、mlp_buf、down_buf、pos_ids 等。
+
+3. **embedding**
+ - `llaisys::ops::embedding(hidden, token_tensor, in_embed)`：用 token id 查表得到 hidden。
+
+4. **逐层 Transformer Block**（`forward_layer`，每层调用一次）
+ - **Attention**：rms_norm → linear 得到 Q/K/V → RoPE → 把本步 K/V 写入 KV cache → 用「当前 Q」和「cache 里拼好的 K/V」做 self_attention → linear(o_proj) → 残差加回 hidden。
+ - **MLP**：rms_norm → linear(gate/up) → swiglu → linear(down) → 残差加回 hidden。
+
+5. **最后一层之后**
+ - 对 hidden 做 rms_norm，再 linear 得到 logits [seq_len, voc]。
+ - 取最后一个位置的 logits，argmax 得到 next_token，更新 `cache_len`，返回 next_token。
+
+### 单层 forward_layer 里用到的作业 2 算子（对应 qwen2.cc 行号）
+
+- **rms_norm**：attention 前 norm、MLP 前 norm、最后输出前 norm。
+- **linear**：Q/K/V 投影、o_proj、gate/up/down、输出层。
+- **rope**：对 Q、K 做旋转。
+- **self_attention**：causal attention。
+- **swiglu**：MLP 激活。
+- **add**：两次残差加（attention 后、MLP 后）。
+- **embedding**：仅最前一步。
+- **argmax**：仅最后取 next token。
+
+KV cache 的写入在 `forward_layer` 里：把本步的 `k_rope`、`v_buf` 按时间步写入 `k_caches`、`v_caches` 的 [cache_start, cache_start+seq_len)；attention 时用 `k_cache->slice(0, 0, kv_len)` 和 `v_cache->slice(...)` 取「当前有效长度」的 K/V。
+
+### 建议看哪些文件、看什么（按顺序）
+
+1. **README**：`README.md` “Assignment #3: Large Language Model Inference”，看目标、约束、测试命令。
+2. **入口与对比**：`test/test_infer.py`
+ - 先 PyTorch 跑一遍得到参考 token 序列；
+ - 再 LLAISYS 加载模型、`llaisys_infer` 里循环 `model.generate(...)`；
+ - 若加 `--test`，会断言两边 token 序列一致。
+3. **Python 封装**：`python/llaisys_py/models/qwen2.py`
+ - `__init__`：读 config、拼 Meta、调 Create、用 Weights 遍历 safetensors 并 tensorLoad；
+ - `generate`：循环里把当前 tokens 传给 `llaisysQwen2ModelInfer`，拿 next_tok，追加到 tokens，直到 EOS 或 max_new_tokens。
+4. **C 接口**：`include/llaisys/models/qwen2.h` —— LlaisysQwen2Meta、LlaisysQwen2Weights 的字段，以及 Create/Destroy/Weights/Infer 的声明。
+5. **C++ 实现**：`src/llaisys/qwen2.cc`
+ - `LlaisysQwen2Model` 结构体（meta、weights、k_caches、v_caches、cache_len）；
+ - `create_weight_tensors`、`llaisysQwen2ModelCreate`（分配权重与 KV cache）；
+ - `llaisysQwen2ModelInfer`（prefill/decode、embed、逐层 forward_layer、最后 norm+linear+argmax）；
+ - `forward_layer`（attention 分支 + MLP 分支，以及写 KV cache 的循环）。
+6. **整体数据流**：`docs/代码导读-推理流程.md` —— 从 test_infer 到 qwen2.py 到 C API 到 qwen2.cc 的完整调用链说明。
+
+### 如何验证看懂
+
+- 运行 `python test/test_infer.py --model <路径> --test`，LLAISYS 与 PyTorch 的 token 序列一致。
+- 能说出：prefill 与 decode 的区别；单层里 attention 和 MLP 各用了哪些算子；KV cache 在哪个函数里写入、在哪个函数里被 slice 出来做 attention。
+
+---
+
+## 总体阅读顺序与“如何验证看懂”（小结）
+
+1. **作业 0**：按 README 装环境、构建、跑 runtime 和 test_infer（仅 PyTorch），知道「C++ 编成 dll，Python 调它」。
+2. **作业 1**：看 `tensor.hpp` 成员、`tensor.cpp` 里 create/load/isContiguous/view/permute/slice，再跑 `test/test_tensor.py`，能解释 view 为何要求连续、slice 为何只改 offset。
+3. **作业 2**：看 add 的 op.cpp → cpu/add_cpu.cpp → ops.cc 的 llaisysAdd → test/ops/add.py；再看 linear 的公式与实现；其余算子按需看 README 与对应 `src/ops//`。
+4. **作业 3**：按 test_infer.py → qwen2.py → qwen2.h → qwen2.cc 走一遍，重点看 Infer 里 prefill/decode、embed、forward_layer、最后 argmax，以及 KV cache 的写入与使用；配合 `代码导读-推理流程.md` 串起来。
+
+若某一作业里某一块想再细看（例如 view 的兼容性检查、或 self_attention 的 causal mask），可以指定文件名或函数名继续问。
diff --git "a/docs/\344\273\243\347\240\201\345\257\274\350\257\273-\346\216\250\347\220\206\346\265\201\347\250\213.md" "b/docs/\344\273\243\347\240\201\345\257\274\350\257\273-\346\216\250\347\220\206\346\265\201\347\250\213.md"
new file mode 100644
index 000000000..3f18328ed
--- /dev/null
+++ "b/docs/\344\273\243\347\240\201\345\257\274\350\257\273-\346\216\250\347\220\206\346\265\201\347\250\213.md"
@@ -0,0 +1,257 @@
+# LLAISYS 推理流程代码导读（零基础版）
+
+本文按「你运行 `python test/test_infer.py --model ...` 时，代码的执行顺序」带你看一遍涉及到的代码，方便理解整条链路。
+
+---
+
+## 一、入口：测试脚本在做什么
+
+**文件：`test/test_infer.py`**
+
+你运行的命令会执行这个脚本的 `if __name__ == "__main__":` 这一段。
+
+### 1.1 整体流程（简化）
+
+```text
+1. 解析命令行参数（--model、--prompt 等）
+2. 用 HuggingFace 加载 tokenizer + PyTorch 模型，跑一遍推理 → 得到「标准答案」
+3. 打印 PyTorch 的 Tokens 和 Contents（=== Answer ===）
+4. 用 LLAISYS 加载同一份模型，再跑一遍推理 → 得到「你的实现的结果」
+5. 打印 LLAISYS 的 Tokens 和 Contents（=== Your Result ===）
+6. 若加了 --test，会比对两边的 token 序列是否一致
+```
+
+所以脚本做了两件事：**用 PyTorch 跑一遍**、**用你的 C++ 后端（LLAISYS）跑一遍**，并对比结果。
+
+### 1.2 关键代码位置
+
+- **加载 LLAISYS 模型**（第 146 行附近）：
+
+```python
+model = load_llaisys_model(model_path, args.device)
+```
+
+- **LLAISYS 推理**（第 155 行附近）：
+
+```python
+llaisys_tokens, llaisys_output = llaisys_infer(
+ args.prompt, tokenizer, model, max_new_tokens=..., ...
+)
+```
+
+- **`load_llaisys_model`**（第 61 行）：只是调用了 `llaisys.models.Qwen2(model_path, device)`，即你项目里的 Python 封装类。
+- **`llaisys_infer`**（第 66 行起）：
+ - 用 tokenizer 把 prompt 转成 token id 列表 `inputs`
+ - 调用 `model.generate(inputs, max_new_tokens=..., ...)`，得到 `outputs`（token 列表）
+ - 用 tokenizer 把 `outputs` 解码成字符串返回
+
+所以：**测试脚本只负责「加载模型」和「调用 generate」**，真正的模型和生成逻辑在 `llaisys.models.Qwen2` 里。
+
+---
+
+## 二、Python 模型封装：Qwen2 类
+
+**文件：`python/llaisys_py/models/qwen2.py`**
+
+这是你在 Python 里直接用的「Qwen2 模型」：它内部会调 C 接口（通过 ctypes 调用的 dll），**不**用 PyTorch 做推理。
+
+### 2.1 初始化 `__init__`：创建模型 + 加载权重
+
+1. **读 config.json**
+ 从模型目录读 `config.json`，取出 `hidden_size`、`num_hidden_layers`、`num_attention_heads`、`vocab_size` 等，用来填「模型元信息」。
+
+2. **拼出 C 的「元信息」结构体**
+ 用这些配置构造 `LlaisysQwen2Meta`（dtype、层数、头数、词表大小、norm 的 eps、RoPE 的 theta、eos_token_id 等）。
+
+3. **调用 C 接口创建模型**
+ ```python
+ self._model = LIB_LLAISYS.llaisysQwen2ModelCreate(byref(meta), device, None, 0)
+ ```
+ 这里 `LIB_LLAISYS` 就是加载好的 `llaisys.dll`（或 .so），`llaisysQwen2ModelCreate` 是 dll 里导出的 C 函数。C 侧会分配模型结构、权重张量和 KV Cache。
+
+4. **加载权重**
+ - 用 `LIB_LLAISYS.llaisysQwen2ModelWeights(self._model)` 拿到「权重句柄」。
+ - 遍历目录下所有 `*.safetensors`，对每个文件里的每个 key，若在「权重名 → 句柄」的映射里，就：
+ - 用 safetensors 读出张量（numpy 或 torch，bf16 时用 torch 再转 float32）；
+ - 调用 `_numpy_to_backend(arr, handle)` 把数据拷进 C 侧对应的张量（内部会调 `LIB_LLAISYS.tensorLoad`）。
+
+所以：**Python 只负责「读 config + 读权重文件 + 把数据灌进 C 侧」**，模型结构和算子在 C++ 里。
+
+### 2.2 生成 `generate`
+
+```python
+def generate(self, inputs, max_new_tokens=128, ...):
+ tokens = list(inputs) # 例如 [151646, 151644, 15191, 525, ...]
+ for _ in range(max_new_tokens):
+ n = len(tokens)
+ token_arr = (c_int64 * n)(*tokens)
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInfer(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ n,
+ )
+ if next_tok == -1:
+ raise RuntimeError(...)
+ tokens.append(next_tok)
+ if next_tok == self._end_token: # 遇到 EOS 就停
+ break
+ return tokens
+```
+
+含义：
+
+- **每次循环**：把当前的 `tokens`（整段序列）交给 C 的 `llaisysQwen2ModelInfer`，C 会做一次「前向」（见下节），返回**下一个 token 的 id**（`next_tok`）。
+- Python 把 `next_tok` 追加到 `tokens`，若是 EOS 就结束循环，否则继续。
+- 所以：**生成是一个「循环：整段序列 → C 前向 → 取最后一个位置的 next token」** 的过程；C 侧内部会维护 KV Cache，所以不用每次传整段历史，但接口上目前是「传整段、C 内部自己用 cache」。
+
+小结：**Python 的 Qwen2 类 = 配置 + 权重加载 + 循环调用 C 的 Infer，得到 token 序列。**
+
+---
+
+## 三、Python 如何调 C：libllaisys
+
+**目录：`python/llaisys_py/libllaisys/`**
+
+### 3.1 加载 dll
+
+**文件：`python/llaisys_py/libllaisys/__init__.py`**
+
+```python
+def load_shared_library():
+ lib_dir = Path(__file__).parent # 即 libllaisys 目录
+ # Windows: llaisys.dll, Linux: libllaisys.so
+ lib_path = os.path.join(lib_dir, libname)
+ return ctypes.CDLL(str(lib_path))
+
+LIB_LLAISYS = load_shared_library()
+```
+
+之后所有「调 C」都是通过 `LIB_LLAISYS.函数名(...)` 完成。`load_qwen2(LIB_LLAISYS)` 会给这些函数声明参数类型和返回类型，这样 ctypes 才能正确把 Python 对象转成 C 的指针和整数。
+
+### 3.2 Qwen2 的 C 接口声明
+
+**文件：`python/llaisys_py/libllaisys/qwen2.py`**
+
+- **LlaisysQwen2Meta / LlaisysQwen2Weights**：用 ctypes 的 `Structure` 定义成和 C 头文件里一样的结构体，这样 `byref(meta)` 传进去 C 能正确读到。
+- **load_qwen2(lib)**：
+ - `lib.llaisysQwen2ModelCreate.argtypes = [POINTER(LlaisysQwen2Meta), llaisysDeviceType_t, ...]`
+ - `lib.llaisysQwen2ModelInfer.argtypes = [LlaisysQwen2Model_t, POINTER(c_int64), c_size_t]`
+ - `lib.llaisysQwen2ModelInfer.restype = c_int64`
+ 这样 Python 传列表、C 拿到的就是 `int64_t*` 和长度，返回的就是下一个 token 的 id。
+
+所以：**libllaisys 的作用 = 加载 dll + 用 ctypes 声明 C 的接口（结构体、函数签名）**，让 `models/qwen2.py` 能无脑调 `LIB_LLAISYS.llaisysQwen2ModelCreate / llaisysQwen2ModelInfer` 等。
+
+---
+
+## 四、C 接口长什么样（头文件）
+
+**文件：`include/llaisys/models/qwen2.h`**
+
+C 侧只暴露 4 个函数和 2 个结构体（给 Python 用）：
+
+- **LlaisysQwen2Meta**：模型超参（dtype、nlayer、hs、nh、nkvh、dh、di、maxseq、voc、epsilon、theta、end_token）。
+- **LlaisysQwen2Weights**：各权重的「句柄」指针（in_embed、out_embed、各层的 attn/mlp 的 weight/bias）。
+- **llaisysQwen2ModelCreate(meta, device, ...)**：创建模型；内部会分配所有权重张量和 KV Cache。
+- **llaisysQwen2ModelDestroy(model)**：释放模型。
+- **llaisysQwen2ModelWeights(model)**：返回指向「权重结构体」的指针，Python 用来往每个句柄里灌数据。
+- **llaisysQwen2ModelInfer(model, token_ids, ntoken)**：输入当前 token 序列和长度，做一次前向，返回**下一个 token 的 id**（int64）。
+
+所以：**头文件定义的是「模型创建 / 销毁 / 拿权重 / 单步推理」的契约**，实现都在 `src/llaisys/qwen2.cc`。
+
+---
+
+## 五、C++ 实现：模型创建与单步推理
+
+**文件：`src/llaisys/qwen2.cc`**
+
+### 5.1 模型里有什么（LlaisysQwen2Model）
+
+```cpp
+struct LlaisysQwen2Model {
+ LlaisysQwen2Meta meta;
+ LlaisysQwen2Weights weights;
+ std::vector k_caches;
+ std::vector v_caches;
+ size_t cache_len;
+ llaisysDeviceType_t device_type;
+ int device_id;
+};
+```
+
+- **meta**：上面说的超参。
+- **weights**：各层权重的张量（C++ 里的 tensor 对象）。
+- **k_caches / v_caches**：每一层的 K、V 的 cache，解码时复用，避免重复算历史。
+- **cache_len**：当前已经填了多少个时间步的 cache（prefill 后或 decode 每步 +1）。
+
+### 5.2 创建模型：llaisysQwen2ModelCreate
+
+1. 根据 `meta` 里的 nlayer、hs、nh、nkvh、dh、di、voc 等，**创建所有权重张量**（embed、norm、每层的 q/k/v/o、mlp 的 gate/up/down），并挂到 `weights` 里。
+2. 为每一层分配 **KV Cache** 张量（形状约 `[maxseq, nkvh, dh]`），用于后续推理时拼 K、V。
+3. 返回 `LlaisysQwen2Model*`。
+Python 拿到这个指针后，通过 `llaisysQwen2ModelWeights` 拿到每个权重句柄，再用 `tensorLoad` 把 safetensors 里的数据拷进去。
+
+### 5.3 单步推理：llaisysQwen2ModelInfer
+
+这是「你问的推理」的核心：**给定当前 token 序列，算下一个 token**。
+
+1. **判断是 prefill 还是 decode**
+ - `cache_len == 0` → prefill：一次性喂入整段 `token_ids`（例如整段 prompt）。
+ - 否则 → decode：只喂入「最后一个 token」（当前要预测的下一个位置），序列长度 `seq_len = 1`。
+
+2. **准备输入**
+ - 把 token id 拷到设备上的 `token_tensor`（长度 `seq_len`）。
+ - 分配这一轮前向用的临时张量：hidden、normed、q_buf、k_buf、v_buf、q_rope、k_rope、attn_val、o_proj_out、res_buf、gate_buf、up_buf、mlp_buf、down_buf、pos_ids 等。
+
+3. **Embedding**
+ - `llaisys::ops::embedding(hidden, token_tensor, in_embed)`：用 token id 查表得到 hidden state。
+
+4. **逐层 Transformer Block**（`forward_layer`）
+ 对每一层做：
+ - **Attention 分支**
+ - RMSNorm → `linear` 得到 Q、K、V → RoPE → 把当前步的 K、V 写入 KV Cache → 用「当前步的 Q」和「cache 里拼好的 K、V」做 `self_attention` → 再 `linear`（o_proj）→ 残差加回 hidden。
+ - **MLP 分支**
+ - RMSNorm → gate/up 两个 linear → SwiGLU → down linear → 残差加回 hidden。
+ 这样 `hidden` 就更新成这一层的输出。
+
+5. **最后一层之后**
+ - 对 hidden 做 RMSNorm，再用 `out_embed` 做一次 linear，得到 **logits**（形状 `[seq_len, voc]`）。
+ - 取**最后一个位置**的 logits（`last_logit_1d`），做 **argmax**，得到 `next_token`（int64）。
+ - 把 `cache_len` 加上本轮的 `seq_len`（prefill 加整段长度，decode 加 1）。
+ - 返回 `next_token`。
+
+所以：**llaisysQwen2ModelInfer = 一次完整的前向（embed + 所有 layer + 最后一层 norm + 输出层 logits + argmax）**，返回「下一个 token 的 id」。Python 的 `generate` 就是反复调这个函数，直到遇到 EOS 或达到 max_new_tokens。
+
+### 5.4 forward_layer 在做什么（一层 Transformer）
+
+- **Attention**：norm → Q/K/V 投影 → RoPE → 更新 KV Cache → causal self-attention → O 投影 → 残差。
+- **MLP**：norm → gate/up（SwiGLU）→ down → 残差。
+这里用到的都是你在作业 2 里实现的算子：`rms_norm`、`linear`、`rope`、`self_attention`、`swiglu`、`add` 等。
+
+---
+
+## 六、数据流小结（从你运行命令到下一个 token）
+
+1. **test/test_infer.py**
+ 解析参数 → 用 HuggingFace 跑一遍（PyTorch）→ 再 `load_llaisys_model` + `llaisys_infer`。
+
+2. **llaisys_infer**
+ 用 tokenizer 把 prompt 转成 `inputs` → 调用 `model.generate(inputs, ...)`。
+
+3. **python/llaisys_py/models/qwen2.py**
+ - `Qwen2.__init__`：读 config、调 `llaisysQwen2ModelCreate`、再读 safetensors 调 `tensorLoad` 灌权重。
+ - `Qwen2.generate`：循环里每次把当前 `tokens` 传给 `llaisysQwen2ModelInfer`，拿到 `next_tok`，追加到 `tokens`，直到 EOS 或达到 max_new_tokens。
+
+4. **python/llaisys_py/libllaisys/**
+ 加载 `llaisys.dll`，用 ctypes 声明 Qwen2 的 C API（Meta、Weights、Create、Destroy、Weights、Infer），供 qwen2.py 调用。
+
+5. **include/llaisys/models/qwen2.h**
+ 定义 C 的「模型元信息、权重句柄、Create/Destroy/Weights/Infer」接口。
+
+6. **src/llaisys/qwen2.cc**
+ - **Create**：按 meta 分配所有权重张量和 KV Cache。
+ - **Infer**：prefill/decode 分支 → embed → 逐层 forward_layer（attention + mlp）→ 最后一层 norm + 输出层 linear → argmax → 返回 next token id。
+
+7. **算子**
+ embedding、linear、rms_norm、rope、self_attention、swiglu、add、argmax 等都在 `src/ops/` 下，由 `qwen2.cc` 的 `forward_layer` 和 `llaisysQwen2ModelInfer` 调用。
+
+如果你愿意，下一步可以单独挑「某一段」（例如只看 Python 的 generate，或只看 C++ 的 Infer）再逐行讲；或者指定一个文件/函数，我按行号带你看。
diff --git "a/docs/\347\256\200\345\216\206-\351\241\271\347\233\2563-AI\350\201\212\345\244\251Agent.md" "b/docs/\347\256\200\345\216\206-\351\241\271\347\233\2563-AI\350\201\212\345\244\251Agent.md"
new file mode 100644
index 000000000..01cc982e1
--- /dev/null
+++ "b/docs/\347\256\200\345\216\206-\351\241\271\347\233\2563-AI\350\201\212\345\244\251Agent.md"
@@ -0,0 +1,43 @@
+# 简历描述：LLAISYS AI 聊天 Agent 项目
+
+以下内容可直接或稍作修改后写入简历「项目经历」一栏，突出 **Agent / 对话式 AI** 相关能力。
+
+---
+
+## 项目名称（建议）
+
+**基于自研推理引擎的对话式 AI Agent 系统**
+或：**LLM 推理引擎与对话 Agent 服务端到端实现**
+
+---
+
+## 项目描述（一段话版）
+
+参与教育型 AI 系统 LLAISYS 的推理与服务层开发，**独立完成对话式 AI Agent 的完整链路**：从底层随机采样算子（Temperature / Top-K / Top-P）实现与 API 打通，到 HTTP 聊天服务（OpenAI chat-completion 兼容）、再到多轮对话 CLI 与 Web 前端。Agent 基于自研 C++ 推理引擎，支持单用户实时多轮对话与流式输出，为后续多用户推理服务与连续批处理奠定基础。
+
+---
+
+## bullet 点（任选 4～6 条）
+
+- **对话 Agent 采样与推理**：设计并实现随机采样算子（Temperature / Top-K / Top-P），替代原有 argmax，使 Agent 回复更自然、可调；在 C++ 推理管线中集成采样逻辑，经 C API 与 Python 封装贯通至上层调用。
+- **Agent 服务端**：使用 FastAPI 实现 OpenAI chat-completion 风格的 HTTP 接口，支持非流式与 SSE 流式响应；单用户阻塞式处理，保证对话上下文一致，便于后续扩展为多 Agent/多用户服务。
+- **多轮对话与上下文管理**：在 CLI 与 Web 客户端维护完整 `messages` 历史，每次请求将整段对话上下文发给服务端，实现**多轮连续对话**，体现 Agent 的会话记忆与上下文理解能力。
+- **端到端对话体验**：实现命令行聊天客户端（Python + requests）与内嵌 Web 聊天页（HTML/JS + Fetch），用户可连续发消息、收回复，形成完整「人机对话 Agent」闭环。
+- **技术栈**：C++ 推理引擎（自研算子、张量、设备抽象）、Python 模型封装与 HTTP 服务、transformers 分词与对话模板，接口设计兼容 OpenAI，便于与现有 Agent 框架对接。
+- **工程实践**：完成从算子实现、C/Python API 打通、服务端到前端的全链路开发；编写项目总结文档与运行说明，便于复现与后续迭代（多用户、KV-Cache 池等）。
+
+---
+
+## 关键词（便于简历筛选）
+
+对话式 AI · Agent · 多轮对话 · LLM 推理 · 随机采样（Temperature / Top-K / Top-P）· FastAPI · OpenAI API 兼容 · 流式输出（SSE）· C++ / Python · 自研推理引擎
+
+---
+
+## 简短版（空间有限时用）
+
+**LLAISYS 对话 Agent**：实现随机采样算子（Temperature / Top-K / Top-P）并打通 C/Python API；基于 FastAPI 提供 OpenAI 兼容的 chat-completion 服务，支持流式输出；开发 CLI 与 Web 多轮对话前端，完成单用户对话式 AI Agent 的端到端链路。
+
+---
+
+按需选用「一段话版」+ 部分 bullet，或仅用「简短版」即可突出 Agent 与对话式 AI 能力。
diff --git "a/docs/\351\241\271\347\233\2563-\351\232\217\346\234\272\351\207\207\346\240\267\344\270\216API\346\211\223\351\200\232\346\200\273\347\273\223.md" "b/docs/\351\241\271\347\233\2563-\351\232\217\346\234\272\351\207\207\346\240\267\344\270\216API\346\211\223\351\200\232\346\200\273\347\273\223.md"
new file mode 100644
index 000000000..25ed304c0
--- /dev/null
+++ "b/docs/\351\241\271\347\233\2563-\351\232\217\346\234\272\351\207\207\346\240\267\344\270\216API\346\211\223\351\200\232\346\200\273\347\273\223.md"
@@ -0,0 +1,272 @@
+# 项目 #3 随机采样与 API 打通 — 实现总结
+
+本文档说明：为支持「构建 AI 聊天机器人」而实现的**随机采样算子**、**采样参数（Temperature / Top-K / Top-P）**以及**从 C 到 Python 的 API 打通**的完整流程、做了哪些修改、以及为什么这样做。
+
+---
+
+## 一、背景与目标
+
+### 1.1 项目 #3 的要求
+
+项目 #3 是「构建 AI 聊天机器人」。要实现能与用户实时对话的聊天机器人，需要：
+
+- **更自然的回复**：不能总是选「概率最高的那个词」（argmax），否则生成会过于死板、重复。
+- **随机采样**：按概率从候选词中抽样，使每次回复有一定随机性。
+- **可调参数**：用 Temperature、Top-K、Top-P 控制随机程度和候选范围。
+
+因此，需要：
+
+1. 实现一个**随机采样算子**，支持 Temperature、Top-K、Top-P。
+2. 在推理路径中**用该算子替代原来的 argmax**（在需要随机时）。
+3. **打通 API**：从 Python 的 `generate(top_k, top_p, temperature, seed)` 一路传到 C++ 的采样算子。
+
+### 1.2 原先的流程（仅 argmax）
+
+在实现前，推理流程是：
+
+```
+Python: model.generate(inputs, max_new_tokens, top_k=1, top_p=0.8, temperature=0.8)
+ ↓ 循环
+ llaisysQwen2ModelInfer(model, token_ids, ntoken) // 只有 3 个参数
+ ↓
+C++: 前向 → 得到 logits → argmax(logits) → 返回 next_token
+```
+
+- `top_k`、`top_p`、`temperature` 在 Python 有参数，但**没有传给 C**，C 侧只做 argmax，所以「未在 C 侧使用」。
+- 要支持随机采样，就必须：**扩展 C 的 Infer 接口** → **在 C++ 里根据参数选择 sample 或 argmax** → **实现 sample 算子**。
+
+---
+
+## 二、整体流程（实现后）
+
+实现后的数据流如下。
+
+```
+用户 / 测试脚本
+ ↓
+Python: model.generate(inputs, max_new_tokens, top_k=50, top_p=0.8, temperature=0.8, seed=0)
+ ↓ 每个 token 循环
+ llaisysQwen2ModelInfer(model, token_ids, ntoken, temperature, top_k, top_p, seed)
+ ↓
+C API (qwen2.h)
+ ↓
+qwen2.cc: llaisysQwen2ModelInfer(...)
+ ├─ 前向：embed → layers → norm → linear → 得到 logits
+ ├─ 取最后一个位置的 logits → last_logit_1d
+ ├─ 判断：use_sampling = (temperature 有效 且 (top_k>1 或 0 1` 或 `0 < top_p < 1`）时用采样，否则用 argmax；这样 Python 侧 `top_k=1` 或 `temperature≈0` 即退化为原来的贪心行为，兼容现有测试脚本（如 `--test`）。
+
+**为什么 GPU 时要把 logits 拷到 CPU 再 sample？**
+
+- 当前只实现了 CPU 版 sample；若模型在 GPU，logits 在显存，不能直接在 GPU 上调用现有 sample 实现。因此先 D2H 拷贝到 CPU 的临时 tensor，在 CPU 上 sample，再把得到的索引（一个 int64）返回，这样无需改 Python 接口即可支持 GPU 模型 + 随机采样。
+
+---
+
+### 3.4 Python 绑定与 generate 传参
+
+| 文件 | 修改内容 |
+|------|----------|
+| `python/llaisys_py/libllaisys/ops.py` | 为 `llaisysSample` 声明 argtypes（out_idx, logits, c_float, c_int, c_float, c_ulonglong），并加入 `load_ops` |
+| `python/llaisys_py/ops.py` | 增加 `Ops.sample(out_idx, logits, temperature=1.0, top_k=0, top_p=0.0, seed=0)`，内部调 `llaisysSample` |
+| `python/llaisys_py/libllaisys/qwen2.py` | `llaisysQwen2ModelInfer` 的 argtypes 增加 `c_float, c_int, c_float, c_ulonglong`（temperature, top_k, top_p, seed） |
+| `python/llaisys_py/models/qwen2.py` | ① `generate` 增加参数 `seed=0`；② 调用 `llaisysQwen2ModelInfer` 时传入 `c_float(temperature), c_int(top_k), c_float(top_p), c_ulonglong(seed)`；③ 文档字符串改为说明「temperature、top_k、top_p 会传入 C 侧」，并说明 seed 含义 |
+
+**为什么要传 seed？**
+
+- 可复现：同一 prompt、同一组参数下，相同 seed 得到相同序列，便于调试和测试。
+- `seed=0` 表示「每次用随机设备」，不保证复现；非 0 则用该种子初始化 `std::mt19937`。
+
+**为什么用 c_ulonglong 表示 seed？**
+
+- C 侧用 `unsigned long long`，与 64 位种子一致；Python 侧用 `c_ulonglong` 和 `c_ulonglong(seed)` 与之对应，避免跨平台位数问题。
+
+---
+
+### 3.5 单元测试（可选）
+
+| 文件 | 作用 |
+|------|------|
+| `test/ops/sample.py` | 不依赖完整模型：① 用 numpy 构造 logits，拷贝到 llaisys Tensor，调 `Ops.sample`，检查返回索引在 [0, voc) 内；② 固定 logits、极小 temperature，检查退化为 argmax（返回最大 logit 的下标）。若环境无 torch 或 DLL 问题，可单独用该脚本验证 sample 绑定与行为。 |
+
+---
+
+## 四、为什么这样设计（简要）
+
+1. **算子独立（sample op）**：采样逻辑集中在一个 op 里，支持多种 dtype、Temperature/Top-K/Top-P，将来加 GPU 或其它采样方式（如 beam）只需改/加算子，不动 qwen2 前向大逻辑。
+2. **一个 Infer 接口**：通过参数控制「采样 vs 贪心」，调用方简单；测试用 `top_k=1` 即与原来行为一致。
+3. **参数从 Python 直通 C++**：`generate(...)` 的 `top_k/top_p/temperature/seed` 原样传到 C，C 侧真正使用，注释中「未在 C 侧使用」的问题被消除。
+4. **GPU 兼容**：在未实现 GPU sample 前，用「logits 拷到 CPU → CPU sample」保证 GPU 模型也能做随机采样，为后续优化留空间。
+
+---
+
+## 五、涉及文件一览
+
+| 类型 | 路径 |
+|------|------|
+| 新增 | `src/ops/sample/op.hpp`, `src/ops/sample/op.cpp` |
+| 修改 | `include/llaisys/ops.h`, `include/llaisys/models/qwen2.h` |
+| 修改 | `src/llaisys/ops.cc`, `src/llaisys/qwen2.cc` |
+| 修改 | `python/llaisys_py/libllaisys/ops.py`, `python/llaisys_py/libllaisys/qwen2.py`, `python/llaisys_py/ops.py`, `python/llaisys_py/models/qwen2.py` |
+| 可选 | `test/ops/sample.py` |
+
+构建时 `xmake` 会扫描 `src/ops/*/op.cpp`，因此无需改 `xmake.lua`，sample 会自动参与编译。
+
+---
+
+## 六、如何运行并检查效果
+
+在项目根目录（即包含 `xmake.lua` 的目录）下按顺序执行即可。
+
+### 6.1 编译并安装到 Python 包
+
+```bash
+# 编译 C++ 与动态库
+xmake build
+
+# 将生成的 llaisys.dll（或 libllaisys.so）复制到 python/llaisys_py/libllaisys/
+# 这样 Python 的 import llaisys_py 会用到刚编译的版本
+xmake install
+```
+
+### 6.2 方式一：完整推理测试（推荐，需模型）
+
+依赖：Python 环境已安装 `torch`、`transformers`、`huggingface_hub` 等（见项目 README）。
+
+**不指定 `--model` 时**：会自动从 Hugging Face 下载 `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B`（约数 GB），首次较慢。
+
+```bash
+# 进入项目根目录后，让 Python 找到本项目的 test 和 llaisys（若未 pip install -e python）
+# Windows PowerShell 示例：
+$env:PYTHONPATH = "python;test"
+# Linux/macOS：
+# export PYTHONPATH=python:test
+
+# 1）贪心模式：与 PyTorch 对齐，用于验证实现正确性
+python test/test_infer.py --test
+
+# 若已有本地模型目录，可指定路径避免重复下载：
+# python test/test_infer.py --model D:/models/DeepSeek-R1-Distill-Qwen-1.5B --test
+```
+
+- 通过条件：终端打印的「Your Result」与「Answer」的 **Tokens 序列完全一致**，且无 AssertionError。
+- 含义：LLAISYS 在 `top_k=1, temperature=1.0` 下走 argmax 分支，与 PyTorch 一致。
+
+**随机采样模式**：同一命令去掉 `--test`，会使用默认 `top_k=50, top_p=0.8, temperature=1.0`，走采样分支。
+
+```bash
+# 2）随机采样：每次运行结果一般不同
+python test/test_infer.py
+
+# 自定义 prompt、生成长度、采样参数示例：
+# python test/test_infer.py --prompt "你好，请介绍一下自己" --max_steps 64 --temperature 0.8 --top_k 40 --top_p 0.9
+```
+
+- 检查方式：多运行几次，观察「Your Result」的文本是否有所变化；或修改 `test_infer.py` 里传给 `generate` 的 `seed` 固定，两次用同一 seed 应得到相同结果。
+
+### 6.3 方式二：仅测 sample 算子（不跑完整模型）
+
+不加载大模型，只验证「从 logits 按概率采样一个 token」的算子和 Python 绑定是否正确。
+
+```bash
+# 在项目根目录，确保使用本项目的 llaisys（例如设置 PYTHONPATH=python）
+# Windows PowerShell：
+$env:PYTHONPATH = "python"
+python test/ops/sample.py
+```
+
+- 通过条件：终端输出 `Sample op tests passed!`，且无报错。
+- 说明：该脚本会 `import llaisys_py`，若项目配置为通过 `llaisys` 包拉取 `models`，则会间接导入 `torch`；若本机 torch 有 DLL 等问题，可能在此报错，此时以方式一在能跑通的环境中验证即可。
+
+### 6.4 可选：用 pip 安装本项目后再跑
+
+若已在本机用 pip 安装过本项目的 Python 包，可直接用：
+
+```bash
+pip install -e python
+xmake build
+xmake install
+python test/test_infer.py --test
+```
+
+这样无需每次设置 `PYTHONPATH`，`import llaisys_py` 会使用当前项目下的包和刚安装的 DLL。
+
+### 6.5 小结
+
+| 目的 | 命令示例 |
+|----------------|----------|
+| 验证贪心对齐 | `python test/test_infer.py --test`（可加 `--model `） |
+| 看随机采样效果 | `python test/test_infer.py`（可加 `--prompt`、`--temperature` 等） |
+| 只测 sample op | `PYTHONPATH=python python test/ops/sample.py` |
+
+以上即为「随机采样 + Temperature / Top-K / Top-P + API 打通」的完整流程、实现内容、设计原因与运行检查方式。
diff --git "a/docs/\351\241\271\347\233\2563\345\256\214\346\210\220\346\240\270\345\257\271\344\270\216\351\241\271\347\233\2564\344\273\273\345\212\241\346\213\206\350\247\243.md" "b/docs/\351\241\271\347\233\2563\345\256\214\346\210\220\346\240\270\345\257\271\344\270\216\351\241\271\347\233\2564\344\273\273\345\212\241\346\213\206\350\247\243.md"
new file mode 100644
index 000000000..afb212bd1
--- /dev/null
+++ "b/docs/\351\241\271\347\233\2563\345\256\214\346\210\220\346\240\270\345\257\271\344\270\216\351\241\271\347\233\2564\344\273\273\345\212\241\346\213\206\350\247\243.md"
@@ -0,0 +1,268 @@
+# 项目#3 完成核对 & 项目#4 多用户推理服务 — 任务拆解
+
+## 一、项目#3：是否全部完成？
+
+按 README_ZN.md「项目#3：构建 AI 聊天机器人」逐条核对如下。
+
+| 条款 | README 要求 | 当前实现 | 结论 |
+|------|-------------|----------|------|
+| **随机采样** | 实现随机采样算子，尽量支持 **Temperature**、**Top-K**、**Top-P** | `src/ops/sample/`：`sample(out_idx, logits, temperature, top_k, top_p, seed)`；qwen2 推理与 `/v1/chat/completions` 请求体均支持 `temperature` / `top_k` / `top_p` | ✅ 完成 |
+| **搭建聊天服务器** | Python 前端实现 HTTP 服务器（如 FastAPI）；接口遵循 **OpenAI chat-completion API**；尽量支持**流式输出**；可先假设单用户、请求阻塞直到处理完成 | FastAPI；`POST /v1/chat/completions`（messages / stream / max_tokens 等）；`stream: true` 时 SSE 流式返回 | ✅ 完成 |
+| **交互式聊天 UI** | 实现 UI，向服务器发请求并接收回复；可为**命令行**或 **Web**；能**连续发送消息**与机器人保持对话 | **Web**：`/chat` 内嵌 HTML+JS，连续发消息、流式显示回复；**CLI**：`chat_cli.py` 多轮对话循环 | ✅ 完成 |
+| **（可选）会话管理** | 多对话、切换；修改历史问题让 AI 重新生成；扩展 UI 支持；实现**支持前缀匹配的 KV-Cache 池**复用 | `/v1/sessions` CRUD；`/v1/sessions/{id}/regenerate`（`from_message_index`）；Web 侧栏会话列表、新建/删除/切换、编辑某条消息后从此处重新生成；`_kv_pool` 以 `(session_id, user_message_index)` 为 key，按 prefix_len 匹配复用并 LRU 淘汰 | ✅ 完成 |
+
+**结论：项目#3 已全部完成（含可选「会话管理」）。**
+
+---
+
+## 二、项目#4：按 README 逐条拆成的设计与实现任务
+
+README 要求（摘要）：
+
+- **前置**：完成项目#3 并实现流式输出。（✅ 已满足）
+- **支持多用户**：推理服务同时为多用户服务，请求随时到来；请求加入**请求池/队列**，用**单独的循环线程/进程**处理。
+- **连续批处理**：为最大化吞吐做**批处理**；每个请求长度不同，需**连续迭代级批处理**——每轮从池中取若干请求组成 batch，执行一次**批量推理**，再把未完成请求放回池；推理时尽量用**批量矩阵乘法**；每个请求绑定**不同 KV-Cache**，实现**支持前缀匹配的 KV-Cache 池**复用。
+
+下面按「设计」与「实现」拆成可执行任务。
+
+---
+
+### 2.1 支持多用户
+
+| 序号 | 类型 | 任务描述 |
+|------|------|----------|
+| 4.1.1 | 设计 | **请求抽象**：定义「推理请求」结构（如 request_id、session_id、messages、stream 与否、callback/queue 用于回写结果），与当前「单请求直接调 model」解耦。 |
+| 4.1.2 | 设计 | **请求入口**：规定所有到达的 chat-completion 请求（含流式/非流式）先**不入队即处理**，而是生成上述请求对象并放入**全局请求队列**，立即返回 202 或通过异步句柄等待结果（或保持当前「同步等待」但改为从队列取结果的语义，二选一需定）。 |
+| 4.1.3 | 设计 | **队列与并发**：选定「有界/无界队列」、是否按 session 或用户做公平性（可选）；确定**单个 worker 线程/进程**从队列取请求并调用当前推理逻辑，避免多线程同时写同一模型实例（除非后续改为多实例或批处理引擎）。 |
+| 4.1.4 | 实现 | **队列与 Worker**：实现请求队列（如 `queue.Queue` 或 `asyncio.Queue`）；实现一个或多个 **worker 线程/协程**，循环：取请求 → 解析 session/messages → 调现有 `_stream_response` / 非流式生成 → 将结果写回请求关联的 response 或 SSE 通道。 |
+| 4.1.5 | 实现 | **API 接入**：将 `POST /v1/chat/completions` 改为：把请求入队 + 阻塞或异步等待 worker 完成并返回/流式写出；保证现有 `/v1/sessions`、regenerate 等仍能与该队列模型兼容（如 session 与 request 的对应关系）。 |
+| 4.1.6 | 实现 | **背压与超时**：队列满时返回 503 或 429；请求等待超时（可选）返回 504；必要时限制单用户/单 session 的并发请求数。 |
+
+---
+
+### 2.2 连续批处理
+
+| 序号 | 类型 | 任务描述 |
+|------|------|----------|
+| 4.2.1 | 设计 | **迭代级调度**：定义「一轮」= 从请求池中取出若干**未完成**请求，组成当前 batch；对 batch 执行**一次**迭代（一次 prefill 或一次 decode）；把本轮新生成的 token 写回各请求，未到 EOS 的请求放回池，下一轮再参与组 batch。 |
+| 4.2.2 | 设计 | **batch 组成规则**：规定每轮最大 batch_size（如 4/8）、最大 total tokens 或最大序列长度和，避免 OOM；规定如何从池中**选择**请求（FIFO、按已生成 token 数等）；不同请求的 prompt 长度与当前生成长度不同，需在数据结构上能表达「每个请求当前 state」。 |
+| 4.2.3 | 设计 | **请求状态**：每个请求维护：input_ids 或已编码的 prompt、当前生成中的 token 序列、KV-Cache 句柄（或指向 KV 池中某条的 key）、是否已 EOS、stream 回调等；池中只放「未完成」请求。 |
+| 4.2.4 | 设计 | **批量推理接口**：当前引擎为「单序列 in → 单 token out」。需设计：要么在 **C++/Python 层** 实现「多序列 batch in → 多 token out」（每序列一个 next token），要么用多线程/多进程每序列调一次现有接口并在上层做「一轮」的同步（后者不是真正批量矩阵乘，但可先实现调度与池化）。 |
+| 4.2.5 | 实现 | **批量 Prefill**：若支持 batch prefill，在引擎侧实现多序列一次 prefill（padding 或 packed），得到每个请求的 KV 写入各自 slot；若不支持，则用「多请求轮流 prefill + 共用或分片 KV 池」并记录每请求的 prefix_len。 |
+| 4.2.6 | 实现 | **批量 Decode**：每轮将当前 batch 内每个请求的「最后一个 token」组成一个 batch 输入，调用批量 next_token（若引擎支持）；将返回的 token 列表按请求写回，更新各请求状态，EOS 的请求移出池并回调完成。 |
+| 4.2.7 | 实现 | **批量矩阵乘法**：若 C++ 侧有 batch 维度的 linear/attention，在 batch 维度上一次算多个序列，以利用批量矩阵乘加速；若暂无，4.2.5/4.2.6 可先按「多序列轮流调用现有单序列 API」实现，再在后续迭代中替换为真正的 batch 算子。 |
+
+**说明**：4.2 的落地不能仅靠 Python 层调度；必须配合 C++ 的 batch 算子与多 slot KV-Cache。详见第四节「4.2 落地瓶颈与改造路径」。
+
+---
+
+### 2.3 KV-Cache 池（多用户/批处理场景）
+
+| 序号 | 类型 | 任务描述 |
+|------|------|----------|
+| 4.3.1 | 设计 | **池 key 与粒度**：当前池 key = `(session_id, user_message_index)`。多用户/多请求下，每个**请求**对应一个逻辑「会话前缀」；若一个 session 可能被多个请求交错（如同一会话的 regenerate 与续写），需规定 key 是否仍为 (session_id, user_msg_idx)，或增加 request_id 以区分同一前缀的不同请求。 |
+| 4.3.2 | 设计 | **前缀匹配**：与项目#3 一致，命中条件为「当前请求的 messages 前缀」对应的 token 长度与池中某条的 prefix_len 一致；多请求共享同一 session 时，需保证并发安全（读 blob、写回新 blob 的原子性/锁）。 |
+| 4.3.3 | 设计 | **每请求独立 KV**：批处理时每个请求有独立 KV-Cache（或独立 slot）；池中存的是「某 (session_id, user_idx) 的 blob」，被某请求复用时，将该 blob 导入到**该请求当前占用的 KV slot**，再继续 decode。 |
+| 4.3.4 | 实现 | **池的并发安全**：对 `_kv_pool` 的 get/put 加锁（或使用线程安全结构），避免 worker 与 API 线程并发写同一池。 |
+| 4.3.5 | 实现 | **batch 与池的配合**：prefill 前查池；若命中则 import_kv_cache 再只对 suffix 做 prefill；未命中则全量 prefill；每请求完成后按 (session_id, user_message_index) 写回池，并支持 LRU 淘汰（沿用当前逻辑即可）。 |
+
+---
+
+### 2.4 流式与 API 兼容
+
+| 序号 | 类型 | 任务描述 |
+|------|------|----------|
+| 4.4.1 | 设计 | **流式在 batch 下的行为**：同一轮 decode 可能推进多个请求各 1 个 token；每个 token 需写回对应请求的 SSE 通道；保证「每个请求的 SSE 顺序」与生成顺序一致，且不互相串线。 |
+| 4.4.2 | 实现 | **流式回写**：worker 或 batch 循环内，每生成一个请求的一个 token，即向该请求的 `StreamingResponse` 或等价通道 push 一块 SSE；若使用队列，需在请求对象上挂接「可写流」句柄。 |
+| 4.4.3 | 实现 | **OpenAI 兼容**：保持现有 `/v1/chat/completions` 请求/响应格式（含 stream、choices[].delta 等），确保多用户 + 队列/批处理后，每个客户端仍收到与自己 request 对应的完整流或非流响应。 |
+
+---
+
+### 2.5 测试与可观测性
+
+| 序号 | 类型 | 任务描述 |
+|------|------|----------|
+| 4.5.1 | 实现 | **多用户压测**：编写脚本或用例：并发多个客户端（或多次异步请求）同时发 chat-completion，验证响应正确归属、无串线、无崩溃；验证队列满或超时时的行为。 |
+| 4.5.2 | 实现 | **批处理正确性**：在 batch_size>1 时，对比「同一组请求分别单序列推理」与「同一组请求进 batch 推理」的输出是否一致（在相同 seed 下）；验证 KV 池命中率与 prefix_len 匹配逻辑。 |
+| 4.5.3 | 可选 | **监控**：队列长度、平均等待时间、每轮 batch 大小、KV 池命中率等指标，便于调优与排障。 |
+
+---
+
+## 三、任务依赖关系（建议实现顺序）
+
+```
+4.1.1 → 4.1.2 → 4.1.3 → 4.1.4 → 4.1.5 → 4.1.6 （先打通「多用户 + 队列 + 单 worker」）
+ ↓
+4.3.1 → 4.3.2 → 4.3.3 → 4.3.4 （池的并发与多请求语义）
+ ↓
+4.2.1 → 4.2.2 → 4.2.3 → 4.2.4 → 4.2.5 → 4.2.6 （连续批处理调度与状态）
+ ↓
+4.2.7 （批量矩阵乘，可选后续优化）
+ ↓
+4.4.1 → 4.4.2 → 4.4.3 （流式在 batch 下的回写与兼容）
+ ↓
+4.3.5 （batch 与 KV 池的完整配合）
+ ↓
+4.5.1 → 4.5.2 → 4.5.3 （测试与可观测性）
+```
+
+可先完成 4.1.x 和 4.3.4，使「多用户排队、单序列逐个推理、KV 池并发安全」上线，再迭代 4.2.x 的批处理与 4.4.x 的流式回写。
+
+---
+
+## 四、项目#4 首阶段实现记录（当前已完成）
+
+已在 `python/llaisys_py/server/app.py` 中实现：
+
+| 任务 | 实现要点 |
+|------|----------|
+| **4.1 支持多用户** | 请求入队 `_request_queue`（容量由 `LLAISYS_REQUEST_QUEUE_MAX` 控制，默认 64）；单 worker 线程 `_worker_loop` 在 `startup` 时启动，循环取任务并执行推理；`POST /v1/chat/completions` 先入队再阻塞/流式从 `response_queue` 取结果；队列满时返回 503。 |
+| **4.3.4 池的并发安全** | `_sessions_lock`、`_kv_pool_lock` 保护会话与 KV 池的读写；所有 `_sessions` / `_kv_pool` 访问均在对应锁内。 |
+| **推理互斥** | `_inference_lock`：worker 处理 chat 与 regenerate 端点使用模型时均先获取该锁，保证同一时刻仅一处跑推理。 |
+
+**环境变量**：`LLAISYS_REQUEST_QUEUE_MAX`（默认 64）、`LLAISYS_KV_POOL_MAX`（沿用项目#3）。
+
+**未实现**：连续批处理（4.2）、迭代级 batch 与批量矩阵乘（4.2.5–4.2.7）；当前为「多用户排队 + 单请求逐个推理」。
+
+### 4.2 落地瓶颈与改造路径（连续批处理必须攻克的部分）
+
+当前 4.1 的实现**本质上是串行处理（Serial Processing）**：虽然外部有队列管理多用户，但底层模型同一时刻只为**一个**请求计算，算力利用率低，无法达到 README 中「最大化吞吐量」的要求。2.2 节规划的 4.2.1–4.2.7 思路与业界主流推理框架（vLLM、TGI）的 **Continuous Batching（连续批处理 / 迭代级调度）** 一致；要将其落地，**仅改 Python 层队列不够，必须修改 C++ 底层算子和 KV-Cache 架构**。4.1 可保留作容灾降级（请求多到 batch 塞不下时排队），但 4.2 及相关 C++ 改造需尽快启动。
+
+以下为必须攻克的**三个核心技术点**与改造路径。
+
+---
+
+#### 1. C++ 底层算子支持 Batch 维度（批量矩阵乘法）
+
+| 现状 | 改造目标 |
+|------|----------|
+| 当前 C++ 算子（如 `linear`、`self_attention`）输入多为单序列形状：`[seq_len, d]` 或 `[1, d]`。 | 将算子输入扩展为 **`[batch_size, seq_len, d]`**（或等价 packed 表示）。 |
+
+**为何关键**：CPU/GPU 做矩阵乘法时，算 $1 \times d$ 与算 $8 \times d$ 的耗时差距很小（Memory Bound）。把多个请求的 token 拼成 batch 一起算，是提升吞吐的根本。
+**改造点**：在 `src/ops/` 下为 linear、self_attention、rms_norm、rope、swiglu 等增加 batch 维或批量 GEMM 路径；Qwen2 模型前向接口支持「多序列一次 forward」而非单序列。
+
+---
+
+#### 2. 状态隔离的批量 KV-Cache（Batched KV-Cache）
+
+| 现状 | 改造目标 |
+|------|----------|
+| `_kv_pool` 在 **Python** 层存序列化 blob，每次推理前全量 `import_kv_cache` / `export_kv_cache`；底层 C++ 模型为**单状态**，一次只服务一个请求。 | **C++** 侧维护支持多 **Slot（槽位）** 的 KV-Cache 池，例如一块形状为 `[max_batch_size, max_seq_len, n_head, head_dim]` 的连续内存（或等价分 slot 管理）。 |
+
+**改造点**：
+
+- Python 层不再在请求间传递巨大 blob；改为向 C++ 引擎**申请 `slot_id`**。
+- 推理时由 Python 告诉 C++：「本轮用第 0、2、5 号 slot 的 KV-Cache 做一次 Batch Compute」。
+- 请求结束时释放对应 slot，供新请求复用。
+
+---
+
+#### 3. Python 层的迭代级调度器（The Scheduler）
+
+| 现状 | 改造目标 |
+|------|----------|
+| 单 worker 循环：`while True: req = queue.get(); process(req)`，一次只处理一个请求。 | Worker 改为持续运行的 **Engine Loop（引擎心跳）**，每轮做三件事：**收集 → 生成 → 分发与清理**。 |
+
+**执行逻辑**：
+
+1. **收集阶段**：若 `running_requests` 未满且内存/显存允许，从 `waiting_requests` 拉新请求进 batch，执行一次 **Batched Prefill**（多序列一起 prefill，KV 写入各自 slot）。
+2. **生成阶段**：取出所有 `running_requests` 的「当前最后一个 token」，拼成 batch，调用 C++ 的 **Batched Decode**（如 `model.batched_next_token([t1, t2, t3, ...])`），得到本轮每个请求的一个新 token。
+3. **分发与清理**：把生成 token 按请求写回对应 SSE 流；若某请求生成了 EOS，将其移出 `running_requests` 并释放其 KV slot。
+
+这样每一轮都是一次「多请求同时推进一个 token」，才能真正提高吞吐。
+
+---
+
+#### 小结
+
+| 层级 | 必须改造内容 |
+|------|--------------|
+| **C++ 算子** | 支持 batch 维或批量 GEMM（linear、attention 等）。 |
+| **C++ KV-Cache** | 多 slot、按 slot_id 读写；Python 只传 slot_id，不传 blob。 |
+| **Python 调度** | 单 worker 改为 Engine Loop：Batched Prefill → Batched Decode → 按请求回写 SSE 并释放 EOS 请求的 slot。 |
+
+4.1 的队列与 503 背压可保留；4.2 落地依赖上述三层改造，不能仅靠 Python 队列改造完成。
+
+---
+
+## 五、多用户含义、是否需要登录、如何测试
+
+### 5.1 多用户是如何实现的？
+
+当前实现的是**请求级排队**，而不是「用户账号 + 登录」：
+
+- **任意客户端**（浏览器、`chat_cli`、curl、其他程序）都可以直接调用 `POST /v1/chat/completions`，无需登录、无需 token。
+- 每个请求进入**全局请求队列** `_request_queue`，由**唯一 worker 线程**按 FIFO 顺序逐个处理。
+- 因此「多用户」= **多个并发请求**：多人在不同浏览器/终端同时发问时，请求会排队，worker 一次只处理一个，处理完再处理下一个；其他请求在队列里等待，不会抢占模型，也不会报错。
+
+会话用 **session_id** 区分：前端先 `POST /v1/sessions` 拿到一个 UUID，之后在 chat 请求里带 `session_id`，服务端把该轮对话追到对应会话并更新 KV 池。**没有用户账号体系**，不区分「谁」在请求，只区分「哪个会话」。
+
+### 5.2 需要登录吗？
+
+**不需要。** 没有登录、鉴权或 API Key；只要网络能访问服务，任何人都可以发请求。适合内网/本机演示；若要对公网开放，需要自己在前面加反向代理鉴权（如 Nginx + 认证、或 API Key 中间件）。
+
+### 5.3 如何测试多用户/排队？
+
+1. **多端同时发请求（推荐）**
+ - 浏览器：开两个标签页，都打开 `/chat`，在 A 里发一条长问题（如「写一篇 200 字短文」），立刻在 B 里也发一条。B 会等 A 的回复流结束后才开始收到自己的流，说明请求在排队。
+ - 命令行：开两个终端都运行 `python -m llaisys_py.server.chat_cli --base-url http://127.0.0.1:8002`，在一个里发消息后马上在另一个里发，现象同上。
+
+2. **用脚本压测并发与队列满**
+ 项目内提供了 `test/test_multi_user_chat.py`（见下），可并发发 N 个非流式请求，看是否都成功、响应是否串行；并可测队列满时是否返回 503。
+
+---
+
+## 六、当前实现核对（与文档/代码逐项对照）
+
+以下为对「文档中描述的内容」与「当前代码」的逐项核对结果。
+
+### 6.1 项目#3（文档第一节）
+
+| 文档条款 | 核对结果 |
+|----------|----------|
+| 随机采样 Temperature / Top-K / Top-P | ✅ `src/ops/sample/` 与请求体参数已支持 |
+| FastAPI + OpenAI chat-completion + 流式 | ✅ `POST /v1/chat/completions`，`stream: true` 走 SSE |
+| Web `/chat` + CLI `chat_cli.py` 连续对话 | ✅ 均已实现 |
+| 会话 CRUD、regenerate、KV 池前缀匹配 + LRU | ✅ `/v1/sessions`、`/regenerate`、`_kv_pool` 与锁一致 |
+
+**结论：项目#3 与文档一致，已全部实现。**
+
+### 6.2 项目#4 首阶段（文档第四节 + 2.1 / 2.3 部分）
+
+| 文档任务 | 核对结果 |
+|----------|----------|
+| 4.1.1 请求抽象 | ✅ 入队结构含 `response_queue`、`session_id`、`messages`、`stream`、`max_tokens` 等 |
+| 4.1.2 请求入口入队 | ✅ chat 请求先 `put_nowait` 入队，再同步等待 `response_queue` 或流式迭代 |
+| 4.1.3 有界队列 + 单 worker | ✅ `queue.Queue(maxsize=_REQUEST_QUEUE_MAX)`，单线程 `_worker_loop` |
+| 4.1.4 队列与 Worker 实现 | ✅ `_request_queue.get()` → 推理 → 结果写入 `response_queue` |
+| 4.1.5 API 接入 | ✅ `POST /v1/chat/completions` 入队并等待；sessions / regenerate 仍可用 |
+| 4.1.6 背压 | ✅ 队列满时 `put_nowait` 抛 `queue.Full`，返回 503；非流式有 `get(timeout=300)`，文档中 504 为可选，当前未单独返回 504 |
+| 4.3.4 池的并发安全 | ✅ `_sessions_lock`、`_kv_pool_lock` 保护所有读写 |
+| 推理互斥（worker + regenerate） | ✅ `_inference_lock` 在 worker 与 regenerate 流式/非流式路径中统一使用 |
+| 4.5.1 多用户压测脚本 | ✅ `test/test_multi_user_chat.py` 存在，支持并发与 `--test-queue-full` |
+
+### 6.3 连续批处理与相关任务（已实现）
+
+以下任务已实现（环境变量 `LLAISYS_USE_ENGINE_LOOP=1` 且 `LLAISYS_MAX_BATCH_SIZE>=2` 时启用）：
+
+| 任务 | 实现要点 |
+|------|----------|
+| **4.2** 连续批处理 | `python/llaisys_py/server/engine.py`：Engine 主循环 Prefill（含 KV 池命中时 suffix prefill）→ Batched Decode（C++ `llaisysQwen2ModelBatchedDecode`）→ 按请求写回 out_queue；多 slot KV-Cache 由 C++ 支持（`qwen2.h` / `qwen2.cc`）。 |
+| **4.3.5** batch 与池配合 | Engine 构造时注入 `get_kv` / `put_kv`；prefill 前查池，命中则 `import_kv_cache_slot` 后只对 suffix prefill；请求完成后 `export_kv_cache_slot` 写回池；C API `ExportKVCacheSlot` / `ImportKVCacheSlot` 已实现。 |
+| **4.4** 流式在 batch 下回写 | 每轮 decode 后按请求将 next_token 放入该请求的 `out_queue`；`/v1/chat/completions` 流式时从 `req_state.out_queue` 读取并推送 SSE，与生成顺序一致、不串线。 |
+| **4.5.2** 批处理正确性测试 | `test/test_batch_correctness.py`：同一组 prompt、相同 seed、确定性采样下，顺序 `model.generate` 与 Engine 批量调度输出 token 序列一致。 |
+| **4.5.3** 监控 | `GET /v1/metrics`：返回 `request_queue_size`、`request_queue_max`、`kv_pool_size`、`kv_pool_max`；启用 Engine 时含 `engine.pending_queue_size`、`running_count`、`free_slots_count`、`max_batch_size`。 |
+
+**说明**：4.2.7 批量矩阵乘（C++ 侧真正的 batch GEMM）仍为可选后续优化；当前 Batched Decode 为多 slot 循环调用单 slot 推理。
+
+### 6.4 文档未写但已实现的修复
+
+| 项 | 说明 |
+|----|------|
+| 会话切换竞态 | `switchSession` 中增加 `if (currentSessionId !== sessionId) return`，避免晚到的会话响应覆盖当前选中会话内容。 |
+
+### 6.5 总结
+
+- **项目#3**：与文档第一节描述一致，**已全部实现**。
+- **项目#4**：4.1（多用户队列）、4.2（连续批处理 Engine）、4.3.4（池并发安全）、4.3.5（batch 与 KV 池配合）、4.4（流式 batch 回写）、4.5.1（多用户压测）、4.5.2（批处理正确性测试）、4.5.3（监控端点）**均已实现**。启用方式：`LLAISYS_USE_ENGINE_LOOP=1`、`LLAISYS_MAX_BATCH_SIZE=4`（或其它 ≥2）。
+- 额外已做：会话页切换时的竞态修复，文档在「六」中已补录。
diff --git a/include/llaisys.h b/include/llaisys.h
index 73ca7eead..96dcbeac1 100644
--- a/include/llaisys.h
+++ b/include/llaisys.h
@@ -10,11 +10,11 @@
#endif
#ifdef __cplusplus
-#define __C extern "C"
+#define LLAISYS_EXTERN_C extern "C"
#include
#include
#else
-#define __C
+#define LLAISYS_EXTERN_C
#include
#include
#endif
diff --git a/include/llaisys/models/qwen2.h b/include/llaisys/models/qwen2.h
index 7054626d4..9df91db8b 100644
--- a/include/llaisys/models/qwen2.h
+++ b/include/llaisys/models/qwen2.h
@@ -3,10 +3,13 @@
#include "../tensor.h"
-__C {
+LLAISYS_EXTERN_C {
struct LlaisysQwen2Meta {
llaisysDataType_t dtype;
size_t nlayer, hs, nh, nkvh, dh, di, maxseq, voc;
+ size_t max_batch_size; /* 连续批处理：KV-Cache 槽位数，1=单序列（默认） */
+ int tp_rank; /* 张量并行 rank，0..tp_world_size-1；默认 0 */
+ int tp_world_size; /* 张量并行 world size，1=非分布式；默认 1 */
float epsilon, theta;
int64_t end_token;
};
@@ -35,8 +38,71 @@ __C {
__export void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model * model);
+ /** 将输出层权重（out_norm_w、out_embed）拷到 CPU 并缓存；GPU 推理时最后一层在 CPU 上算以规避 GPU 输出层异常。应在权重加载完成后调用一次。 */
+ __export void llaisysQwen2ModelCacheOutputLayerOnCPU(struct LlaisysQwen2Model * model);
+
+ /** 将所有权重与 KV cache 拷到 CPU 并缓存；GPU 推理时整次前向在 CPU 上执行以规避 GPU 算子异常。应在权重加载完成后调用一次。 */
+ __export void llaisysQwen2ModelCacheAllWeightsOnCPU(struct LlaisysQwen2Model * model);
+
__export struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model * model);
- __export int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken);
+ /** 返回当前已写入 cache 的长度（prefill/suffix prefill/decode 后更新） */
+ __export size_t llaisysQwen2ModelGetCacheLen(struct LlaisysQwen2Model * model);
+
+ /** 返回存储前缀长度为 prefix_len 的 KV cache 所需字节数（供 Export/Import 分配缓冲区） */
+ __export size_t llaisysQwen2ModelGetKVCacheBytes(struct LlaisysQwen2Model * model, size_t prefix_len);
+
+ /** 将当前 cache 内容导出到 ptr_out（调用方需预先分配 GetKVCacheBytes(model, cache_len) 字节） */
+ __export void llaisysQwen2ModelExportKVCache(struct LlaisysQwen2Model * model, void * ptr_out);
+
+ /** 从 ptr_in 导入前缀长度为 prefix_len 的 KV，并设 cache_len = prefix_len；之后可做 suffix prefill */
+ __export void llaisysQwen2ModelImportKVCache(struct LlaisysQwen2Model * model, const void * ptr_in, size_t prefix_len);
+
+ /** 将 cache_len 置 0，用于新请求全量 prefill 前清掉上一轮状态（单 slot 或 slot_id=0） */
+ __export void llaisysQwen2ModelResetKVCache(struct LlaisysQwen2Model * model);
+
+ /** 将指定 slot 的 cache_len 置 0；仅当 meta.max_batch_size > 1 时有效 */
+ __export void llaisysQwen2ModelResetKVCacheSlot(struct LlaisysQwen2Model * model, size_t slot_id);
+
+ /** 将指定 slot 的 KV cache 导出到 ptr_out（调用方需分配 GetKVCacheBytes(model, GetCacheLenSlot(model, slot_id)) 字节） */
+ __export void llaisysQwen2ModelExportKVCacheSlot(struct LlaisysQwen2Model * model, size_t slot_id, void * ptr_out);
+
+ /** 从 ptr_in 导入前缀长度为 prefix_len 的 KV 到指定 slot，并设该 slot 的 cache_len = prefix_len；之后可做 suffix prefill */
+ __export void llaisysQwen2ModelImportKVCacheSlot(struct LlaisysQwen2Model * model, size_t slot_id, const void * ptr_in, size_t prefix_len);
+
+ /** 返回指定 slot 的 cache_len；当 max_batch_size==1 时 slot_id 忽略，返回当前唯一 cache_len */
+ __export size_t llaisysQwen2ModelGetCacheLenSlot(struct LlaisysQwen2Model * model, size_t slot_id);
+
+ /**
+ * 单步推理（支持多 slot）。
+ * 当 meta.max_batch_size==1 时 slot_id 被忽略，行为与 llaisysQwen2ModelInfer 一致。
+ * 当 max_batch_size>1 时，使用指定 slot 的 KV-Cache 进行 prefill/decode，并更新该 slot 的 cache_len。
+ */
+ __export int64_t llaisysQwen2ModelInferWithSlot(struct LlaisysQwen2Model * model, size_t slot_id, int64_t * token_ids, size_t ntoken, float temperature, int top_k, float top_p, unsigned long long seed);
+
+ __export int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken, float temperature, int top_k, float top_p, unsigned long long seed);
+
+ /**
+ * 诊断用：前 (gpu_up_to_layer+1) 层在 GPU 上跑，其余层与输出在 CPU 上跑；需已调用 CacheAllWeightsOnCPU。
+ * gpu_up_to_layer < 0：整次前向在 CPU；=0：仅 embedding 在 GPU；=1：embedding+layer0 在 GPU；依此类推。
+ * 返回 next_token。用于逐层对比找出首个产生错误结果的 GPU 层。
+ */
+ __export int64_t llaisysQwen2ModelInferHybrid(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken, float temperature, int top_k, float top_p, unsigned long long seed, int gpu_up_to_layer);
+
+ /**
+ * 批量 Decode：一次传入多 slot 的当前 token，返回多个 next token。
+ * 用于连续批处理调度器；过渡期内部以 for 循环调用单 slot 单 token 推理，后续可替换为真正的 Batched 算子。
+ * 要求：n_batch <= meta.max_batch_size，且每个 slot_id 有效；每个 slot 此时应处于 decode 阶段（cache_len > 0）。
+ */
+ __export void llaisysQwen2ModelBatchedDecode(
+ struct LlaisysQwen2Model * model,
+ const size_t * slot_ids, /* 长度为 n_batch，例如 [0, 2, 5] */
+ const int64_t * token_ids, /* 长度为 n_batch，每个 slot 的当前 token */
+ size_t n_batch,
+ int64_t * out_next_tokens, /* 长度为 n_batch 的输出 */
+ float temperature,
+ int top_k,
+ float top_p,
+ unsigned long long seed);
}
#endif // LLAISYS_MODELS_QWEN2_H
diff --git a/include/llaisys/nccl_comm.h b/include/llaisys/nccl_comm.h
new file mode 100644
index 000000000..05f54d96e
--- /dev/null
+++ b/include/llaisys/nccl_comm.h
@@ -0,0 +1,53 @@
+/**
+ * NCCL 通信封装：用于张量并行下的 AllReduce / AllGather。
+ * 仅当 ENABLE_NVIDIA_API 且 ENABLE_NCCL 时有效；否则接口为空实现或返回错误。
+ */
+#ifndef LLAISYS_NCCL_COMM_H
+#define LLAISYS_NCCL_COMM_H
+
+#include "../llaisys.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/** NCCL 唯一 ID 字节数（用于进程间广播，再调用 llaisysNcclInitRank） */
+#define LLAISYS_NCCL_UNIQUE_ID_BYTES 128
+
+/**
+ * 在 rank 0 上调用，将唯一 ID 写入 buffer（至少 LLAISYS_NCCL_UNIQUE_ID_BYTES 字节），
+ * 然后通过文件/MPI/等广播给其他 rank，供 llaisysNcclInitRank 使用。
+ */
+__export void llaisysNcclGetUniqueId(void *buffer);
+
+/**
+ * 每个进程调用一次：用 rank 0 广播得到的 unique_id 初始化本进程的 NCCL 通信器。
+ * rank in [0, world_size), world_size >= 1。GPU 由调用方在此前通过 setDevice 等设定。
+ */
+__export int llaisysNcclInitRank(int rank, int world_size, const void *unique_id);
+
+/**
+ * AllReduce：所有 rank 的 sendbuf 做 sum，结果写入各 rank 的 recvbuf。
+ * count 为元素个数，dtype 为元素类型。stream 为 CUDA 流（void*），CPU 推理时传 NULL。
+ */
+__export int llaisysNcclAllReduce(const void *sendbuf, void *recvbuf, size_t count,
+ llaisysDataType_t dtype, void *stream);
+
+/**
+ * AllGather：每个 rank 提供 sendbuf（count_per_rank 个元素），
+ * 结果 recvbuf 为所有 rank 的 sendbuf 按 rank 顺序拼接（总长 count_per_rank * world_size）。
+ */
+__export int llaisysNcclAllGather(const void *sendbuf, void *recvbuf, size_t count_per_rank,
+ llaisysDataType_t dtype, void *stream);
+
+/** 释放 NCCL 通信器，进程退出前调用。 */
+__export void llaisysNcclDestroy(void);
+
+/** 返回最近一次 NCCL/CUDA 错误的描述（静态缓冲区，仅用于调试）。 */
+__export const char *llaisysNcclGetLastError(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* LLAISYS_NCCL_COMM_H */
diff --git a/include/llaisys/ops.h b/include/llaisys/ops.h
index ddb3be246..4d9a0631d 100644
--- a/include/llaisys/ops.h
+++ b/include/llaisys/ops.h
@@ -3,7 +3,7 @@
#include "tensor.h"
-__C {
+LLAISYS_EXTERN_C {
__export void llaisysAdd(llaisysTensor_t c, llaisysTensor_t a, llaisysTensor_t b);
__export void llaisysArgmax(llaisysTensor_t max_idx, llaisysTensor_t max_val, llaisysTensor_t vals);
__export void llaisysEmbedding(llaisysTensor_t out, llaisysTensor_t index, llaisysTensor_t weight);
@@ -13,6 +13,7 @@ __C {
__export void llaisysROPE(llaisysTensor_t out, llaisysTensor_t in, llaisysTensor_t pos_ids, float theta);
__export void llaisysSelfAttention(llaisysTensor_t attn_val, llaisysTensor_t q, llaisysTensor_t k, llaisysTensor_t v, float scale);
__export void llaisysSwiGLU(llaisysTensor_t out, llaisysTensor_t gate, llaisysTensor_t up);
+ __export void llaisysSample(llaisysTensor_t out_idx, llaisysTensor_t logits, float temperature, int top_k, float top_p, unsigned long long seed);
}
#endif
diff --git a/include/llaisys/ops_nvidia.h b/include/llaisys/ops_nvidia.h
new file mode 100644
index 000000000..bc22e8c72
--- /dev/null
+++ b/include/llaisys/ops_nvidia.h
@@ -0,0 +1,35 @@
+/**
+ * NVIDIA CUDA 算子声明，供 op.cpp 在 LLAISYS_DEVICE_NVIDIA 分支调用。
+ * 仅在使用 ENABLE_NVIDIA_API 编译时由 op 引用；实现位于 src/ops/nvidia/ops_nvidia.cu。
+ */
+#ifndef LLAISYS_OPS_NVIDIA_H
+#define LLAISYS_OPS_NVIDIA_H
+
+#include "../llaisys.h"
+#include
+
+#ifdef ENABLE_NVIDIA_API
+
+namespace llaisys::ops::nvidia {
+
+void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t dtype, size_t numel);
+
+void embedding(std::byte *out, const std::byte *weight, const int64_t *index, size_t num_index, size_t embed_dim, size_t vocab_size, size_t elem_size);
+
+void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, llaisysDataType_t dtype, size_t B, size_t M, size_t K);
+
+void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t vals_dtype, size_t numel);
+
+void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, llaisysDataType_t dtype, size_t rows, size_t dim, float eps);
+
+void rope(std::byte *out, const std::byte *in, const int64_t *pos_ids, llaisysDataType_t dtype, size_t seq_len, size_t num_heads, size_t head_dim, float theta);
+
+void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t dtype, size_t numel);
+
+void self_attention(std::byte *out, const std::byte *q, const std::byte *k, const std::byte *v, llaisysDataType_t dtype, size_t qlen, size_t kvlen, size_t num_heads, size_t nkvh, size_t head_dim, float scale);
+
+} // namespace llaisys::ops::nvidia
+
+#endif // ENABLE_NVIDIA_API
+
+#endif // LLAISYS_OPS_NVIDIA_H
diff --git a/include/llaisys/runtime.h b/include/llaisys/runtime.h
index d8e6f66f1..b79bd3c3e 100644
--- a/include/llaisys/runtime.h
+++ b/include/llaisys/runtime.h
@@ -3,7 +3,7 @@
#include "../llaisys.h"
-__C {
+LLAISYS_EXTERN_C {
// Runtime API Functions
// Device
typedef int (*get_device_count_api)();
diff --git a/include/llaisys/tensor.h b/include/llaisys/tensor.h
index 76f13fbc3..dfca9c6be 100644
--- a/include/llaisys/tensor.h
+++ b/include/llaisys/tensor.h
@@ -3,7 +3,7 @@
#include "../llaisys.h"
-__C {
+LLAISYS_EXTERN_C {
typedef struct LlaisysTensor *llaisysTensor_t;
__export llaisysTensor_t tensorCreate(
diff --git a/pic/83627ff8dcdc8d366b068858b07ca2c0.png b/pic/83627ff8dcdc8d366b068858b07ca2c0.png
new file mode 100644
index 000000000..34ca0126a
Binary files /dev/null and b/pic/83627ff8dcdc8d366b068858b07ca2c0.png differ
diff --git a/python/llaisys/models/qwen2.py b/python/llaisys/models/qwen2.py
deleted file mode 100644
index 0d07b0b21..000000000
--- a/python/llaisys/models/qwen2.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from typing import Sequence
-from ..libllaisys import LIB_LLAISYS
-from ..libllaisys import DeviceType
-
-from pathlib import Path
-import safetensors
-
-
-class Qwen2:
-
- def __init__(self, model_path, device: DeviceType = DeviceType.CPU):
- # TODO: Implement model constructor
-
- model_path = Path(model_path)
-
- for file in sorted(model_path.glob("*.safetensors")):
- data_ = safetensors.safe_open(file, framework="numpy", device="cpu")
- for name_ in data_.keys():
- ## TODO: load the model weights
- pass
-
- def generate(
- self,
- inputs: Sequence[int],
- max_new_tokens: int = None,
- top_k: int = 1,
- top_p: float = 0.8,
- temperature: float = 0.8,
- ):
-
- # TODO: Implement generate function
-
- return []
diff --git a/python/llaisys/__init__.py b/python/llaisys_py/__init__.py
similarity index 61%
rename from python/llaisys/__init__.py
rename to python/llaisys_py/__init__.py
index de8d99f48..9a83c711a 100644
--- a/python/llaisys/__init__.py
+++ b/python/llaisys_py/__init__.py
@@ -1,20 +1,25 @@
-from .runtime import RuntimeAPI
-from .libllaisys import DeviceType
-from .libllaisys import DataType
-from .libllaisys import MemcpyKind
-from .libllaisys import llaisysStream_t as Stream
-from .tensor import Tensor
-from .ops import Ops
-from . import models
-from .models import *
-
-__all__ = [
- "RuntimeAPI",
- "DeviceType",
- "DataType",
- "MemcpyKind",
- "Stream",
- "Tensor",
- "Ops",
- "models",
-]
+# 在任何会加载 libllaisys/CUDA 的 import 之前修正：CUDA_VISIBLE_DEVICES="" 会导致 cudaGetDeviceCount() 一直为 0
+import os
+if os.environ.get("CUDA_VISIBLE_DEVICES") == "":
+ os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+from .runtime import RuntimeAPI
+from .libllaisys import DeviceType
+from .libllaisys import DataType
+from .libllaisys import MemcpyKind
+from .libllaisys import llaisysStream_t as Stream
+from .tensor import Tensor
+from .ops import Ops
+from . import models
+from .models import *
+
+__all__ = [
+ "RuntimeAPI",
+ "DeviceType",
+ "DataType",
+ "MemcpyKind",
+ "Stream",
+ "Tensor",
+ "Ops",
+ "models",
+]
diff --git a/python/llaisys/libllaisys/__init__.py b/python/llaisys_py/libllaisys/__init__.py
similarity index 64%
rename from python/llaisys/libllaisys/__init__.py
rename to python/llaisys_py/libllaisys/__init__.py
index f536fb527..3edb90847 100644
--- a/python/llaisys/libllaisys/__init__.py
+++ b/python/llaisys_py/libllaisys/__init__.py
@@ -12,6 +12,9 @@
from .tensor import llaisysTensor_t
from .tensor import load_tensor
from .ops import load_ops
+from .qwen2 import load_qwen2
+from .qwen2 import LlaisysQwen2Meta, LlaisysQwen2Weights, LlaisysQwen2Model_t
+from . import nccl_comm
def load_shared_library():
@@ -31,6 +34,16 @@ def load_shared_library():
if not os.path.isfile(lib_path):
raise FileNotFoundError(f"Shared library not found: {lib_path}")
+ # 预加载 OpenMP 运行时，避免 libllaisys.so 出现 undefined symbol: omp_get_thread_num
+ if sys.platform.startswith("linux"):
+ try:
+ ctypes.CDLL("libgomp.so.1", mode=ctypes.RTLD_GLOBAL)
+ except OSError:
+ pass # 若系统无 libgomp 或已链接进 .so，忽略
+
+ lib_path_abs = os.path.abspath(lib_path)
+ if os.environ.get("LLAISYS_DEBUG"):
+ print(f"[LLAISYS] Loading shared library: {lib_path_abs}")
return ctypes.CDLL(str(lib_path))
@@ -38,7 +51,8 @@ def load_shared_library():
load_runtime(LIB_LLAISYS)
load_tensor(LIB_LLAISYS)
load_ops(LIB_LLAISYS)
-
+load_qwen2(LIB_LLAISYS)
+nccl_comm.load_nccl(LIB_LLAISYS)
__all__ = [
"LIB_LLAISYS",
@@ -52,4 +66,7 @@ def load_shared_library():
"llaisysMemcpyKind_t",
"MemcpyKind",
"llaisysStream_t",
+ "LlaisysQwen2Meta",
+ "LlaisysQwen2Weights",
+ "LlaisysQwen2Model_t",
]
diff --git a/python/llaisys/libllaisys/llaisys_types.py b/python/llaisys_py/libllaisys/llaisys_types.py
similarity index 93%
rename from python/llaisys/libllaisys/llaisys_types.py
rename to python/llaisys_py/libllaisys/llaisys_types.py
index c5a0b4679..60e136dbf 100644
--- a/python/llaisys/libllaisys/llaisys_types.py
+++ b/python/llaisys_py/libllaisys/llaisys_types.py
@@ -1,63 +1,63 @@
-import ctypes
-from enum import IntEnum
-
-
-# Device Type enum
-class DeviceType(IntEnum):
- CPU = 0
- NVIDIA = 1
- COUNT = 2
-
-
-llaisysDeviceType_t = ctypes.c_int
-
-
-# Data Type enum
-class DataType(IntEnum):
- INVALID = 0
- BYTE = 1
- BOOL = 2
- I8 = 3
- I16 = 4
- I32 = 5
- I64 = 6
- U8 = 7
- U16 = 8
- U32 = 9
- U64 = 10
- F8 = 11
- F16 = 12
- F32 = 13
- F64 = 14
- C16 = 15
- C32 = 16
- C64 = 17
- C128 = 18
- BF16 = 19
-
-
-llaisysDataType_t = ctypes.c_int
-
-
-# Memory Copy Kind enum
-class MemcpyKind(IntEnum):
- H2H = 0
- H2D = 1
- D2H = 2
- D2D = 3
-
-
-llaisysMemcpyKind_t = ctypes.c_int
-
-# Stream type (opaque pointer)
-llaisysStream_t = ctypes.c_void_p
-
-__all__ = [
- "llaisysDeviceType_t",
- "DeviceType",
- "llaisysDataType_t",
- "DataType",
- "llaisysMemcpyKind_t",
- "MemcpyKind",
- "llaisysStream_t",
-]
+import ctypes
+from enum import IntEnum
+
+
+# Device Type enum
+class DeviceType(IntEnum):
+ CPU = 0
+ NVIDIA = 1
+ COUNT = 2
+
+
+llaisysDeviceType_t = ctypes.c_int
+
+
+# Data Type enum
+class DataType(IntEnum):
+ INVALID = 0
+ BYTE = 1
+ BOOL = 2
+ I8 = 3
+ I16 = 4
+ I32 = 5
+ I64 = 6
+ U8 = 7
+ U16 = 8
+ U32 = 9
+ U64 = 10
+ F8 = 11
+ F16 = 12
+ F32 = 13
+ F64 = 14
+ C16 = 15
+ C32 = 16
+ C64 = 17
+ C128 = 18
+ BF16 = 19
+
+
+llaisysDataType_t = ctypes.c_int
+
+
+# Memory Copy Kind enum
+class MemcpyKind(IntEnum):
+ H2H = 0
+ H2D = 1
+ D2H = 2
+ D2D = 3
+
+
+llaisysMemcpyKind_t = ctypes.c_int
+
+# Stream type (opaque pointer)
+llaisysStream_t = ctypes.c_void_p
+
+__all__ = [
+ "llaisysDeviceType_t",
+ "DeviceType",
+ "llaisysDataType_t",
+ "DataType",
+ "llaisysMemcpyKind_t",
+ "MemcpyKind",
+ "llaisysStream_t",
+]
diff --git a/python/llaisys_py/libllaisys/nccl_comm.py b/python/llaisys_py/libllaisys/nccl_comm.py
new file mode 100644
index 000000000..31348c5af
--- /dev/null
+++ b/python/llaisys_py/libllaisys/nccl_comm.py
@@ -0,0 +1,73 @@
+"""ctypes 绑定：NCCL 通信（项目#5 张量并行）。仅当编译时启用 ENABLE_NCCL 时符号存在。"""
+from ctypes import c_int, c_void_p, c_size_t, c_char_p, create_string_buffer
+import os
+
+# 与 include/llaisys/nccl_comm.h 一致
+LLAISYS_NCCL_UNIQUE_ID_BYTES = 128
+
+
+def load_nccl(lib):
+ """为 lib 绑定 NCCL 相关符号；若未编译 NCCL 则部分可能缺失。"""
+ try:
+ lib.llaisysNcclGetUniqueId.argtypes = [c_void_p]
+ lib.llaisysNcclGetUniqueId.restype = None
+
+ lib.llaisysNcclInitRank.argtypes = [c_int, c_int, c_void_p]
+ lib.llaisysNcclInitRank.restype = c_int
+
+ lib.llaisysNcclAllReduce.argtypes = [
+ c_void_p, c_void_p, c_size_t,
+ c_int, # llaisysDataType_t
+ c_void_p, # stream
+ ]
+ lib.llaisysNcclAllReduce.restype = c_int
+
+ lib.llaisysNcclAllGather.argtypes = [
+ c_void_p, c_void_p, c_size_t,
+ c_int,
+ c_void_p,
+ ]
+ lib.llaisysNcclAllGather.restype = c_int
+
+ lib.llaisysNcclDestroy.argtypes = []
+ lib.llaisysNcclDestroy.restype = None
+
+ lib.llaisysNcclGetLastError.argtypes = []
+ lib.llaisysNcclGetLastError.restype = c_char_p
+ return True
+ except AttributeError:
+ return False
+
+
+def get_unique_id(lib):
+ """在 rank 0 上调用，返回 LLAISYS_NCCL_UNIQUE_ID_BYTES 字节的 bytes，供广播给其他 rank。"""
+ buf = create_string_buffer(LLAISYS_NCCL_UNIQUE_ID_BYTES)
+ lib.llaisysNcclGetUniqueId(buf)
+ return bytes(buf.raw)
+
+
+def init_rank(lib, rank: int, world_size: int, unique_id: bytes) -> int:
+ """每个进程调用一次；unique_id 来自 rank 0 的 get_unique_id()。返回 0 成功，-1 失败。"""
+ if len(unique_id) < LLAISYS_NCCL_UNIQUE_ID_BYTES:
+ return -1
+ # Python 3 中 create_string_buffer(n).raw 为不可变 bytes，不能切片赋值；用内容初始化 buffer
+ data = (unique_id[:LLAISYS_NCCL_UNIQUE_ID_BYTES]).ljust(LLAISYS_NCCL_UNIQUE_ID_BYTES, b"\x00")
+ buf = create_string_buffer(data)
+ return lib.llaisysNcclInitRank(rank, world_size, buf)
+
+
+def get_last_error(lib):
+ """返回 C 端记录的最近一次 NCCL/CUDA 错误（调试用）。"""
+ try:
+ p = lib.llaisysNcclGetLastError()
+ return p.decode("utf-8") if p else ""
+ except AttributeError:
+ return ""
+
+
+def destroy(lib):
+ """进程退出前释放 NCCL 通信器。"""
+ try:
+ lib.llaisysNcclDestroy()
+ except AttributeError:
+ pass
diff --git a/python/llaisys/libllaisys/ops.py b/python/llaisys_py/libllaisys/ops.py
similarity index 79%
rename from python/llaisys/libllaisys/ops.py
rename to python/llaisys_py/libllaisys/ops.py
index 5be095eff..8049f88f2 100644
--- a/python/llaisys/libllaisys/ops.py
+++ b/python/llaisys_py/libllaisys/ops.py
@@ -1,5 +1,5 @@
from .tensor import llaisysTensor_t
-from ctypes import c_float
+from ctypes import c_float, c_int, c_ulonglong
def load_ops(lib):
lib.llaisysAdd.argtypes = [llaisysTensor_t, llaisysTensor_t, llaisysTensor_t]
@@ -34,3 +34,13 @@ def load_ops(lib):
lib.llaisysSwiGLU.argtypes = [llaisysTensor_t, llaisysTensor_t, llaisysTensor_t]
lib.llaisysSwiGLU.restype = None
+
+ lib.llaisysSample.argtypes = [
+ llaisysTensor_t, # out_idx
+ llaisysTensor_t, # logits
+ c_float, # temperature
+ c_int, # top_k
+ c_float, # top_p
+ c_ulonglong, # seed
+ ]
+ lib.llaisysSample.restype = None
diff --git a/python/llaisys_py/libllaisys/qwen2.py b/python/llaisys_py/libllaisys/qwen2.py
new file mode 100644
index 000000000..e70eda042
--- /dev/null
+++ b/python/llaisys_py/libllaisys/qwen2.py
@@ -0,0 +1,162 @@
+"""ctypes bindings for Qwen2 C API."""
+from ctypes import (
+ POINTER,
+ Structure,
+ c_float,
+ c_int,
+ c_int64,
+ c_size_t,
+ c_ulonglong,
+ c_void_p,
+)
+from .llaisys_types import llaisysDataType_t, llaisysDeviceType_t
+from .tensor import llaisysTensor_t
+
+
+class LlaisysQwen2Meta(Structure):
+ _fields_ = [
+ ("dtype", llaisysDataType_t),
+ ("nlayer", c_size_t),
+ ("hs", c_size_t),
+ ("nh", c_size_t),
+ ("nkvh", c_size_t),
+ ("dh", c_size_t),
+ ("di", c_size_t),
+ ("maxseq", c_size_t),
+ ("voc", c_size_t),
+ ("max_batch_size", c_size_t), # 连续批处理槽位数，1=单序列
+ ("tp_rank", c_int), # 张量并行 rank，默认 0
+ ("tp_world_size", c_int), # 张量并行 world size，1=非分布式
+ ("epsilon", c_float),
+ ("theta", c_float),
+ ("end_token", c_int64),
+ ]
+
+
+class LlaisysQwen2Weights(Structure):
+ _fields_ = [
+ ("in_embed", llaisysTensor_t),
+ ("out_embed", llaisysTensor_t),
+ ("out_norm_w", llaisysTensor_t),
+ ("attn_norm_w", POINTER(llaisysTensor_t)),
+ ("attn_q_w", POINTER(llaisysTensor_t)),
+ ("attn_q_b", POINTER(llaisysTensor_t)),
+ ("attn_k_w", POINTER(llaisysTensor_t)),
+ ("attn_k_b", POINTER(llaisysTensor_t)),
+ ("attn_v_w", POINTER(llaisysTensor_t)),
+ ("attn_v_b", POINTER(llaisysTensor_t)),
+ ("attn_o_w", POINTER(llaisysTensor_t)),
+ ("mlp_norm_w", POINTER(llaisysTensor_t)),
+ ("mlp_gate_w", POINTER(llaisysTensor_t)),
+ ("mlp_up_w", POINTER(llaisysTensor_t)),
+ ("mlp_down_w", POINTER(llaisysTensor_t)),
+ ]
+
+
+LlaisysQwen2Model_t = c_void_p
+
+
+def load_qwen2(lib):
+ lib.llaisysQwen2ModelCreate.argtypes = [
+ POINTER(LlaisysQwen2Meta),
+ llaisysDeviceType_t,
+ POINTER(c_int),
+ c_int,
+ ]
+ lib.llaisysQwen2ModelCreate.restype = LlaisysQwen2Model_t
+
+ lib.llaisysQwen2ModelDestroy.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelDestroy.restype = None
+
+ lib.llaisysQwen2ModelCacheOutputLayerOnCPU.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelCacheOutputLayerOnCPU.restype = None
+
+ lib.llaisysQwen2ModelCacheAllWeightsOnCPU.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelCacheAllWeightsOnCPU.restype = None
+
+ lib.llaisysQwen2ModelWeights.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelWeights.restype = POINTER(LlaisysQwen2Weights)
+
+ lib.llaisysQwen2ModelGetCacheLen.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelGetCacheLen.restype = c_size_t
+
+ lib.llaisysQwen2ModelGetKVCacheBytes.argtypes = [LlaisysQwen2Model_t, c_size_t]
+ lib.llaisysQwen2ModelGetKVCacheBytes.restype = c_size_t
+
+ lib.llaisysQwen2ModelExportKVCache.argtypes = [LlaisysQwen2Model_t, c_void_p]
+ lib.llaisysQwen2ModelExportKVCache.restype = None
+
+ lib.llaisysQwen2ModelImportKVCache.argtypes = [
+ LlaisysQwen2Model_t,
+ c_void_p,
+ c_size_t,
+ ]
+ lib.llaisysQwen2ModelImportKVCache.restype = None
+
+ lib.llaisysQwen2ModelResetKVCache.argtypes = [LlaisysQwen2Model_t]
+ lib.llaisysQwen2ModelResetKVCache.restype = None
+
+ lib.llaisysQwen2ModelResetKVCacheSlot.argtypes = [LlaisysQwen2Model_t, c_size_t]
+ lib.llaisysQwen2ModelResetKVCacheSlot.restype = None
+
+ lib.llaisysQwen2ModelExportKVCacheSlot.argtypes = [LlaisysQwen2Model_t, c_size_t, c_void_p]
+ lib.llaisysQwen2ModelExportKVCacheSlot.restype = None
+
+ lib.llaisysQwen2ModelImportKVCacheSlot.argtypes = [
+ LlaisysQwen2Model_t,
+ c_size_t,
+ c_void_p,
+ c_size_t,
+ ]
+ lib.llaisysQwen2ModelImportKVCacheSlot.restype = None
+
+ lib.llaisysQwen2ModelGetCacheLenSlot.argtypes = [LlaisysQwen2Model_t, c_size_t]
+ lib.llaisysQwen2ModelGetCacheLenSlot.restype = c_size_t
+
+ lib.llaisysQwen2ModelInferWithSlot.argtypes = [
+ LlaisysQwen2Model_t,
+ c_size_t, # slot_id
+ POINTER(c_int64),
+ c_size_t,
+ c_float,
+ c_int,
+ c_float,
+ c_ulonglong,
+ ]
+ lib.llaisysQwen2ModelInferWithSlot.restype = c_int64
+
+ lib.llaisysQwen2ModelInfer.argtypes = [
+ LlaisysQwen2Model_t,
+ POINTER(c_int64),
+ c_size_t,
+ c_float, # temperature
+ c_int, # top_k
+ c_float, # top_p
+ c_ulonglong, # seed
+ ]
+ lib.llaisysQwen2ModelInfer.restype = c_int64
+
+ lib.llaisysQwen2ModelInferHybrid.argtypes = [
+ LlaisysQwen2Model_t,
+ POINTER(c_int64),
+ c_size_t,
+ c_float,
+ c_int,
+ c_float,
+ c_ulonglong,
+ c_int, # gpu_up_to_layer: -1=all CPU, 0=only embed GPU, 1=embed+layer0 GPU, ...
+ ]
+ lib.llaisysQwen2ModelInferHybrid.restype = c_int64
+
+ lib.llaisysQwen2ModelBatchedDecode.argtypes = [
+ LlaisysQwen2Model_t,
+ POINTER(c_size_t), # slot_ids
+ POINTER(c_int64), # token_ids
+ c_size_t, # n_batch
+ POINTER(c_int64), # out_next_tokens
+ c_float,
+ c_int,
+ c_float,
+ c_ulonglong,
+ ]
+ lib.llaisysQwen2ModelBatchedDecode.restype = None
diff --git a/python/llaisys/libllaisys/runtime.py b/python/llaisys_py/libllaisys/runtime.py
similarity index 97%
rename from python/llaisys/libllaisys/runtime.py
rename to python/llaisys_py/libllaisys/runtime.py
index 3e5b8be5b..d69ecec5f 100644
--- a/python/llaisys/libllaisys/runtime.py
+++ b/python/llaisys_py/libllaisys/runtime.py
@@ -1,48 +1,48 @@
-import ctypes
-from ctypes import c_void_p, c_size_t, c_int, Structure, CFUNCTYPE
-from .llaisys_types import *
-
-# Define function pointer types
-get_device_count_api = CFUNCTYPE(c_int)
-set_device_api = CFUNCTYPE(None, c_int)
-device_synchronize_api = CFUNCTYPE(None)
-
-create_stream_api = CFUNCTYPE(llaisysStream_t)
-destroy_stream_api = CFUNCTYPE(None, llaisysStream_t)
-stream_synchronize_api = CFUNCTYPE(None, llaisysStream_t)
-
-malloc_device_api = CFUNCTYPE(c_void_p, c_size_t)
-free_device_api = CFUNCTYPE(None, c_void_p)
-malloc_host_api = CFUNCTYPE(c_void_p, c_size_t)
-free_host_api = CFUNCTYPE(None, c_void_p)
-
-memcpy_sync_api = CFUNCTYPE(None, c_void_p, c_void_p, c_size_t, llaisysMemcpyKind_t)
-memcpy_async_api = CFUNCTYPE(None, c_void_p, c_void_p, c_size_t, llaisysMemcpyKind_t, llaisysStream_t)
-
-
-# Define the struct matching LlaisysRuntimeAPI
-class LlaisysRuntimeAPI(Structure):
- _fields_ = [
- ("get_device_count", get_device_count_api),
- ("set_device", set_device_api),
- ("device_synchronize", device_synchronize_api),
- ("create_stream", create_stream_api),
- ("destroy_stream", destroy_stream_api),
- ("stream_synchronize", stream_synchronize_api),
- ("malloc_device", malloc_device_api),
- ("free_device", free_device_api),
- ("malloc_host", malloc_host_api),
- ("free_host", free_host_api),
- ("memcpy_sync", memcpy_sync_api),
- ("memcpy_async", memcpy_async_api),
- ]
-
-
-# Load shared library
-def load_runtime(lib):
- # Declare API function prototypes
- lib.llaisysGetRuntimeAPI.argtypes = [llaisysDeviceType_t]
- lib.llaisysGetRuntimeAPI.restype = ctypes.POINTER(LlaisysRuntimeAPI)
-
- lib.llaisysSetContextRuntime.argtypes = [llaisysDeviceType_t, c_int]
- lib.llaisysSetContextRuntime.restype = None
+import ctypes
+from ctypes import c_void_p, c_size_t, c_int, Structure, CFUNCTYPE
+from .llaisys_types import *
+
+# Define function pointer types
+get_device_count_api = CFUNCTYPE(c_int)
+set_device_api = CFUNCTYPE(None, c_int)
+device_synchronize_api = CFUNCTYPE(None)
+
+create_stream_api = CFUNCTYPE(llaisysStream_t)
+destroy_stream_api = CFUNCTYPE(None, llaisysStream_t)
+stream_synchronize_api = CFUNCTYPE(None, llaisysStream_t)
+
+malloc_device_api = CFUNCTYPE(c_void_p, c_size_t)
+free_device_api = CFUNCTYPE(None, c_void_p)
+malloc_host_api = CFUNCTYPE(c_void_p, c_size_t)
+free_host_api = CFUNCTYPE(None, c_void_p)
+
+memcpy_sync_api = CFUNCTYPE(None, c_void_p, c_void_p, c_size_t, llaisysMemcpyKind_t)
+memcpy_async_api = CFUNCTYPE(None, c_void_p, c_void_p, c_size_t, llaisysMemcpyKind_t, llaisysStream_t)
+
+
+# Define the struct matching LlaisysRuntimeAPI
+class LlaisysRuntimeAPI(Structure):
+ _fields_ = [
+ ("get_device_count", get_device_count_api),
+ ("set_device", set_device_api),
+ ("device_synchronize", device_synchronize_api),
+ ("create_stream", create_stream_api),
+ ("destroy_stream", destroy_stream_api),
+ ("stream_synchronize", stream_synchronize_api),
+ ("malloc_device", malloc_device_api),
+ ("free_device", free_device_api),
+ ("malloc_host", malloc_host_api),
+ ("free_host", free_host_api),
+ ("memcpy_sync", memcpy_sync_api),
+ ("memcpy_async", memcpy_async_api),
+ ]
+
+
+# Load shared library
+def load_runtime(lib):
+ # Declare API function prototypes
+ lib.llaisysGetRuntimeAPI.argtypes = [llaisysDeviceType_t]
+ lib.llaisysGetRuntimeAPI.restype = ctypes.POINTER(LlaisysRuntimeAPI)
+
+ lib.llaisysSetContextRuntime.argtypes = [llaisysDeviceType_t, c_int]
+ lib.llaisysSetContextRuntime.restype = None
diff --git a/python/llaisys/libllaisys/tensor.py b/python/llaisys_py/libllaisys/tensor.py
similarity index 97%
rename from python/llaisys/libllaisys/tensor.py
rename to python/llaisys_py/libllaisys/tensor.py
index b58057883..cef02243a 100644
--- a/python/llaisys/libllaisys/tensor.py
+++ b/python/llaisys_py/libllaisys/tensor.py
@@ -1,78 +1,78 @@
-from ctypes import POINTER, c_uint8, c_void_p, c_size_t, c_ssize_t, c_int
-from .llaisys_types import llaisysDataType_t, llaisysDeviceType_t
-
-# Handle type
-llaisysTensor_t = c_void_p
-
-
-def load_tensor(lib):
- lib.tensorCreate.argtypes = [
- POINTER(c_size_t), # shape
- c_size_t, # ndim
- llaisysDataType_t, # dtype
- llaisysDeviceType_t, # device_type
- c_int, # device_id
- ]
- lib.tensorCreate.restype = llaisysTensor_t
-
- # Function: tensorDestroy
- lib.tensorDestroy.argtypes = [llaisysTensor_t]
- lib.tensorDestroy.restype = None
-
- # Function: tensorGetData
- lib.tensorGetData.argtypes = [llaisysTensor_t]
- lib.tensorGetData.restype = c_void_p
-
- # Function: tensorGetNdim
- lib.tensorGetNdim.argtypes = [llaisysTensor_t]
- lib.tensorGetNdim.restype = c_size_t
-
- # Function: tensorGetShape
- lib.tensorGetShape.argtypes = [llaisysTensor_t, POINTER(c_size_t)]
- lib.tensorGetShape.restype = None
-
- # Function: tensorGetStrides
- lib.tensorGetStrides.argtypes = [llaisysTensor_t, POINTER(c_ssize_t)]
- lib.tensorGetStrides.restype = None
-
- # Function: tensorGetDataType
- lib.tensorGetDataType.argtypes = [llaisysTensor_t]
- lib.tensorGetDataType.restype = llaisysDataType_t
-
- # Function: tensorGetDeviceType
- lib.tensorGetDeviceType.argtypes = [llaisysTensor_t]
- lib.tensorGetDeviceType.restype = llaisysDeviceType_t
-
- # Function: tensorGetDeviceId
- lib.tensorGetDeviceId.argtypes = [llaisysTensor_t]
- lib.tensorGetDeviceId.restype = c_int
-
- # Function: tensorDebug
- lib.tensorDebug.argtypes = [llaisysTensor_t]
- lib.tensorDebug.restype = None
-
- # Function: tensorIsContiguous
- lib.tensorIsContiguous.argtypes = [llaisysTensor_t]
- lib.tensorIsContiguous.restype = c_uint8
-
- # Function: tensorLoad
- lib.tensorLoad.argtypes = [llaisysTensor_t, c_void_p]
- lib.tensorLoad.restype = None
-
- # Function: tensorView(llaisysTensor_t tensor, size_t *shape);
- lib.tensorView.argtypes = [llaisysTensor_t, POINTER(c_size_t), c_size_t]
- lib.tensorView.restype = llaisysTensor_t
-
- # Function: tensorPermute(llaisysTensor_t tensor, size_t *order);
- lib.tensorPermute.argtypes = [llaisysTensor_t, POINTER(c_size_t)]
- lib.tensorPermute.restype = llaisysTensor_t
-
- # Function: tensorSlice(llaisysTensor_t tensor,
- # size_t dim, size_t start, size_t end);
- lib.tensorSlice.argtypes = [
- llaisysTensor_t, # tensor handle
- c_size_t, # dim : which axis to slice
- c_size_t, # start: inclusive
- c_size_t, # end : exclusive
- ]
- lib.tensorSlice.restype = llaisysTensor_t
+from ctypes import POINTER, c_uint8, c_void_p, c_size_t, c_ssize_t, c_int
+from .llaisys_types import llaisysDataType_t, llaisysDeviceType_t
+
+# Handle type
+llaisysTensor_t = c_void_p
+
+
+def load_tensor(lib):
+ lib.tensorCreate.argtypes = [
+ POINTER(c_size_t), # shape
+ c_size_t, # ndim
+ llaisysDataType_t, # dtype
+ llaisysDeviceType_t, # device_type
+ c_int, # device_id
+ ]
+ lib.tensorCreate.restype = llaisysTensor_t
+
+ # Function: tensorDestroy
+ lib.tensorDestroy.argtypes = [llaisysTensor_t]
+ lib.tensorDestroy.restype = None
+
+ # Function: tensorGetData
+ lib.tensorGetData.argtypes = [llaisysTensor_t]
+ lib.tensorGetData.restype = c_void_p
+
+ # Function: tensorGetNdim
+ lib.tensorGetNdim.argtypes = [llaisysTensor_t]
+ lib.tensorGetNdim.restype = c_size_t
+
+ # Function: tensorGetShape
+ lib.tensorGetShape.argtypes = [llaisysTensor_t, POINTER(c_size_t)]
+ lib.tensorGetShape.restype = None
+
+ # Function: tensorGetStrides
+ lib.tensorGetStrides.argtypes = [llaisysTensor_t, POINTER(c_ssize_t)]
+ lib.tensorGetStrides.restype = None
+
+ # Function: tensorGetDataType
+ lib.tensorGetDataType.argtypes = [llaisysTensor_t]
+ lib.tensorGetDataType.restype = llaisysDataType_t
+
+ # Function: tensorGetDeviceType
+ lib.tensorGetDeviceType.argtypes = [llaisysTensor_t]
+ lib.tensorGetDeviceType.restype = llaisysDeviceType_t
+
+ # Function: tensorGetDeviceId
+ lib.tensorGetDeviceId.argtypes = [llaisysTensor_t]
+ lib.tensorGetDeviceId.restype = c_int
+
+ # Function: tensorDebug
+ lib.tensorDebug.argtypes = [llaisysTensor_t]
+ lib.tensorDebug.restype = None
+
+ # Function: tensorIsContiguous
+ lib.tensorIsContiguous.argtypes = [llaisysTensor_t]
+ lib.tensorIsContiguous.restype = c_uint8
+
+ # Function: tensorLoad
+ lib.tensorLoad.argtypes = [llaisysTensor_t, c_void_p]
+ lib.tensorLoad.restype = None
+
+ # Function: tensorView(llaisysTensor_t tensor, size_t *shape);
+ lib.tensorView.argtypes = [llaisysTensor_t, POINTER(c_size_t), c_size_t]
+ lib.tensorView.restype = llaisysTensor_t
+
+ # Function: tensorPermute(llaisysTensor_t tensor, size_t *order);
+ lib.tensorPermute.argtypes = [llaisysTensor_t, POINTER(c_size_t)]
+ lib.tensorPermute.restype = llaisysTensor_t
+
+ # Function: tensorSlice(llaisysTensor_t tensor,
+ # size_t dim, size_t start, size_t end);
+ lib.tensorSlice.argtypes = [
+ llaisysTensor_t, # tensor handle
+ c_size_t, # dim : which axis to slice
+ c_size_t, # start: inclusive
+ c_size_t, # end : exclusive
+ ]
+ lib.tensorSlice.restype = llaisysTensor_t
diff --git a/python/llaisys/models/__init__.py b/python/llaisys_py/models/__init__.py
similarity index 96%
rename from python/llaisys/models/__init__.py
rename to python/llaisys_py/models/__init__.py
index af9918b0d..b84deabb8 100644
--- a/python/llaisys/models/__init__.py
+++ b/python/llaisys_py/models/__init__.py
@@ -1 +1 @@
-from .qwen2 import Qwen2
+from .qwen2 import Qwen2
diff --git a/python/llaisys_py/models/qwen2.py b/python/llaisys_py/models/qwen2.py
new file mode 100644
index 000000000..daa63e4f4
--- /dev/null
+++ b/python/llaisys_py/models/qwen2.py
@@ -0,0 +1,599 @@
+"""
+Qwen2 模型的 Python 封装：通过 LLAISYS C++ 后端进行推理。
+
+本模块不依赖 PyTorch 做推理，仅用 C 动态库（llaisys.dll / libllaisys.so）实现
+前向计算。权重从 safetensors 文件加载；若为 bfloat16，需用 PyTorch 读取后转成
+numpy 再灌入后端。
+"""
+from typing import Sequence
+import ctypes
+from ctypes import byref, cast, c_float, c_int, c_int64, c_size_t, c_ulonglong, POINTER
+
+# 张量并行时传 device_ids
+
+from ..libllaisys import LIB_LLAISYS
+from ..libllaisys import DeviceType
+from ..libllaisys import LlaisysQwen2Meta, LlaisysQwen2Weights, LlaisysQwen2Model_t
+from ..libllaisys.llaisys_types import DataType
+
+from pathlib import Path
+import json
+import numpy as np
+import safetensors
+
+try:
+ import torch
+except ImportError:
+ torch = None
+
+
+def _weight_key_to_handle(weights_ptr, nlayer: int):
+ """
+ 生成「safetensors 中的权重 key」到「C 侧权重句柄」的映射。
+
+ 权重文件名（key）与 Qwen2/HuggingFace 命名一致；部分 key 有别名（如 model.norm.w）
+ 以兼容 ModelScope 等不同来源的 checkpoint。
+
+ Yields:
+ (key, handle): key 为 safetensors 中的张量名，handle 为 C 侧 LlaisysQwen2Weights 中对应成员的指针。
+ """
+ w = weights_ptr.contents
+ # 词嵌入与输出层
+ yield "model.embed_tokens.weight", w.in_embed
+ yield "model.norm.weight", w.out_norm_w
+ yield "model.norm.w", w.out_norm_w # 部分 checkpoint 用 model.norm.w
+ yield "lm_head.weight", w.out_embed
+ # 每一层的 attention / MLP 权重与 bias
+ for i in range(nlayer):
+ yield f"model.layers.{i}.input_layernorm.weight", w.attn_norm_w[i]
+ yield f"model.layers.{i}.self_attn.q_proj.weight", w.attn_q_w[i]
+ yield f"model.layers.{i}.self_attn.q_proj.bias", w.attn_q_b[i]
+ yield f"model.layers.{i}.self_attn.k_proj.weight", w.attn_k_w[i]
+ yield f"model.layers.{i}.self_attn.k_proj.bias", w.attn_k_b[i]
+ yield f"model.layers.{i}.self_attn.v_proj.weight", w.attn_v_w[i]
+ yield f"model.layers.{i}.self_attn.v_proj.bias", w.attn_v_b[i]
+ yield f"model.layers.{i}.self_attn.o_proj.weight", w.attn_o_w[i]
+ yield f"model.layers.{i}.post_attention_layernorm.weight", w.mlp_norm_w[i]
+ yield f"model.layers.{i}.mlp.gate_proj.weight", w.mlp_gate_w[i]
+ yield f"model.layers.{i}.mlp.up_proj.weight", w.mlp_up_w[i]
+ yield f"model.layers.{i}.mlp.down_proj.weight", w.mlp_down_w[i]
+
+
+def _shard_weight_for_tp(key: str, arr: np.ndarray, tp_rank: int, tp_world_size: int,
+ nh: int, nkvh: int, dh: int, di: int, hs: int, nlayer: int):
+ """
+ 张量并行：按 key 对权重做行/列切分，返回当前 rank 应加载的 shard。
+ 列并行（输出维切分）：q/k/v_proj, gate/up_proj -> 切行。
+ 行并行（输入维切分）：o_proj, down_proj -> 切列。
+ """
+ if tp_world_size <= 1:
+ return arr
+ w = tp_world_size
+ r = tp_rank
+ arr = np.ascontiguousarray(arr)
+ if "q_proj.weight" in key or "q_proj.bias" in key:
+ # [nh*dh, hs] or [nh*dh]
+ size = nh * dh
+ step = size // w
+ if "bias" in key:
+ return arr[r * step : (r + 1) * step].copy()
+ return arr[r * step : (r + 1) * step, :].copy()
+ if "k_proj.weight" in key or "k_proj.bias" in key or "v_proj.weight" in key or "v_proj.bias" in key:
+ size = nkvh * dh
+ step = size // w
+ if "bias" in key:
+ return arr[r * step : (r + 1) * step].copy()
+ return arr[r * step : (r + 1) * step, :].copy()
+ if "o_proj.weight" in key:
+ # [hs, nh*dh] 行并行
+ size = nh * dh
+ step = size // w
+ return arr[:, r * step : (r + 1) * step].copy()
+ if "gate_proj.weight" in key or "up_proj.weight" in key:
+ step = di // w
+ return arr[r * step : (r + 1) * step, :].copy()
+ if "down_proj.weight" in key:
+ step = di // w
+ return arr[:, r * step : (r + 1) * step].copy()
+ return arr
+
+
+def _bf16_bytes_to_float32(raw: bytes) -> np.ndarray:
+ """将 safetensors 中的 bf16 原始字节转为 float32 numpy（无 torch 依赖）。"""
+ n = len(raw) // 2
+ u16 = np.frombuffer(raw, dtype=np.uint16)
+ u32 = (u16.astype(np.uint32) << 16)
+ return np.frombuffer(u32.tobytes(), dtype=np.float32).copy()
+
+
+def _read_safetensors_header(fpath: Path):
+ """返回 (data_start_offset, key -> {dtype, shape, data_offsets})，便于无 torch 时读 bf16。"""
+ with open(fpath, "rb") as f:
+ header_len = int.from_bytes(f.read(8), "little")
+ header_json = f.read(header_len).decode("utf-8")
+ data_start = 8 + (header_len + 7) // 8 * 8
+ header = json.loads(header_json)
+ key_to_meta = {}
+ for k, v in header.items():
+ if k == "__metadata__":
+ continue
+ if isinstance(v, dict) and "dtype" in v and "shape" in v and "data_offsets" in v:
+ key_to_meta[k] = v
+ return data_start, key_to_meta
+
+
+def _load_bf16_weights_from_safetensors_no_torch(
+ fpath: Path,
+ key_to_handle: dict,
+ key_to_meta: dict,
+ data_start: int,
+ tp_world_size: int,
+ tp_rank: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ dh: int,
+ intermediate_size: int,
+ hidden_size: int,
+ num_hidden_layers: int,
+) -> set:
+ """从 safetensors 文件按 key 读 bf16 原始数据，转 float32 后灌入后端（不依赖 torch）。"""
+ loaded = set()
+ with open(fpath, "rb") as f:
+ for key, handle in key_to_handle.items():
+ if key not in key_to_meta:
+ continue
+ meta = key_to_meta[key]
+ if meta.get("dtype") != "BF16":
+ continue
+ shape = meta["shape"]
+ start, end = meta["data_offsets"]
+ offset_in_file = data_start + start
+ size_bytes = end - start
+ f.seek(offset_in_file)
+ raw = f.read(size_bytes)
+ arr = _bf16_bytes_to_float32(raw).reshape(shape)
+ arr = np.ascontiguousarray(arr)
+ if tp_world_size > 1:
+ arr = _shard_weight_for_tp(
+ key, arr, tp_rank, tp_world_size,
+ num_attention_heads, num_key_value_heads, dh,
+ intermediate_size, hidden_size, num_hidden_layers,
+ )
+ _numpy_to_backend(arr, handle)
+ loaded.add(key)
+ return loaded
+
+
+def _numpy_to_backend(arr: np.ndarray, tensor_handle) -> None:
+ """
+ 将 numpy 数组拷贝到 LLAISYS 后端张量（CPU 或设备内存）。
+
+ 若数组为 float32，会按 bfloat16 的「高 16 位」方式截断后传入后端，
+ 以兼容从 bfloat16 转成 float32 再传过来的权重。
+
+ Args:
+ arr: 主机侧 numpy 数组，需为连续内存。
+ tensor_handle: C 侧张量句柄（LlaisysTensor*），由 tensorLoad 写入。
+ """
+ arr = np.ascontiguousarray(arr)
+ if arr.dtype == np.float32:
+ # float32 视为“从 bf16 转来的”，取高 16 位作为 bf16 比特表示
+ arr_bf16 = (arr.view(np.uint32) >> 16).astype(np.uint16)
+ LIB_LLAISYS.tensorLoad(tensor_handle, arr_bf16.ctypes.data)
+ elif arr.dtype == np.uint16 or arr.dtype == np.float16:
+ LIB_LLAISYS.tensorLoad(tensor_handle, arr.ctypes.data)
+ else:
+ LIB_LLAISYS.tensorLoad(tensor_handle, arr.ctypes.data)
+
+
+class Qwen2:
+ """
+ Qwen2 模型的 Python 封装类。
+
+ 通过 LLAISYS C 接口创建模型、加载 safetensors 权重，并对外提供 generate()
+ 做自回归生成。推理全程在 C++ 后端执行，Python 只做配置、权重加载和循环调用 Infer。
+ """
+
+ def __init__(
+ self,
+ model_path,
+ device: DeviceType = DeviceType.CPU,
+ max_batch_size: int = 1,
+ tp_rank: int = 0,
+ tp_world_size: int = 1,
+ ):
+ """
+ 从本地目录加载 Qwen2 模型：读 config、创建 C 模型、灌入权重。
+
+ Args:
+ model_path: 模型目录路径（需含 config.json 和 *.safetensors）。
+ device: 运行设备，如 DeviceType.CPU 或 DeviceType.NVIDIA。
+ max_batch_size: KV-Cache 槽位数，用于连续批处理；1 为单序列（默认）。
+ tp_rank: 张量并行 rank（0..tp_world_size-1），默认 0。
+ tp_world_size: 张量并行 world size，1 表示非分布式，默认 1。
+ """
+ model_path = Path(model_path)
+
+ # ---------- 1. 读取 config.json ----------
+ config_path = model_path / "config.json"
+ if not config_path.exists():
+ raise FileNotFoundError(f"config.json not found in {model_path}")
+ with open(config_path, "r", encoding="utf-8") as f:
+ config = json.load(f)
+
+ hidden_size = config["hidden_size"]
+ num_hidden_layers = config["num_hidden_layers"]
+ num_attention_heads = config["num_attention_heads"]
+ num_key_value_heads = config.get("num_key_value_heads", num_attention_heads)
+ intermediate_size = config["intermediate_size"]
+ vocab_size = config["vocab_size"]
+ rms_norm_eps = float(config.get("rms_norm_eps", 1e-6))
+ rope_theta = float(config.get("rope_theta", 10000.0))
+ eos_id = config.get("eos_token_id", config.get("bos_token_id", 151643))
+ max_position = config.get("max_position_embeddings", 131072)
+ maxseq = min(4096, max_position)
+
+ # 解析 dtype：支持 config 里 "dtype" 或 "torch_dtype"（如 "bfloat16" / "float16"）
+ cfg_dtype = config.get("dtype", config.get("torch_dtype", "bfloat16"))
+ if isinstance(cfg_dtype, str) and "bfloat" in cfg_dtype.lower():
+ dtype = DataType.BF16
+ elif isinstance(cfg_dtype, str) and "float16" in cfg_dtype.lower():
+ dtype = DataType.F16
+ else:
+ dtype = DataType.BF16
+
+ # 每个 attention 头的维度
+ dh = hidden_size // num_attention_heads
+
+ # ---------- 2. 组装 C 侧元信息并创建模型 ----------
+ # max_batch_size：连续批处理时 KV 槽位数，默认 1 保持单序列行为
+ # tp_rank / tp_world_size：张量并行（项目#5），默认 0/1 表示单卡
+ meta = LlaisysQwen2Meta(
+ dtype=dtype,
+ nlayer=num_hidden_layers,
+ hs=hidden_size,
+ nh=num_attention_heads,
+ nkvh=num_key_value_heads,
+ dh=dh,
+ di=intermediate_size,
+ maxseq=maxseq,
+ voc=vocab_size,
+ max_batch_size=max_batch_size,
+ tp_rank=tp_rank,
+ tp_world_size=tp_world_size,
+ epsilon=rms_norm_eps,
+ theta=rope_theta,
+ end_token=eos_id,
+ )
+
+ # TP 时每进程通常通过 CUDA_VISIBLE_DEVICES 只暴露一张卡，故传 device_ids=None 用默认 0 即可
+ self._model = LIB_LLAISYS.llaisysQwen2ModelCreate(
+ byref(meta),
+ device,
+ None,
+ 0,
+ )
+ if not self._model:
+ raise RuntimeError("llaisysQwen2ModelCreate failed")
+
+ self._end_token = eos_id
+ self._nlayer = num_hidden_layers
+ self._max_batch_size = max_batch_size
+
+ # ---------- 3. 从 safetensors 加载权重到 C 侧 ----------
+ weights_ptr = LIB_LLAISYS.llaisysQwen2ModelWeights(self._model)
+ key_to_handle = dict(_weight_key_to_handle(weights_ptr, num_hidden_layers))
+ loaded_keys = set()
+
+ # bfloat16 时可用 PyTorch 读；若无 torch（如 spawn 子进程避免 nccl 冲突）则走纯 Python 解析 safetensors
+ use_pt = torch is not None
+ use_bf16_no_torch = (dtype == DataType.BF16 and torch is None)
+
+ safetensor_files = sorted(model_path.glob("*.safetensors"))
+ for idx, fpath in enumerate(safetensor_files):
+ if torch is not None and idx == 0:
+ pass # 首次用 torch 打开文件时可能触发 c10 等，便于定位崩溃
+ print(f" Loading weights: {fpath.name} ({idx + 1}/{len(safetensor_files)})", flush=True)
+ if use_bf16_no_torch:
+ data_start, key_to_meta = _read_safetensors_header(fpath)
+ loaded = _load_bf16_weights_from_safetensors_no_torch(
+ fpath, key_to_handle, key_to_meta, data_start,
+ tp_world_size, tp_rank,
+ num_attention_heads, num_key_value_heads, dh,
+ intermediate_size, hidden_size, num_hidden_layers,
+ )
+ loaded_keys.update(loaded)
+ continue
+ with safetensors.safe_open(
+ fpath, framework="pt" if use_pt else "numpy", device="cpu"
+ ) as data:
+ for key in data.keys():
+ if key not in key_to_handle:
+ continue
+ handle = key_to_handle[key]
+ t = data.get_tensor(key)
+ if use_pt:
+ if t.dtype == torch.bfloat16:
+ arr = t.float().numpy()
+ else:
+ arr = t.numpy()
+ else:
+ arr = np.ascontiguousarray(t)
+ arr = np.ascontiguousarray(arr)
+ if tp_world_size > 1:
+ arr = _shard_weight_for_tp(
+ key, arr, tp_rank, tp_world_size,
+ num_attention_heads, num_key_value_heads, dh,
+ intermediate_size, hidden_size, num_hidden_layers,
+ )
+ _numpy_to_backend(arr, handle)
+ loaded_keys.add(key)
+
+ # 至少需要嵌入层权重，否则说明 key 对不上
+ embed_loaded = "model.embed_tokens.weight" in loaded_keys
+ if not embed_loaded:
+ sample_keys = []
+ for fpath in sorted(model_path.glob("*.safetensors")):
+ with safetensors.safe_open(
+ fpath, framework="pt" if use_pt else "numpy", device="cpu"
+ ) as data:
+ sample_keys.extend(list(data.keys())[:40])
+ break
+ raise RuntimeError(
+ "No embedding weights loaded. Loaded %d keys; sample keys from file: %s"
+ % (len(loaded_keys), sample_keys[:30])
+ )
+
+ # GPU 时：默认只缓存输出层（推理主体在 GPU，快）；若环境变量 LLAISYS_GPU_FULL_CPU=1 则全量缓存，整次前向在 CPU（慢但可规避其它 GPU 算子问题）
+ if device == DeviceType.NVIDIA:
+ import os
+ if os.environ.get("LLAISYS_GPU_FULL_CPU") == "1":
+ LIB_LLAISYS.llaisysQwen2ModelCacheAllWeightsOnCPU(self._model)
+ print(" [Qwen2] GPU 已缓存全量权重到 CPU，推理全程在 CPU 上执行（LLAISYS_GPU_FULL_CPU=1）。", flush=True)
+ else:
+ LIB_LLAISYS.llaisysQwen2ModelCacheOutputLayerOnCPU(self._model)
+ print(" [Qwen2] GPU：embedding 与输出层在 CPU，其余层在 GPU。", flush=True)
+
+ def kv_cache_bytes(self, prefix_len: int) -> int:
+ """存储前缀长度为 prefix_len 的 KV cache 所需字节数。"""
+ return LIB_LLAISYS.llaisysQwen2ModelGetKVCacheBytes(self._model, prefix_len)
+
+ def export_kv_cache(self) -> bytes:
+ """导出当前 KV cache 到字节串（当前 cache_len 由 C 侧维护）。"""
+ n = LIB_LLAISYS.llaisysQwen2ModelGetCacheLen(self._model)
+ if n == 0:
+ return b""
+ size = LIB_LLAISYS.llaisysQwen2ModelGetKVCacheBytes(self._model, n)
+ buf = (ctypes.c_byte * size)()
+ LIB_LLAISYS.llaisysQwen2ModelExportKVCache(self._model, ctypes.cast(buf, ctypes.c_void_p))
+ return bytes(buf)
+
+ def reset_kv_cache(self) -> None:
+ """将 KV cache 长度置 0，新请求全量 prefill 前调用，避免沿用上一轮状态。"""
+ LIB_LLAISYS.llaisysQwen2ModelResetKVCache(self._model)
+
+ def import_kv_cache(self, data: bytes, prefix_len: int) -> None:
+ """从字节串导入前缀长度为 prefix_len 的 KV cache；之后可对 suffix 做 prefill。"""
+ if prefix_len == 0 or not data:
+ return
+ expected = self.kv_cache_bytes(prefix_len)
+ if len(data) < expected:
+ raise ValueError(f"import_kv_cache: need {expected} bytes, got {len(data)}")
+ buf = (ctypes.c_byte * len(data))()
+ ctypes.memmove(ctypes.addressof(buf), data, len(data))
+ LIB_LLAISYS.llaisysQwen2ModelImportKVCache(
+ self._model, ctypes.cast(buf, ctypes.c_void_p), prefix_len
+ )
+
+ def export_kv_cache_slot(self, slot_id: int) -> bytes:
+ """导出指定 slot 的 KV cache 到字节串（用于连续批处理 + KV 池）。"""
+ n = LIB_LLAISYS.llaisysQwen2ModelGetCacheLenSlot(self._model, slot_id)
+ if n == 0:
+ return b""
+ size = LIB_LLAISYS.llaisysQwen2ModelGetKVCacheBytes(self._model, n)
+ buf = (ctypes.c_byte * size)()
+ LIB_LLAISYS.llaisysQwen2ModelExportKVCacheSlot(
+ self._model, slot_id, ctypes.cast(buf, ctypes.c_void_p)
+ )
+ return bytes(buf)
+
+ def import_kv_cache_slot(self, slot_id: int, data: bytes, prefix_len: int) -> None:
+ """将字节串导入到指定 slot 的前缀 KV cache；之后可对该 slot 做 suffix prefill。"""
+ if prefix_len == 0 or not data:
+ return
+ expected = self.kv_cache_bytes(prefix_len)
+ if len(data) < expected:
+ raise ValueError(
+ f"import_kv_cache_slot: need {expected} bytes, got {len(data)}"
+ )
+ buf = (ctypes.c_byte * len(data))()
+ ctypes.memmove(ctypes.addressof(buf), data, len(data))
+ LIB_LLAISYS.llaisysQwen2ModelImportKVCacheSlot(
+ self._model,
+ slot_id,
+ ctypes.cast(buf, ctypes.c_void_p),
+ prefix_len,
+ )
+
+ @property
+ def cache_len(self) -> int:
+ """当前已写入 KV cache 的长度。"""
+ return LIB_LLAISYS.llaisysQwen2ModelGetCacheLen(self._model)
+
+ def generate(
+ self,
+ inputs: Sequence[int],
+ max_new_tokens: int = 128,
+ top_k: int = 1,
+ top_p: float = 0.8,
+ temperature: float = 0.8,
+ seed: int = 0,
+ prefix_len: int = 0,
+ ):
+ """
+ 自回归生成：从当前 token 序列出发，每次调用 C 的 Infer 得到下一个 token，直到 EOS 或达到 max_new_tokens。
+ 支持随机采样：temperature、top_k、top_p 会传入 C 侧；top_k=1 且 temperature 接近 0 时为 argmax 贪心。
+ prefix_len>0 时表示已通过 import_kv_cache 导入前缀，仅对 inputs[prefix_len:] 做 suffix prefill 再 decode。
+
+ Args:
+ inputs: 初始 token id 序列（如 prompt 经 tokenizer 编码后的列表）。
+ max_new_tokens: 最多新生成多少个 token。
+ top_k: 保留概率最高的 k 个 token，<=0 表示不限制。
+ top_p: nucleus 采样阈值，<=0 或 >=1 表示不限制。
+ temperature: 温度，<=0 或极小为贪心。
+ seed: 随机种子，0 表示每次随机。
+ prefix_len: 若已 import_kv_cache，则为前缀长度；0 表示全量 prefill。
+
+ Returns:
+ 完整 token 序列（inputs + 新生成的 token），包含 EOS 在内。
+ """
+ # 检查输入是否为空，如果连问题都没有，AI 没法往下接话
+ if not inputs:
+ raise ValueError("generate() called with empty input token list")
+
+ import os
+ if os.environ.get("LLAISYS_DEBUG"):
+ print(f"[LLAISYS] Qwen2.generate() n_inputs={len(inputs)} prefix_len={prefix_len} max_new_tokens={max_new_tokens}")
+
+ tokens = list(inputs)
+ if prefix_len == 0:
+ self.reset_kv_cache()
+ # 首步：全量 prefill 或 suffix prefill（需先 import_kv_cache）
+ if prefix_len > 0:
+ if prefix_len >= len(tokens):
+ raise ValueError("prefix_len must be < len(inputs)")
+ suffix = tokens[prefix_len:]
+ n = len(suffix)
+ token_arr = (c_int64 * n)(*suffix)
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInfer(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ n,
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ )
+ if next_tok == -1:
+ raise RuntimeError("llaisysQwen2ModelInfer failed (returned -1)")
+ tokens.append(next_tok)
+ if next_tok == self._end_token:
+ return tokens
+ n_decoded = 1
+ else:
+ n = len(tokens)
+ token_arr = (c_int64 * n)(*tokens)
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInfer(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ n,
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ )
+ if next_tok == -1:
+ raise RuntimeError("llaisysQwen2ModelInfer failed (returned -1)")
+ tokens.append(next_tok)
+ if next_tok == self._end_token:
+ return tokens
+ n_decoded = 1
+
+ # 后续为 decode 步：每次只传最后一个 token（ntoken=1）
+ for _ in range(max_new_tokens - n_decoded):
+ token_arr = (c_int64 * 1)(tokens[-1])
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInfer(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ 1,
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ )
+ if next_tok == -1:
+ raise RuntimeError("llaisysQwen2ModelInfer failed (returned -1)")
+ tokens.append(next_tok)
+ if next_tok == self._end_token:
+ break
+ return tokens
+
+ @property
+ def end_token(self) -> int:
+ """EOS token id，用于流式生成时判断结束。"""
+ return self._end_token
+
+ @property
+ def max_batch_size(self) -> int:
+ """KV-Cache 槽位数，供连续批处理 Engine 使用。"""
+ return getattr(self, "_max_batch_size", 1)
+
+ def next_token(
+ self,
+ token_ids: Sequence[int],
+ temperature: float = 0.8,
+ top_k: int = 50,
+ top_p: float = 0.8,
+ seed: int = 0,
+ ) -> int:
+ """
+ 单步推理：给定当前 token 序列，返回下一个 token id。
+ 供流式输出或外部自回归循环使用。
+ """
+ if not token_ids:
+ raise ValueError("next_token() requires non-empty token_ids")
+ n = len(token_ids)
+ token_arr = (c_int64 * n)(*token_ids)
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInfer(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ n,
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ )
+ if next_tok == -1:
+ raise RuntimeError("llaisysQwen2ModelInfer failed (returned -1)")
+ return next_tok
+
+ def infer_hybrid(
+ self,
+ token_ids: Sequence[int],
+ temperature: float = 0.0,
+ top_k: int = 1,
+ top_p: float = 1.0,
+ seed: int = 0,
+ gpu_up_to_layer: int = -1,
+ ) -> int:
+ """
+ 诊断用：前 (gpu_up_to_layer+1) 层在 GPU 上跑，其余在 CPU。需已调用 CacheAllWeightsOnCPU。
+ gpu_up_to_layer=-1：全 CPU；=0：仅 embedding 在 GPU；=1：embedding+layer0 在 GPU；依此类推。
+ 调用前会 reset_kv_cache。
+ """
+ self.reset_kv_cache()
+ n = len(token_ids)
+ token_arr = (c_int64 * n)(*token_ids)
+ next_tok = LIB_LLAISYS.llaisysQwen2ModelInferHybrid(
+ self._model,
+ cast(token_arr, POINTER(c_int64)),
+ n,
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ c_int(gpu_up_to_layer),
+ )
+ if next_tok == -1:
+ raise RuntimeError("llaisysQwen2ModelInferHybrid failed (returned -1)")
+ return next_tok
+
+ # ===== Python 的魔法方法：析构函数 =====
+ def __del__(self):
+ """析构时释放 C 侧模型，避免泄漏。""" #
+ # 当 Python 里的 Qwen2 对象不再被使用，准备被垃圾回收时，会自动触发这个函数
+ # 它负责打电话通知 C++ 侧：“我要下线了，你把那些占了几个 G 内存的模型张量（Tensor）全删了吧！”
+ if getattr(self, "_model", None) is not None:
+ LIB_LLAISYS.llaisysQwen2ModelDestroy(self._model) # 调用 C++ 的销毁接口
+ self._model = None # 清空指针，防止重复释放报错
+
\ No newline at end of file
diff --git a/python/llaisys/ops.py b/python/llaisys_py/ops.py
similarity index 70%
rename from python/llaisys/ops.py
rename to python/llaisys_py/ops.py
index ed0180bc8..6878342da 100644
--- a/python/llaisys/ops.py
+++ b/python/llaisys_py/ops.py
@@ -1,6 +1,6 @@
from .libllaisys import LIB_LLAISYS
from .tensor import Tensor
-from ctypes import c_float, c_int
+from ctypes import c_float, c_int, c_ulonglong
class Ops:
@@ -19,9 +19,10 @@ def embedding(out: Tensor, index: Tensor, weight: Tensor):
)
@staticmethod
- def linear(out: Tensor, inp: Tensor, weight: Tensor, bias: Tensor):
+ def linear(out: Tensor, inp: Tensor, weight: Tensor, bias=None):
LIB_LLAISYS.llaisysLinear(
- out.lib_tensor(), inp.lib_tensor(), weight.lib_tensor(), bias.lib_tensor()
+ out.lib_tensor(), inp.lib_tensor(), weight.lib_tensor(),
+ bias.lib_tensor() if bias is not None else None,
)
@staticmethod
@@ -53,3 +54,22 @@ def self_attention(attn_val: Tensor, q: Tensor, k: Tensor, v: Tensor, scale: flo
@staticmethod
def swiglu(out: Tensor, gate: Tensor, up: Tensor):
LIB_LLAISYS.llaisysSwiGLU(out.lib_tensor(), gate.lib_tensor(), up.lib_tensor())
+
+ @staticmethod
+ def sample(
+ out_idx: Tensor,
+ logits: Tensor,
+ temperature: float = 1.0,
+ top_k: int = 0,
+ top_p: float = 0.0,
+ seed: int = 0,
+ ):
+ """从 logits 按概率采样一个 token 索引。支持 Temperature、Top-K、Top-P。seed=0 表示随机。"""
+ LIB_LLAISYS.llaisysSample(
+ out_idx.lib_tensor(),
+ logits.lib_tensor(),
+ c_float(temperature),
+ c_int(top_k),
+ c_float(top_p),
+ c_ulonglong(seed),
+ )
diff --git a/python/llaisys/runtime.py b/python/llaisys_py/runtime.py
similarity index 96%
rename from python/llaisys/runtime.py
rename to python/llaisys_py/runtime.py
index 15be1aa17..15ea752be 100644
--- a/python/llaisys/runtime.py
+++ b/python/llaisys_py/runtime.py
@@ -1,68 +1,68 @@
-from . import libllaisys
-from .libllaisys import LIB_LLAISYS
-from ctypes import c_void_p
-
-
-class RuntimeAPI:
- def __init__(self, device_type: libllaisys.DeviceType):
- self._api = LIB_LLAISYS.llaisysGetRuntimeAPI(
- libllaisys.llaisysDeviceType_t(device_type)
- )
-
- def get_device_count(self) -> int:
- result = self._api.contents.get_device_count()
- return result
-
- def set_device(self, device_id: int) -> None:
- self._api.contents.set_device(device_id)
-
- def device_synchronize(self) -> None:
- self._api.contents.device_synchronize()
-
- def create_stream(self) -> libllaisys.llaisysStream_t:
- stream = self._api.contents.create_stream()
- return stream
-
- def destroy_stream(self, stream: libllaisys.llaisysStream_t) -> None:
- self._api.contents.destroy_stream(stream)
-
- def stream_synchronize(self, stream: libllaisys.llaisysStream_t) -> None:
- self._api.contents.stream_synchronize(stream)
-
- def malloc_device(self, size: int) -> c_void_p:
- ptr = self._api.contents.malloc_device(size)
- return ptr
-
- def free_device(self, ptr: c_void_p) -> None:
- print(f"[llaisys] free_device({ptr})")
- self._api.contents.free_device(ptr)
-
- def malloc_host(self, size: int) -> c_void_p:
- ptr = self._api.contents.malloc_host(size)
- return ptr
-
- def free_host(self, ptr: c_void_p) -> None:
- self._api.contents.free_host(ptr)
-
- def memcpy_sync(
- self,
- dst: c_void_p,
- src: c_void_p,
- size: int,
- kind: libllaisys.MemcpyKind,
- ) -> None:
- self._api.contents.memcpy_sync(
- dst, src, size, libllaisys.llaisysMemcpyKind_t(kind)
- )
-
- def memcpy_async(
- self,
- dst: c_void_p,
- src: c_void_p,
- size: int,
- kind: libllaisys.MemcpyKind,
- stream: libllaisys.llaisysStream_t,
- ) -> None:
- self._api.contents.memcpy_async(
- dst, src, size, libllaisys.llaisysMemcpyKind_t(kind), stream
- )
+from . import libllaisys
+from .libllaisys import LIB_LLAISYS
+from ctypes import c_void_p
+
+
+class RuntimeAPI:
+ def __init__(self, device_type: libllaisys.DeviceType):
+ self._api = LIB_LLAISYS.llaisysGetRuntimeAPI(
+ libllaisys.llaisysDeviceType_t(device_type)
+ )
+
+ def get_device_count(self) -> int:
+ result = self._api.contents.get_device_count()
+ return result
+
+ def set_device(self, device_id: int) -> None:
+ self._api.contents.set_device(device_id)
+
+ def device_synchronize(self) -> None:
+ self._api.contents.device_synchronize()
+
+ def create_stream(self) -> libllaisys.llaisysStream_t:
+ stream = self._api.contents.create_stream()
+ return stream
+
+ def destroy_stream(self, stream: libllaisys.llaisysStream_t) -> None:
+ self._api.contents.destroy_stream(stream)
+
+ def stream_synchronize(self, stream: libllaisys.llaisysStream_t) -> None:
+ self._api.contents.stream_synchronize(stream)
+
+ def malloc_device(self, size: int) -> c_void_p:
+ ptr = self._api.contents.malloc_device(size)
+ return ptr
+
+ def free_device(self, ptr: c_void_p) -> None:
+ print(f"[llaisys] free_device({ptr})")
+ self._api.contents.free_device(ptr)
+
+ def malloc_host(self, size: int) -> c_void_p:
+ ptr = self._api.contents.malloc_host(size)
+ return ptr
+
+ def free_host(self, ptr: c_void_p) -> None:
+ self._api.contents.free_host(ptr)
+
+ def memcpy_sync(
+ self,
+ dst: c_void_p,
+ src: c_void_p,
+ size: int,
+ kind: libllaisys.MemcpyKind,
+ ) -> None:
+ self._api.contents.memcpy_sync(
+ dst, src, size, libllaisys.llaisysMemcpyKind_t(kind)
+ )
+
+ def memcpy_async(
+ self,
+ dst: c_void_p,
+ src: c_void_p,
+ size: int,
+ kind: libllaisys.MemcpyKind,
+ stream: libllaisys.llaisysStream_t,
+ ) -> None:
+ self._api.contents.memcpy_async(
+ dst, src, size, libllaisys.llaisysMemcpyKind_t(kind), stream
+ )
diff --git a/python/llaisys_py/server/README.md b/python/llaisys_py/server/README.md
new file mode 100644
index 000000000..b22be7ba9
--- /dev/null
+++ b/python/llaisys_py/server/README.md
@@ -0,0 +1,90 @@
+# LLAISYS Chatbot Server
+
+OpenAI chat-completion 风格的 HTTP 服务，单用户、支持流式 (SSE)。
+
+## 依赖
+
+```bash
+pip install fastapi uvicorn
+```
+
+（若已安装 `transformers` 用于分词器则无需额外依赖。）
+
+## 启动
+
+指定模型目录（与 `test_infer.py` 使用的 Qwen2 模型一致）：
+
+```bash
+# 方式一：环境变量
+set MODEL_PATH=C:\path\to\DeepSeek-R1-Distill-Qwen-1.5B
+python -m llaisys_py.server
+
+# 方式二：命令行参数
+python -m llaisys_py.server --model "C:\path\to\DeepSeek-R1-Distill-Qwen-1.5B" --port 8000
+```
+
+可选参数：`--host`, `--port`, `--device`（cpu / nvidia）。
+
+## 接口
+
+- `GET /health`：健康检查，返回是否已加载模型。
+- `POST /v1/chat/completions`：与 OpenAI 兼容的对话补全。
+
+请求体示例：
+
+```json
+{
+ "model": "default",
+ "messages": [{"role": "user", "content": "你好"}],
+ "max_tokens": 128,
+ "temperature": 0.8,
+ "top_p": 0.9,
+ "top_k": 50,
+ "stream": false,
+ "seed": null
+}
+```
+
+- `stream: true` 时返回 SSE 流式响应。
+
+## 交互式聊天 UI（多轮对话）
+
+### 方式一：Web 页面
+
+服务启动后，在浏览器打开：
+
+- **http://127.0.0.1:8000/chat**
+
+即可使用网页聊天：输入框输入内容，点击发送或回车，对话历史会保留在页面上，支持连续多轮对话。
+
+### 方式二：命令行（CLI）
+
+先启动服务，再在**另一个终端**运行：
+
+```bash
+pip install requests
+python -m llaisys_py.server.chat_cli
+```
+
+默认连到 `http://127.0.0.1:8000`。输入一句话回车发送，收到回复后继续输入下一句；输入 `quit` 或 `q` 退出。
+
+可选参数：`--base-url`, `--max-tokens`, `--temperature`, `--top-k`, `--top-p`。
+
+## 从项目根目录运行
+
+若未 `pip install -e python`，需把 `python` 加入 PYTHONPATH：
+
+```bash
+set PYTHONPATH=python
+python -m llaisys_py.server --model "C:\path\to\model"
+```
+
+## 答非所问时
+
+若模型经常跑题、只回复客套话或固定“我是 DeepSeek-R1…”自我介绍，可尝试：
+
+1. **默认不加系统提示**：服务端已改为不注入任何 system 内容（避免触发模型自报家门）。若需要自定义系统人设，可设环境变量 `LLAISYS_SYSTEM_PROMPT="你的说明"` 再启动。
+2. **降低 temperature**：请求里传 `"temperature": 0.3` 或 `0.5`，回答会更聚焦。
+3. **新开会话**：在 `/sessions` 里新建对话再问，避免被之前的长回复干扰。
+4. **看实际发给模型的 prompt**：启动前设 `LLAISYS_DEBUG=1`，控制台会打印每轮 prompt 长度和末尾 300 字符，便于排查。
+5. **模型能力**：DeepSeek-R1-Distill-Qwen-1.5B 为 1.5B 小模型，知识型问题可能表现有限，可换更大模型或仅作演示。
diff --git a/python/llaisys_py/server/__init__.py b/python/llaisys_py/server/__init__.py
new file mode 100644
index 000000000..1d70c6b73
--- /dev/null
+++ b/python/llaisys_py/server/__init__.py
@@ -0,0 +1,8 @@
+"""
+Chatbot HTTP server (OpenAI chat-completion style).
+Run with: python -m llaisys_py.server
+"""
+
+from .app import create_app
+
+__all__ = ["create_app"]
diff --git a/python/llaisys_py/server/__main__.py b/python/llaisys_py/server/__main__.py
new file mode 100644
index 000000000..839bdd019
--- /dev/null
+++ b/python/llaisys_py/server/__main__.py
@@ -0,0 +1,89 @@
+"""
+启动 Chatbot Server。
+用法:
+ python -m llaisys_py.server # 使用下方 DEFAULT_MODEL_PATH 或环境变量 MODEL_PATH
+ python -m llaisys_py.server --port 8000 --device nvidia
+"""
+# 最先修正 CUDA_VISIBLE_DEVICES：若为空串，CUDA 会认为 0 张卡且之后无法更改。
+# 必须在 import 任何会间接加载 torch/CUDA 的模块之前执行（如 create_app -> transformers -> torch）。
+import os
+import sys
+# 发生 segfault 时打印 Python 栈，便于定位是否在 C 扩展内崩溃
+try:
+ import faulthandler
+ faulthandler.enable(all_threads=True)
+except Exception:
+ pass
+if os.environ.get("CUDA_VISIBLE_DEVICES") == "":
+ os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+import argparse
+
+# 未传 --model 且未设置 MODEL_PATH 时使用的默认模型目录（请按本机实际路径修改）
+DEFAULT_MODEL_PATH = "/home/chenncy/llaisys/DeepSeek-R1-Distill-Qwen-1___5B"
+
+
+def _log(msg: str) -> None:
+ print(msg, flush=True)
+
+
+def main():
+ parser = argparse.ArgumentParser(description="LLAISYS Chatbot Server (OpenAI chat-completion API)")
+ parser.add_argument("--host", default="127.0.0.1", help="bind host")
+ parser.add_argument("--port", type=int, default=8000, help="bind port")
+ parser.add_argument("--model", default=None, help="model path (overrides MODEL_PATH env)")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], help="device")
+ args = parser.parse_args()
+
+ model_path = args.model or os.environ.get("MODEL_PATH") or DEFAULT_MODEL_PATH
+ if not model_path or not os.path.isdir(model_path):
+ _log("Warning: MODEL_PATH not set or not a directory. Set MODEL_PATH or use --model. Requests will return 503 until model is loaded.")
+ else:
+ _log(f"Model path: {model_path}")
+
+ _log("Importing uvicorn...")
+ try:
+ import uvicorn
+ except ImportError:
+ _log("Install uvicorn and fastapi: pip install uvicorn fastapi")
+ sys.exit(1)
+
+ _log("Importing create_app...")
+ try:
+ from .app import create_app
+ except Exception as e:
+ _log(f"Import failed (e.g. torch/llaisys): {e}")
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
+
+ _log("Loading tokenizer and model (may take 1-2 minutes)...")
+ try:
+ app = create_app(model_path=model_path, device=args.device)
+ except Exception as e:
+ _log(f"Failed to load model: {e}")
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
+
+ _log("Model ready. Starting server...")
+ try:
+ uvicorn.run(app, host=args.host, port=args.port)
+ except Exception as e:
+ _log(f"Server exited: {e}")
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
+
+
+if __name__ == "__main__":
+ try:
+ main()
+ except KeyboardInterrupt:
+ _log("\nStopped by user (Ctrl+C).")
+ sys.exit(0)
+ except Exception as e:
+ _log(f"Unexpected error: {e}")
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
diff --git a/python/llaisys_py/server/app.py b/python/llaisys_py/server/app.py
new file mode 100644
index 000000000..d0e5701f7
--- /dev/null
+++ b/python/llaisys_py/server/app.py
@@ -0,0 +1,1993 @@
+"""
+LLAISYS 聊天机器人服务端（FastAPI）。
+
+- 提供 OpenAI 风格的 Chat Completions API（POST /v1/chat/completions）。
+- 项目#4：多用户请求入队，单 worker 线程顺序处理；支持流式响应（SSE）与非流式。
+- 会话管理：GET/POST/PATCH/DELETE /v1/sessions、GET /v1/sessions/{id}、POST /v1/sessions/{id}/regenerate。
+- 模型与分词器在 create_app() 中按 MODEL_PATH 或参数加载，未加载时请求返回 503。
+"""
+import os
+import json
+import queue
+import threading
+import uuid
+from datetime import datetime, timezone
+from typing import Callable, Optional
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
+from pydantic import BaseModel, Field
+
+# ---------- 全局状态（在 create_app 中按需加载） ----------
+_tokenizer = None # HuggingFace AutoTokenizer，用于编码/解码与对话模板
+_model = None # LLAISYS Qwen2 模型实例
+_device_type = None # 当前设备类型（如 "cpu"），供后续扩展用
+_engine = None # 连续批处理 Engine（当 LLAISYS_USE_ENGINE_LOOP=1 且 max_batch_size>1 时创建）
+
+# ---------- 多轮对话 Web 页面（内联 HTML+JS，由 GET /chat 直接返回） ----------
+# 包含：样式（气泡、滚动、状态栏）、对话历史 DOM、流式请求与 SSE 解析逻辑
+_CHAT_HTML = """
+
+
+ LLAISYS 聊天
+
+
+
+
+
+
+
+"""
+
+# ---------- Agent 风格页面（新路由 /agent，同 API，多块展示：思考 + 回答） ----------
+_AGENT_HTML = """
+
+
+ LLAISYS Agent
+
+
+
+
+
+
+
+"""
+
+# ---------- 多会话管理页面：侧栏列表、新建/删除/切换、编辑消息并从此处重新生成 ----------
+_SESSIONS_HTML = """
+
+
+ LLAISYS 会话
+
+
+
+
+
+ 选择或新建一个对话
+
+
+ 发送
+
+
+
+
+
+
+"""
+
+
+def _get_model():
+ """返回已加载的 LLAISYS 模型；未加载时抛出 503。"""
+ if _model is None:
+ raise HTTPException(status_code=503, detail="Model not loaded. Set MODEL_PATH and restart.")
+ return _model
+
+
+def _get_tokenizer():
+ """返回已加载的分词器；未加载时抛出 503。"""
+ if _tokenizer is None:
+ raise HTTPException(status_code=503, detail="Tokenizer not loaded. Set MODEL_PATH and restart.")
+ return _tokenizer
+
+
+# ---------- 请求/响应模型（与 OpenAI Chat Completions 对齐） ----------
+
+class ChatMessage(BaseModel):
+ """单条对话消息。"""
+ role: str = Field(..., description="user | assistant | system")
+ content: str = Field(default="", description="message content")
+
+
+class ChatCompletionRequest(BaseModel):
+ """POST /v1/chat/completions 的请求体。"""
+ model: str = Field(default="default", description="模型名（当前忽略，使用服务端加载的模型）")
+ messages: list[ChatMessage] = Field(..., description="多轮对话历史，最后一条一般为 user")
+ session_id: Optional[str] = Field(default=None, description="可选；若提供，完成后将本轮 user+assistant 追加到该会话")
+ max_tokens: int = Field(default=512, ge=1, le=2048, description="本次最多生成的新 token 数")
+ temperature: float = Field(default=0.3, ge=0.0, le=2.0, description="默认 0.3 减少胡言乱语；若回复太死板可试 0.5")
+ top_p: float = Field(default=0.9, ge=0.0, le=1.0)
+ top_k: int = Field(default=40, ge=0, le=100)
+ stream: bool = Field(default=False, description="是否以 SSE 流式返回")
+ seed: Optional[int] = Field(default=None, description="随机种子；None 表示非确定性")
+
+
+# ---------- 会话存储与模型（内存存储，重启清空） ----------
+_sessions: dict[str, dict] = {} # session_id -> { id, title, messages, created_at, updated_at }
+_sessions_lock = threading.Lock()
+
+# ---------- KV-Cache 池（Phase 3）：key=(session_id, user_message_index), value={blob, prefix_len, last_used}, LRU ----------
+_KV_POOL_MAX_ENTRIES = int(os.environ.get("LLAISYS_KV_POOL_MAX", "16"))
+_kv_pool: dict[tuple[str, int], dict] = {} # (session_id, user_idx) -> {"blob": bytes, "prefix_len": int, "last_used": float}
+_kv_pool_lock = threading.Lock()
+import time as _time_module
+
+# ---------- 项目#4：请求队列与 worker ----------
+_REQUEST_QUEUE_MAX = int(os.environ.get("LLAISYS_REQUEST_QUEUE_MAX", "64"))
+_request_queue: queue.Queue = queue.Queue(maxsize=_REQUEST_QUEUE_MAX)
+_inference_lock = threading.Lock() # 推理互斥，worker 与 regenerate 共用模型时串行化
+_STREAM_SENTINEL = None # 流式结束标记，worker 放入 response_queue 表示结束
+
+
+def _kv_pool_get(session_id: str, user_message_index: int):
+ """若存在则返回 (blob, prefix_len) 并更新 last_used；否则返回 None。"""
+ with _kv_pool_lock:
+ key = (session_id, user_message_index)
+ if key not in _kv_pool:
+ return None
+ entry = _kv_pool[key]
+ entry["last_used"] = _time_module.perf_counter()
+ return (entry["blob"], entry["prefix_len"])
+
+
+def _kv_pool_put(session_id: str, user_message_index: int, blob: bytes, prefix_len: int) -> None:
+ """写入池；若超过容量则 LRU 淘汰。"""
+ with _kv_pool_lock:
+ while len(_kv_pool) >= _KV_POOL_MAX_ENTRIES and _kv_pool:
+ oldest_key = min(_kv_pool, key=lambda k: _kv_pool[k]["last_used"])
+ del _kv_pool[oldest_key]
+ key = (session_id, user_message_index)
+ _kv_pool[key] = {"blob": blob, "prefix_len": prefix_len, "last_used": _time_module.perf_counter()}
+
+
+def resolve_kv_prefix(session_id: Optional[str], request_messages: list, tokenizer, input_ids: list) -> tuple[int, Optional[bytes]]:
+ """若 session 且可命中 KV 池则返回 (prefix_len, blob)；否则返回 (0, None)。供 Engine 与 worker 共用。"""
+ if not session_id or not request_messages or tokenizer is None:
+ return (0, None)
+ user_count = _count_user_messages(request_messages)
+ if user_count == 0:
+ return (0, None)
+ hit = _kv_pool_get(session_id, user_count - 1)
+ if hit is None:
+ return (0, None)
+ blob, stored_prefix_len = hit
+ need_prefix_len = _prefix_len_for_messages(request_messages, tokenizer)
+ if stored_prefix_len != need_prefix_len or need_prefix_len <= 0 or need_prefix_len >= len(input_ids):
+ return (0, None)
+ return (stored_prefix_len, blob)
+
+
+def _now_iso() -> str:
+ return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _get_session(session_id: str) -> dict:
+ with _sessions_lock:
+ if session_id not in _sessions:
+ raise HTTPException(status_code=404, detail="Session not found")
+ return _sessions[session_id]
+
+
+class SessionCreate(BaseModel):
+ """创建会话请求体。"""
+ title: Optional[str] = Field(default=None, description="可选标题，默认用首条用户消息摘要")
+
+
+class SessionUpdate(BaseModel):
+ """更新会话请求体（PATCH）。"""
+ title: Optional[str] = Field(default=None, description="新标题")
+
+
+class SessionOut(BaseModel):
+ """会话响应。"""
+ id: str
+ title: Optional[str]
+ messages: list[dict] # [{ role, content }]
+ created_at: str
+ updated_at: str
+
+
+class RegenerateRequest(BaseModel):
+ """从某条消息后重新生成请求体。"""
+ from_message_index: int = Field(..., ge=0, description="保留到该条用户消息（含），之后全部删除并重新生成")
+ new_content: Optional[str] = Field(default=None, description="若提供，替换该条用户消息内容")
+ max_tokens: int = Field(default=512, ge=1, le=2048)
+ temperature: float = Field(default=0.3, ge=0.0, le=2.0)
+ top_k: int = Field(default=40, ge=0, le=100)
+ top_p: float = Field(default=0.9, ge=0.0, le=1.0)
+ stream: bool = Field(default=True)
+ seed: Optional[int] = Field(default=None)
+
+
+def _strip_think_tags(content: str) -> str:
+ """去掉助手回复中的 ... 推理块，只保留实际回答，避免存入会话后干扰后续对话。"""
+ if not content or not isinstance(content, str):
+ return content or ""
+ import re
+ content = re.sub(r"[\s\S]*?<\s*/\s*think\s*>", "", content, flags=re.IGNORECASE)
+ content = re.sub(r"\u003cthink\u003e[\s\S]*?\u003c/think\u003e", "", content)
+ return content.strip() or content
+
+
+# 退化输出时返回给用户的提示（避免界面只显示 "1 1 1" 等）
+_DEGENERATE_FALLBACK = "（回复异常，请重试。）"
+
+
+def _is_degenerate_output(full_content: list[str], last_token_ids: list[int], max_recent: int = 20) -> bool:
+ """
+ 检测是否陷入退化输出（如重复的 "1\\n0\\n0\\n"、纯数字、或同一词/ token 大量重复）。
+ 若最近内容仅包含数字/换行/空格且有一定长度，或同一词在最近片段中出现过多，则视为退化并建议停止。
+ """
+ if len(full_content) < 5:
+ return False
+ recent_text = "".join(full_content[-max_recent:])
+ if not recent_text.strip():
+ return True
+ allowed = set(" \n\t\r0123456789")
+ if not all(c in allowed for c in recent_text):
+ # 非纯数字：检查是否同一词重复过多（如 "regards regards regards"）
+ words = recent_text.split()
+ if len(words) >= 5:
+ from collections import Counter
+ cnt = Counter(words)
+ if cnt and cnt.most_common(1)[0][1] >= 4:
+ return True
+ return False
+ # 纯数字/换行且长度>=4 即视为退化（如 "1\\n0\\n" 或 "12\\n13"）
+ if len(recent_text.strip()) >= 4:
+ return True
+ return False
+
+
+def _is_content_only_digits_and_whitespace(full_content: list[str]) -> bool:
+ """判断整段内容是否仅包含数字、空格、换行（用于退化时是否用 fallback 替换）。"""
+ if not full_content:
+ return True
+ text = "".join(full_content).strip()
+ if not text:
+ return True
+ return all(c in " \n\t\r0123456789" for c in text)
+
+
+# 系统提示：留空避免触发 DeepSeek-R1 的固定“自我介绍”；仅当明确需要时可设环境变量
+def _get_system_prompt() -> str:
+ return os.environ.get("LLAISYS_SYSTEM_PROMPT", "").strip()
+
+
+def _messages_to_prompt(messages: list[ChatMessage], tokenizer) -> str:
+ """将 OpenAI 风格 messages 转为带对话模板的输入文本（含 system/user/assistant 格式）。"""
+ conversation = [{"role": m.role, "content": m.content} for m in messages]
+ system = _get_system_prompt()
+ if system and not any(m.get("role") == "system" for m in conversation):
+ conversation.insert(0, {"role": "system", "content": system})
+ prompt = tokenizer.apply_chat_template(
+ conversation,
+ add_generation_prompt=True,
+ tokenize=False,
+ )
+ if os.environ.get("LLAISYS_DEBUG"):
+ import sys
+ print(f"[LLAISYS] prompt len={len(prompt)} last_300={repr(prompt[-300:])}", file=sys.stderr)
+ return prompt
+
+
+def _count_user_messages(messages: list[ChatMessage]) -> int:
+ return sum(1 for m in messages if m.role == "user")
+
+
+def _prefix_len_for_messages(messages: list[ChatMessage], tokenizer) -> int:
+ """编码「去掉最后一条」的 messages 得到的 token 长度，用于池 key 的 prefix_len 校验。"""
+ if not messages:
+ return 0
+ conv = [{"role": m.role, "content": m.content} for m in messages[:-1]]
+ system = _get_system_prompt()
+ if system and not any(c.get("role") == "system" for c in conv):
+ conv.insert(0, {"role": "system", "content": system})
+ prompt = tokenizer.apply_chat_template(conv, add_generation_prompt=True, tokenize=False)
+ return len(tokenizer.encode(prompt))
+
+
+# 当传入路径无效时尝试的备用模型目录（与 __main__.py 中 DEFAULT_MODEL_PATH 保持一致）
+_FALLBACK_MODEL_PATH = "/home/chenncy/llaisys/DeepSeek-R1-Distill-Qwen-1___5B"
+
+
+def create_app(model_path: Optional[str] = None, device: str = "cpu"):
+ """
+ 创建 FastAPI 应用并（在路径有效时）加载模型与分词器。
+ model_path 为空时从环境变量 MODEL_PATH 读取；路径无效时再尝试 _FALLBACK_MODEL_PATH；仍无效则不加载，请求返回 503。
+ """
+ global _tokenizer, _model, _device_type
+ path = model_path or os.environ.get("MODEL_PATH")
+ if not path or not os.path.isdir(path):
+ path = _FALLBACK_MODEL_PATH if os.path.isdir(_FALLBACK_MODEL_PATH) else None
+ if not path or not os.path.isdir(path):
+ _tokenizer = None
+ _model = None
+ _device_type = device
+ app = FastAPI(title="LLAISYS Chatbot", description="OpenAI chat-completion style API")
+ _register_routes(app)
+ return app
+
+ # 在导入 transformers（会拉取 torch）之前检测 GPU；CUDA_VISIBLE_DEVICES="" 会导致 cudaGetDeviceCount() 返回 0
+ if device == "nvidia":
+ import sys
+ _cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
+ if _cvd == "" or _cvd is None:
+ # 空或未设置时强制至少可见 0 号卡，避免容器/DSW 默认传 "" 导致看不到 GPU
+ os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+ if _cvd == "":
+ print("Info: CUDA_VISIBLE_DEVICES was empty, set to 0 for this process.", file=sys.stderr)
+ try:
+ from ..libllaisys import DeviceType, LIB_LLAISYS
+ nvidia_api = LIB_LLAISYS.llaisysGetRuntimeAPI(DeviceType.NVIDIA)
+ _count = nvidia_api.contents.get_device_count()
+ if _count == 0:
+ print("Warning: no NVIDIA GPUs available (get_device_count()=0), using CPU. (CUDA_VISIBLE_DEVICES=%r)" % os.environ.get("CUDA_VISIBLE_DEVICES"), file=sys.stderr)
+ device = "cpu"
+ else:
+ print("Info: using NVIDIA GPU(s), count=%s." % _count, file=sys.stderr)
+ except Exception as e:
+ print("Warning: NVIDIA check failed (%s), using CPU." % e, file=sys.stderr)
+ device = "cpu"
+
+ from transformers import AutoTokenizer
+ from ..libllaisys import DeviceType
+ from ..models.qwen2 import Qwen2
+
+ _device_type = device
+ _tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+ dev = DeviceType.NVIDIA if device == "nvidia" else DeviceType.CPU
+ use_engine = os.environ.get("LLAISYS_USE_ENGINE_LOOP", "").strip().lower() in ("1", "true", "yes")
+ max_batch_size = int(os.environ.get("LLAISYS_MAX_BATCH_SIZE", "4" if use_engine else "1"))
+ _model = Qwen2(path, device=dev, max_batch_size=max_batch_size)
+
+ global _engine
+ _engine = None
+ if use_engine and max_batch_size >= 1:
+ from .engine import Engine
+ def _engine_get_kv(session_id, request_messages, input_ids):
+ return resolve_kv_prefix(session_id, request_messages, _tokenizer, input_ids)
+ def _engine_put_kv(session_id, request_messages, blob, prefix_len):
+ _kv_pool_put(session_id, _count_user_messages(request_messages), blob, prefix_len)
+ _engine = Engine(
+ _model,
+ max_batch_size,
+ pending_maxsize=_REQUEST_QUEUE_MAX,
+ get_kv=_engine_get_kv,
+ put_kv=_engine_put_kv,
+ )
+
+ app = FastAPI(title="LLAISYS Chatbot", description="OpenAI chat-completion style API")
+ _register_routes(app)
+ return app
+
+
+def _register_routes(app: FastAPI):
+ """注册所有 HTTP 路由。"""
+
+ @app.get("/")
+ def root():
+ """根路径：返回服务说明与常用链接。"""
+ return {
+ "message": "LLAISYS Chatbot Server",
+ "docs": "/docs",
+ "health": "/health",
+ "chat_ui": "/chat",
+ "agent_ui": "/agent",
+ "sessions_ui": "/sessions",
+ "chat_api": "POST /v1/chat/completions",
+ }
+
+ @app.get("/health")
+ def health():
+ """健康检查：是否存活及模型是否已加载。"""
+ return {"status": "ok", "model_loaded": _model is not None}
+
+ @app.get("/v1/metrics")
+ def metrics():
+ """项目#4 监控指标（4.5.3）：队列长度、Engine 状态、KV 池大小等。"""
+ out = {
+ "request_queue_size": _request_queue.qsize(),
+ "request_queue_max": _REQUEST_QUEUE_MAX,
+ "kv_pool_size": len(_kv_pool),
+ "kv_pool_max": _KV_POOL_MAX_ENTRIES,
+ }
+ if _engine is not None:
+ out["engine"] = _engine.get_metrics()
+ return out
+
+ @app.get("/chat", response_class=HTMLResponse)
+ def chat_page():
+ """返回内联的 Web 聊天页（HTML+JS），支持多轮对话与流式显示。"""
+ return _CHAT_HTML
+
+ @app.get("/agent", response_class=HTMLResponse)
+ def agent_page():
+ """返回 Agent 风格页面：思考 + 回答分块展示，仍使用同一流式 API。"""
+ return _AGENT_HTML
+
+ @app.get("/sessions", response_class=HTMLResponse)
+ def sessions_page():
+ """多会话管理 UI：列表、新建/删除/切换、编辑消息并从此处重新生成。"""
+ return _SESSIONS_HTML
+
+ # ---------- 会话 API ----------
+ @app.get("/v1/sessions")
+ def list_sessions():
+ """列出所有会话，按 updated_at 倒序。"""
+ with _sessions_lock:
+ out = []
+ for sid, s in _sessions.items():
+ out.append({
+ "id": s["id"],
+ "title": s.get("title"),
+ "updated_at": s["updated_at"],
+ "message_count": len(s.get("messages", [])),
+ })
+ out.sort(key=lambda x: x["updated_at"], reverse=True)
+ return {"sessions": out}
+
+ @app.post("/v1/sessions")
+ def create_session(body: Optional[SessionCreate] = None):
+ """创建新会话。"""
+ sid = str(uuid.uuid4())
+ now = _now_iso()
+ with _sessions_lock:
+ _sessions[sid] = {
+ "id": sid,
+ "title": (body.title if body else None) or "新对话",
+ "messages": [],
+ "created_at": now,
+ "updated_at": now,
+ }
+ return _sessions[sid]
+
+ @app.get("/v1/sessions/{session_id}")
+ def get_session(session_id: str):
+ """获取会话详情。"""
+ s = _get_session(session_id)
+ return SessionOut(
+ id=s["id"],
+ title=s.get("title"),
+ messages=s.get("messages", []),
+ created_at=s["created_at"],
+ updated_at=s["updated_at"],
+ )
+
+ @app.patch("/v1/sessions/{session_id}")
+ def update_session(session_id: str, body: SessionUpdate):
+ """更新会话（如标题）。"""
+ with _sessions_lock:
+ if session_id not in _sessions:
+ raise HTTPException(status_code=404, detail="Session not found")
+ s = _sessions[session_id]
+ if body.title is not None:
+ s["title"] = body.title
+ s["updated_at"] = _now_iso()
+ return dict(s)
+
+ @app.delete("/v1/sessions/{session_id}")
+ def delete_session(session_id: str):
+ """删除会话。"""
+ with _sessions_lock:
+ if session_id not in _sessions:
+ raise HTTPException(status_code=404, detail="Session not found")
+ del _sessions[session_id]
+ return {"ok": True}
+
+ def _resolve_kv_prefix_regenerate(session_id: str, from_message_index: int, request_messages: list[ChatMessage], tokenizer, input_ids: list[int]):
+ """Regenerate 专用：查池 key=(session_id, from_message_index)，命中则返回 (prefix_len, blob)。"""
+ hit = _kv_pool_get(session_id, from_message_index)
+ if hit is None:
+ return (0, None)
+ blob, stored_prefix_len = hit
+ need_prefix_len = _prefix_len_for_messages(request_messages, tokenizer)
+ if stored_prefix_len != need_prefix_len or need_prefix_len <= 0 or need_prefix_len >= len(input_ids):
+ return (0, None)
+ return (stored_prefix_len, blob)
+
+ @app.post("/v1/sessions/{session_id}/regenerate")
+ def regenerate_session(session_id: str, body: RegenerateRequest):
+ """从某条用户消息后截断并重新生成；可选替换该条内容；命中 KV 池则 suffix prefill，并写回池。"""
+ s = _get_session(session_id)
+ messages = list(s.get("messages", []))
+ user_indices = [i for i, m in enumerate(messages) if m.get("role") == "user"]
+ if body.from_message_index >= len(user_indices):
+ raise HTTPException(
+ status_code=400,
+ detail="from_message_index out of range (no such user message)",
+ )
+ cut_at = user_indices[body.from_message_index]
+ messages = messages[: cut_at + 1]
+ if body.new_content is not None:
+ messages[-1] = {"role": "user", "content": body.new_content}
+ s["messages"] = messages
+ s["updated_at"] = _now_iso()
+
+ tokenizer = _get_tokenizer()
+ model = _get_model()
+ seed = body.seed if body.seed is not None else 0
+ request_messages = [ChatMessage(role=m["role"], content=m["content"]) for m in messages]
+ prompt = _messages_to_prompt(request_messages, tokenizer)
+ input_ids = tokenizer.encode(prompt)
+ if not input_ids:
+ raise HTTPException(status_code=400, detail="Empty input after encoding")
+
+ prefix_len, kv_blob = _resolve_kv_prefix_regenerate(session_id, body.from_message_index, request_messages, tokenizer, input_ids)
+
+ if body.stream:
+ def gen():
+ with _inference_lock:
+ if kv_blob is not None and prefix_len > 0:
+ model.import_kv_cache(kv_blob, prefix_len)
+ full_content = []
+ tokens = list(input_ids)
+ next_id = None
+ n_remaining = body.max_tokens
+ if prefix_len == 0:
+ model.reset_kv_cache()
+ if prefix_len > 0 and prefix_len < len(input_ids):
+ suffix = input_ids[prefix_len:]
+ next_id = model.next_token(
+ suffix,
+ temperature=body.temperature,
+ top_k=body.top_k,
+ top_p=body.top_p,
+ seed=seed,
+ )
+ tokens.append(next_id)
+ if next_id != model.end_token:
+ delta_text = tokenizer.decode([next_id], skip_special_tokens=True)
+ if delta_text:
+ full_content.append(delta_text)
+ yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:24]}', 'choices': [{'index': 0, 'delta': {'content': delta_text}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
+ n_remaining = body.max_tokens - 1
+ for _ in range(n_remaining):
+ if next_id == model.end_token:
+ break
+ next_id = model.next_token(
+ tokens[-1:] if len(tokens) > 1 else tokens,
+ temperature=body.temperature,
+ top_k=body.top_k,
+ top_p=body.top_p,
+ seed=seed,
+ )
+ tokens.append(next_id)
+ if next_id == model.end_token:
+ break
+ delta_text = tokenizer.decode([next_id], skip_special_tokens=True)
+ if not delta_text:
+ continue
+ full_content.append(delta_text)
+ yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:24]}', 'choices': [{'index': 0, 'delta': {'content': delta_text}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
+ yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:24]}', 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
+ yield "data: [DONE]\n\n"
+ content = "".join(full_content) or "(无回复)"
+ s["messages"].append({"role": "assistant", "content": _strip_think_tags(content)})
+ s["updated_at"] = _now_iso()
+ _kv_pool_put(session_id, body.from_message_index + 1, model.export_kv_cache(), len(input_ids))
+
+ return StreamingResponse(
+ gen(),
+ media_type="text/event-stream",
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+ )
+ with _inference_lock:
+ if kv_blob is not None and prefix_len > 0:
+ model.import_kv_cache(kv_blob, prefix_len)
+ full_tokens = model.generate(
+ input_ids,
+ max_new_tokens=body.max_tokens,
+ temperature=body.temperature,
+ top_k=body.top_k,
+ top_p=body.top_p,
+ seed=seed,
+ prefix_len=prefix_len,
+ )
+ new_tokens = full_tokens[len(input_ids):]
+ content = tokenizer.decode(new_tokens, skip_special_tokens=True) or "(无回复)"
+ content_clean = _strip_think_tags(content)
+ s["messages"].append({"role": "assistant", "content": content_clean})
+ s["updated_at"] = _now_iso()
+ _kv_pool_put(session_id, body.from_message_index + 1, model.export_kv_cache(), len(input_ids))
+ return {
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "object": "chat.completion",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": content_clean}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": len(input_ids), "completion_tokens": len(new_tokens), "total_tokens": len(full_tokens)},
+ }
+
+ def _resolve_kv_prefix(session_id: Optional[str], request_messages: list[ChatMessage], tokenizer, input_ids: list[int]):
+ """若 session 且可命中池则返回 (prefix_len, blob)；否则返回 (0, None)。"""
+ return resolve_kv_prefix(session_id, request_messages, tokenizer, input_ids)
+
+ def _worker_loop():
+ """项目#4：单 worker 线程，从请求队列取任务并执行推理，结果放入各请求的 response_queue。"""
+ while True:
+ item = _request_queue.get()
+ resp_queue = item["response_queue"]
+ cancel_event = item.get("cancel_event")
+ try:
+ with _inference_lock:
+ tokenizer = _get_tokenizer()
+ model = _get_model()
+ if tokenizer is None or model is None:
+ resp_queue.put({"__error": "Model not loaded", "status_code": 503})
+ continue
+ session_id = item.get("session_id")
+ messages = item["messages"]
+ stream = item["stream"]
+ seed = item["seed"]
+ prompt = _messages_to_prompt(messages, tokenizer)
+ input_ids = tokenizer.encode(prompt)
+ if not input_ids:
+ resp_queue.put({"__error": "Empty input after encoding", "status_code": 400})
+ continue
+ prefix_len, kv_blob = _resolve_kv_prefix(session_id, messages, tokenizer, input_ids)
+ if kv_blob is not None and prefix_len > 0:
+ model.import_kv_cache(kv_blob, prefix_len)
+ if stream:
+ if session_id:
+ with _sessions_lock:
+ s = _sessions.get(session_id)
+ if s is not None:
+ s["messages"] = [{"role": m.role, "content": m.content} for m in messages]
+ s["messages"].append({"role": "assistant", "content": ""})
+ s["updated_at"] = _now_iso()
+ cancel_check = (lambda: cancel_event.is_set()) if cancel_event else None
+ for chunk in _stream_chunks_generator(
+ model, tokenizer, input_ids, item["max_tokens"],
+ item["temperature"], item["top_k"], item["top_p"], seed,
+ prompt=prompt, session_id=session_id, request_messages=messages, prefix_len=prefix_len,
+ cancel_check=cancel_check,
+ ):
+ if cancel_event and cancel_event.is_set():
+ break
+ resp_queue.put(chunk)
+ if session_id and isinstance(chunk, str) and chunk.startswith("data: "):
+ try:
+ payload = chunk[6:].strip().strip("\n")
+ if payload and payload != "[DONE]":
+ data = json.loads(payload)
+ delta = (data.get("choices") or [{}])[0].get("delta") or {}
+ delta_content = delta.get("content")
+ if isinstance(delta_content, str) and delta_content:
+ with _sessions_lock:
+ s = _sessions.get(session_id)
+ if s and s["messages"] and s["messages"][-1]["role"] == "assistant":
+ s["messages"][-1]["content"] += delta_content
+ s["updated_at"] = _now_iso()
+ except (json.JSONDecodeError, IndexError, KeyError, TypeError):
+ pass
+ resp_queue.put(_STREAM_SENTINEL)
+ else:
+ full_tokens = model.generate(
+ input_ids,
+ max_new_tokens=item["max_tokens"],
+ temperature=item["temperature"],
+ top_k=item["top_k"],
+ top_p=item["top_p"],
+ seed=seed,
+ prefix_len=prefix_len,
+ )
+ new_tokens = full_tokens[len(input_ids):]
+ content = tokenizer.decode(new_tokens, skip_special_tokens=True)
+ content_clean = _strip_think_tags(content)
+ if session_id:
+ with _sessions_lock:
+ s = _sessions.get(session_id)
+ if s is not None:
+ s["messages"] = [{"role": m.role, "content": m.content} for m in messages]
+ s["messages"].append({"role": "assistant", "content": content_clean})
+ s["updated_at"] = _now_iso()
+ if s.get("title") == "新对话" and messages:
+ first = messages[0].content.strip()[:30]
+ if first:
+ s["title"] = first + ("…" if len(messages[0].content) > 30 else "")
+ _kv_pool_put(session_id, _count_user_messages(messages), model.export_kv_cache(), len(input_ids))
+ resp_queue.put({
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "object": "chat.completion",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": content_clean}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": len(input_ids), "completion_tokens": len(new_tokens), "total_tokens": len(full_tokens)},
+ })
+ except Exception as e:
+ resp_queue.put({"__error": str(e), "status_code": 500})
+
+ @app.post("/v1/chat/completions")
+ def chat_completions(req: ChatCompletionRequest):
+ """OpenAI 风格对话补全。启用 Engine 时走连续批处理；否则入队由单 worker 顺序处理。"""
+ if req.session_id:
+ _get_session(req.session_id) # 404 if not found
+ seed = req.seed if req.seed is not None else 0
+
+ # ---------- Engine 路径：连续批处理，Prefill + Batched Decode ----------
+ if _engine is not None:
+ from .engine import RequestState, _StreamError
+ tokenizer = _get_tokenizer()
+ model = _get_model()
+ if tokenizer is None or model is None:
+ raise HTTPException(status_code=503, detail="Model not loaded")
+ prompt = _messages_to_prompt(req.messages, tokenizer)
+ input_ids = tokenizer.encode(prompt)
+ if not input_ids:
+ raise HTTPException(status_code=400, detail="Empty input after encoding")
+ out_queue = queue.Queue()
+ req_state = RequestState(
+ request_id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ prompt_tokens=input_ids,
+ max_tokens=req.max_tokens,
+ out_queue=out_queue,
+ temperature=req.temperature,
+ top_k=req.top_k,
+ top_p=req.top_p,
+ seed=seed,
+ session_id=req.session_id,
+ request_messages=req.messages,
+ )
+ try:
+ _engine.submit_request(req_state)
+ except queue.Full:
+ raise HTTPException(status_code=503, detail="Request queue full, try again later")
+
+ if req.stream:
+ def stream_from_engine():
+ full_content = []
+ while True:
+ item = req_state.out_queue.get(timeout=300)
+ if item is None:
+ break
+ if isinstance(item, _StreamError):
+ yield f"data: {json.dumps({'error': item.message})}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+ delta_text = tokenizer.decode([item], skip_special_tokens=True)
+ if delta_text:
+ full_content.append(delta_text)
+ yield f"data: {json.dumps({'id': req_state.request_id, 'choices': [{'index': 0, 'delta': {'content': delta_text}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
+ yield f"data: {json.dumps({'id': req_state.request_id, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
+ yield "data: [DONE]\n\n"
+ content = "".join(full_content) or "(无回复)"
+ if req.session_id and req.messages:
+ with _sessions_lock:
+ s = _sessions.get(req.session_id)
+ if s is not None:
+ s["messages"] = [{"role": m.role, "content": m.content} for m in req.messages]
+ s["messages"].append({"role": "assistant", "content": _strip_think_tags(content)})
+ s["updated_at"] = _now_iso()
+ if s.get("title") == "新对话" and req.messages:
+ first = req.messages[0].content.strip()[:30]
+ if first:
+ s["title"] = first + ("…" if len(req.messages[0].content) > 30 else "")
+
+ return StreamingResponse(
+ stream_from_engine(),
+ media_type="text/event-stream",
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+ )
+
+ collected = []
+ while True:
+ item = req_state.out_queue.get(timeout=300)
+ if item is None:
+ break
+ if isinstance(item, _StreamError):
+ raise HTTPException(status_code=500, detail=item.message)
+ collected.append(item)
+ new_tokens = collected
+ content = tokenizer.decode(new_tokens, skip_special_tokens=True) or "(无回复)"
+ content_clean = _strip_think_tags(content)
+ if req.session_id and req.messages:
+ with _sessions_lock:
+ s = _sessions.get(req.session_id)
+ if s is not None:
+ s["messages"] = [{"role": m.role, "content": m.content} for m in req.messages]
+ s["messages"].append({"role": "assistant", "content": content_clean})
+ s["updated_at"] = _now_iso()
+ if s.get("title") == "新对话" and req.messages:
+ first = req.messages[0].content.strip()[:30]
+ if first:
+ s["title"] = first + ("…" if len(req.messages[0].content) > 30 else "")
+ return {
+ "id": req_state.request_id,
+ "object": "chat.completion",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": content_clean}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": len(input_ids), "completion_tokens": len(new_tokens), "total_tokens": len(input_ids) + len(new_tokens)},
+ }
+
+ # ---------- 原有 worker 路径 ----------
+ response_queue = queue.Queue()
+ cancel_event = threading.Event()
+ try:
+ _request_queue.put_nowait({
+ "response_queue": response_queue,
+ "cancel_event": cancel_event,
+ "session_id": req.session_id,
+ "messages": req.messages,
+ "stream": req.stream,
+ "max_tokens": req.max_tokens,
+ "temperature": req.temperature,
+ "top_k": req.top_k,
+ "top_p": req.top_p,
+ "seed": seed,
+ })
+ except queue.Full:
+ raise HTTPException(status_code=503, detail="Request queue full, try again later")
+ if req.stream:
+ response_queue.put("data: " + json.dumps(
+ {"choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": None}]},
+ ensure_ascii=False,
+ ) + "\n\n")
+ def stream_from_queue():
+ try:
+ while True:
+ chunk = response_queue.get()
+ if chunk is _STREAM_SENTINEL:
+ return
+ if isinstance(chunk, dict):
+ yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+ yield chunk
+ finally:
+ cancel_event.set()
+ return StreamingResponse(
+ stream_from_queue(),
+ media_type="text/event-stream",
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+ )
+ resp = response_queue.get(timeout=300)
+ if isinstance(resp, dict) and resp.get("__error"):
+ raise HTTPException(status_code=resp.get("status_code", 500), detail=resp["__error"])
+ return resp
+
+ def _stream_chunks_generator(
+ model, tokenizer, input_ids, max_tokens, temperature, top_k, top_p, seed,
+ prompt: str = "",
+ session_id: Optional[str] = None,
+ request_messages: Optional[list[ChatMessage]] = None,
+ prefix_len: int = 0,
+ cancel_check: Optional[Callable[[], bool]] = None,
+ ):
+ """
+ 流式生成 SSE 块；供 _stream_response 与 worker 共用。yield 若干 "data: {...}\\n\\n" 字符串。
+ cancel_check: 若提供且返回 True 则提前结束生成（用于客户端断开时取消推理）。
+ """
+ import sys
+ _chat_debug = os.environ.get("LLAISYS_CHAT_DEBUG", "").strip() in ("1", "true", "yes")
+
+ def generate():
+ full_content = []
+ recent_token_ids = [] # 用于退化检测
+ stopped_degenerate = False # 是否因退化输出而提前停止
+ tokens = list(input_ids)
+ next_id = None
+ n_remaining = max_tokens
+ if _chat_debug and prompt:
+ eos_id = getattr(tokenizer, "eos_token_id", None)
+ print(f"[CHAT_DEBUG] prompt_len={len(prompt)} input_ids_len={len(input_ids)} prefix_len={prefix_len} tokenizer.eos_token_id={eos_id} model.end_token={model.end_token}", file=sys.stderr)
+ print(f"[CHAT_DEBUG] prompt_tail(400)={repr(prompt[-400:])}", file=sys.stderr)
+ try:
+ decoded_prompt = tokenizer.decode(input_ids, skip_special_tokens=False)
+ print(f"[CHAT_DEBUG] decoded_input_ids={repr(decoded_prompt)}", file=sys.stderr)
+ except Exception as e:
+ print(f"[CHAT_DEBUG] decode err: {e}", file=sys.stderr)
+ if len(input_ids) <= 20:
+ print(f"[CHAT_DEBUG] input_ids={input_ids}", file=sys.stderr)
+ else:
+ print(f"[CHAT_DEBUG] input_ids[:15]={input_ids[:15]} ... input_ids[-5:]={input_ids[-5:]}", file=sys.stderr)
+ if prefix_len == 0:
+ model.reset_kv_cache()
+ # 每次采样使用不同 seed，避免 C 层每步用同一 seed 重设 RNG 导致重复采样同一 token（如出现 "0" 后立刻退化）
+ sampling_step = 0
+ if prefix_len > 0 and prefix_len < len(input_ids):
+ suffix = input_ids[prefix_len:]
+ next_id = model.next_token(
+ suffix,
+ temperature=temperature,
+ top_k=top_k,
+ top_p=top_p,
+ seed=seed + sampling_step,
+ )
+ sampling_step += 1
+ tokens.append(next_id)
+ recent_token_ids.append(next_id)
+ if _chat_debug:
+ print(f"[CHAT_DEBUG] token_id={next_id} end_token={model.end_token} delta={repr(tokenizer.decode([next_id], skip_special_tokens=True))}", file=sys.stderr)
+ if next_id != model.end_token:
+ delta_text = tokenizer.decode([next_id], skip_special_tokens=True)
+ if delta_text:
+ full_content.append(delta_text)
+ chunk = {
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "choices": [{"index": 0, "delta": {"content": delta_text}, "finish_reason": None}],
+ }
+ yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
+ n_remaining = max_tokens - 1
+ step = 0
+ for _ in range(n_remaining):
+ if cancel_check and cancel_check():
+ break
+ if next_id == model.end_token:
+ break
+ # 首步且 prefix_len==0 时必须传入完整 prompt 做 prefill，否则只传最后一个 token 做 decode
+ if next_id is None and len(tokens) > 1:
+ next_id = model.next_token(
+ tokens,
+ temperature=temperature,
+ top_k=top_k,
+ top_p=top_p,
+ seed=seed + sampling_step,
+ )
+ else:
+ next_id = model.next_token(
+ tokens[-1:] if len(tokens) > 1 else tokens,
+ temperature=temperature,
+ top_k=top_k,
+ top_p=top_p,
+ seed=seed + sampling_step,
+ )
+ sampling_step += 1
+ tokens.append(next_id)
+ recent_token_ids.append(next_id)
+ if len(recent_token_ids) > 30:
+ recent_token_ids.pop(0)
+ if _chat_debug:
+ step += 1
+ dt = tokenizer.decode([next_id], skip_special_tokens=True)
+ print(f"[CHAT_DEBUG] step={step} token_id={next_id} end={next_id == model.end_token} delta={repr(dt)}", file=sys.stderr)
+ if next_id == model.end_token:
+ break
+ # 尽早检测：同一 token 连续或多次重复则视为退化，不输出当前 token 直接停
+ if len(recent_token_ids) >= 2 and next_id == recent_token_ids[-2]:
+ stopped_degenerate = True
+ break
+ if len(recent_token_ids) >= 5:
+ from collections import Counter
+ cnt = Counter(recent_token_ids[-6:])
+ if cnt and cnt.most_common(1)[0][1] >= 3:
+ stopped_degenerate = True
+ break
+ delta_text = tokenizer.decode([next_id], skip_special_tokens=True)
+ if not delta_text:
+ continue
+ full_content.append(delta_text)
+ # 退化输出检测：同一词重复、纯数字/换行等
+ if _is_degenerate_output(full_content, recent_token_ids):
+ if _chat_debug:
+ print(f"[CHAT_DEBUG] stop: degenerate output detected", file=sys.stderr)
+ stopped_degenerate = True
+ break
+ chunk = {
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "choices": [{"index": 0, "delta": {"content": delta_text}, "finish_reason": None}],
+ }
+ yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
+ if _chat_debug:
+ full_text = "".join(full_content)
+ print(f"[CHAT_DEBUG] done steps={step} full_text_len={len(full_text)} full_text={repr(full_text[:500])}", file=sys.stderr)
+ # 若因退化停止：仅当整段内容纯数字/空白等明显垃圾时才用 fallback 替换；否则保留已生成内容，只停止续写
+ if stopped_degenerate:
+ if _is_content_only_digits_and_whitespace(full_content):
+ fallback_chunk = {
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "choices": [{"index": 0, "delta": {"content": _DEGENERATE_FALLBACK}, "finish_reason": None}],
+ }
+ yield f"data: {json.dumps(fallback_chunk, ensure_ascii=False)}\n\n"
+ full_content = [_DEGENERATE_FALLBACK]
+ # 否则 full_content 保持不变，会话保存已生成的部分内容
+ chunk = {
+ "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }
+ yield f"data: {json.dumps(chunk)}\n\n"
+ yield "data: [DONE]\n\n"
+ if session_id and request_messages is not None:
+ content = "".join(full_content) or "(无回复)"
+ with _sessions_lock:
+ s = _sessions.get(session_id)
+ if s is not None:
+ s["messages"] = [{"role": m.role, "content": m.content} for m in request_messages]
+ s["messages"].append({"role": "assistant", "content": _strip_think_tags(content)})
+ s["updated_at"] = _now_iso()
+ if s.get("title") == "新对话" and request_messages:
+ first = request_messages[0].content.strip()[:30]
+ if first:
+ s["title"] = first + ("…" if len(request_messages[0].content) > 30 else "")
+ _kv_pool_put(session_id, _count_user_messages(request_messages), model.export_kv_cache(), len(input_ids))
+
+ yield from generate()
+
+ def _stream_response(
+ model, tokenizer, input_ids, max_tokens, temperature, top_k, top_p, seed,
+ prompt: str = "",
+ session_id: Optional[str] = None,
+ request_messages: Optional[list[ChatMessage]] = None,
+ prefix_len: int = 0,
+ ):
+ """SSE 流式响应；委托 _stream_chunks_generator 生成块。"""
+ return StreamingResponse(
+ _stream_chunks_generator(
+ model, tokenizer, input_ids, max_tokens, temperature, top_k, top_p, seed,
+ prompt=prompt,
+ session_id=session_id,
+ request_messages=request_messages,
+ prefix_len=prefix_len,
+ ),
+ media_type="text/event-stream",
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+ )
+
+ @app.on_event("startup")
+ def _start_request_worker():
+ """启动请求处理：启用 Engine 时由其自带线程负责；否则启动单 worker 线程。"""
+ if _engine is not None:
+ return # Engine 在 create_app 中已启动 _step_loop 线程
+ t = threading.Thread(target=_worker_loop, daemon=True)
+ t.start()
diff --git a/python/llaisys_py/server/chat_cli.py b/python/llaisys_py/server/chat_cli.py
new file mode 100644
index 000000000..0737bc27b
--- /dev/null
+++ b/python/llaisys_py/server/chat_cli.py
@@ -0,0 +1,92 @@
+"""
+命令行聊天客户端：多轮与 LLAISYS Chatbot Server 对话。
+用法：先启动 server，再在本机运行
+ python -m llaisys_py.server.chat_cli
+ python -m llaisys_py.server.chat_cli --base-url http://127.0.0.1:8000 --max-tokens 128
+"""
+import argparse
+import json
+import sys
+
+try:
+ import requests
+except ImportError:
+ print("请安装 requests: pip install requests")
+ sys.exit(1)
+
+
+def main():
+ parser = argparse.ArgumentParser(description="LLAISYS Chatbot 命令行客户端（多轮对话）")
+ parser.add_argument("--base-url", default="http://127.0.0.1:8000", help="服务器地址")
+ parser.add_argument("--max-tokens", type=int, default=128, help="每轮最多生成 token 数")
+ parser.add_argument("--temperature", type=float, default=0.8)
+ parser.add_argument("--top-k", type=int, default=50)
+ parser.add_argument("--top-p", type=float, default=0.9)
+ args = parser.parse_args()
+
+ url = args.base_url.rstrip("/") + "/v1/chat/completions"
+ health_url = args.base_url.rstrip("/") + "/health"
+
+ print("LLAISYS 命令行聊天（多轮对话）")
+ print(f"服务器: {args.base_url}")
+ try:
+ r = requests.get(health_url, timeout=5)
+ if r.status_code != 200:
+ print("警告: /health 返回非 200，请确认服务已启动且已加载模型")
+ else:
+ d = r.json()
+ print("模型已加载" if d.get("model_loaded") else "模型未加载，请求可能返回 503")
+ except requests.RequestException as e:
+ print(f"无法连接服务器: {e}")
+ print("请先启动: python -m llaisys_py.server --model --port 8000")
+ sys.exit(1)
+
+ print("输入内容后回车发送；输入 quit / exit / q 退出。\n")
+ messages = []
+
+ while True:
+ try:
+ user_input = input("你: ").strip()
+ except (EOFError, KeyboardInterrupt):
+ print("\n再见。")
+ break
+ if not user_input:
+ continue
+ if user_input.lower() in ("quit", "exit", "q"):
+ print("再见。")
+ break
+
+ messages.append({"role": "user", "content": user_input})
+ payload = {
+ "messages": messages,
+ "max_tokens": args.max_tokens,
+ "temperature": args.temperature,
+ "top_k": args.top_k,
+ "top_p": args.top_p,
+ "stream": False,
+ }
+
+ try:
+ r = requests.post(url, json=payload, timeout=300)
+ r.raise_for_status()
+ data = r.json()
+ except requests.RequestException as e:
+ print(f"请求失败: {e}")
+ if hasattr(e, "response") and e.response is not None and e.response.text:
+ print(e.response.text[:500])
+ messages.pop()
+ continue
+
+ choice = data.get("choices", [{}])[0]
+ msg = choice.get("message", {})
+ content = msg.get("content", "")
+ if not content:
+ print("(无回复内容)")
+ else:
+ print("助手:", content)
+ messages.append({"role": "assistant", "content": content})
+ print()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/python/llaisys_py/server/engine.py b/python/llaisys_py/server/engine.py
new file mode 100644
index 000000000..3ac667f8f
--- /dev/null
+++ b/python/llaisys_py/server/engine.py
@@ -0,0 +1,263 @@
+"""
+连续批处理引擎：按迭代为周期的状态机循环。
+
+- SlotManager：管理 KV-Cache 槽位分配与回收。
+- RequestState：单个请求在生命周期内的状态。
+- Engine：后台 _step_loop 执行 Prefill + Batched Decode，与 FastAPI 通过 out_queue 解耦。
+"""
+from __future__ import annotations
+
+import queue
+import threading
+import ctypes
+from ctypes import POINTER, cast, c_float, c_int, c_int64, c_size_t, c_ulonglong
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+from ..libllaisys import LIB_LLAISYS
+
+if TYPE_CHECKING:
+ from ..models.qwen2 import Qwen2
+
+
+class _StreamError:
+ """流式输出错误标记，放入 out_queue 后消费者应处理并结束流。"""
+ def __init__(self, message: str):
+ self.message = message
+
+
+class SlotManager:
+ """
+ 管理 KV-Cache 的槽位分配。
+ 使用列表存储空闲的 slot_id，确保不会超过 max_batch_size。
+ """
+ def __init__(self, max_batch_size: int):
+ self.max_batch_size = max_batch_size
+ self.free_slots = list(range(max_batch_size))
+ self.used_slots: set[int] = set()
+
+ def allocate(self) -> int:
+ """弹出并返回一个空闲槽位；无空闲时抛出 RuntimeError。"""
+ if not self.free_slots:
+ raise RuntimeError("No free slots available")
+ slot_id = self.free_slots.pop(0)
+ self.used_slots.add(slot_id)
+ return slot_id
+
+ def free(self, slot_id: int) -> None:
+ """将使用完毕的槽位回收。"""
+ if slot_id in self.used_slots:
+ self.used_slots.discard(slot_id)
+ self.free_slots.append(slot_id)
+
+
+class RequestState:
+ """
+ 单个请求在生命周期内的状态。
+ """
+ def __init__(
+ self,
+ request_id: str,
+ prompt_tokens: list[int],
+ max_tokens: int,
+ out_queue: queue.Queue,
+ *,
+ temperature: float = 0.7,
+ top_k: int = 50,
+ top_p: float = 0.9,
+ seed: int = 0,
+ session_id: Optional[str] = None,
+ request_messages: Optional[list] = None,
+ ):
+ self.request_id = request_id
+ self.prompt_tokens = prompt_tokens
+ self.max_tokens = max_tokens
+ self.out_queue = out_queue
+
+ self.temperature = temperature
+ self.top_k = top_k
+ self.top_p = top_p
+ self.seed = seed
+
+ self.slot_id: int = -1
+ self.generated_tokens: list[int] = []
+ self.is_finished: bool = False
+ self.last_token_id: int = -1
+
+ self.session_id = session_id
+ self.request_messages = request_messages
+
+
+class Engine:
+ """
+ 连续批处理引擎：pending_queue 接收新请求，_step_loop 中 Prefill + Batched Decode，
+ 结果通过 RequestState.out_queue 回写，与网络 I/O 解耦。
+ 支持 KV 池：get_kv/put_kv 可选，prefill 前查池命中则 suffix prefill，请求完成后写回池。
+ """
+ def __init__(
+ self,
+ model: "Qwen2",
+ max_batch_size: int,
+ pending_maxsize: int = 64,
+ *,
+ get_kv: Optional[Callable[..., Any]] = None,
+ put_kv: Optional[Callable[..., Any]] = None,
+ ):
+ self.model = model
+ self._c_model = model._model
+ self._end_token = model._end_token
+ self.max_batch_size = max_batch_size
+ self.slot_manager = SlotManager(max_batch_size)
+ self._get_kv = get_kv # (session_id, request_messages, input_ids) -> (prefix_len, blob)
+ self._put_kv = put_kv # (session_id, request_messages, blob, prefix_len) -> None
+
+ self.pending_queue: queue.Queue = queue.Queue(maxsize=pending_maxsize)
+ self.running_requests: list[RequestState] = []
+
+ self._engine_thread = threading.Thread(target=self._step_loop, daemon=True)
+ self._engine_thread.start()
+
+ def submit_request(self, req_state: RequestState) -> None:
+ """供 FastAPI 路由调用，将新请求加入等待队列。队列满时抛出 queue.Full。"""
+ self.pending_queue.put_nowait(req_state)
+
+ def _do_prefill(self, req: RequestState) -> int:
+ """
+ 对指定请求做 Prefill：若 KV 池命中则 import 到该 slot 后只对 suffix 做 prefill；
+ 否则重置该 slot 并传入完整 prompt。返回首个生成的 token_id。
+ """
+ if not req.prompt_tokens:
+ return self._end_token
+ input_ids = req.prompt_tokens
+ prefix_len = 0
+ blob = None
+ if self._get_kv and req.session_id and req.request_messages:
+ prefix_len, blob = self._get_kv(req.session_id, req.request_messages, input_ids)
+ if prefix_len > 0 and blob:
+ self.model.import_kv_cache_slot(req.slot_id, blob, prefix_len)
+ suffix = input_ids[prefix_len:]
+ if not suffix:
+ return self._end_token
+ n = len(suffix)
+ token_arr = (c_int64 * n)(*suffix)
+ first_token = LIB_LLAISYS.llaisysQwen2ModelInferWithSlot(
+ self._c_model,
+ c_size_t(req.slot_id),
+ cast(token_arr, POINTER(c_int64)),
+ c_size_t(n),
+ c_float(req.temperature),
+ c_int(req.top_k),
+ c_float(req.top_p),
+ c_ulonglong(req.seed),
+ )
+ return int(first_token)
+ LIB_LLAISYS.llaisysQwen2ModelResetKVCacheSlot(self._c_model, req.slot_id)
+ n = len(input_ids)
+ token_arr = (c_int64 * n)(*input_ids)
+ first_token = LIB_LLAISYS.llaisysQwen2ModelInferWithSlot(
+ self._c_model,
+ c_size_t(req.slot_id),
+ cast(token_arr, POINTER(c_int64)),
+ c_size_t(n),
+ c_float(req.temperature),
+ c_int(req.top_k),
+ c_float(req.top_p),
+ c_ulonglong(req.seed),
+ )
+ return int(first_token)
+
+ def _step_loop(self) -> None:
+ """
+ 引擎主循环：Prefill 新请求 → Batched Decode 当前请求 → 状态更新与槽位回收。
+ """
+ while True:
+ # 阶段 1：Prefill（有空闲槽位且有待处理请求时）
+ while self.free_slots and not self.pending_queue.empty():
+ try:
+ req = self.pending_queue.get_nowait()
+ except queue.Empty:
+ break
+ try:
+ req.slot_id = self.slot_manager.allocate()
+ except RuntimeError:
+ self.pending_queue.put(req)
+ break
+ try:
+ first_token = self._do_prefill(req)
+ except Exception as e:
+ self.slot_manager.free(req.slot_id)
+ LIB_LLAISYS.llaisysQwen2ModelResetKVCacheSlot(self._c_model, req.slot_id)
+ req.out_queue.put(_StreamError(str(e)))
+ req.out_queue.put(None)
+ continue
+ req.last_token_id = first_token
+ req.generated_tokens.append(first_token)
+ req.out_queue.put(first_token)
+ self.running_requests.append(req)
+
+ # 无运行中请求时，阻塞等待新请求以降低 CPU 占用
+ if not self.running_requests:
+ req = self.pending_queue.get()
+ self.pending_queue.put(req)
+ continue
+
+ # 阶段 2：Batched Decode
+ n_batch = len(self.running_requests)
+ slot_ids_array = (c_size_t * n_batch)()
+ token_ids_array = (c_int64 * n_batch)()
+ out_next_tokens = (c_int64 * n_batch)()
+
+ first = self.running_requests[0]
+ for i, r in enumerate(self.running_requests):
+ slot_ids_array[i] = r.slot_id
+ token_ids_array[i] = r.last_token_id
+
+ LIB_LLAISYS.llaisysQwen2ModelBatchedDecode(
+ self._c_model,
+ cast(slot_ids_array, POINTER(c_size_t)),
+ cast(token_ids_array, POINTER(c_int64)),
+ c_size_t(n_batch),
+ cast(out_next_tokens, POINTER(c_int64)),
+ c_float(first.temperature),
+ c_int(first.top_k),
+ c_float(first.top_p),
+ c_ulonglong(first.seed),
+ )
+
+ # 阶段 3：状态更新与清理（含 4.3.5 完成后写回 KV 池）
+ active_requests: list[RequestState] = []
+ for i, req in enumerate(self.running_requests):
+ next_token = int(out_next_tokens[i])
+ req.last_token_id = next_token
+ req.generated_tokens.append(next_token)
+ req.out_queue.put(next_token)
+
+ if next_token == self._end_token or len(req.generated_tokens) >= req.max_tokens:
+ req.is_finished = True
+ req.out_queue.put(None)
+ if self._put_kv and req.session_id and req.request_messages:
+ try:
+ blob = self.model.export_kv_cache_slot(req.slot_id)
+ prefix_len = len(req.prompt_tokens)
+ self._put_kv(req.session_id, req.request_messages, blob, prefix_len)
+ except Exception:
+ pass
+ self.slot_manager.free(req.slot_id)
+ LIB_LLAISYS.llaisysQwen2ModelResetKVCacheSlot(self._c_model, req.slot_id)
+ else:
+ active_requests.append(req)
+
+ self.running_requests = active_requests
+
+ @property
+ def free_slots(self) -> list[int]:
+ """当前空闲槽位列表（只读视图，用于判断是否有空位）。"""
+ return self.slot_manager.free_slots
+
+ def get_metrics(self) -> dict:
+ """返回引擎监控指标（4.5.3），供 /v1/metrics 等使用。"""
+ return {
+ "pending_queue_size": self.pending_queue.qsize(),
+ "running_count": len(self.running_requests),
+ "free_slots_count": len(self.slot_manager.free_slots),
+ "max_batch_size": self.max_batch_size,
+ }
diff --git a/python/llaisys/tensor.py b/python/llaisys_py/tensor.py
similarity index 96%
rename from python/llaisys/tensor.py
rename to python/llaisys_py/tensor.py
index 1466d851e..919f886e3 100644
--- a/python/llaisys/tensor.py
+++ b/python/llaisys_py/tensor.py
@@ -1,97 +1,97 @@
-from typing import Sequence, Tuple
-
-from .libllaisys import (
- LIB_LLAISYS,
- llaisysTensor_t,
- llaisysDeviceType_t,
- DeviceType,
- llaisysDataType_t,
- DataType,
-)
-from ctypes import c_size_t, c_int, c_ssize_t, c_void_p
-
-
-class Tensor:
- def __init__(
- self,
- shape: Sequence[int] = None,
- dtype: DataType = DataType.F32,
- device: DeviceType = DeviceType.CPU,
- device_id: int = 0,
- tensor: llaisysTensor_t = None,
- ):
- if tensor:
- self._tensor = tensor
- else:
- _ndim = 0 if shape is None else len(shape)
- _shape = None if shape is None else (c_size_t * len(shape))(*shape)
- self._tensor: llaisysTensor_t = LIB_LLAISYS.tensorCreate(
- _shape,
- c_size_t(_ndim),
- llaisysDataType_t(dtype),
- llaisysDeviceType_t(device),
- c_int(device_id),
- )
-
- def __del__(self):
- if hasattr(self, "_tensor") and self._tensor is not None:
- LIB_LLAISYS.tensorDestroy(self._tensor)
- self._tensor = None
-
- def shape(self) -> Tuple[int]:
- buf = (c_size_t * self.ndim())()
- LIB_LLAISYS.tensorGetShape(self._tensor, buf)
- return tuple(buf[i] for i in range(self.ndim()))
-
- def strides(self) -> Tuple[int]:
- buf = (c_ssize_t * self.ndim())()
- LIB_LLAISYS.tensorGetStrides(self._tensor, buf)
- return tuple(buf[i] for i in range(self.ndim()))
-
- def ndim(self) -> int:
- return int(LIB_LLAISYS.tensorGetNdim(self._tensor))
-
- def dtype(self) -> DataType:
- return DataType(LIB_LLAISYS.tensorGetDataType(self._tensor))
-
- def device_type(self) -> DeviceType:
- return DeviceType(LIB_LLAISYS.tensorGetDeviceType(self._tensor))
-
- def device_id(self) -> int:
- return int(LIB_LLAISYS.tensorGetDeviceId(self._tensor))
-
- def data_ptr(self) -> c_void_p:
- return LIB_LLAISYS.tensorGetData(self._tensor)
-
- def lib_tensor(self) -> llaisysTensor_t:
- return self._tensor
-
- def debug(self):
- LIB_LLAISYS.tensorDebug(self._tensor)
-
- def __repr__(self):
- return f""
-
- def load(self, data: c_void_p):
- LIB_LLAISYS.tensorLoad(self._tensor, data)
-
- def is_contiguous(self) -> bool:
- return bool(LIB_LLAISYS.tensorIsContiguous(self._tensor))
-
- def view(self, *shape: int) -> llaisysTensor_t:
- _shape = (c_size_t * len(shape))(*shape)
- return Tensor(
- tensor=LIB_LLAISYS.tensorView(self._tensor, _shape, c_size_t(len(shape)))
- )
-
- def permute(self, *perm: int) -> llaisysTensor_t:
- assert len(perm) == self.ndim()
- _perm = (c_size_t * len(perm))(*perm)
- return Tensor(tensor=LIB_LLAISYS.tensorPermute(self._tensor, _perm))
-
- def slice(self, dim: int, start: int, end: int):
- return Tensor(
- tensor=LIB_LLAISYS.tensorSlice(
- self._tensor, c_size_t(dim), c_size_t(start), c_size_t(end)
- )
- )
+from typing import Sequence, Tuple
+
+from .libllaisys import (
+ LIB_LLAISYS,
+ llaisysTensor_t,
+ llaisysDeviceType_t,
+ DeviceType,
+ llaisysDataType_t,
+ DataType,
+)
+from ctypes import c_size_t, c_int, c_ssize_t, c_void_p
+
+
+class Tensor:
+ def __init__(
+ self,
+ shape: Sequence[int] = None,
+ dtype: DataType = DataType.F32,
+ device: DeviceType = DeviceType.CPU,
+ device_id: int = 0,
+ tensor: llaisysTensor_t = None,
+ ):
+ if tensor:
+ self._tensor = tensor
+ else:
+ _ndim = 0 if shape is None else len(shape)
+ _shape = None if shape is None else (c_size_t * len(shape))(*shape)
+ self._tensor: llaisysTensor_t = LIB_LLAISYS.tensorCreate(
+ _shape,
+ c_size_t(_ndim),
+ llaisysDataType_t(dtype),
+ llaisysDeviceType_t(device),
+ c_int(device_id),
+ )
+
+ def __del__(self):
+ if hasattr(self, "_tensor") and self._tensor is not None:
+ LIB_LLAISYS.tensorDestroy(self._tensor)
+ self._tensor = None
+
+ def shape(self) -> Tuple[int]:
+ buf = (c_size_t * self.ndim())()
+ LIB_LLAISYS.tensorGetShape(self._tensor, buf)
+ return tuple(buf[i] for i in range(self.ndim()))
+
+ def strides(self) -> Tuple[int]:
+ buf = (c_ssize_t * self.ndim())()
+ LIB_LLAISYS.tensorGetStrides(self._tensor, buf)
+ return tuple(buf[i] for i in range(self.ndim()))
+
+ def ndim(self) -> int:
+ return int(LIB_LLAISYS.tensorGetNdim(self._tensor))
+
+ def dtype(self) -> DataType:
+ return DataType(LIB_LLAISYS.tensorGetDataType(self._tensor))
+
+ def device_type(self) -> DeviceType:
+ return DeviceType(LIB_LLAISYS.tensorGetDeviceType(self._tensor))
+
+ def device_id(self) -> int:
+ return int(LIB_LLAISYS.tensorGetDeviceId(self._tensor))
+
+ def data_ptr(self) -> c_void_p:
+ return LIB_LLAISYS.tensorGetData(self._tensor)
+
+ def lib_tensor(self) -> llaisysTensor_t:
+ return self._tensor
+
+ def debug(self):
+ LIB_LLAISYS.tensorDebug(self._tensor)
+
+ def __repr__(self):
+ return f""
+
+ def load(self, data: c_void_p):
+ LIB_LLAISYS.tensorLoad(self._tensor, data)
+
+ def is_contiguous(self) -> bool:
+ return bool(LIB_LLAISYS.tensorIsContiguous(self._tensor))
+
+ def view(self, *shape: int) -> llaisysTensor_t:
+ _shape = (c_size_t * len(shape))(*shape)
+ return Tensor(
+ tensor=LIB_LLAISYS.tensorView(self._tensor, _shape, c_size_t(len(shape)))
+ )
+
+ def permute(self, *perm: int) -> llaisysTensor_t:
+ assert len(perm) == self.ndim()
+ _perm = (c_size_t * len(perm))(*perm)
+ return Tensor(tensor=LIB_LLAISYS.tensorPermute(self._tensor, _perm))
+
+ def slice(self, dim: int, start: int, end: int):
+ return Tensor(
+ tensor=LIB_LLAISYS.tensorSlice(
+ self._tensor, c_size_t(dim), c_size_t(start), c_size_t(end)
+ )
+ )
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 8fe2f47af..022472c3d 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -1,3 +1,3 @@
-[build-system]
-requires = ["setuptools>=42", "wheel"]
-build-backend = "setuptools.build_meta"
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
diff --git a/python/setup.cfg b/python/setup.cfg
index b35fc65f7..10cbf84a7 100644
--- a/python/setup.cfg
+++ b/python/setup.cfg
@@ -1,21 +1,21 @@
-[metadata]
-name = llaisys
-version = 0.1.0
-description = Python APIs for llaisys
-author = Pan Zezhong
-license = MIT
-
-[options]
-packages = find:
-include_package_data = True
-zip_safe = False
-install_requires =
- torch>=2.4.0
- transformers
- accelerate
-
-[options.package_data]
-llaisys =
- libllaisys/*.so
- libllaisys/*.dll
- libllaisys/*.dylib
+[metadata]
+name = llaisys-py
+version = 0.1.0
+description = Python APIs for llaisys
+author = Pan Zezhong
+license = MIT
+
+[options]
+packages = find:
+include_package_data = True
+zip_safe = False
+install_requires =
+ torch>=2.4.0
+ transformers
+ accelerate
+
+[options.package_data]
+llaisys_py =
+ libllaisys/*.so
+ libllaisys/*.dll
+ libllaisys/*.dylib
diff --git a/scripts/download_model.py b/scripts/download_model.py
new file mode 100644
index 000000000..1eddb57ce
--- /dev/null
+++ b/scripts/download_model.py
@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+从 Hugging Face 下载 LLAISYS 使用的模型（默认 DeepSeek-R1-Distill-Qwen-1.5B）。
+
+用法:
+ python scripts/download_model.py
+ python scripts/download_model.py --dir /自定义/保存目录
+ python scripts/download_model.py --repo 其他组织/其他模型名
+
+依赖: 在虚拟环境中安装 huggingface_hub，例如:
+ python3 -m venv .venv && .venv/bin/pip install huggingface_hub
+ 然后运行: .venv/bin/python scripts/download_model.py
+"""
+import argparse
+import os
+import sys
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="从 Hugging Face 下载模型到本地目录",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+示例:
+ python scripts/download_model.py
+ python scripts/download_model.py --dir ./my_models/DeepSeek-R1-Distill-Qwen-1.5B
+ python scripts/download_model.py --repo deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --dir ./models
+ """,
+ )
+ parser.add_argument(
+ "--repo",
+ default="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+ help="Hugging Face 仓库 ID，默认: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+ )
+ parser.add_argument(
+ "--dir",
+ default=None,
+ help="保存目录。不指定时使用: <项目根>/models/<仓库名>",
+ )
+ parser.add_argument(
+ "--resume",
+ action="store_true",
+ help="断点续传（已存在的文件会跳过）",
+ )
+ args = parser.parse_args()
+
+ try:
+ from huggingface_hub import snapshot_download
+ except ImportError:
+ print("请先安装 huggingface_hub:", file=sys.stderr)
+ print(" pip install huggingface_hub", file=sys.stderr)
+ print("或在项目 venv 中: .venv/bin/pip install huggingface_hub", file=sys.stderr)
+ sys.exit(1)
+
+ repo_id = args.repo
+ if args.dir:
+ local_dir = os.path.abspath(os.path.expanduser(args.dir))
+ else:
+ script_dir = os.path.dirname(os.path.abspath(__file__))
+ root = os.path.dirname(script_dir)
+ # 用仓库名最后一段作为子目录名，如 DeepSeek-R1-Distill-Qwen-1.5B
+ repo_name = repo_id.split("/")[-1]
+ local_dir = os.path.join(root, "models", repo_name)
+
+ parent = os.path.dirname(local_dir)
+ if parent:
+ os.makedirs(parent, exist_ok=True)
+
+ print(f"仓库: {repo_id}")
+ print(f"目标: {local_dir}")
+ if args.resume:
+ print("模式: 断点续传（已存在文件将跳过）")
+ print()
+
+ try:
+ path = snapshot_download(
+ repo_id,
+ local_dir=local_dir,
+ local_dir_use_symlinks=False,
+ resume_download=args.resume,
+ )
+ print(f"下载完成: {path}")
+ print(f"\n启动 Project3 服务示例:")
+ print(f" python -m llaisys_py.server --model {path}")
+ except Exception as e:
+ print(f"下载失败: {e}", file=sys.stderr)
+ sys.exit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/format.py b/scripts/format.py
index 376eaf233..eaf322b40 100644
--- a/scripts/format.py
+++ b/scripts/format.py
@@ -1,204 +1,204 @@
-import argparse
-import subprocess
-import os
-from pathlib import Path
-from colorama import Fore, Style
-
-# 支持的文件类型
-SUPPORTED_FILES = {
- ".h": "c",
- ".hh": "c",
- ".hpp": "c",
- ".c": "c",
- ".cc": "c",
- ".cpp": "c",
- ".cxx": "c",
- ".cu": "c",
- ".cuh": "c",
- ".mlu": "c",
- ".cl": "c",
- ".py": "py",
-}
-
-
-def format_file(file: Path, check: bool, formatter) -> bool:
- formatter = formatter.get(SUPPORTED_FILES.get(file.suffix, None), None)
- if not formatter:
- return True # 文件类型不支持，跳过
-
- try:
- cmd = []
- if formatter.startswith("clang-format"):
- cmd = [formatter, "-style=file", "-i", file]
- if check:
- cmd.insert(2, "-dry-run")
- process = subprocess.run(
- cmd,
- capture_output=True,
- text=True,
- check=True,
- )
- if process.stderr:
- print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}")
- print(
- f"Use {Fore.CYAN}{formatter} -style=file -i {file}{Style.RESET_ALL} to format it."
- )
- return False
- else:
- subprocess.run(
- cmd,
- capture_output=True,
- text=True,
- check=True,
- )
- print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}")
- elif formatter == "black":
- cmd = [formatter, file]
- if check:
- cmd.insert(1, "--check")
- process = subprocess.run(
- cmd,
- capture_output=True,
- text=True,
- check=True,
- )
- if process.returncode != 0:
- print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}")
- print(
- f"Use {Fore.CYAN}{formatter} {file}{Style.RESET_ALL} to format it."
- )
- return False
- else:
- subprocess.run(
- cmd,
- capture_output=True,
- text=True,
- check=True,
- )
- print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}")
- except FileNotFoundError:
- print(
- f"{Fore.RED}Formatter {formatter} not found, {file} skipped.{Style.RESET_ALL}"
- )
- except subprocess.CalledProcessError as e:
- print(f"{Fore.RED}Formatter {formatter} failed: {e}{Style.RESET_ALL}")
-
- return True
-
-
-def git_added_files():
- """获取所有已暂存更改的文件"""
- try:
- # 使用 git diff --cached --name-only 获取所有已添加到暂存区的文件
- result = subprocess.run(
- ["git", "diff", "--cached", "--diff-filter=AMR", "--name-only"],
- capture_output=True,
- text=True,
- check=True,
- )
- for file in result.stdout.splitlines():
- yield Path(file.strip())
- except subprocess.CalledProcessError as e:
- print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}")
-
-
-def git_modified_since_ref(ref):
- """获取从指定的 Git 引用到当前状态的修改文件列表"""
- try:
- result = subprocess.run(
- ["git", "diff", f"{ref}..", "--diff-filter=AMR", "--name-only"],
- capture_output=True,
- text=True,
- check=True,
- )
- for file in result.stdout.splitlines():
- yield Path(file.strip())
- except subprocess.CalledProcessError as e:
- print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}")
-
-
-def list_files(paths):
- """递归获取指定路径下的所有文件"""
- files = []
- for path in paths:
- if path.is_file():
- yield path
- elif path.is_dir():
- for dirpath, _, filenames in os.walk(path):
- for name in filenames:
- yield Path(dirpath) / name
- else:
- print(
- f"{Fore.RED}Error: {path} is not a file or directory.{Style.RESET_ALL}"
- )
-
-
-def filter_in_path(file: Path, path) -> bool:
- """判断文件是否在指定路径下"""
- for p in path:
- if file.is_relative_to(p):
- return True
- return False
-
-
-def main():
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--ref", type=str, help="Git reference (commit hash) to compare against."
- )
- parser.add_argument(
- "--path", nargs="*", type=Path, help="Files to format or check."
- )
- parser.add_argument(
- "--check", action="store_true", help="Check files without modifying them."
- )
- parser.add_argument(
- "--c", default="clang-format-16", help="C formatter (default: clang-format-16)"
- )
- parser.add_argument(
- "--py", default="black", help="Python formatter (default: black)"
- )
- args = parser.parse_args()
-
- if args.ref is None and args.path is None:
- # Last commit.
- print(f"{Fore.GREEN}Formating git added files.{Style.RESET_ALL}")
- files = git_added_files()
-
- else:
- if args.ref is None:
- print(f"{Fore.GREEN}Formating files in {args.path}.{Style.RESET_ALL}")
- files = list_files(args.path)
- elif args.path is None:
- print(
- f"{Fore.GREEN}Formating git modified files from {args.ref}.{Style.RESET_ALL}"
- )
- files = git_modified_since_ref(args.ref)
- else:
- print(
- f"{Fore.GREEN}Formating git modified files from {args.ref} in {args.path}.{Style.RESET_ALL}"
- )
- files = (
- file
- for file in git_modified_since_ref(args.ref)
- if filter_in_path(file, args.path)
- )
-
- formatted = True
- for file in files:
- if not format_file(
- file,
- args.check,
- {
- "c": args.c,
- "py": args.py,
- },
- ):
- formatted = False
-
- if not formatted:
- exit(1)
-
-
-if __name__ == "__main__":
- main()
+import argparse
+import subprocess
+import os
+from pathlib import Path
+from colorama import Fore, Style
+
+# 支持的文件类型
+SUPPORTED_FILES = {
+ ".h": "c",
+ ".hh": "c",
+ ".hpp": "c",
+ ".c": "c",
+ ".cc": "c",
+ ".cpp": "c",
+ ".cxx": "c",
+ ".cu": "c",
+ ".cuh": "c",
+ ".mlu": "c",
+ ".cl": "c",
+ ".py": "py",
+}
+
+
+def format_file(file: Path, check: bool, formatter) -> bool:
+ formatter = formatter.get(SUPPORTED_FILES.get(file.suffix, None), None)
+ if not formatter:
+ return True # 文件类型不支持，跳过
+
+ try:
+ cmd = []
+ if formatter.startswith("clang-format"):
+ cmd = [formatter, "-style=file", "-i", file]
+ if check:
+ cmd.insert(2, "-dry-run")
+ process = subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ if process.stderr:
+ print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}")
+ print(
+ f"Use {Fore.CYAN}{formatter} -style=file -i {file}{Style.RESET_ALL} to format it."
+ )
+ return False
+ else:
+ subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}")
+ elif formatter == "black":
+ cmd = [formatter, file]
+ if check:
+ cmd.insert(1, "--check")
+ process = subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ if process.returncode != 0:
+ print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}")
+ print(
+ f"Use {Fore.CYAN}{formatter} {file}{Style.RESET_ALL} to format it."
+ )
+ return False
+ else:
+ subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}")
+ except FileNotFoundError:
+ print(
+ f"{Fore.RED}Formatter {formatter} not found, {file} skipped.{Style.RESET_ALL}"
+ )
+ except subprocess.CalledProcessError as e:
+ print(f"{Fore.RED}Formatter {formatter} failed: {e}{Style.RESET_ALL}")
+
+ return True
+
+
+def git_added_files():
+ """获取所有已暂存更改的文件"""
+ try:
+ # 使用 git diff --cached --name-only 获取所有已添加到暂存区的文件
+ result = subprocess.run(
+ ["git", "diff", "--cached", "--diff-filter=AMR", "--name-only"],
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ for file in result.stdout.splitlines():
+ yield Path(file.strip())
+ except subprocess.CalledProcessError as e:
+ print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}")
+
+
+def git_modified_since_ref(ref):
+ """获取从指定的 Git 引用到当前状态的修改文件列表"""
+ try:
+ result = subprocess.run(
+ ["git", "diff", f"{ref}..", "--diff-filter=AMR", "--name-only"],
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ for file in result.stdout.splitlines():
+ yield Path(file.strip())
+ except subprocess.CalledProcessError as e:
+ print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}")
+
+
+def list_files(paths):
+ """递归获取指定路径下的所有文件"""
+ files = []
+ for path in paths:
+ if path.is_file():
+ yield path
+ elif path.is_dir():
+ for dirpath, _, filenames in os.walk(path):
+ for name in filenames:
+ yield Path(dirpath) / name
+ else:
+ print(
+ f"{Fore.RED}Error: {path} is not a file or directory.{Style.RESET_ALL}"
+ )
+
+
+def filter_in_path(file: Path, path) -> bool:
+ """判断文件是否在指定路径下"""
+ for p in path:
+ if file.is_relative_to(p):
+ return True
+ return False
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--ref", type=str, help="Git reference (commit hash) to compare against."
+ )
+ parser.add_argument(
+ "--path", nargs="*", type=Path, help="Files to format or check."
+ )
+ parser.add_argument(
+ "--check", action="store_true", help="Check files without modifying them."
+ )
+ parser.add_argument(
+ "--c", default="clang-format-16", help="C formatter (default: clang-format-16)"
+ )
+ parser.add_argument(
+ "--py", default="black", help="Python formatter (default: black)"
+ )
+ args = parser.parse_args()
+
+ if args.ref is None and args.path is None:
+ # Last commit.
+ print(f"{Fore.GREEN}Formating git added files.{Style.RESET_ALL}")
+ files = git_added_files()
+
+ else:
+ if args.ref is None:
+ print(f"{Fore.GREEN}Formating files in {args.path}.{Style.RESET_ALL}")
+ files = list_files(args.path)
+ elif args.path is None:
+ print(
+ f"{Fore.GREEN}Formating git modified files from {args.ref}.{Style.RESET_ALL}"
+ )
+ files = git_modified_since_ref(args.ref)
+ else:
+ print(
+ f"{Fore.GREEN}Formating git modified files from {args.ref} in {args.path}.{Style.RESET_ALL}"
+ )
+ files = (
+ file
+ for file in git_modified_since_ref(args.ref)
+ if filter_in_path(file, args.path)
+ )
+
+ formatted = True
+ for file in files:
+ if not format_file(
+ file,
+ args.check,
+ {
+ "c": args.c,
+ "py": args.py,
+ },
+ ):
+ formatted = False
+
+ if not formatted:
+ exit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/list_safetensors_keys.py b/scripts/list_safetensors_keys.py
new file mode 100644
index 000000000..1b8330905
--- /dev/null
+++ b/scripts/list_safetensors_keys.py
@@ -0,0 +1,26 @@
+"""List keys in .safetensors files without loading tensor data (metadata only)."""
+import sys
+from pathlib import Path
+
+import safetensors
+
+
+def main():
+ if len(sys.argv) < 2:
+ print("Usage: python list_safetensors_keys.py ")
+ sys.exit(1)
+ model_dir = Path(sys.argv[1])
+ if not model_dir.is_dir():
+ print("Not a directory:", model_dir)
+ sys.exit(1)
+ for fpath in sorted(model_dir.glob("*.safetensors")):
+ print("\n---", fpath.name, "---")
+ with safetensors.safe_open(fpath, framework="numpy", device="cpu") as f:
+ keys = list(f.keys())
+ for k in keys:
+ print(k)
+ print("Total keys:", len(keys))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/run_server.sh b/scripts/run_server.sh
new file mode 100644
index 000000000..87d561db5
--- /dev/null
+++ b/scripts/run_server.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+# 不依赖「pip install -e ./python/」也能启动服务：用 PYTHONPATH 找到 llaisys 包。
+# 前提：已执行 xmake && xmake install（动态库在 python/llaisys_py/libllaisys/ 下）。
+#
+# 用法: ./scripts/run_server.sh [模型目录]
+# 示例: ./scripts/run_server.sh /home/chenncy/llaisys/DeepSeek-R1-Distill-Qwen-1___5B
+
+set -e
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+MODEL="${1:-}"
+PORT="${2:-8000}"
+
+if [ -z "$MODEL" ]; then
+ echo "用法: $0 <模型目录> [端口]"
+ echo "示例: $0 $ROOT/DeepSeek-R1-Distill-Qwen-1___5B 8000"
+ exit 1
+fi
+
+if [ ! -d "$MODEL" ]; then
+ echo "错误: 模型目录不存在: $MODEL"
+ exit 1
+fi
+
+# 用项目里的 python 包，不要求 pip install
+export PYTHONPATH="${ROOT}/python${PYTHONPATH:+:$PYTHONPATH}"
+
+# 优先用 venv 的 Python（里面可能有 torch、transformers 等）
+if [ -x "${ROOT}/.venv/bin/python" ]; then
+ PY="${ROOT}/.venv/bin/python"
+else
+ PY="python3"
+fi
+
+echo "PYTHONPATH=$PYTHONPATH"
+echo "Python: $PY"
+echo "模型: $MODEL"
+echo "端口: $PORT"
+exec "$PY" -m llaisys_py.server --model "$MODEL" --port "$PORT"
diff --git a/src/core/allocator/allocator.hpp b/src/core/allocator/allocator.hpp
index 2388927e4..69907e32a 100644
--- a/src/core/allocator/allocator.hpp
+++ b/src/core/allocator/allocator.hpp
@@ -1,19 +1,19 @@
-#pragma once
-
-#include "llaisys/runtime.h"
-
-#include "../storage/storage.hpp"
-
-namespace llaisys::core {
-class MemoryAllocator {
-protected:
- const LlaisysRuntimeAPI *_api;
- MemoryAllocator(const LlaisysRuntimeAPI *runtime_api) : _api(runtime_api){};
-
-public:
- virtual ~MemoryAllocator() = default;
- virtual std::byte *allocate(size_t size) = 0;
- virtual void release(std::byte *memory) = 0;
-};
-
-} // namespace llaisys::core
+#pragma once
+
+#include "llaisys/runtime.h"
+
+#include "../storage/storage.hpp"
+
+namespace llaisys::core {
+class MemoryAllocator {
+protected:
+ const LlaisysRuntimeAPI *_api;
+ MemoryAllocator(const LlaisysRuntimeAPI *runtime_api) : _api(runtime_api){};
+
+public:
+ virtual ~MemoryAllocator() = default;
+ virtual std::byte *allocate(size_t size) = 0;
+ virtual void release(std::byte *memory) = 0;
+};
+
+} // namespace llaisys::core
diff --git a/src/core/allocator/naive_allocator.cpp b/src/core/allocator/naive_allocator.cpp
index 723f2975c..09eff1699 100644
--- a/src/core/allocator/naive_allocator.cpp
+++ b/src/core/allocator/naive_allocator.cpp
@@ -1,16 +1,16 @@
-#include "naive_allocator.hpp"
-
-#include "../runtime/runtime.hpp"
-
-namespace llaisys::core::allocators {
-NaiveAllocator::NaiveAllocator(const LlaisysRuntimeAPI *runtime_api) : MemoryAllocator(runtime_api) {
-}
-
-std::byte *NaiveAllocator::allocate(size_t size) {
- return static_cast(_api->malloc_device(size));
-}
-
-void NaiveAllocator::release(std::byte *memory) {
- _api->free_device(memory);
-}
+#include "naive_allocator.hpp"
+
+#include "../runtime/runtime.hpp"
+
+namespace llaisys::core::allocators {
+NaiveAllocator::NaiveAllocator(const LlaisysRuntimeAPI *runtime_api) : MemoryAllocator(runtime_api) {
+}
+
+std::byte *NaiveAllocator::allocate(size_t size) {
+ return static_cast(_api->malloc_device(size));
+}
+
+void NaiveAllocator::release(std::byte *memory) {
+ _api->free_device(memory);
+}
} // namespace llaisys::core::allocators
\ No newline at end of file
diff --git a/src/core/allocator/naive_allocator.hpp b/src/core/allocator/naive_allocator.hpp
index e93cb5303..91e88303e 100644
--- a/src/core/allocator/naive_allocator.hpp
+++ b/src/core/allocator/naive_allocator.hpp
@@ -1,13 +1,13 @@
-#pragma once
-
-#include "allocator.hpp"
-
-namespace llaisys::core::allocators {
-class NaiveAllocator : public MemoryAllocator {
-public:
- NaiveAllocator(const LlaisysRuntimeAPI *runtime_api);
- ~NaiveAllocator() = default;
- std::byte *allocate(size_t size) override;
- void release(std::byte *memory) override;
-};
+#pragma once
+
+#include "allocator.hpp"
+
+namespace llaisys::core::allocators {
+class NaiveAllocator : public MemoryAllocator {
+public:
+ NaiveAllocator(const LlaisysRuntimeAPI *runtime_api);
+ ~NaiveAllocator() = default;
+ std::byte *allocate(size_t size) override;
+ void release(std::byte *memory) override;
+};
} // namespace llaisys::core::allocators
\ No newline at end of file
diff --git a/src/core/context/context.cpp b/src/core/context/context.cpp
index 44894b9e7..c619d4209 100644
--- a/src/core/context/context.cpp
+++ b/src/core/context/context.cpp
@@ -1,50 +1,68 @@
#include "context.hpp"
#include "../../utils.hpp"
+#include
#include
namespace llaisys::core {
-Context::Context() {
- // All device types, put CPU at the end
- std::vector device_typs;
- for (int i = 1; i < LLAISYS_DEVICE_TYPE_COUNT; i++) {
- device_typs.push_back(static_cast(i));
- }
- device_typs.push_back(LLAISYS_DEVICE_CPU);
+namespace {
+// 进程内共享的 Runtime 池，保证多线程（如 worker）与主线程使用同一 CUDA 上下文，避免模型在主线程分配、在 worker 访问时 segfault。
+struct GlobalRuntimePool {
+ std::mutex mutex;
+ std::unordered_map> pool;
+ bool initialized = false;
- // Create runtimes for each device type.
- // Activate the first available device. If no other device is available, activate CPU runtime.
- for (auto device_type : device_typs) {
- const LlaisysRuntimeAPI *api_ = llaisysGetRuntimeAPI(device_type);
- int device_count = api_->get_device_count();
- std::vector runtimes_(device_count);
- for (int device_id = 0; device_id < device_count; device_id++) {
-
- if (_current_runtime == nullptr) {
- auto runtime = new Runtime(device_type, device_id);
- runtime->_activate();
- runtimes_[device_id] = runtime;
- _current_runtime = runtime;
+ ~GlobalRuntimePool() {
+ for (auto &entry : pool) {
+ for (Runtime *r : entry.second) {
+ if (r != nullptr) {
+ r->deactivateForShutdown();
+ delete r;
+ }
}
}
- _runtime_map[device_type] = runtimes_;
}
-}
+} g_runtime_pool;
+} // namespace
-Context::~Context() {
- // Destroy current runtime first.
- delete _current_runtime;
+Context::Context() {
+ std::lock_guard lock(g_runtime_pool.mutex);
+ if (!g_runtime_pool.initialized) {
+ std::vector device_typs;
+ for (int i = 1; i < LLAISYS_DEVICE_TYPE_COUNT; i++)
+ device_typs.push_back(static_cast(i));
+ device_typs.push_back(LLAISYS_DEVICE_CPU);
- for (auto &runtime_entry : _runtime_map) {
- std::vector runtimes = runtime_entry.second;
- for (auto runtime : runtimes) {
- if (runtime != nullptr && runtime != _current_runtime) {
- runtime->_activate();
- delete runtime;
+ Runtime *first = nullptr;
+ for (auto device_type : device_typs) {
+ const LlaisysRuntimeAPI *api_ = llaisysGetRuntimeAPI(device_type);
+ int device_count = api_->get_device_count();
+ std::vector runtimes_(device_count, nullptr);
+ for (int device_id = 0; device_id < device_count; device_id++) {
+ auto *r = new Runtime(device_type, device_id);
+ runtimes_[device_id] = r;
+ if (first == nullptr) {
+ r->_activate();
+ first = r;
+ }
}
+ g_runtime_pool.pool[device_type] = runtimes_;
}
- runtimes.clear();
+ g_runtime_pool.initialized = true;
}
+ _runtime_map = g_runtime_pool.pool;
+ _current_runtime = nullptr;
+ // 默认激活 CPU runtime，避免未调用 setDevice 的代码路径（如首次 Tensor::create）触发 runtime() 断言
+ auto it = _runtime_map.find(LLAISYS_DEVICE_CPU);
+ if (it != _runtime_map.end() && !it->second.empty() && it->second[0] != nullptr) {
+ it->second[0]->_activate();
+ _current_runtime = it->second[0];
+ }
+}
+
+Context::~Context() {
+ if (_current_runtime != nullptr)
+ _current_runtime->_deactivate();
_current_runtime = nullptr;
_runtime_map.clear();
}
@@ -53,13 +71,18 @@ void Context::setDevice(llaisysDeviceType_t device_type, int device_id) {
// If doest not match the current runtime.
if (_current_runtime == nullptr || _current_runtime->deviceType() != device_type || _current_runtime->deviceId() != device_id) {
auto runtimes = _runtime_map[device_type];
+ if (runtimes.empty()) {
+ if (device_type == LLAISYS_DEVICE_NVIDIA)
+ CHECK_ARGUMENT(false, "no NVIDIA GPUs available (get_device_count() returned 0). Use --device cpu or check CUDA/driver.");
+ else
+ CHECK_ARGUMENT(false, "no devices available for this device type (get_device_count() returned 0).");
+ }
CHECK_ARGUMENT((size_t)device_id < runtimes.size() && device_id >= 0, "invalid device id");
if (_current_runtime != nullptr) {
_current_runtime->_deactivate();
}
- if (runtimes[device_id] == nullptr) {
- runtimes[device_id] = new Runtime(device_type, device_id);
- }
+ // Runtime 必须来自全局池，保证多线程与主线程使用同一 CUDA 上下文；不得在此处 new 仅写入本线程 _runtime_map 副本，否则会泄漏且破坏设备一致性。
+ CHECK_ARGUMENT(runtimes[device_id] != nullptr, "runtime for device type/id not found in pool; ensure Context is used after global pool init.");
runtimes[device_id]->_activate();
_current_runtime = runtimes[device_id];
}
diff --git a/src/core/context/context.hpp b/src/core/context/context.hpp
index a3ebcdecf..e74f2771b 100644
--- a/src/core/context/context.hpp
+++ b/src/core/context/context.hpp
@@ -1,35 +1,35 @@
-#pragma once
-
-#include "llaisys.h"
-
-#include "../core.hpp"
-
-#include "../runtime/runtime.hpp"
-
-#include
-#include
-
-namespace llaisys::core {
-class Context {
-private:
- std::unordered_map> _runtime_map;
- Runtime *_current_runtime;
- Context();
-
-public:
- ~Context();
-
- // Prevent copy
- Context(const Context &) = delete;
- Context &operator=(const Context &) = delete;
-
- // Prevent move
- Context(Context &&) = delete;
- Context &operator=(Context &&) = delete;
-
- void setDevice(llaisysDeviceType_t device_type, int device_id);
- Runtime &runtime();
-
- friend Context &context();
-};
-} // namespace llaisys::core
+#pragma once
+
+#include "llaisys.h"
+
+#include "../core.hpp"
+
+#include "../runtime/runtime.hpp"
+
+#include
+#include
+
+namespace llaisys::core {
+class Context {
+private:
+ std::unordered_map> _runtime_map;
+ Runtime *_current_runtime;
+ Context();
+
+public:
+ ~Context();
+
+ // Prevent copy
+ Context(const Context &) = delete;
+ Context &operator=(const Context &) = delete;
+
+ // Prevent move
+ Context(Context &&) = delete;
+ Context &operator=(Context &&) = delete;
+
+ void setDevice(llaisysDeviceType_t device_type, int device_id);
+ Runtime &runtime();
+
+ friend Context &context();
+};
+} // namespace llaisys::core
diff --git a/src/core/core.hpp b/src/core/core.hpp
index 2eed7bbfb..7258030f5 100644
--- a/src/core/core.hpp
+++ b/src/core/core.hpp
@@ -1,18 +1,18 @@
-#pragma once
-#include
-
-namespace llaisys {
-namespace core {
-class Storage;
-using storage_t = std::shared_ptr;
-
-class MemoryAllocator;
-
-class Runtime;
-class Context;
-
-// Global function to get thread local context
-Context &context();
-} // namespace core
-
+#pragma once
+#include
+
+namespace llaisys {
+namespace core {
+class Storage;
+using storage_t = std::shared_ptr;
+
+class MemoryAllocator;
+
+class Runtime;
+class Context;
+
+// Global function to get thread local context
+Context &context();
+} // namespace core
+
} // namespace llaisys
\ No newline at end of file
diff --git a/src/core/llaisys_core.hpp b/src/core/llaisys_core.hpp
index 8d30b9427..33aa2a4ea 100644
--- a/src/core/llaisys_core.hpp
+++ b/src/core/llaisys_core.hpp
@@ -1,9 +1,9 @@
-#pragma once
-
-// Header file for using llaisys core functionalities.
-
-#include "core.hpp"
-
-#include "context/context.hpp"
-#include "runtime/runtime.hpp"
-#include "storage/storage.hpp"
+#pragma once
+
+// Header file for using llaisys core functionalities.
+
+#include "core.hpp"
+
+#include "context/context.hpp"
+#include "runtime/runtime.hpp"
+#include "storage/storage.hpp"
diff --git a/src/core/runtime/runtime.cpp b/src/core/runtime/runtime.cpp
index 7f03a8622..a0567334f 100644
--- a/src/core/runtime/runtime.cpp
+++ b/src/core/runtime/runtime.cpp
@@ -5,14 +5,14 @@
namespace llaisys::core {
Runtime::Runtime(llaisysDeviceType_t device_type, int device_id)
- : _device_type(device_type), _device_id(device_id), _is_active(false) {
+ : _device_type(device_type), _device_id(device_id), _is_active(false), _deactivated_for_shutdown(false) {
_api = llaisys::device::getRuntimeAPI(_device_type);
_stream = _api->create_stream();
_allocator = new allocators::NaiveAllocator(_api);
}
Runtime::~Runtime() {
- if (!_is_active) {
+ if (!_is_active && !_deactivated_for_shutdown) {
std::cerr << "Mallicious destruction of inactive runtime." << std::endl;
}
delete _allocator;
@@ -30,6 +30,11 @@ void Runtime::_deactivate() {
_is_active = false;
}
+void Runtime::deactivateForShutdown() {
+ _is_active = false;
+ _deactivated_for_shutdown = true;
+}
+
bool Runtime::isActive() const {
return _is_active;
}
diff --git a/src/core/runtime/runtime.hpp b/src/core/runtime/runtime.hpp
index 43235824e..4025b112a 100644
--- a/src/core/runtime/runtime.hpp
+++ b/src/core/runtime/runtime.hpp
@@ -12,6 +12,7 @@ class Runtime {
const LlaisysRuntimeAPI *_api;
MemoryAllocator *_allocator;
bool _is_active;
+ bool _deactivated_for_shutdown;
void _activate();
void _deactivate();
llaisysStream_t _stream;
@@ -43,5 +44,7 @@ class Runtime {
llaisysStream_t stream() const;
void synchronize() const;
+ /// 进程退出时由全局 Runtime 池调用，避免析构时误报
+ void deactivateForShutdown();
};
} // namespace llaisys::core
diff --git a/src/core/storage/storage.cpp b/src/core/storage/storage.cpp
index f131111c7..054e5cd3a 100644
--- a/src/core/storage/storage.cpp
+++ b/src/core/storage/storage.cpp
@@ -1,40 +1,40 @@
-#include "storage.hpp"
-
-#include "../runtime/runtime.hpp"
-
-namespace llaisys::core {
-Storage::Storage(std::byte *memory, size_t size, Runtime &runtime, bool is_host)
- : _memory(memory), _size(size), _runtime(runtime), _is_host(is_host) {}
-
-Storage::~Storage() {
- _runtime.freeStorage(this);
-}
-
-std::byte *Storage::memory() const {
- return _memory;
-}
-
-size_t Storage::size() const {
- return _size;
-}
-
-llaisysDeviceType_t Storage::deviceType() const {
- if (isHost()) {
- return LLAISYS_DEVICE_CPU;
- } else {
- return _runtime.deviceType();
- }
-}
-
-int Storage::deviceId() const {
- if (isHost()) {
- return 0;
- } else {
- return _runtime.deviceId();
- }
-}
-
-bool Storage::isHost() const {
- return _is_host;
-}
+#include "storage.hpp"
+
+#include "../runtime/runtime.hpp"
+
+namespace llaisys::core {
+Storage::Storage(std::byte *memory, size_t size, Runtime &runtime, bool is_host)
+ : _memory(memory), _size(size), _runtime(runtime), _is_host(is_host) {}
+
+Storage::~Storage() {
+ _runtime.freeStorage(this);
+}
+
+std::byte *Storage::memory() const {
+ return _memory;
+}
+
+size_t Storage::size() const {
+ return _size;
+}
+
+llaisysDeviceType_t Storage::deviceType() const {
+ if (isHost()) {
+ return LLAISYS_DEVICE_CPU;
+ } else {
+ return _runtime.deviceType();
+ }
+}
+
+int Storage::deviceId() const {
+ if (isHost()) {
+ return 0;
+ } else {
+ return _runtime.deviceId();
+ }
+}
+
+bool Storage::isHost() const {
+ return _is_host;
+}
} // namespace llaisys::core
\ No newline at end of file
diff --git a/src/core/storage/storage.hpp b/src/core/storage/storage.hpp
index 7260e30a2..f7427aaf3 100644
--- a/src/core/storage/storage.hpp
+++ b/src/core/storage/storage.hpp
@@ -1,28 +1,28 @@
-#pragma once
-#include "llaisys.h"
-
-#include "../core.hpp"
-
-#include
-
-namespace llaisys::core {
-class Storage {
-private:
- std::byte *_memory;
- size_t _size;
- Runtime &_runtime;
- bool _is_host;
- Storage(std::byte *memory, size_t size, Runtime &runtime, bool is_host);
-
-public:
- friend class Runtime;
- ~Storage();
-
- std::byte *memory() const;
- size_t size() const;
- llaisysDeviceType_t deviceType() const;
- int deviceId() const;
- bool isHost() const;
-};
-
-}; // namespace llaisys::core
+#pragma once
+#include "llaisys.h"
+
+#include "../core.hpp"
+
+#include
+
+namespace llaisys::core {
+class Storage {
+private:
+ std::byte *_memory;
+ size_t _size;
+ Runtime &_runtime;
+ bool _is_host;
+ Storage(std::byte *memory, size_t size, Runtime &runtime, bool is_host);
+
+public:
+ friend class Runtime;
+ ~Storage();
+
+ std::byte *memory() const;
+ size_t size() const;
+ llaisysDeviceType_t deviceType() const;
+ int deviceId() const;
+ bool isHost() const;
+};
+
+}; // namespace llaisys::core
diff --git a/src/device/cpu/cpu_resource.cpp b/src/device/cpu/cpu_resource.cpp
index 4fb28bd06..7e0580a8d 100644
--- a/src/device/cpu/cpu_resource.cpp
+++ b/src/device/cpu/cpu_resource.cpp
@@ -1,5 +1,5 @@
-#include "cpu_resource.hpp"
-
-namespace llaisys::device::cpu {
-Resource::Resource() : llaisys::device::DeviceResource(LLAISYS_DEVICE_CPU, 0) {}
-} // namespace llaisys::device::cpu
+#include "cpu_resource.hpp"
+
+namespace llaisys::device::cpu {
+Resource::Resource() : llaisys::device::DeviceResource(LLAISYS_DEVICE_CPU, 0) {}
+} // namespace llaisys::device::cpu
diff --git a/src/device/cpu/cpu_resource.hpp b/src/device/cpu/cpu_resource.hpp
index a99a67391..fc1f784ae 100644
--- a/src/device/cpu/cpu_resource.hpp
+++ b/src/device/cpu/cpu_resource.hpp
@@ -1,11 +1,11 @@
-#pragma once
-
-#include "../device_resource.hpp"
-
-namespace llaisys::device::cpu {
-class Resource : public llaisys::device::DeviceResource {
-public:
- Resource();
- ~Resource() = default;
-};
+#pragma once
+
+#include "../device_resource.hpp"
+
+namespace llaisys::device::cpu {
+class Resource : public llaisys::device::DeviceResource {
+public:
+ Resource();
+ ~Resource() = default;
+};
} // namespace llaisys::device::cpu
\ No newline at end of file
diff --git a/src/device/cpu/cpu_runtime_api.cpp b/src/device/cpu/cpu_runtime_api.cpp
index 8d57cc402..1149bc980 100644
--- a/src/device/cpu/cpu_runtime_api.cpp
+++ b/src/device/cpu/cpu_runtime_api.cpp
@@ -1,75 +1,75 @@
-#include "../runtime_api.hpp"
-
-#include
-#include
-
-namespace llaisys::device::cpu {
-
-namespace runtime_api {
-int getDeviceCount() {
- return 1;
-}
-
-void setDevice(int) {
- // do nothing
-}
-
-void deviceSynchronize() {
- // do nothing
-}
-
-llaisysStream_t createStream() {
- return (llaisysStream_t)0; // null stream
-}
-
-void destroyStream(llaisysStream_t stream) {
- // do nothing
-}
-void streamSynchronize(llaisysStream_t stream) {
- // do nothing
-}
-
-void *mallocDevice(size_t size) {
- return std::malloc(size);
-}
-
-void freeDevice(void *ptr) {
- std::free(ptr);
-}
-
-void *mallocHost(size_t size) {
- return mallocDevice(size);
-}
-
-void freeHost(void *ptr) {
- freeDevice(ptr);
-}
-
-void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
- std::memcpy(dst, src, size);
-}
-
-void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
- memcpySync(dst, src, size, kind);
-}
-
-static const LlaisysRuntimeAPI RUNTIME_API = {
- &getDeviceCount,
- &setDevice,
- &deviceSynchronize,
- &createStream,
- &destroyStream,
- &streamSynchronize,
- &mallocDevice,
- &freeDevice,
- &mallocHost,
- &freeHost,
- &memcpySync,
- &memcpyAsync};
-
-} // namespace runtime_api
-
-const LlaisysRuntimeAPI *getRuntimeAPI() {
- return &runtime_api::RUNTIME_API;
-}
-} // namespace llaisys::device::cpu
+#include "../runtime_api.hpp"
+
+#include
+#include
+
+namespace llaisys::device::cpu {
+
+namespace runtime_api {
+int getDeviceCount() {
+ return 1;
+}
+
+void setDevice(int) {
+ // do nothing
+}
+
+void deviceSynchronize() {
+ // do nothing
+}
+
+llaisysStream_t createStream() {
+ return (llaisysStream_t)0; // null stream
+}
+
+void destroyStream(llaisysStream_t stream) {
+ // do nothing
+}
+void streamSynchronize(llaisysStream_t stream) {
+ // do nothing
+}
+
+void *mallocDevice(size_t size) {
+ return std::malloc(size);
+}
+
+void freeDevice(void *ptr) {
+ std::free(ptr);
+}
+
+void *mallocHost(size_t size) {
+ return mallocDevice(size);
+}
+
+void freeHost(void *ptr) {
+ freeDevice(ptr);
+}
+
+void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
+ std::memcpy(dst, src, size);
+}
+
+void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
+ memcpySync(dst, src, size, kind);
+}
+
+static const LlaisysRuntimeAPI RUNTIME_API = {
+ &getDeviceCount,
+ &setDevice,
+ &deviceSynchronize,
+ &createStream,
+ &destroyStream,
+ &streamSynchronize,
+ &mallocDevice,
+ &freeDevice,
+ &mallocHost,
+ &freeHost,
+ &memcpySync,
+ &memcpyAsync};
+
+} // namespace runtime_api
+
+const LlaisysRuntimeAPI *getRuntimeAPI() {
+ return &runtime_api::RUNTIME_API;
+}
+} // namespace llaisys::device::cpu
diff --git a/src/device/device_resource.hpp b/src/device/device_resource.hpp
index e9062e510..f6283b45c 100644
--- a/src/device/device_resource.hpp
+++ b/src/device/device_resource.hpp
@@ -1,22 +1,22 @@
-#pragma once
-#include "llaisys.h"
-
-#include "../utils.hpp"
-
-namespace llaisys::device {
-class DeviceResource {
-private:
- llaisysDeviceType_t _device_type;
- int _device_id;
-
-public:
- DeviceResource(llaisysDeviceType_t device_type, int device_id)
- : _device_type(device_type),
- _device_id(device_id) {
- }
- ~DeviceResource() = default;
-
- llaisysDeviceType_t getDeviceType() const { return _device_type; }
- int getDeviceId() const { return _device_id; };
-};
-} // namespace llaisys::device
+#pragma once
+#include "llaisys.h"
+
+#include "../utils.hpp"
+
+namespace llaisys::device {
+class DeviceResource {
+private:
+ llaisysDeviceType_t _device_type;
+ int _device_id;
+
+public:
+ DeviceResource(llaisysDeviceType_t device_type, int device_id)
+ : _device_type(device_type),
+ _device_id(device_id) {
+ }
+ ~DeviceResource() = default;
+
+ llaisysDeviceType_t getDeviceType() const { return _device_type; }
+ int getDeviceId() const { return _device_id; };
+};
+} // namespace llaisys::device
diff --git a/src/device/nvidia/nvidia_resource.cu b/src/device/nvidia/nvidia_resource.cu
index 2e63647e5..66d56f131 100644
--- a/src/device/nvidia/nvidia_resource.cu
+++ b/src/device/nvidia/nvidia_resource.cu
@@ -1,7 +1,7 @@
-#include "nvidia_resource.cuh"
-
-namespace llaisys::device::nvidia {
-
-Resource::Resource(int device_id) : llaisys::device::DeviceResource(LLAISYS_DEVICE_NVIDIA, device_id) {}
-
-} // namespace llaisys::device::nvidia
+#include "nvidia_resource.cuh"
+
+namespace llaisys::device::nvidia {
+
+Resource::Resource(int device_id) : llaisys::device::DeviceResource(LLAISYS_DEVICE_NVIDIA, device_id) {}
+
+} // namespace llaisys::device::nvidia
diff --git a/src/device/nvidia/nvidia_resource.cuh b/src/device/nvidia/nvidia_resource.cuh
index a3002170b..42065beb9 100644
--- a/src/device/nvidia/nvidia_resource.cuh
+++ b/src/device/nvidia/nvidia_resource.cuh
@@ -1,11 +1,11 @@
-#pragma once
-
-#include "../device_resource.hpp"
-
-namespace llaisys::device::nvidia {
-class Resource : public llaisys::device::DeviceResource {
-public:
- Resource(int device_id);
- ~Resource();
-};
-} // namespace llaisys::device::nvidia
+#pragma once
+
+#include "../device_resource.hpp"
+
+namespace llaisys::device::nvidia {
+class Resource : public llaisys::device::DeviceResource {
+public:
+ Resource(int device_id);
+ ~Resource();
+};
+} // namespace llaisys::device::nvidia
diff --git a/src/device/nvidia/nvidia_runtime_api.cu b/src/device/nvidia/nvidia_runtime_api.cu
index cab928261..f8737f2f5 100644
--- a/src/device/nvidia/nvidia_runtime_api.cu
+++ b/src/device/nvidia/nvidia_runtime_api.cu
@@ -1,56 +1,95 @@
+/**
+ * NVIDIA CUDA Runtime API 实现：设备管理、显存分配、同步与拷贝。
+ * 对应 include/llaisys/runtime.h 中的 LlaisysRuntimeAPI。
+ */
#include "../runtime_api.hpp"
+#include "../../../include/llaisys.h"
+#include
#include
#include
+#include
namespace llaisys::device::nvidia {
namespace runtime_api {
+
+static cudaMemcpyKind toCudaMemcpyKind(llaisysMemcpyKind_t kind) {
+ switch (kind) {
+ case LLAISYS_MEMCPY_H2H: return cudaMemcpyHostToHost;
+ case LLAISYS_MEMCPY_H2D: return cudaMemcpyHostToDevice;
+ case LLAISYS_MEMCPY_D2H: return cudaMemcpyDeviceToHost;
+ case LLAISYS_MEMCPY_D2D: return cudaMemcpyDeviceToDevice;
+ default: return cudaMemcpyHostToHost;
+ }
+}
+
int getDeviceCount() {
- TO_BE_IMPLEMENTED();
+ int n = 0;
+ cudaError_t e = cudaGetDeviceCount(&n);
+ if (e != cudaSuccess) {
+ std::cerr << "[llaisys/nvidia] cudaGetDeviceCount failed: " << cudaGetErrorString(e)
+ << " (" << e << "). Check CUDA_VISIBLE_DEVICES and driver." << std::endl;
+ return 0;
+ }
+ return n;
}
-void setDevice(int) {
- TO_BE_IMPLEMENTED();
+void setDevice(int id) {
+ cudaSetDevice(id);
}
void deviceSynchronize() {
- TO_BE_IMPLEMENTED();
+ cudaDeviceSynchronize();
}
llaisysStream_t createStream() {
- TO_BE_IMPLEMENTED();
+ cudaStream_t s = nullptr;
+ cudaStreamCreate(&s);
+ return (llaisysStream_t)s;
}
void destroyStream(llaisysStream_t stream) {
- TO_BE_IMPLEMENTED();
+ if (stream)
+ cudaStreamDestroy((cudaStream_t)stream);
}
+
void streamSynchronize(llaisysStream_t stream) {
- TO_BE_IMPLEMENTED();
+ if (stream)
+ cudaStreamSynchronize((cudaStream_t)stream);
+ else
+ cudaDeviceSynchronize();
}
void *mallocDevice(size_t size) {
- TO_BE_IMPLEMENTED();
+ void *ptr = nullptr;
+ cudaMalloc(&ptr, size);
+ return ptr;
}
void freeDevice(void *ptr) {
- TO_BE_IMPLEMENTED();
+ if (ptr)
+ cudaFree(ptr);
}
void *mallocHost(size_t size) {
- TO_BE_IMPLEMENTED();
+ void *ptr = nullptr;
+ cudaMallocHost(&ptr, size);
+ return ptr;
}
void freeHost(void *ptr) {
- TO_BE_IMPLEMENTED();
+ if (ptr)
+ cudaFreeHost(ptr);
}
void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
- TO_BE_IMPLEMENTED();
+ cudaMemcpy(dst, src, size, toCudaMemcpyKind(kind));
}
-void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
- TO_BE_IMPLEMENTED();
+void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
+ cudaStream_t s = stream ? (cudaStream_t)stream : (cudaStream_t)0;
+ cudaMemcpyAsync(dst, src, size, toCudaMemcpyKind(kind), s);
}
static const LlaisysRuntimeAPI RUNTIME_API = {
@@ -65,11 +104,13 @@ static const LlaisysRuntimeAPI RUNTIME_API = {
&mallocHost,
&freeHost,
&memcpySync,
- &memcpyAsync};
+ &memcpyAsync,
+};
} // namespace runtime_api
const LlaisysRuntimeAPI *getRuntimeAPI() {
return &runtime_api::RUNTIME_API;
}
+
} // namespace llaisys::device::nvidia
diff --git a/src/device/runtime_api.cpp b/src/device/runtime_api.cpp
index 2de3eca02..1bb394989 100644
--- a/src/device/runtime_api.cpp
+++ b/src/device/runtime_api.cpp
@@ -1,89 +1,89 @@
-#include "runtime_api.hpp"
-
-namespace llaisys::device {
-
-int getDeviceCount() {
- return 0;
-}
-
-void setDevice(int) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-void deviceSynchronize() {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-llaisysStream_t createStream() {
- EXCEPTION_UNSUPPORTED_DEVICE;
- return nullptr;
-}
-
-void destroyStream(llaisysStream_t stream) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-void streamSynchronize(llaisysStream_t stream) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-void *mallocDevice(size_t size) {
- EXCEPTION_UNSUPPORTED_DEVICE;
- return nullptr;
-}
-
-void freeDevice(void *ptr) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-void *mallocHost(size_t size) {
- EXCEPTION_UNSUPPORTED_DEVICE;
- return nullptr;
-}
-
-void freeHost(void *ptr) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
- EXCEPTION_UNSUPPORTED_DEVICE;
-}
-
-static const LlaisysRuntimeAPI NOOP_RUNTIME_API = {
- &getDeviceCount,
- &setDevice,
- &deviceSynchronize,
- &createStream,
- &destroyStream,
- &streamSynchronize,
- &mallocDevice,
- &freeDevice,
- &mallocHost,
- &freeHost,
- &memcpySync,
- &memcpyAsync};
-
-const LlaisysRuntimeAPI *getUnsupportedRuntimeAPI() {
- return &NOOP_RUNTIME_API;
-}
-
-const LlaisysRuntimeAPI *getRuntimeAPI(llaisysDeviceType_t device_type) {
- // Implement for all device types
- switch (device_type) {
- case LLAISYS_DEVICE_CPU:
- return llaisys::device::cpu::getRuntimeAPI();
- case LLAISYS_DEVICE_NVIDIA:
-#ifdef ENABLE_NVIDIA_API
- return llaisys::device::nvidia::getRuntimeAPI();
-#else
- return getUnsupportedRuntimeAPI();
-#endif
- default:
- EXCEPTION_UNSUPPORTED_DEVICE;
- return nullptr;
- }
-}
-} // namespace llaisys::device
+#include "runtime_api.hpp"
+
+namespace llaisys::device {
+
+int getDeviceCount() {
+ return 0;
+}
+
+void setDevice(int) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+void deviceSynchronize() {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+llaisysStream_t createStream() {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ return nullptr;
+}
+
+void destroyStream(llaisysStream_t stream) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+void streamSynchronize(llaisysStream_t stream) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+void *mallocDevice(size_t size) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ return nullptr;
+}
+
+void freeDevice(void *ptr) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+void *mallocHost(size_t size) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ return nullptr;
+}
+
+void freeHost(void *ptr) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) {
+ EXCEPTION_UNSUPPORTED_DEVICE;
+}
+
+static const LlaisysRuntimeAPI NOOP_RUNTIME_API = {
+ &getDeviceCount,
+ &setDevice,
+ &deviceSynchronize,
+ &createStream,
+ &destroyStream,
+ &streamSynchronize,
+ &mallocDevice,
+ &freeDevice,
+ &mallocHost,
+ &freeHost,
+ &memcpySync,
+ &memcpyAsync};
+
+const LlaisysRuntimeAPI *getUnsupportedRuntimeAPI() {
+ return &NOOP_RUNTIME_API;
+}
+
+const LlaisysRuntimeAPI *getRuntimeAPI(llaisysDeviceType_t device_type) {
+ // Implement for all device types
+ switch (device_type) {
+ case LLAISYS_DEVICE_CPU:
+ return llaisys::device::cpu::getRuntimeAPI();
+ case LLAISYS_DEVICE_NVIDIA:
+#ifdef ENABLE_NVIDIA_API
+ return llaisys::device::nvidia::getRuntimeAPI();
+#else
+ return getUnsupportedRuntimeAPI();
+#endif
+ default:
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ return nullptr;
+ }
+}
+} // namespace llaisys::device
diff --git a/src/device/runtime_api.hpp b/src/device/runtime_api.hpp
index e6b9f80d6..a522d0f7b 100644
--- a/src/device/runtime_api.hpp
+++ b/src/device/runtime_api.hpp
@@ -1,20 +1,20 @@
-#pragma once
-#include "llaisys/runtime.h"
-
-#include "../utils.hpp"
-
-namespace llaisys::device {
-const LlaisysRuntimeAPI *getRuntimeAPI(llaisysDeviceType_t device_type);
-
-const LlaisysRuntimeAPI *getUnsupportedRuntimeAPI();
-
-namespace cpu {
-const LlaisysRuntimeAPI *getRuntimeAPI();
-}
-
-#ifdef ENABLE_NVIDIA_API
-namespace nvidia {
-const LlaisysRuntimeAPI *getRuntimeAPI();
-}
-#endif
-} // namespace llaisys::device
+#pragma once
+#include "llaisys/runtime.h"
+
+#include "../utils.hpp"
+
+namespace llaisys::device {
+const LlaisysRuntimeAPI *getRuntimeAPI(llaisysDeviceType_t device_type);
+
+const LlaisysRuntimeAPI *getUnsupportedRuntimeAPI();
+
+namespace cpu {
+const LlaisysRuntimeAPI *getRuntimeAPI();
+}
+
+#ifdef ENABLE_NVIDIA_API
+namespace nvidia {
+const LlaisysRuntimeAPI *getRuntimeAPI();
+}
+#endif
+} // namespace llaisys::device
diff --git a/src/llaisys/llaisys_tensor.hpp b/src/llaisys/llaisys_tensor.hpp
index d1274ca5a..4828398fb 100644
--- a/src/llaisys/llaisys_tensor.hpp
+++ b/src/llaisys/llaisys_tensor.hpp
@@ -3,7 +3,7 @@
#include "../tensor/tensor.hpp"
-__C {
+LLAISYS_EXTERN_C {
typedef struct LlaisysTensor {
llaisys::tensor_t tensor;
} LlaisysTensor;
diff --git a/src/llaisys/nccl_comm.cu b/src/llaisys/nccl_comm.cu
new file mode 100644
index 000000000..e2e4e6b34
--- /dev/null
+++ b/src/llaisys/nccl_comm.cu
@@ -0,0 +1,91 @@
+/**
+ * NCCL 通信实现（仅当 ENABLE_NCCL 且 ENABLE_NVIDIA_API 时编译）。
+ */
+#if defined(ENABLE_NCCL) && defined(ENABLE_NVIDIA_API)
+
+#include "llaisys/nccl_comm.h"
+#include
+#include
+#include
+#include
+
+static ncclComm_t g_nccl_comm = nullptr;
+static thread_local char g_last_error[256] = "";
+
+static ncclDataType_t to_nccl_dtype(llaisysDataType_t dtype) {
+ switch (dtype) {
+ case LLAISYS_DTYPE_F32: return ncclFloat32;
+ case LLAISYS_DTYPE_F16:
+ case LLAISYS_DTYPE_BF16: return ncclFloat16;
+ case LLAISYS_DTYPE_I64: return ncclInt64;
+ default: return ncclFloat32;
+ }
+}
+
+extern "C" {
+
+void llaisysNcclGetUniqueId(void *buffer) {
+ if (!buffer) return;
+ ncclUniqueId id;
+ ncclGetUniqueId(&id);
+ std::memcpy(buffer, &id, LLAISYS_NCCL_UNIQUE_ID_BYTES);
+}
+
+int llaisysNcclInitRank(int rank, int world_size, const void *unique_id) {
+ if (g_nccl_comm != nullptr) return 0;
+ if (!unique_id || world_size < 1 || rank < 0 || rank >= world_size) return -1;
+ // NCCL 要求每进程在 ncclCommInitRank 前已初始化 CUDA；测试中每进程通过 CUDA_VISIBLE_DEVICES 仅见一卡，设备号为 0
+ cudaError_t ce = cudaSetDevice(0);
+ if (ce != cudaSuccess) {
+ std::snprintf(g_last_error, sizeof(g_last_error), "cudaSetDevice(0) failed: %s", cudaGetErrorString(ce));
+ return -1;
+ }
+ ncclUniqueId id;
+ std::memcpy(&id, unique_id, LLAISYS_NCCL_UNIQUE_ID_BYTES);
+ ncclResult_t r = ncclCommInitRank(&g_nccl_comm, world_size, id, rank);
+ if (r != ncclSuccess) {
+ std::snprintf(g_last_error, sizeof(g_last_error), "ncclCommInitRank: %s", ncclGetErrorString(r));
+ return -1;
+ }
+ g_last_error[0] = '\0';
+ return 0;
+}
+
+int llaisysNcclAllReduce(const void *sendbuf, void *recvbuf, size_t count,
+ llaisysDataType_t dtype, void *stream) {
+ if (!g_nccl_comm || !sendbuf || !recvbuf) return -1;
+ cudaStream_t s = stream ? (cudaStream_t)stream : (cudaStream_t)0;
+ ncclResult_t r = ncclAllReduce(sendbuf, recvbuf, count, to_nccl_dtype(dtype), ncclSum, g_nccl_comm, s);
+ if (r != ncclSuccess) {
+ std::fprintf(stderr, "[llaisys] ncclAllReduce failed: %s\n", ncclGetErrorString(r));
+ return -1;
+ }
+ return 0;
+}
+
+int llaisysNcclAllGather(const void *sendbuf, void *recvbuf, size_t count_per_rank,
+ llaisysDataType_t dtype, void *stream) {
+ if (!g_nccl_comm || !sendbuf || !recvbuf) return -1;
+ cudaStream_t s = stream ? (cudaStream_t)stream : (cudaStream_t)0;
+ ncclResult_t r = ncclAllGather(sendbuf, recvbuf, count_per_rank, to_nccl_dtype(dtype), g_nccl_comm, s);
+ if (r != ncclSuccess) {
+ std::fprintf(stderr, "[llaisys] ncclAllGather failed: %s\n", ncclGetErrorString(r));
+ return -1;
+ }
+ return 0;
+}
+
+void llaisysNcclDestroy(void) {
+ if (g_nccl_comm != nullptr) {
+ ncclCommDestroy(g_nccl_comm);
+ g_nccl_comm = nullptr;
+ }
+}
+
+const char *llaisysNcclGetLastError(void) {
+ return g_last_error[0] ? g_last_error : "(no error)";
+}
+
+} // extern "C"
+
+#endif /* ENABLE_NCCL && ENABLE_NVIDIA_API */
diff --git a/src/llaisys/nccl_comm_stub.cc b/src/llaisys/nccl_comm_stub.cc
new file mode 100644
index 000000000..c09626588
--- /dev/null
+++ b/src/llaisys/nccl_comm_stub.cc
@@ -0,0 +1,47 @@
+/**
+ * NCCL 接口空实现：未启用 ENABLE_NCCL 时提供符号，避免链接失败。
+ * 调用方不应在未启用 NCCL 时使用张量并行。
+ */
+#include "llaisys/nccl_comm.h"
+#include
+
+extern "C" {
+
+void llaisysNcclGetUniqueId(void *buffer) {
+ if (buffer) std::memset(buffer, 0, LLAISYS_NCCL_UNIQUE_ID_BYTES);
+}
+
+int llaisysNcclInitRank(int rank, int world_size, const void *unique_id) {
+ (void)rank;
+ (void)world_size;
+ (void)unique_id;
+ return -1; /* not supported */
+}
+
+int llaisysNcclAllReduce(const void *sendbuf, void *recvbuf, size_t count,
+ llaisysDataType_t dtype, void *stream) {
+ (void)sendbuf;
+ (void)recvbuf;
+ (void)count;
+ (void)dtype;
+ (void)stream;
+ return -1;
+}
+
+int llaisysNcclAllGather(const void *sendbuf, void *recvbuf, size_t count_per_rank,
+ llaisysDataType_t dtype, void *stream) {
+ (void)sendbuf;
+ (void)recvbuf;
+ (void)count_per_rank;
+ (void)dtype;
+ (void)stream;
+ return -1;
+}
+
+void llaisysNcclDestroy(void) {}
+
+const char *llaisysNcclGetLastError(void) {
+ return "(NCCL not compiled)";
+}
+
+} // extern "C"
diff --git a/src/llaisys/ops.cc b/src/llaisys/ops.cc
index c99fbc32f..fe0575919 100644
--- a/src/llaisys/ops.cc
+++ b/src/llaisys/ops.cc
@@ -11,8 +11,9 @@
#include "../ops/rope/op.hpp"
#include "../ops/self_attention/op.hpp"
#include "../ops/swiglu/op.hpp"
+#include "../ops/sample/op.hpp"
-__C {
+LLAISYS_EXTERN_C {
void llaisysAdd(llaisysTensor_t c, llaisysTensor_t a, llaisysTensor_t b) {
llaisys::ops::add(c->tensor, a->tensor, b->tensor);
}
@@ -23,7 +24,7 @@ __C {
llaisys::ops::embedding(out->tensor, index->tensor, weight->tensor);
}
void llaisysLinear(llaisysTensor_t out, llaisysTensor_t in, llaisysTensor_t weight, llaisysTensor_t bias) {
- llaisys::ops::linear(out->tensor, in->tensor, weight->tensor, bias->tensor);
+ llaisys::ops::linear(out->tensor, in->tensor, weight->tensor, bias ? bias->tensor : nullptr);
}
void llaisysRearrange(llaisysTensor_t out, llaisysTensor_t in) {
llaisys::ops::rearrange(out->tensor, in->tensor);
@@ -40,4 +41,7 @@ __C {
void llaisysSwiGLU(llaisysTensor_t out, llaisysTensor_t gate, llaisysTensor_t up) {
llaisys::ops::swiglu(out->tensor, gate->tensor, up->tensor);
}
+ void llaisysSample(llaisysTensor_t out_idx, llaisysTensor_t logits, float temperature, int top_k, float top_p, unsigned long long seed) {
+ llaisys::ops::sample(out_idx->tensor, logits->tensor, temperature, top_k, top_p, static_cast(seed));
+ }
}
diff --git a/src/llaisys/qwen2.cc b/src/llaisys/qwen2.cc
new file mode 100644
index 000000000..e92042844
--- /dev/null
+++ b/src/llaisys/qwen2.cc
@@ -0,0 +1,1561 @@
+/**
+ * Qwen2 模型推理的 C++ 实现：模型创建、权重分配、单步前向与 KV Cache 管理。
+ *
+ * 对外暴露 C 接口（见 include/llaisys/models/qwen2.h）：
+ * - llaisysQwen2ModelCreate / Destroy / Weights：创建、销毁、获取权重句柄；
+ * - llaisysQwen2ModelInfer：给定当前 token 序列，执行一次前向，返回下一个 token id。
+ *
+ * 前向流程：embedding -> 逐层 Transformer Block（attention + MLP）-> 最后一层 norm -> 输出层 linear -> argmax。
+ */
+#include "llaisys/llaisys_tensor.hpp"
+#include "llaisys/models/qwen2.h"
+#include "llaisys/ops.h"
+#ifdef ENABLE_NCCL
+#include "llaisys/nccl_comm.h"
+#endif
+
+#include "../core/llaisys_core.hpp"
+#include "../ops/add/op.hpp"
+#include "../ops/argmax/op.hpp"
+#include "../ops/embedding/op.hpp"
+#include "../ops/linear/op.hpp"
+#include "../ops/rms_norm/op.hpp"
+#include "../ops/rope/op.hpp"
+#include "../ops/self_attention/op.hpp"
+#include "../ops/swiglu/op.hpp"
+#include "../ops/sample/op.hpp"
+#include "../tensor/tensor.hpp"
+#include "../utils.hpp"
+
+#include
+#include
+#include
+#include
+#include
+
+namespace {
+
+using namespace llaisys;
+using tensor_t = llaisys::tensor_t;
+
+/// 从 C 接口的 LlaisysTensor 包装中取出内部 tensor_t 指针
+inline tensor_t get_t(llaisysTensor_t t) { return t->tensor; }
+
+/// 按当前设备做同步内存拷贝（H2H 或 D2D），供写 KV cache 和写回 hidden 用
+void copy_sync(void *dst, const void *src, size_t bytes, llaisysDeviceType_t dev) {
+ llaisys::core::context().setDevice(dev, 0);
+ llaisysMemcpyKind_t kind = (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_D2D;
+ llaisys::core::context().runtime().api()->memcpy_sync(dst, src, bytes, kind);
+}
+
+/// 设备 -> 主机拷贝（Export 用）；device_id 默认 0，多卡时由调用方传入 model->device_id
+void copy_device_to_host(void *host_dst, const void *dev_src, size_t bytes, llaisysDeviceType_t dev, int device_id = 0) {
+ if (dev == LLAISYS_DEVICE_CPU) {
+ std::memcpy(host_dst, dev_src, bytes);
+ return;
+ }
+ llaisys::core::context().setDevice(dev, device_id);
+ llaisys::core::context().runtime().api()->memcpy_sync(host_dst, dev_src, bytes, LLAISYS_MEMCPY_D2H);
+}
+
+/// 主机 -> 设备拷贝（Import 用）；device_id 默认 0，多卡时由调用方传入 model->device_id
+void copy_host_to_device(void *dev_dst, const void *host_src, size_t bytes, llaisysDeviceType_t dev, int device_id = 0) {
+ if (dev == LLAISYS_DEVICE_CPU) {
+ std::memcpy(dev_dst, host_src, bytes);
+ return;
+ }
+ llaisys::core::context().setDevice(dev, device_id);
+ llaisys::core::context().runtime().api()->memcpy_sync(dev_dst, host_src, bytes, LLAISYS_MEMCPY_H2D);
+}
+
+} // namespace
+
+/// Qwen2 模型内部表示：元信息、权重张量、每层 K/V cache 及当前已填充长度
+/// 当 meta.max_batch_size > 1 时：k/v_caches 每层形状 [max_batch_size, maxseq, nkvh, dh]，cache_lens 长度为 max_batch_size
+/// 当 max_batch_size == 1：k/v_caches 每层形状 [maxseq, nkvh, dh]，cache_lens.size()==1
+struct LlaisysQwen2Model {
+ LlaisysQwen2Meta meta;
+ LlaisysQwen2Weights weights;
+ std::vector k_caches; /// 每层 key cache；单 slot [maxseq,nkvh,dh]，多 slot [max_batch,maxseq,nkvh,dh]
+ std::vector v_caches; /// 每层 value cache，同上
+ std::vector cache_lens; /// 每 slot 已写入长度；单 slot 时 size()==1
+ llaisysDeviceType_t device_type;
+ int device_id;
+ tensor_t out_norm_w_cpu; /// 可选：输出层 norm 权重的 CPU 副本，GPU 时最后一层在 CPU 算用
+ tensor_t out_embed_cpu; /// 可选：输出层 linear 权重的 CPU 副本
+ /// 全量 CPU 缓存：若 in_embed_cpu 非空则推理时整次前向在 CPU 上执行
+ tensor_t in_embed_cpu;
+ std::vector attn_norm_w_cpu, attn_q_w_cpu, attn_q_b_cpu, attn_k_w_cpu, attn_k_b_cpu;
+ std::vector attn_v_w_cpu, attn_v_b_cpu, attn_o_w_cpu;
+ std::vector mlp_norm_w_cpu, mlp_gate_w_cpu, mlp_up_w_cpu, mlp_down_w_cpu;
+ std::vector k_caches_cpu, v_caches_cpu;
+ /// 张量并行：AllGather 结果缓冲区，仅 tp_world_size>1 且 GPU 时非空
+ tensor_t tp_gather_q; /// [maxseq, nh*dh]
+ tensor_t tp_gather_k; /// [maxseq, nkvh*dh]
+ tensor_t tp_gather_v; /// [maxseq, nkvh*dh]
+ tensor_t tp_gather_gate; /// [maxseq, di]
+ tensor_t tp_gather_up; /// [maxseq, di]
+};
+
+LLAISYS_EXTERN_C {
+
+/// 根据 meta 为模型分配所有权重张量；tp_world_size>1 时按张量并行分片（列并行：输出维切分；行并行：输入维切分）
+static void create_weight_tensors(LlaisysQwen2Model *m) {
+ const LlaisysQwen2Meta *meta = &m->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t hs = meta->hs;
+ const size_t nh = meta->nh;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t di = meta->di;
+ const size_t voc = meta->voc;
+ const int tp_world = meta->tp_world_size > 0 ? meta->tp_world_size : 1;
+ const llaisysDataType_t dtype = meta->dtype;
+ const llaisysDeviceType_t dev = m->device_type;
+ const int dev_id = m->device_id;
+
+ auto mk = [&](const std::vector &shape) {
+ return LlaisysTensor{llaisys::Tensor::create(shape, dtype, dev, dev_id)};
+ };
+
+ const bool use_tp = (tp_world > 1);
+ const size_t W = use_tp ? static_cast(tp_world) : 1u;
+ const size_t nhdh = nh * dh;
+ const size_t nkvhdh = nkvh * dh;
+
+ // 词嵌入与输出层：非 TP 时全量，TP 时也全量复制（每 rank 一份，避免 token 路由）
+ m->weights.in_embed = new LlaisysTensor(mk({voc, hs}));
+ m->weights.out_embed = new LlaisysTensor(mk({voc, hs}));
+ m->weights.out_norm_w = new LlaisysTensor(mk({hs}));
+
+ m->weights.attn_norm_w = new llaisysTensor_t[nlayer];
+ m->weights.attn_q_w = new llaisysTensor_t[nlayer];
+ m->weights.attn_q_b = new llaisysTensor_t[nlayer];
+ m->weights.attn_k_w = new llaisysTensor_t[nlayer];
+ m->weights.attn_k_b = new llaisysTensor_t[nlayer];
+ m->weights.attn_v_w = new llaisysTensor_t[nlayer];
+ m->weights.attn_v_b = new llaisysTensor_t[nlayer];
+ m->weights.attn_o_w = new llaisysTensor_t[nlayer];
+ m->weights.mlp_norm_w = new llaisysTensor_t[nlayer];
+ m->weights.mlp_gate_w = new llaisysTensor_t[nlayer];
+ m->weights.mlp_up_w = new llaisysTensor_t[nlayer];
+ m->weights.mlp_down_w = new llaisysTensor_t[nlayer];
+
+ for (size_t i = 0; i < nlayer; i++) {
+ m->weights.attn_norm_w[i] = new LlaisysTensor(mk({hs}));
+ if (use_tp) {
+ const size_t nhdh_l = nhdh / W;
+ const size_t nkvhdh_l = nkvhdh / W;
+ const size_t di_l = di / W;
+ m->weights.attn_q_w[i] = new LlaisysTensor(mk({nhdh_l, hs}));
+ m->weights.attn_q_b[i] = new LlaisysTensor(mk({nhdh_l}));
+ m->weights.attn_k_w[i] = new LlaisysTensor(mk({nkvhdh_l, hs}));
+ m->weights.attn_k_b[i] = new LlaisysTensor(mk({nkvhdh_l}));
+ m->weights.attn_v_w[i] = new LlaisysTensor(mk({nkvhdh_l, hs}));
+ m->weights.attn_v_b[i] = new LlaisysTensor(mk({nkvhdh_l}));
+ m->weights.attn_o_w[i] = new LlaisysTensor(mk({hs, nhdh_l}));
+ m->weights.mlp_norm_w[i] = new LlaisysTensor(mk({hs}));
+ m->weights.mlp_gate_w[i] = new LlaisysTensor(mk({di_l, hs}));
+ m->weights.mlp_up_w[i] = new LlaisysTensor(mk({di_l, hs}));
+ m->weights.mlp_down_w[i] = new LlaisysTensor(mk({hs, di_l}));
+ } else {
+ m->weights.attn_q_w[i] = new LlaisysTensor(mk({nhdh, hs}));
+ m->weights.attn_q_b[i] = new LlaisysTensor(mk({nhdh}));
+ m->weights.attn_k_w[i] = new LlaisysTensor(mk({nkvhdh, hs}));
+ m->weights.attn_k_b[i] = new LlaisysTensor(mk({nkvhdh}));
+ m->weights.attn_v_w[i] = new LlaisysTensor(mk({nkvhdh, hs}));
+ m->weights.attn_v_b[i] = new LlaisysTensor(mk({nkvhdh}));
+ m->weights.attn_o_w[i] = new LlaisysTensor(mk({hs, nhdh}));
+ m->weights.mlp_norm_w[i] = new LlaisysTensor(mk({hs}));
+ m->weights.mlp_gate_w[i] = new LlaisysTensor(mk({di, hs}));
+ m->weights.mlp_up_w[i] = new LlaisysTensor(mk({di, hs}));
+ m->weights.mlp_down_w[i] = new LlaisysTensor(mk({hs, di}));
+ }
+ }
+}
+
+/// 创建模型：分配 meta、权重张量、每层 KV cache，返回模型指针；权重数据由调用方通过 llaisysQwen2ModelWeights + tensorLoad 写入
+struct LlaisysQwen2Model *llaisysQwen2ModelCreate(const LlaisysQwen2Meta *meta,
+ llaisysDeviceType_t device,
+ int *device_ids,
+ int ndevice) {
+ (void)device_ids;
+ (void)ndevice;
+ LlaisysQwen2Model *m = new LlaisysQwen2Model();
+ m->meta = *meta;
+ if (m->meta.max_batch_size == 0)
+ m->meta.max_batch_size = 1;
+ const size_t max_batch = m->meta.max_batch_size;
+ m->cache_lens.assign(max_batch, 0);
+ m->device_type = device;
+ m->device_id = (ndevice > 0 && device_ids) ? device_ids[0] : 0;
+
+ llaisys::core::context().setDevice(device, m->device_id);
+
+ create_weight_tensors(m);
+
+ const size_t nlayer = meta->nlayer;
+ const size_t maxseq = meta->maxseq;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+
+ m->out_norm_w_cpu = nullptr;
+ m->out_embed_cpu = nullptr;
+ m->in_embed_cpu = nullptr;
+
+ m->k_caches.resize(nlayer);
+ m->v_caches.resize(nlayer);
+ if (max_batch > 1) {
+ for (size_t i = 0; i < nlayer; i++) {
+ m->k_caches[i] = llaisys::Tensor::create(
+ {max_batch, maxseq, nkvh, dh}, meta->dtype, device, m->device_id);
+ m->v_caches[i] = llaisys::Tensor::create(
+ {max_batch, maxseq, nkvh, dh}, meta->dtype, device, m->device_id);
+ }
+ } else {
+ for (size_t i = 0; i < nlayer; i++) {
+ m->k_caches[i] = llaisys::Tensor::create(
+ {maxseq, nkvh, dh}, meta->dtype, device, m->device_id);
+ m->v_caches[i] = llaisys::Tensor::create(
+ {maxseq, nkvh, dh}, meta->dtype, device, m->device_id);
+ }
+ }
+
+ if (meta->tp_world_size > 1 && device == LLAISYS_DEVICE_NVIDIA) {
+ const size_t nhdh = meta->nh * meta->dh;
+ const size_t nkvhdh = meta->nkvh * meta->dh;
+ m->tp_gather_q = llaisys::Tensor::create(
+ {maxseq, nhdh}, meta->dtype, device, m->device_id);
+ m->tp_gather_k = llaisys::Tensor::create(
+ {maxseq, nkvhdh}, meta->dtype, device, m->device_id);
+ m->tp_gather_v = llaisys::Tensor::create(
+ {maxseq, nkvhdh}, meta->dtype, device, m->device_id);
+ m->tp_gather_gate = llaisys::Tensor::create(
+ {maxseq, meta->di}, meta->dtype, device, m->device_id);
+ m->tp_gather_up = llaisys::Tensor::create(
+ {maxseq, meta->di}, meta->dtype, device, m->device_id);
+ } else {
+ m->tp_gather_q = nullptr;
+ m->tp_gather_k = nullptr;
+ m->tp_gather_v = nullptr;
+ m->tp_gather_gate = nullptr;
+ m->tp_gather_up = nullptr;
+ }
+ return m;
+}
+
+/// 释放模型：先销毁所有权重张量和指针数组，再 delete 模型本体
+void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model *model) {
+ if (!model) return;
+ const size_t nlayer = model->meta.nlayer;
+
+ tensorDestroy(model->weights.in_embed);
+ tensorDestroy(model->weights.out_embed);
+ tensorDestroy(model->weights.out_norm_w);
+
+ for (size_t i = 0; i < nlayer; i++) {
+ tensorDestroy(model->weights.attn_norm_w[i]);
+ tensorDestroy(model->weights.attn_q_w[i]);
+ tensorDestroy(model->weights.attn_q_b[i]);
+ tensorDestroy(model->weights.attn_k_w[i]);
+ tensorDestroy(model->weights.attn_k_b[i]);
+ tensorDestroy(model->weights.attn_v_w[i]);
+ tensorDestroy(model->weights.attn_v_b[i]);
+ tensorDestroy(model->weights.attn_o_w[i]);
+ tensorDestroy(model->weights.mlp_norm_w[i]);
+ tensorDestroy(model->weights.mlp_gate_w[i]);
+ tensorDestroy(model->weights.mlp_up_w[i]);
+ tensorDestroy(model->weights.mlp_down_w[i]);
+ }
+ delete[] model->weights.attn_norm_w;
+ delete[] model->weights.attn_q_w;
+ delete[] model->weights.attn_q_b;
+ delete[] model->weights.attn_k_w;
+ delete[] model->weights.attn_k_b;
+ delete[] model->weights.attn_v_w;
+ delete[] model->weights.attn_v_b;
+ delete[] model->weights.attn_o_w;
+ delete[] model->weights.mlp_norm_w;
+ delete[] model->weights.mlp_gate_w;
+ delete[] model->weights.mlp_up_w;
+ delete[] model->weights.mlp_down_w;
+
+ if (model->out_norm_w_cpu) model->out_norm_w_cpu = nullptr;
+ if (model->out_embed_cpu) model->out_embed_cpu = nullptr;
+ if (model->in_embed_cpu) model->in_embed_cpu = nullptr;
+ model->attn_norm_w_cpu.clear();
+ model->attn_q_w_cpu.clear();
+ model->attn_q_b_cpu.clear();
+ model->attn_k_w_cpu.clear();
+ model->attn_k_b_cpu.clear();
+ model->attn_v_w_cpu.clear();
+ model->attn_v_b_cpu.clear();
+ model->attn_o_w_cpu.clear();
+ model->mlp_norm_w_cpu.clear();
+ model->mlp_gate_w_cpu.clear();
+ model->mlp_up_w_cpu.clear();
+ model->mlp_down_w_cpu.clear();
+ model->k_caches_cpu.clear();
+ model->v_caches_cpu.clear();
+ if (model->tp_gather_q) model->tp_gather_q = nullptr;
+ if (model->tp_gather_k) model->tp_gather_k = nullptr;
+ if (model->tp_gather_v) model->tp_gather_v = nullptr;
+ if (model->tp_gather_gate) model->tp_gather_gate = nullptr;
+ if (model->tp_gather_up) model->tp_gather_up = nullptr;
+
+ delete model;
+}
+
+void llaisysQwen2ModelCacheOutputLayerOnCPU(struct LlaisysQwen2Model *model) {
+ if (!model || model->device_type == LLAISYS_DEVICE_CPU) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t hs = meta->hs, voc = meta->voc;
+ const llaisysDataType_t dtype = meta->dtype;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ llaisys::core::context().runtime().api()->device_synchronize();
+
+ tensor_t in_embed_cpu = llaisys::Tensor::create({voc, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(in_embed_cpu->data(), get_t(model->weights.in_embed)->data(), voc * hs * elem_size, model->device_type, model->device_id);
+ model->in_embed_cpu = in_embed_cpu;
+
+ tensor_t norm_cpu = llaisys::Tensor::create({hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t embed_cpu = llaisys::Tensor::create({voc, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(norm_cpu->data(), get_t(model->weights.out_norm_w)->data(), hs * elem_size, model->device_type, model->device_id);
+ copy_device_to_host(embed_cpu->data(), get_t(model->weights.out_embed)->data(), voc * hs * elem_size, model->device_type, model->device_id);
+
+ model->out_norm_w_cpu = norm_cpu;
+ model->out_embed_cpu = embed_cpu;
+}
+
+void llaisysQwen2ModelCacheAllWeightsOnCPU(struct LlaisysQwen2Model *model) {
+ if (!model || model->device_type == LLAISYS_DEVICE_CPU) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, di = meta->di, voc = meta->voc;
+ const llaisysDataType_t dtype = meta->dtype;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ llaisys::core::context().runtime().api()->device_synchronize();
+
+ tensor_t in_embed_cpu = llaisys::Tensor::create({voc, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(in_embed_cpu->data(), get_t(model->weights.in_embed)->data(), voc * hs * elem_size, model->device_type, model->device_id);
+ if (!model->out_norm_w_cpu) {
+ model->out_norm_w_cpu = llaisys::Tensor::create({hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ model->out_embed_cpu = llaisys::Tensor::create({voc, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ }
+ copy_device_to_host(model->out_norm_w_cpu->data(), get_t(model->weights.out_norm_w)->data(), hs * elem_size, model->device_type, model->device_id);
+ copy_device_to_host(model->out_embed_cpu->data(), get_t(model->weights.out_embed)->data(), voc * hs * elem_size, model->device_type, model->device_id);
+
+ model->attn_norm_w_cpu.resize(nlayer);
+ model->attn_q_w_cpu.resize(nlayer);
+ model->attn_q_b_cpu.resize(nlayer);
+ model->attn_k_w_cpu.resize(nlayer);
+ model->attn_k_b_cpu.resize(nlayer);
+ model->attn_v_w_cpu.resize(nlayer);
+ model->attn_v_b_cpu.resize(nlayer);
+ model->attn_o_w_cpu.resize(nlayer);
+ model->mlp_norm_w_cpu.resize(nlayer);
+ model->mlp_gate_w_cpu.resize(nlayer);
+ model->mlp_up_w_cpu.resize(nlayer);
+ model->mlp_down_w_cpu.resize(nlayer);
+ for (size_t i = 0; i < nlayer; i++) {
+ model->attn_norm_w_cpu[i] = llaisys::Tensor::create({hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_norm_w_cpu[i]->data(), get_t(model->weights.attn_norm_w[i])->data(), hs * elem_size, model->device_type, model->device_id);
+ model->attn_q_w_cpu[i] = llaisys::Tensor::create({nh * dh, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_q_w_cpu[i]->data(), get_t(model->weights.attn_q_w[i])->data(), nh * dh * hs * elem_size, model->device_type, model->device_id);
+ model->attn_q_b_cpu[i] = llaisys::Tensor::create({nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_q_b_cpu[i]->data(), get_t(model->weights.attn_q_b[i])->data(), nh * dh * elem_size, model->device_type, model->device_id);
+ model->attn_k_w_cpu[i] = llaisys::Tensor::create({nkvh * dh, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_k_w_cpu[i]->data(), get_t(model->weights.attn_k_w[i])->data(), nkvh * dh * hs * elem_size, model->device_type, model->device_id);
+ model->attn_k_b_cpu[i] = llaisys::Tensor::create({nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_k_b_cpu[i]->data(), get_t(model->weights.attn_k_b[i])->data(), nkvh * dh * elem_size, model->device_type, model->device_id);
+ model->attn_v_w_cpu[i] = llaisys::Tensor::create({nkvh * dh, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_v_w_cpu[i]->data(), get_t(model->weights.attn_v_w[i])->data(), nkvh * dh * hs * elem_size, model->device_type, model->device_id);
+ model->attn_v_b_cpu[i] = llaisys::Tensor::create({nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_v_b_cpu[i]->data(), get_t(model->weights.attn_v_b[i])->data(), nkvh * dh * elem_size, model->device_type, model->device_id);
+ model->attn_o_w_cpu[i] = llaisys::Tensor::create({hs, nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->attn_o_w_cpu[i]->data(), get_t(model->weights.attn_o_w[i])->data(), hs * nh * dh * elem_size, model->device_type, model->device_id);
+ model->mlp_norm_w_cpu[i] = llaisys::Tensor::create({hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->mlp_norm_w_cpu[i]->data(), get_t(model->weights.mlp_norm_w[i])->data(), hs * elem_size, model->device_type, model->device_id);
+ model->mlp_gate_w_cpu[i] = llaisys::Tensor::create({di, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->mlp_gate_w_cpu[i]->data(), get_t(model->weights.mlp_gate_w[i])->data(), di * hs * elem_size, model->device_type, model->device_id);
+ model->mlp_up_w_cpu[i] = llaisys::Tensor::create({di, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->mlp_up_w_cpu[i]->data(), get_t(model->weights.mlp_up_w[i])->data(), di * hs * elem_size, model->device_type, model->device_id);
+ model->mlp_down_w_cpu[i] = llaisys::Tensor::create({hs, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(model->mlp_down_w_cpu[i]->data(), get_t(model->weights.mlp_down_w[i])->data(), hs * di * elem_size, model->device_type, model->device_id);
+ }
+
+ model->k_caches_cpu.resize(nlayer);
+ model->v_caches_cpu.resize(nlayer);
+ const std::vector &k0_shape = model->k_caches[0]->shape();
+ for (size_t i = 0; i < nlayer; i++) {
+ model->k_caches_cpu[i] = llaisys::Tensor::create(k0_shape, dtype, LLAISYS_DEVICE_CPU, 0);
+ model->v_caches_cpu[i] = llaisys::Tensor::create(k0_shape, dtype, LLAISYS_DEVICE_CPU, 0);
+ }
+ model->in_embed_cpu = in_embed_cpu;
+}
+
+/// 返回模型权重结构体指针，供 Python 侧根据 safetensors key 找到对应句柄并 tensorLoad
+struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model *model) {
+ return model ? &model->weights : nullptr;
+}
+
+} // extern "C"
+
+namespace {
+
+/**
+ * 单层 Transformer Block 前向：Attention（norm -> q/k/v -> RoPE -> 写 KV cache -> attention -> o_proj -> 残差）+ MLP（norm -> gate/up -> SwiGLU -> down -> 残差）。
+ * 输入输出通过 hidden 传入并在本函数内原地更新；slot_id 指定使用哪一槽的 KV cache（多 slot 时）。
+ */
+void forward_layer(LlaisysQwen2Model *m, size_t layer_idx, size_t slot_id,
+ tensor_t hidden, // [seq, hs]
+ tensor_t normed, // [seq, hs]
+ tensor_t q_buf, // [seq, nh*dh]
+ tensor_t k_buf, // [seq, nkvh*dh]
+ tensor_t v_buf, // [seq, nkvh*dh]
+ tensor_t q_rope, // [seq, nh, dh]
+ tensor_t k_rope, // [seq, nkvh, dh]
+ tensor_t attn_val, // [seq, nh, dh]
+ tensor_t o_proj_out, // [seq, hs]
+ tensor_t res_buf, // [seq, hs]
+ tensor_t gate_buf, // [seq, di]
+ tensor_t up_buf, // [seq, di]
+ tensor_t mlp_buf, // [seq, di]
+ tensor_t down_buf, // [seq, hs]
+ size_t seq_len, size_t cache_start,
+ tensor_t pos_ids_t) {
+ const LlaisysQwen2Meta *meta = &m->meta;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh;
+ const size_t maxseq = meta->maxseq;
+ const float eps = meta->epsilon, theta = meta->theta;
+ LlaisysQwen2Weights *w = &m->weights;
+ tensor_t wt = get_t(w->attn_norm_w[layer_idx]);
+ const float scale = 1.f / std::sqrt(static_cast(dh));
+
+ tensor_t k_cache_raw = m->k_caches[layer_idx];
+ tensor_t v_cache_raw = m->v_caches[layer_idx];
+ tensor_t k_cache, v_cache;
+ if (meta->max_batch_size > 1) {
+ k_cache = k_cache_raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ v_cache = v_cache_raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ } else {
+ k_cache = k_cache_raw;
+ v_cache = v_cache_raw;
+ }
+
+ // ---------- Attention 分支：norm -> Q/K/V 投影 -> RoPE -> 写 KV cache -> causal attention -> o_proj -> 残差 ----------
+ llaisys::ops::rms_norm(normed, hidden, wt, eps);
+ llaisys::ops::linear(q_buf, normed, get_t(w->attn_q_w[layer_idx]), get_t(w->attn_q_b[layer_idx]));
+ llaisys::ops::linear(k_buf, normed, get_t(w->attn_k_w[layer_idx]), get_t(w->attn_k_b[layer_idx]));
+ llaisys::ops::linear(v_buf, normed, get_t(w->attn_v_w[layer_idx]), get_t(w->attn_v_b[layer_idx]));
+
+ // 展平为 [seq, nh, dh] / [seq, nkvh, dh] 以做 RoPE
+ std::vector shape_q = {seq_len, nh, dh};
+ std::vector shape_kv = {seq_len, nkvh, dh};
+ tensor_t q_view = q_buf->view(shape_q);
+ tensor_t k_view = k_buf->view(shape_kv);
+ tensor_t v_view = v_buf->view(shape_kv);
+
+ llaisys::ops::rope(q_rope, q_view, pos_ids_t, theta);
+ llaisys::ops::rope(k_rope, k_view, pos_ids_t, theta);
+
+ // 将本步的 K/V 写入 cache 的 [cache_start, cache_start+seq_len) 位置，decode 时复用历史
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const size_t kv_row_bytes = nkvh * dh * elem_size;
+ for (size_t s = 0; s < seq_len; s++) {
+ size_t cache_pos = cache_start + s;
+ copy_sync(
+ reinterpret_cast(k_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(k_rope->data()) + s * kv_row_bytes,
+ kv_row_bytes, m->device_type);
+ copy_sync(
+ reinterpret_cast(v_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(v_buf->data()) + s * kv_row_bytes,
+ kv_row_bytes, m->device_type);
+ }
+
+ // 当前有效长度为 cache_start + seq_len；取 cache 的该前缀做 attention
+ size_t kv_len = cache_start + seq_len;
+ tensor_t k_slice = k_cache->slice(0, 0, kv_len);
+ tensor_t v_slice = v_cache->slice(0, 0, kv_len);
+
+ llaisys::ops::self_attention(attn_val, q_rope, k_slice, v_slice, scale);
+
+ std::vector shape_attn_flat = {seq_len, nh * dh};
+ tensor_t attn_flat = attn_val->view(shape_attn_flat);
+ llaisys::ops::linear(o_proj_out, attn_flat, get_t(w->attn_o_w[layer_idx]), nullptr);
+
+ llaisys::ops::add(res_buf, hidden, o_proj_out);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, m->device_type);
+
+ // ---------- MLP 分支：norm -> gate/up -> SwiGLU -> down -> 残差 ----------
+ llaisys::ops::rms_norm(normed, hidden, get_t(w->mlp_norm_w[layer_idx]), eps);
+ llaisys::ops::linear(gate_buf, normed, get_t(w->mlp_gate_w[layer_idx]), nullptr);
+ llaisys::ops::linear(up_buf, normed, get_t(w->mlp_up_w[layer_idx]), nullptr);
+ llaisys::ops::swiglu(mlp_buf, gate_buf, up_buf);
+ llaisys::ops::linear(down_buf, mlp_buf, get_t(w->mlp_down_w[layer_idx]), nullptr);
+ llaisys::ops::add(res_buf, hidden, down_buf);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, m->device_type);
+}
+
+/// 单层 Transformer 前向，使用 CPU 权重与 CPU KV cache（全量 CPU 推理用）
+void forward_layer_cpu(LlaisysQwen2Model *m, size_t layer_idx, size_t slot_id,
+ tensor_t hidden, tensor_t normed, tensor_t q_buf, tensor_t k_buf, tensor_t v_buf,
+ tensor_t q_rope, tensor_t k_rope, tensor_t attn_val, tensor_t o_proj_out, tensor_t res_buf,
+ tensor_t gate_buf, tensor_t up_buf, tensor_t mlp_buf, tensor_t down_buf,
+ size_t seq_len, size_t cache_start, tensor_t pos_ids_t) {
+ const LlaisysQwen2Meta *meta = &m->meta;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, maxseq = meta->maxseq;
+ const float eps = meta->epsilon, theta = meta->theta;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const float scale = 1.f / std::sqrt(static_cast(dh));
+
+ tensor_t k_cache = m->k_caches_cpu[layer_idx];
+ tensor_t v_cache = m->v_caches_cpu[layer_idx];
+ if (meta->max_batch_size > 1) {
+ k_cache = k_cache->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ v_cache = v_cache->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ }
+
+ llaisys::ops::rms_norm(normed, hidden, m->attn_norm_w_cpu[layer_idx], eps);
+ llaisys::ops::linear(q_buf, normed, m->attn_q_w_cpu[layer_idx], m->attn_q_b_cpu[layer_idx]);
+ llaisys::ops::linear(k_buf, normed, m->attn_k_w_cpu[layer_idx], m->attn_k_b_cpu[layer_idx]);
+ llaisys::ops::linear(v_buf, normed, m->attn_v_w_cpu[layer_idx], m->attn_v_b_cpu[layer_idx]);
+
+ std::vector shape_q = {seq_len, nh, dh};
+ std::vector shape_kv = {seq_len, nkvh, dh};
+ tensor_t q_view = q_buf->view(shape_q);
+ tensor_t k_view = k_buf->view(shape_kv);
+ tensor_t v_view = v_buf->view(shape_kv);
+ llaisys::ops::rope(q_rope, q_view, pos_ids_t, theta);
+ llaisys::ops::rope(k_rope, k_view, pos_ids_t, theta);
+
+ const size_t kv_row_bytes = nkvh * dh * elem_size;
+ for (size_t s = 0; s < seq_len; s++) {
+ size_t cache_pos = cache_start + s;
+ copy_sync(
+ reinterpret_cast(k_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(k_rope->data()) + s * kv_row_bytes,
+ kv_row_bytes, LLAISYS_DEVICE_CPU);
+ copy_sync(
+ reinterpret_cast(v_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(v_buf->data()) + s * kv_row_bytes,
+ kv_row_bytes, LLAISYS_DEVICE_CPU);
+ }
+
+ size_t kv_len = cache_start + seq_len;
+ tensor_t k_slice = k_cache->slice(0, 0, kv_len);
+ tensor_t v_slice = v_cache->slice(0, 0, kv_len);
+ llaisys::ops::self_attention(attn_val, q_rope, k_slice, v_slice, scale);
+
+ std::vector shape_attn_flat = {seq_len, nh * dh};
+ tensor_t attn_flat = attn_val->view(shape_attn_flat);
+ llaisys::ops::linear(o_proj_out, attn_flat, m->attn_o_w_cpu[layer_idx], nullptr);
+ llaisys::ops::add(res_buf, hidden, o_proj_out);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, LLAISYS_DEVICE_CPU);
+
+ llaisys::ops::rms_norm(normed, hidden, m->mlp_norm_w_cpu[layer_idx], eps);
+ llaisys::ops::linear(gate_buf, normed, m->mlp_gate_w_cpu[layer_idx], nullptr);
+ llaisys::ops::linear(up_buf, normed, m->mlp_up_w_cpu[layer_idx], nullptr);
+ llaisys::ops::swiglu(mlp_buf, gate_buf, up_buf);
+ llaisys::ops::linear(down_buf, mlp_buf, m->mlp_down_w_cpu[layer_idx], nullptr);
+ llaisys::ops::add(res_buf, hidden, down_buf);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, LLAISYS_DEVICE_CPU);
+}
+
+#ifdef ENABLE_NCCL
+/// 张量并行单层前向：列并行 Q/K/V/Gate/Up + AllGather，行并行 O/Down + AllReduce(Sum)。stream 为 CUDA 流。
+void forward_layer_tp(LlaisysQwen2Model *m, size_t layer_idx, size_t slot_id,
+ tensor_t hidden, tensor_t normed,
+ tensor_t q_buf_local, tensor_t k_buf_local, tensor_t v_buf_local,
+ tensor_t q_rope, tensor_t k_rope, tensor_t attn_val, tensor_t o_proj_out, tensor_t res_buf,
+ tensor_t gate_buf_local, tensor_t up_buf_local, tensor_t mlp_buf, tensor_t down_buf,
+ size_t seq_len, size_t cache_start, tensor_t pos_ids_t, void *stream) {
+ const LlaisysQwen2Meta *meta = &m->meta;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, di = meta->di;
+ const size_t maxseq = meta->maxseq;
+ const float eps = meta->epsilon, theta = meta->theta;
+ const llaisysDataType_t dtype = meta->dtype;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+ const float scale = 1.f / std::sqrt(static_cast(dh));
+ const int W = meta->tp_world_size;
+ const size_t nhdh = nh * dh;
+ const size_t nkvhdh = nkvh * dh;
+ const size_t nhdh_l = nhdh / static_cast(W);
+ const size_t nkvhdh_l = nkvhdh / static_cast(W);
+ const size_t di_l = di / static_cast(W);
+
+ LlaisysQwen2Weights *w = &m->weights;
+ tensor_t wt = get_t(w->attn_norm_w[layer_idx]);
+ tensor_t k_cache_raw = m->k_caches[layer_idx];
+ tensor_t v_cache_raw = m->v_caches[layer_idx];
+ tensor_t k_cache, v_cache;
+ if (meta->max_batch_size > 1) {
+ k_cache = k_cache_raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ v_cache = v_cache_raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ } else {
+ k_cache = k_cache_raw;
+ v_cache = v_cache_raw;
+ }
+
+ llaisys::ops::rms_norm(normed, hidden, wt, eps);
+ llaisys::ops::linear(q_buf_local, normed, get_t(w->attn_q_w[layer_idx]), get_t(w->attn_q_b[layer_idx]));
+ llaisys::ops::linear(k_buf_local, normed, get_t(w->attn_k_w[layer_idx]), get_t(w->attn_k_b[layer_idx]));
+ llaisys::ops::linear(v_buf_local, normed, get_t(w->attn_v_w[layer_idx]), get_t(w->attn_v_b[layer_idx]));
+
+ tensor_t gather_q = m->tp_gather_q->slice(0, 0, seq_len);
+ tensor_t gather_k = m->tp_gather_k->slice(0, 0, seq_len);
+ tensor_t gather_v = m->tp_gather_v->slice(0, 0, seq_len);
+ llaisysNcclAllGather(q_buf_local->data(), gather_q->data(), seq_len * nhdh_l, dtype, stream);
+ llaisysNcclAllGather(k_buf_local->data(), gather_k->data(), seq_len * nkvhdh_l, dtype, stream);
+ llaisysNcclAllGather(v_buf_local->data(), gather_v->data(), seq_len * nkvhdh_l, dtype, stream);
+ llaisys::core::context().runtime().api()->stream_synchronize(stream);
+
+ std::vector shape_q = {seq_len, nh, dh};
+ std::vector shape_kv = {seq_len, nkvh, dh};
+ tensor_t q_view = gather_q->view(shape_q);
+ tensor_t k_view = gather_k->view(shape_kv);
+ tensor_t v_view = gather_v->view(shape_kv);
+ llaisys::ops::rope(q_rope, q_view, pos_ids_t, theta);
+ llaisys::ops::rope(k_rope, k_view, pos_ids_t, theta);
+
+ const size_t kv_row_bytes = nkvh * dh * elem_size;
+ for (size_t s = 0; s < seq_len; s++) {
+ size_t cache_pos = cache_start + s;
+ copy_sync(
+ reinterpret_cast(k_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(k_rope->data()) + s * kv_row_bytes,
+ kv_row_bytes, m->device_type);
+ copy_sync(
+ reinterpret_cast(v_cache->data()) + cache_pos * kv_row_bytes,
+ reinterpret_cast(gather_v->data()) + s * nkvhdh * elem_size,
+ kv_row_bytes, m->device_type);
+ }
+ size_t kv_len = cache_start + seq_len;
+ tensor_t k_slice = k_cache->slice(0, 0, kv_len);
+ tensor_t v_slice = v_cache->slice(0, 0, kv_len);
+ llaisys::ops::self_attention(attn_val, q_rope, k_slice, v_slice, scale);
+
+ std::vector shape_attn_flat = {seq_len, nhdh};
+ tensor_t attn_flat = attn_val->view(shape_attn_flat);
+ llaisys::ops::linear(o_proj_out, attn_flat, get_t(w->attn_o_w[layer_idx]), nullptr);
+ llaisysNcclAllReduce(o_proj_out->data(), o_proj_out->data(), seq_len * hs, dtype, stream);
+ llaisys::core::context().runtime().api()->stream_synchronize(stream);
+ llaisys::ops::add(res_buf, hidden, o_proj_out);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, m->device_type);
+
+ llaisys::ops::rms_norm(normed, hidden, get_t(w->mlp_norm_w[layer_idx]), eps);
+ llaisys::ops::linear(gate_buf_local, normed, get_t(w->mlp_gate_w[layer_idx]), nullptr);
+ llaisys::ops::linear(up_buf_local, normed, get_t(w->mlp_up_w[layer_idx]), nullptr);
+ tensor_t gather_gate = m->tp_gather_gate->slice(0, 0, seq_len);
+ tensor_t gather_up = m->tp_gather_up->slice(0, 0, seq_len);
+ llaisysNcclAllGather(gate_buf_local->data(), gather_gate->data(), seq_len * di_l, dtype, stream);
+ llaisysNcclAllGather(up_buf_local->data(), gather_up->data(), seq_len * di_l, dtype, stream);
+ llaisys::core::context().runtime().api()->stream_synchronize(stream);
+ llaisys::ops::swiglu(mlp_buf, gather_gate, gather_up);
+ llaisys::ops::linear(down_buf, mlp_buf, get_t(w->mlp_down_w[layer_idx]), nullptr);
+ llaisysNcclAllReduce(down_buf->data(), down_buf->data(), seq_len * hs, dtype, stream);
+ llaisys::core::context().runtime().api()->stream_synchronize(stream);
+ llaisys::ops::add(res_buf, hidden, down_buf);
+ copy_sync(hidden->data(), res_buf->data(), seq_len * hs * elem_size, m->device_type);
+}
+#endif
+
+} // namespace
+
+LLAISYS_EXTERN_C {
+
+size_t llaisysQwen2ModelGetCacheLen(struct LlaisysQwen2Model *model) {
+ return model && !model->cache_lens.empty() ? model->cache_lens[0] : 0;
+}
+
+size_t llaisysQwen2ModelGetCacheLenSlot(struct LlaisysQwen2Model *model, size_t slot_id) {
+ if (!model || model->cache_lens.empty()) return 0;
+ if (slot_id >= model->cache_lens.size()) return 0;
+ return model->cache_lens[slot_id];
+}
+
+size_t llaisysQwen2ModelGetKVCacheBytes(struct LlaisysQwen2Model *model, size_t prefix_len) {
+ if (!model || prefix_len == 0) return 0;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ return nlayer * 2 * prefix_len * nkvh * dh * elem_size;
+}
+
+} // extern "C"
+
+// 内部辅助函数：使用 C++ linkage，避免 MSVC C4190（C-linkage 返回 C++ 类型）
+static tensor_t get_slot_k_cache(LlaisysQwen2Model *model, size_t layer_idx, size_t slot_id) {
+ const size_t maxseq = model->meta.maxseq, nkvh = model->meta.nkvh, dh = model->meta.dh;
+ tensor_t raw = model->k_caches[layer_idx];
+ if (model->meta.max_batch_size > 1)
+ return raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ return raw;
+}
+static tensor_t get_slot_v_cache(LlaisysQwen2Model *model, size_t layer_idx, size_t slot_id) {
+ const size_t maxseq = model->meta.maxseq, nkvh = model->meta.nkvh, dh = model->meta.dh;
+ tensor_t raw = model->v_caches[layer_idx];
+ if (model->meta.max_batch_size > 1)
+ return raw->slice(0, slot_id, slot_id + 1)->view({maxseq, nkvh, dh});
+ return raw;
+}
+
+LLAISYS_EXTERN_C {
+
+void llaisysQwen2ModelExportKVCache(struct LlaisysQwen2Model *model, void *ptr_out) {
+ if (!model || !ptr_out || model->cache_lens.empty()) return;
+ const size_t cache_len = model->cache_lens[0];
+ if (cache_len == 0) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const size_t row_bytes = nkvh * dh * elem_size;
+ const size_t layer_bytes = cache_len * row_bytes;
+ std::byte *out = static_cast(ptr_out);
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ for (size_t i = 0; i < nlayer; i++) {
+ tensor_t k_slot = get_slot_k_cache(model, i, 0);
+ tensor_t v_slot = get_slot_v_cache(model, i, 0);
+ copy_device_to_host(out, k_slot->data(), layer_bytes, model->device_type, model->device_id);
+ out += layer_bytes;
+ copy_device_to_host(out, v_slot->data(), layer_bytes, model->device_type, model->device_id);
+ out += layer_bytes;
+ }
+}
+
+void llaisysQwen2ModelImportKVCache(struct LlaisysQwen2Model *model, const void *ptr_in, size_t prefix_len) {
+ if (!model || !ptr_in || prefix_len == 0) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const size_t row_bytes = nkvh * dh * elem_size;
+ const size_t layer_bytes = prefix_len * row_bytes;
+ const std::byte *in = static_cast(ptr_in);
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ for (size_t i = 0; i < nlayer; i++) {
+ tensor_t k_slot = get_slot_k_cache(model, i, 0);
+ tensor_t v_slot = get_slot_v_cache(model, i, 0);
+ copy_host_to_device(k_slot->data(), in, layer_bytes, model->device_type, model->device_id);
+ in += layer_bytes;
+ copy_host_to_device(v_slot->data(), in, layer_bytes, model->device_type, model->device_id);
+ in += layer_bytes;
+ }
+ model->cache_lens[0] = prefix_len;
+}
+
+void llaisysQwen2ModelResetKVCache(struct LlaisysQwen2Model *model) {
+ if (model) {
+ for (size_t i = 0; i < model->cache_lens.size(); i++)
+ model->cache_lens[i] = 0;
+ }
+}
+
+void llaisysQwen2ModelResetKVCacheSlot(struct LlaisysQwen2Model *model, size_t slot_id) {
+ if (!model || model->cache_lens.empty()) return;
+ if (slot_id < model->cache_lens.size())
+ model->cache_lens[slot_id] = 0;
+}
+
+void llaisysQwen2ModelExportKVCacheSlot(struct LlaisysQwen2Model *model, size_t slot_id, void *ptr_out) {
+ if (!model || !ptr_out || model->cache_lens.empty()) return;
+ if (slot_id >= model->cache_lens.size()) return;
+ const size_t cache_len = model->cache_lens[slot_id];
+ if (cache_len == 0) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const size_t row_bytes = nkvh * dh * elem_size;
+ const size_t layer_bytes = cache_len * row_bytes;
+ std::byte *out = static_cast(ptr_out);
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ for (size_t i = 0; i < nlayer; i++) {
+ tensor_t k_slot = get_slot_k_cache(model, i, slot_id);
+ tensor_t v_slot = get_slot_v_cache(model, i, slot_id);
+ copy_device_to_host(out, k_slot->data(), layer_bytes, model->device_type, model->device_id);
+ out += layer_bytes;
+ copy_device_to_host(out, v_slot->data(), layer_bytes, model->device_type, model->device_id);
+ out += layer_bytes;
+ }
+}
+
+void llaisysQwen2ModelImportKVCacheSlot(struct LlaisysQwen2Model *model, size_t slot_id, const void *ptr_in, size_t prefix_len) {
+ if (!model || !ptr_in || prefix_len == 0 || model->cache_lens.empty()) return;
+ if (slot_id >= model->cache_lens.size()) return;
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t nkvh = meta->nkvh;
+ const size_t dh = meta->dh;
+ const size_t elem_size = llaisys::utils::dsize(meta->dtype);
+ const size_t row_bytes = nkvh * dh * elem_size;
+ const size_t layer_bytes = prefix_len * row_bytes;
+ const std::byte *in = static_cast(ptr_in);
+ llaisys::core::context().setDevice(model->device_type, model->device_id);
+ for (size_t i = 0; i < nlayer; i++) {
+ tensor_t k_slot = get_slot_k_cache(model, i, slot_id);
+ tensor_t v_slot = get_slot_v_cache(model, i, slot_id);
+ copy_host_to_device(k_slot->data(), in, layer_bytes, model->device_type, model->device_id);
+ in += layer_bytes;
+ copy_host_to_device(v_slot->data(), in, layer_bytes, model->device_type, model->device_id);
+ in += layer_bytes;
+ }
+ model->cache_lens[slot_id] = prefix_len;
+}
+
+/**
+ * 单步推理：根据当前 token 序列做一次前向，返回下一个 token 的 id。
+ * - 若 cache_len==0（prefill）：传入整段 token_ids，seq_len=ntoken，并填充 KV cache；
+ * - 若 cache_len>0 且 ntoken>1（suffix prefill）：传入后缀 token_ids，seq_len=ntoken，只对后缀做 prefill；
+ * - 否则（decode）：只传入最后一个 token，seq_len=1，用已有 cache 做 attention。
+ * - temperature<=0 或极小且 top_k<=1、top_p>=1 时使用 argmax；否则使用随机采样（Temperature/Top-K/Top-P）。
+ */
+int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model *model,
+ int64_t *token_ids,
+ size_t ntoken,
+ float temperature,
+ int top_k,
+ float top_p,
+ unsigned long long seed) {
+ if (std::getenv("LLAISYS_DEBUG_INFER")) {
+ std::fprintf(stderr, "[DBG] Infer entered model=%p ntoken=%zu\n",
+ (void *)model, (unsigned long)ntoken);
+ std::fflush(stderr);
+ }
+ if (!model || ntoken == 0) return static_cast(-1);
+ if (model->cache_lens.empty()) return static_cast(-1);
+
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, di = meta->di;
+ const size_t voc = meta->voc;
+ const llaisysDataType_t dtype = meta->dtype;
+ const llaisysDeviceType_t dev = model->device_type;
+ const int dev_id = model->device_id;
+
+ const size_t slot0_len = model->cache_lens[0];
+ const bool is_prefill = (slot0_len == 0);
+ const bool is_suffix_prefill = (slot0_len > 0 && ntoken > 1);
+ const size_t seq_len = is_suffix_prefill ? ntoken : (is_prefill ? ntoken : 1);
+ const size_t cache_start = slot0_len;
+
+ if (std::getenv("LLAISYS_DEBUG_INFER")) {
+ std::fprintf(stderr, "[DBG] Infer entry ntoken=%zu seq_len=%zu prefill=%d\n",
+ (unsigned long)ntoken, (unsigned long)seq_len, is_prefill ? 1 : 0);
+ std::fflush(stderr);
+ }
+
+ // GPU 且已全量缓存 CPU 权重：整次前向在 CPU 上执行。使用 k_caches_cpu 非空判断，避免仅 CacheOutputLayerOnCPU 时误入导致越界。
+ if (dev != LLAISYS_DEVICE_CPU && !model->k_caches_cpu.empty()) {
+ const bool use_sampling = (temperature > 1e-6f && (top_k > 1 || (top_p > 0.f && top_p < 1.f)));
+ int64_t next_token = meta->end_token;
+ llaisys::core::context().setDevice(LLAISYS_DEVICE_CPU, 0);
+
+ tensor_t token_tensor = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ if (is_suffix_prefill)
+ std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else if (is_prefill)
+ std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else
+ std::memcpy(token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t));
+
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++)
+ pos_ids_host[s] = static_cast(cache_start + s);
+ std::memcpy(pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t));
+
+ llaisys::ops::embedding(hidden, token_tensor, model->in_embed_cpu);
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer_cpu(model, i, 0, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+ model->cache_lens[0] = cache_start + seq_len;
+
+ tensor_t normed_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_cpu, hidden, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({seq_len, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_cpu, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->slice(0, seq_len - 1, seq_len)->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ std::memcpy(&next_token, sampled_idx->data(), sizeof(int64_t));
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ std::memcpy(&next_token, max_idx_t->data(), sizeof(int64_t));
+ }
+ return next_token;
+ }
+
+ llaisys::core::context().setDevice(dev, dev_id);
+
+ tensor_t token_tensor = llaisys::Tensor::create(
+ {seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ if (is_suffix_prefill) {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids, ntoken * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ } else if (is_prefill) {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids, ntoken * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ } else {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ }
+
+ const bool use_tp_buf = (meta->tp_world_size > 1 && model->tp_gather_q != nullptr && dev == LLAISYS_DEVICE_NVIDIA);
+ const size_t nhdh_buf = use_tp_buf ? (nh * dh / static_cast(meta->tp_world_size)) : (nh * dh);
+ const size_t nkvhdh_buf = use_tp_buf ? (nkvh * dh / static_cast(meta->tp_world_size)) : (nkvh * dh);
+ const size_t di_buf = use_tp_buf ? (di / static_cast(meta->tp_world_size)) : di;
+
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nhdh_buf}, dtype, dev, dev_id);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, dev, dev_id);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, dev, dev_id);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++)
+ pos_ids_host[s] = static_cast(cache_start + s);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+
+ if (dev != LLAISYS_DEVICE_CPU && model->in_embed_cpu) {
+ tensor_t token_cpu = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ if (is_suffix_prefill) std::memcpy(token_cpu->data(), token_ids, ntoken * sizeof(int64_t));
+ else if (is_prefill) std::memcpy(token_cpu->data(), token_ids, ntoken * sizeof(int64_t));
+ else std::memcpy(token_cpu->data(), token_ids + ntoken - 1, sizeof(int64_t));
+ tensor_t hidden_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::embedding(hidden_cpu, token_cpu, model->in_embed_cpu);
+ llaisys::core::context().setDevice(dev, dev_id);
+ copy_host_to_device(hidden->data(), hidden_cpu->data(), seq_len * hs * llaisys::utils::dsize(dtype), dev, dev_id);
+ llaisys::core::context().setDevice(dev, dev_id);
+ } else {
+ llaisys::ops::embedding(hidden, token_tensor, get_t(model->weights.in_embed));
+ }
+
+ const bool use_tp_infer = (meta->tp_world_size > 1 && model->tp_gather_q != nullptr && dev == LLAISYS_DEVICE_NVIDIA);
+ if (use_tp_infer) {
+#ifdef ENABLE_NCCL
+ void *stream = llaisys::core::context().runtime().stream();
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer_tp(model, i, 0, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t, stream);
+ }
+#else
+ (void)seq_len;
+ (void)cache_start;
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer(model, i, 0, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+#endif
+ } else {
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer(model, i, 0, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+ }
+ model->cache_lens[0] = cache_start + seq_len;
+
+ const bool use_sampling = (temperature > 1e-6f && (top_k > 1 || (top_p > 0.f && top_p < 1.f)));
+ int64_t next_token = meta->end_token;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+
+ // GPU 且已缓存输出层 CPU 权重时：最后一层 norm + linear 在 CPU 上算，规避 GPU 输出层异常
+ if (dev != LLAISYS_DEVICE_CPU && model->out_norm_w_cpu && model->out_embed_cpu) {
+ if (std::getenv("LLAISYS_DEBUG_INFER")) { std::fprintf(stderr, "[DBG] Infer before out_cpu sync+d2h\n"); std::fflush(stderr); }
+ llaisys::core::context().setDevice(dev, dev_id);
+ llaisys::core::context().runtime().api()->device_synchronize();
+ // 整块拷贝 hidden 到 CPU，再在 CPU 上取最后一行，避免 D2H 按“最后一行”拷贝时 stride/layout 与 GPU 不一致
+ const size_t full_bytes = seq_len * hs * elem_size;
+ tensor_t hidden_cpu_full = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(hidden_cpu_full->data(), hidden->data(), full_bytes, dev, dev_id);
+ tensor_t hidden_last = hidden_cpu_full->slice(0, seq_len - 1, seq_len);
+ if (std::getenv("LLAISYS_DEBUG_INFER")) { std::fprintf(stderr, "[DBG] Infer after d2h\n"); std::fflush(stderr); }
+ tensor_t normed_cpu = llaisys::Tensor::create({1, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_cpu, hidden_last, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({1, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_cpu, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, sampled_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ }
+ if (std::getenv("LLAISYS_DEBUG_INFER")) {
+ const std::byte *logit_ptr = last_logit_1d->data();
+ auto logit_at = [&](size_t i) -> float {
+ if (i >= voc) return 0.f;
+ const std::byte *p = logit_ptr + i * elem_size;
+ switch (dtype) {
+ case LLAISYS_DTYPE_F32: return *reinterpret_cast(p);
+ case LLAISYS_DTYPE_F16: return llaisys::utils::cast(*reinterpret_cast(p));
+ case LLAISYS_DTYPE_BF16: return llaisys::utils::cast(*reinterpret_cast(p));
+ default: return 0.f;
+ }
+ };
+ std::fprintf(stderr, "[DBG] Infer return(out_cpu) next=%lld logit[0]=%.4f logit[15]=%.4f logit[%lld]=%.4f\n",
+ (long long)next_token, logit_at(0), logit_at(15), (long long)next_token, logit_at(static_cast(next_token)));
+ std::fflush(stderr);
+ }
+ return next_token;
+ }
+
+ // 默认路径：在当前设备执行输出层的 RMSNorm 与 Linear 投影
+ llaisys::ops::rms_norm(normed, hidden, get_t(model->weights.out_norm_w), meta->epsilon);
+
+ tensor_t logits_t = llaisys::Tensor::create({seq_len, voc}, dtype, dev, dev_id);
+ llaisys::ops::linear(logits_t, normed, get_t(model->weights.out_embed), nullptr);
+
+ tensor_t last_logit = logits_t->slice(0, seq_len - 1, seq_len);
+ tensor_t last_logit_1d = last_logit->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t logits_for_sample = last_logit_1d;
+ if (dev != LLAISYS_DEVICE_CPU) {
+ llaisys::core::context().runtime().api()->device_synchronize();
+ tensor_t logits_cpu = llaisys::Tensor::create({voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ const size_t logits_row_bytes = voc * elem_size;
+ const std::byte *src_row = reinterpret_cast(logits_t->data()) + (seq_len - 1) * logits_row_bytes;
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ logits_cpu->data(), src_row, logits_row_bytes, LLAISYS_MEMCPY_D2H);
+ logits_for_sample = logits_cpu;
+ }
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, logits_for_sample, temperature, top_k, top_p, static_cast(seed));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, sampled_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ if (dev != LLAISYS_DEVICE_CPU) {
+ llaisys::core::context().runtime().api()->device_synchronize();
+ tensor_t logits_cpu = llaisys::Tensor::create({voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ const std::byte *src_row = reinterpret_cast(logits_t->data()) + (seq_len - 1) * voc * elem_size;
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ logits_cpu->data(), src_row, voc * elem_size, LLAISYS_MEMCPY_D2H);
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, logits_cpu->view(std::vector{voc}));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, dev, dev_id);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, dev, dev_id);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ }
+ }
+
+ return next_token;
+}
+
+int64_t llaisysQwen2ModelInferHybrid(struct LlaisysQwen2Model *model,
+ int64_t *token_ids,
+ size_t ntoken,
+ float temperature,
+ int top_k,
+ float top_p,
+ unsigned long long seed,
+ int gpu_up_to_layer) {
+ if (!model || ntoken == 0) return static_cast(-1);
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, di = meta->di;
+ const size_t voc = meta->voc;
+ const llaisysDataType_t dtype = meta->dtype;
+ const llaisysDeviceType_t dev = model->device_type;
+ const int dev_id = model->device_id;
+
+ const size_t slot0_len = model->cache_lens[0];
+ const bool is_prefill = (slot0_len == 0);
+ const bool is_suffix_prefill = (slot0_len > 0 && ntoken > 1);
+ const size_t seq_len = is_suffix_prefill ? ntoken : (is_prefill ? ntoken : 1);
+ const size_t cache_start = slot0_len;
+
+ if (!model->in_embed_cpu) return llaisysQwen2ModelInfer(model, token_ids, ntoken, temperature, top_k, top_p, seed);
+ // 仅 embedding+输出层在 CPU（未调用 CacheAllWeightsOnCPU）时不能走全量 CPU 分支，否则 forward_layer_cpu 会访问空的 k_caches_cpu
+ if (model->k_caches_cpu.empty()) return llaisysQwen2ModelInfer(model, token_ids, ntoken, temperature, top_k, top_p, seed);
+
+ const bool use_sampling = (temperature > 1e-6f && (top_k > 1 || (top_p > 0.f && top_p < 1.f)));
+ int64_t next_token = meta->end_token;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+
+ if (gpu_up_to_layer < 0) {
+ llaisys::core::context().setDevice(LLAISYS_DEVICE_CPU, 0);
+ tensor_t token_tensor = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ if (is_suffix_prefill) std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else if (is_prefill) std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else std::memcpy(token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t));
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++) pos_ids_host[s] = static_cast(cache_start + s);
+ std::memcpy(pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t));
+ llaisys::ops::embedding(hidden, token_tensor, model->in_embed_cpu);
+ for (size_t i = 0; i < nlayer; i++)
+ forward_layer_cpu(model, i, 0, hidden, normed, q_buf, k_buf, v_buf, q_rope, k_rope, attn_val, o_proj_out, res_buf, gate_buf, up_buf, mlp_buf, down_buf, seq_len, cache_start, pos_ids_t);
+ model->cache_lens[0] = cache_start + seq_len;
+ tensor_t normed_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_cpu, hidden, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({seq_len, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_cpu, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->slice(0, seq_len - 1, seq_len)->view(std::vector{voc});
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ std::memcpy(&next_token, sampled_idx->data(), sizeof(int64_t));
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ std::memcpy(&next_token, max_idx_t->data(), sizeof(int64_t));
+ }
+ return next_token;
+ }
+
+ int gpu_layers = static_cast(gpu_up_to_layer);
+ if (gpu_layers >= static_cast(nlayer)) gpu_layers = static_cast(nlayer) - 1;
+
+ llaisys::core::context().setDevice(dev, dev_id);
+ tensor_t token_tensor = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ if (is_suffix_prefill) llaisys::core::context().runtime().api()->memcpy_sync(token_tensor->data(), token_ids, ntoken * sizeof(int64_t), LLAISYS_MEMCPY_H2D);
+ else if (is_prefill) llaisys::core::context().runtime().api()->memcpy_sync(token_tensor->data(), token_ids, ntoken * sizeof(int64_t), LLAISYS_MEMCPY_H2D);
+ else llaisys::core::context().runtime().api()->memcpy_sync(token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t), LLAISYS_MEMCPY_H2D);
+
+ const bool use_tp_buf = (meta->tp_world_size > 1 && model->tp_gather_q != nullptr && dev == LLAISYS_DEVICE_NVIDIA);
+ const size_t nhdh_buf = use_tp_buf ? (nh * dh / static_cast(meta->tp_world_size)) : (nh * dh);
+ const size_t nkvhdh_buf = use_tp_buf ? (nkvh * dh / static_cast(meta->tp_world_size)) : (nkvh * dh);
+ const size_t di_buf = use_tp_buf ? (di / static_cast(meta->tp_world_size)) : di;
+
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nhdh_buf}, dtype, dev, dev_id);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, dev, dev_id);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, dev, dev_id);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++) pos_ids_host[s] = static_cast(cache_start + s);
+ llaisys::core::context().runtime().api()->memcpy_sync(pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t), LLAISYS_MEMCPY_H2D);
+
+ llaisys::ops::embedding(hidden, token_tensor, get_t(model->weights.in_embed));
+ for (int i = 0; i <= gpu_layers; i++)
+ forward_layer(model, i, 0, hidden, normed, q_buf, k_buf, v_buf, q_rope, k_rope, attn_val, o_proj_out, res_buf, gate_buf, up_buf, mlp_buf, down_buf, seq_len, cache_start, pos_ids_t);
+
+ llaisys::core::context().runtime().api()->device_synchronize();
+ tensor_t hidden_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(hidden_cpu->data(), hidden->data(), seq_len * hs * elem_size, dev, dev_id);
+
+ llaisys::core::context().setDevice(LLAISYS_DEVICE_CPU, 0);
+ tensor_t normed_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_buf_cpu = llaisys::Tensor::create({seq_len, nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_buf_cpu = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t v_buf_cpu = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_rope_cpu = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_rope_cpu = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t attn_val_cpu = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t o_proj_out_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t res_buf_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t gate_buf_cpu = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t up_buf_cpu = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t mlp_buf_cpu = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t down_buf_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t pos_ids_cpu = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ std::memcpy(pos_ids_cpu->data(), pos_ids_host.data(), seq_len * sizeof(int64_t));
+
+ for (size_t i = static_cast(gpu_layers) + 1; i < nlayer; i++)
+ forward_layer_cpu(model, i, 0, hidden_cpu, normed_cpu, q_buf_cpu, k_buf_cpu, v_buf_cpu, q_rope_cpu, k_rope_cpu, attn_val_cpu, o_proj_out_cpu, res_buf_cpu, gate_buf_cpu, up_buf_cpu, mlp_buf_cpu, down_buf_cpu, seq_len, cache_start, pos_ids_cpu);
+
+ model->cache_lens[0] = cache_start + seq_len;
+ tensor_t normed_out = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_out, hidden_cpu, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({seq_len, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_out, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->slice(0, seq_len - 1, seq_len)->view(std::vector{voc});
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ std::memcpy(&next_token, sampled_idx->data(), sizeof(int64_t));
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ std::memcpy(&next_token, max_idx_t->data(), sizeof(int64_t));
+ }
+ return next_token;
+}
+
+int64_t llaisysQwen2ModelInferWithSlot(struct LlaisysQwen2Model *model,
+ size_t slot_id,
+ int64_t *token_ids,
+ size_t ntoken,
+ float temperature,
+ int top_k,
+ float top_p,
+ unsigned long long seed) {
+ if (!model || ntoken == 0) return static_cast(-1);
+ if (slot_id >= model->cache_lens.size()) return static_cast(-1);
+
+ const LlaisysQwen2Meta *meta = &model->meta;
+ const size_t nlayer = meta->nlayer;
+ const size_t hs = meta->hs, nh = meta->nh, nkvh = meta->nkvh, dh = meta->dh, di = meta->di;
+ const size_t voc = meta->voc;
+ const llaisysDataType_t dtype = meta->dtype;
+ const llaisysDeviceType_t dev = model->device_type;
+ const int dev_id = model->device_id;
+
+ size_t *p_cache_len = &model->cache_lens[slot_id];
+ const bool is_prefill = (*p_cache_len == 0);
+ const bool is_suffix_prefill = (*p_cache_len > 0 && ntoken > 1);
+ const size_t seq_len = is_suffix_prefill ? ntoken : (is_prefill ? ntoken : 1);
+ const size_t cache_start = *p_cache_len;
+
+ // GPU 且已全量缓存 CPU 权重：整次前向在 CPU 上执行。使用 k_caches_cpu 非空判断，避免仅 CacheOutputLayerOnCPU 时误入导致越界。
+ if (dev != LLAISYS_DEVICE_CPU && !model->k_caches_cpu.empty()) {
+ const bool use_sampling = (temperature > 1e-6f && (top_k > 1 || (top_p > 0.f && top_p < 1.f)));
+ int64_t next_token = meta->end_token;
+ llaisys::core::context().setDevice(LLAISYS_DEVICE_CPU, 0);
+
+ tensor_t token_tensor = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ if (is_suffix_prefill)
+ std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else if (is_prefill)
+ std::memcpy(token_tensor->data(), token_ids, ntoken * sizeof(int64_t));
+ else
+ std::memcpy(token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t));
+
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvh * dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++)
+ pos_ids_host[s] = static_cast(cache_start + s);
+ std::memcpy(pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t));
+
+ llaisys::ops::embedding(hidden, token_tensor, model->in_embed_cpu);
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer_cpu(model, i, slot_id, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+ *p_cache_len = cache_start + seq_len;
+
+ tensor_t normed_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_cpu, hidden, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({seq_len, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_cpu, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->slice(0, seq_len - 1, seq_len)->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ std::memcpy(&next_token, sampled_idx->data(), sizeof(int64_t));
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ std::memcpy(&next_token, max_idx_t->data(), sizeof(int64_t));
+ }
+ return next_token;
+ }
+
+ llaisys::core::context().setDevice(dev, dev_id);
+
+ tensor_t token_tensor = llaisys::Tensor::create(
+ {seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ if (is_suffix_prefill) {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids, ntoken * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ } else if (is_prefill) {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids, ntoken * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ } else {
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ token_tensor->data(), token_ids + ntoken - 1, sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+ }
+
+ const bool use_tp_buf = (meta->tp_world_size > 1 && model->tp_gather_q != nullptr && dev == LLAISYS_DEVICE_NVIDIA);
+ const size_t nhdh_buf = use_tp_buf ? (nh * dh / static_cast(meta->tp_world_size)) : (nh * dh);
+ const size_t nkvhdh_buf = use_tp_buf ? (nkvh * dh / static_cast(meta->tp_world_size)) : (nkvh * dh);
+ const size_t di_buf = use_tp_buf ? (di / static_cast(meta->tp_world_size)) : di;
+
+ tensor_t hidden = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t normed = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t q_buf = llaisys::Tensor::create({seq_len, nhdh_buf}, dtype, dev, dev_id);
+ tensor_t k_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t v_buf = llaisys::Tensor::create({seq_len, nkvhdh_buf}, dtype, dev, dev_id);
+ tensor_t q_rope = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t k_rope = llaisys::Tensor::create({seq_len, nkvh, dh}, dtype, dev, dev_id);
+ tensor_t attn_val = llaisys::Tensor::create({seq_len, nh, dh}, dtype, dev, dev_id);
+ tensor_t o_proj_out = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t res_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+ tensor_t gate_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t up_buf = llaisys::Tensor::create({seq_len, di_buf}, dtype, dev, dev_id);
+ tensor_t mlp_buf = llaisys::Tensor::create({seq_len, di}, dtype, dev, dev_id);
+ tensor_t down_buf = llaisys::Tensor::create({seq_len, hs}, dtype, dev, dev_id);
+
+ tensor_t pos_ids_t = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, dev, dev_id);
+ std::vector pos_ids_host(seq_len);
+ for (size_t s = 0; s < seq_len; s++)
+ pos_ids_host[s] = static_cast(cache_start + s);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ pos_ids_t->data(), pos_ids_host.data(), seq_len * sizeof(int64_t),
+ (dev == LLAISYS_DEVICE_CPU) ? LLAISYS_MEMCPY_H2H : LLAISYS_MEMCPY_H2D);
+
+ if (dev != LLAISYS_DEVICE_CPU && model->in_embed_cpu) {
+ tensor_t token_cpu = llaisys::Tensor::create({seq_len}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ if (is_suffix_prefill) std::memcpy(token_cpu->data(), token_ids, ntoken * sizeof(int64_t));
+ else if (is_prefill) std::memcpy(token_cpu->data(), token_ids, ntoken * sizeof(int64_t));
+ else std::memcpy(token_cpu->data(), token_ids + ntoken - 1, sizeof(int64_t));
+ tensor_t hidden_cpu = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::embedding(hidden_cpu, token_cpu, model->in_embed_cpu);
+ llaisys::core::context().setDevice(dev, dev_id);
+ copy_host_to_device(hidden->data(), hidden_cpu->data(), seq_len * hs * llaisys::utils::dsize(dtype), dev, dev_id);
+ llaisys::core::context().setDevice(dev, dev_id);
+ } else {
+ llaisys::ops::embedding(hidden, token_tensor, get_t(model->weights.in_embed));
+ }
+
+ const bool use_tp_slot = (meta->tp_world_size > 1 && model->tp_gather_q != nullptr && dev == LLAISYS_DEVICE_NVIDIA);
+ if (use_tp_slot) {
+#ifdef ENABLE_NCCL
+ void *stream = llaisys::core::context().runtime().stream();
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer_tp(model, i, slot_id, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t, stream);
+ }
+#else
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer(model, i, slot_id, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+#endif
+ } else {
+ for (size_t i = 0; i < nlayer; i++) {
+ forward_layer(model, i, slot_id, hidden, normed, q_buf, k_buf, v_buf,
+ q_rope, k_rope, attn_val, o_proj_out, res_buf,
+ gate_buf, up_buf, mlp_buf, down_buf,
+ seq_len, cache_start, pos_ids_t);
+ }
+ }
+
+ *p_cache_len = cache_start + seq_len;
+
+ const bool use_sampling = (temperature > 1e-6f && (top_k > 1 || (top_p > 0.f && top_p < 1.f)));
+ int64_t next_token = meta->end_token;
+ const size_t elem_size = llaisys::utils::dsize(dtype);
+
+ if (dev != LLAISYS_DEVICE_CPU && model->out_norm_w_cpu && model->out_embed_cpu) {
+ llaisys::core::context().setDevice(dev, dev_id);
+ llaisys::core::context().runtime().api()->device_synchronize();
+ // 整块拷贝 hidden 到 CPU，再在 CPU 上取最后一行
+ const size_t full_bytes = seq_len * hs * elem_size;
+ tensor_t hidden_cpu_full = llaisys::Tensor::create({seq_len, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ copy_device_to_host(hidden_cpu_full->data(), hidden->data(), full_bytes, dev, dev_id);
+ tensor_t hidden_last = hidden_cpu_full->slice(0, seq_len - 1, seq_len);
+ tensor_t normed_cpu = llaisys::Tensor::create({1, hs}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::rms_norm(normed_cpu, hidden_last, model->out_norm_w_cpu, meta->epsilon);
+ tensor_t logits_cpu = llaisys::Tensor::create({1, voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::linear(logits_cpu, normed_cpu, model->out_embed_cpu, nullptr);
+ tensor_t last_logit_1d = logits_cpu->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, last_logit_1d, temperature, top_k, top_p, static_cast(seed));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, sampled_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ }
+ if (std::getenv("LLAISYS_DEBUG_INFER")) {
+ const std::byte *logit_ptr = last_logit_1d->data();
+ auto logit_at = [&](size_t i) -> float {
+ if (i >= voc) return 0.f;
+ const std::byte *p = logit_ptr + i * elem_size;
+ switch (dtype) {
+ case LLAISYS_DTYPE_F32: return *reinterpret_cast(p);
+ case LLAISYS_DTYPE_F16: return llaisys::utils::cast(*reinterpret_cast(p));
+ case LLAISYS_DTYPE_BF16: return llaisys::utils::cast(*reinterpret_cast(p));
+ default: return 0.f;
+ }
+ };
+ std::fprintf(stderr, "[DBG] Infer return(out_cpu slot) next=%lld logit[0]=%.4f logit[15]=%.4f logit[%lld]=%.4f\n",
+ (long long)next_token, logit_at(0), logit_at(15), (long long)next_token, logit_at(static_cast(next_token)));
+ std::fflush(stderr);
+ }
+ return next_token;
+ }
+
+ // 默认路径：在当前设备执行输出层的 RMSNorm 与 Linear 投影
+ llaisys::ops::rms_norm(normed, hidden, get_t(model->weights.out_norm_w), meta->epsilon);
+
+ tensor_t logits_t = llaisys::Tensor::create({seq_len, voc}, dtype, dev, dev_id);
+ llaisys::ops::linear(logits_t, normed, get_t(model->weights.out_embed), nullptr);
+
+ tensor_t last_logit = logits_t->slice(0, seq_len - 1, seq_len);
+ tensor_t last_logit_1d = last_logit->view(std::vector{voc});
+
+ if (use_sampling) {
+ tensor_t logits_for_sample = last_logit_1d;
+ if (dev != LLAISYS_DEVICE_CPU) {
+ llaisys::core::context().runtime().api()->device_synchronize();
+ tensor_t logits_cpu = llaisys::Tensor::create({voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ const std::byte *src_row = reinterpret_cast(logits_t->data()) + (seq_len - 1) * voc * elem_size;
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ logits_cpu->data(), src_row, voc * elem_size, LLAISYS_MEMCPY_D2H);
+ logits_for_sample = logits_cpu;
+ }
+ tensor_t sampled_idx = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::sample(sampled_idx, logits_for_sample, temperature, top_k, top_p, static_cast(seed));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, sampled_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ if (dev != LLAISYS_DEVICE_CPU) {
+ llaisys::core::context().runtime().api()->device_synchronize();
+ tensor_t logits_cpu = llaisys::Tensor::create({voc}, dtype, LLAISYS_DEVICE_CPU, 0);
+ const std::byte *src_row = reinterpret_cast(logits_t->data()) + (seq_len - 1) * voc * elem_size;
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ logits_cpu->data(), src_row, voc * elem_size, LLAISYS_MEMCPY_D2H);
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, LLAISYS_DEVICE_CPU, 0);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, LLAISYS_DEVICE_CPU, 0);
+ llaisys::ops::argmax(max_idx_t, max_val_t, logits_cpu->view(std::vector{voc}));
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ } else {
+ tensor_t max_idx_t = llaisys::Tensor::create({1}, LLAISYS_DTYPE_I64, dev, dev_id);
+ tensor_t max_val_t = llaisys::Tensor::create({1}, dtype, dev, dev_id);
+ llaisys::ops::argmax(max_idx_t, max_val_t, last_logit_1d);
+ llaisys::core::context().runtime().api()->memcpy_sync(
+ &next_token, max_idx_t->data(), sizeof(int64_t), LLAISYS_MEMCPY_H2H);
+ }
+ }
+
+ return next_token;
+}
+
+void llaisysQwen2ModelBatchedDecode(struct LlaisysQwen2Model *model,
+ const size_t *slot_ids,
+ const int64_t *token_ids,
+ size_t n_batch,
+ int64_t *out_next_tokens,
+ float temperature,
+ int top_k,
+ float top_p,
+ unsigned long long seed) {
+ if (!model || !slot_ids || !token_ids || !out_next_tokens || n_batch == 0)
+ return;
+ const size_t max_batch = model->meta.max_batch_size;
+ if (n_batch > max_batch)
+ n_batch = max_batch;
+ for (size_t i = 0; i < n_batch; i++) {
+ int64_t one_token = token_ids[i];
+ out_next_tokens[i] = llaisysQwen2ModelInferWithSlot(
+ model, slot_ids[i], &one_token, 1,
+ temperature, top_k, top_p, seed);
+ }
+}
+
+} // extern "C"
diff --git a/src/llaisys/runtime.cc b/src/llaisys/runtime.cc
index 7b00ff1bb..703131ea0 100644
--- a/src/llaisys/runtime.cc
+++ b/src/llaisys/runtime.cc
@@ -3,11 +3,11 @@
#include "../device/runtime_api.hpp"
// Llaisys API for setting context runtime.
-__C void llaisysSetContextRuntime(llaisysDeviceType_t device_type, int device_id) {
+LLAISYS_EXTERN_C void llaisysSetContextRuntime(llaisysDeviceType_t device_type, int device_id) {
llaisys::core::context().setDevice(device_type, device_id);
}
// Llaisys API for getting the runtime APIs
-__C const LlaisysRuntimeAPI *llaisysGetRuntimeAPI(llaisysDeviceType_t device_type) {
+LLAISYS_EXTERN_C const LlaisysRuntimeAPI *llaisysGetRuntimeAPI(llaisysDeviceType_t device_type) {
return llaisys::device::getRuntimeAPI(device_type);
}
\ No newline at end of file
diff --git a/src/llaisys/tensor.cc b/src/llaisys/tensor.cc
index 5e6e50124..294fa787b 100644
--- a/src/llaisys/tensor.cc
+++ b/src/llaisys/tensor.cc
@@ -2,7 +2,7 @@
#include
-__C {
+LLAISYS_EXTERN_C {
llaisysTensor_t tensorCreate(
size_t * shape,
size_t ndim,
diff --git a/src/ops/add/cpu/add_cpu.cpp b/src/ops/add/cpu/add_cpu.cpp
index 47f6a3d49..e2584870b 100644
--- a/src/ops/add/cpu/add_cpu.cpp
+++ b/src/ops/add/cpu/add_cpu.cpp
@@ -1,33 +1,75 @@
-#include "add_cpu.hpp"
-
-#include "../../../utils.hpp"
-
-#include
-
-template
-void add_(T *c, const T *a, const T *b, size_t numel) {
- for (size_t i = 0; i < numel; i++) {
- if constexpr (std::is_same_v || std::is_same_v) {
- c[i] = llaisys::utils::cast(llaisys::utils::cast(a[i]) + llaisys::utils::cast(b[i]));
- } else {
- c[i] = a[i] + b[i];
- }
- }
-}
-
-namespace llaisys::ops::cpu {
-void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t numel) {
- switch (type) {
- case LLAISYS_DTYPE_F32:
- return add_(reinterpret_cast(c), reinterpret_cast(a), reinterpret_cast(b), numel);
- case LLAISYS_DTYPE_BF16:
- return add_(reinterpret_cast(c), reinterpret_cast(a),
- reinterpret_cast(b), numel);
- case LLAISYS_DTYPE_F16:
- return add_(reinterpret_cast(c), reinterpret_cast(a),
- reinterpret_cast(b), numel);
- default:
- EXCEPTION_UNSUPPORTED_DATATYPE(type);
- }
-}
-} // namespace llaisys::ops::cpu
+/**
+ * Add 算子的 CPU 具体实现：按 dtype 分支，逐元素 c[i] = a[i] + b[i]。
+ *
+ * F16/BF16 先转 float 再相加再转回，避免半精度运算的精度与溢出问题；
+ * F32 直接相加。其它 dtype 通过 EXCEPTION_UNSUPPORTED_DATATYPE 报错。
+ */
+ #include "add_cpu.hpp" // 【大白话】：引入我们刚才说的“任务派发单”
+
+ #include "../../../utils.hpp" // 【大白话】：引入车间里的通用工具箱，比如这里的 cast（类型转换工具）
+
+ #include // 【大白话】：引入 C++ 标准的数学库（虽然这里加法没直接用到，但算子文件一般都会备着）
+
+// 当这段模板函数被调用时，假设传入的类型 T 是 llaisys::bf16_t（一种 16 位浮点数）：
+// 编译器看到 if constexpr，检查条件。
+// 因为 T 是 bf16_t，std::is_same_v 返回 true。
+// 编译器只保留 if 里面的代码进行编译，把 a[i] 和 b[i] 强制转换成 float 类型相加，然后把结果再次强制转换回 bf16_t 类型，赋值给 c[i]。
+//
+// 假设传入的类型 T 是 float：
+// 编译器检查条件。
+// 两个 is_same_v 都返回 false。
+// 编译器直接把 if 块里的代码删掉，只编译 else 里面的代码：c[i] = a[i] + b[i];。
+// 这种写法保证了底层数学运算的代码既精简，又能针对不同的数据类型生成最高效的底层机器指令。
+ template
+ void add_(T *c, const T *a, const T *b, size_t numel) {
+ // 【大白话】：流水线开启！numel 就是包裹里总共有多少个数字。循环一次，处理一个数字。
+ for (size_t i = 0; i < numel; i++) {
+
+ // 与普通的 if（在程序运行时判断）不同，if constexpr 是在程序编译时执行的判断。
+ // 编译器在编译阶段会计算括号里的条件。如果条件为 true，编译器就把大括号 {} 里的代码编译进最终的机器码；
+ // 如果为 false，编译器会直接丢弃这块代码，就像它从来没写过一样。
+ if constexpr (std::is_same_v || std::is_same_v) { // 如果 T 是 bf16_t 或 fp16_t
+
+ // llaisys::utils::cast 表示我们要调用的是 llaisys 命名空间下、utils 子命名空间里的 cast 函数，而不是其他地方可能存在的同名 cast 函数。
+ // 这里的 cast 是项目代码库中自定义的一个类型转换函数。尖括号 明确指示该函数：无论输入参数 a[i] 的原始类型是什么，必须将其转换为 float（32位浮点数）类型并返回。
+ c[i] = llaisys::utils::cast(llaisys::utils::cast(a[i]) + llaisys::utils::cast(b[i]));
+
+ } else {
+ // 【大白话】：如果本来就是大尺寸（float），那就别折腾了，直接简单粗暴地相加即可。
+ c[i] = a[i] + b[i];
+ }
+ }
+ }
+
+ namespace llaisys::ops::cpu {
+
+ // 【流水线入口】：这是车间主任实际调用的地方。
+ // 注意看参数：进来的 a, b, c 全是 std::byte *。
+ // 在电脑底层，其实根本没有所谓的“数字”，全是一堆毫无意义的二进制乱码（字节 byte）。
+ void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t numel) {
+
+ // 【大白话】：调度员根据标签（type）来看看这批乱码到底该按什么规格处理。
+ switch (type) {
+ case LLAISYS_DTYPE_F32:
+ // 【魔法核心 3：强制透视镜 reinterpret_cast】
+ // reinterpret_cast(c) 的意思就是：给工人戴上一副“浮点数透视镜”。
+ // 工人戴上后，看那些原本的字节乱码，就会自动把每 4 个字节当成一个 32位小数(float) 来理解！
+ // 然后调用我们上面写的 add_ 通用模具开始干活。
+ return add_(reinterpret_cast(c), reinterpret_cast(a), reinterpret_cast(b), numel);
+
+ case LLAISYS_DTYPE_BF16:
+ // 【大白话】：如果是 BF16，就戴上 BF16 的透视镜（每 2 个字节看成一个数字）。
+ return add_(reinterpret_cast(c), reinterpret_cast(a),
+ reinterpret_cast(b), numel);
+
+ case LLAISYS_DTYPE_F16:
+ // 【大白话】：如果是 F16，就戴上 F16 的透视镜。
+ return add_(reinterpret_cast(c), reinterpret_cast(a),
+ reinterpret_cast(b), numel);
+
+ default:
+ // 【大白话】：如果送来了不认识的材料（比如整数格式），直接报警停工。
+ EXCEPTION_UNSUPPORTED_DATATYPE(type);
+ }
+ }
+ } // namespace llaisys::ops::cpu
\ No newline at end of file
diff --git a/src/ops/add/cpu/add_cpu.hpp b/src/ops/add/cpu/add_cpu.hpp
index 34d809a11..29b464b85 100644
--- a/src/ops/add/cpu/add_cpu.hpp
+++ b/src/ops/add/cpu/add_cpu.hpp
@@ -1,8 +1,16 @@
-#pragma once
-#include "llaisys.h"
-
-#include
-
-namespace llaisys::ops::cpu {
-void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t size);
-}
\ No newline at end of file
+/**
+ * Add 算子的 CPU 实现声明。
+ *
+ * 接口使用裸指针 + dtype + 元素个数，由 op.cpp 在通过张量合法性检查后调用；
+ * 不关心张量形状布局，假定内存连续、按 numel 逐元素计算即可。
+ */
+#pragma once
+#include "llaisys.h"
+
+#include
+
+namespace llaisys::ops::cpu {
+
+/// 逐元素加法 c[i]=a[i]+b[i]，c/a/b 为同一 dtype 的连续内存，size 为元素个数
+void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t size);
+}
diff --git a/src/ops/add/op.cpp b/src/ops/add/op.cpp
index a057330d7..40b36d4c3 100644
--- a/src/ops/add/op.cpp
+++ b/src/ops/add/op.cpp
@@ -1,23 +1,33 @@
+/**
+ * Add 算子实现：校验输入合法性，按设备分发到具体实现（当前仅 CPU）。
+ *
+ * 被 src/llaisys/ops.cc 的 llaisysAdd 调用，供 Python 与 qwen2 等 C++ 代码使用。
+ */
#include "op.hpp"
#include "../../core/llaisys_core.hpp"
#include "../../utils.hpp"
#include "cpu/add_cpu.hpp"
+#ifdef ENABLE_NVIDIA_API
+#include "llaisys/ops_nvidia.h"
+#endif
namespace llaisys::ops {
+
void add(tensor_t c, tensor_t a, tensor_t b) {
+ // ---------- 合法性检查：同设备、同形状、同 dtype、三者皆连续 ----------
CHECK_SAME_DEVICE(c, a, b);
- // Only support contiguous inputs with same shape for now.
CHECK_SAME_SHAPE(c->shape(), a->shape(), b->shape());
CHECK_SAME_DTYPE(c->dtype(), a->dtype(), b->dtype());
ASSERT(c->isContiguous() && a->isContiguous() && b->isContiguous(), "Add: all tensors must be contiguous.");
- // always support cpu calculation
+ // CPU 分支：直接调 CPU 实现，无需切换 Context
if (c->deviceType() == LLAISYS_DEVICE_CPU) {
return cpu::add(c->data(), a->data(), b->data(), c->dtype(), c->numel());
}
+ // 非 CPU 时先切到当前张量所在设备，再按设备类型分发
llaisys::core::context().setDevice(c->deviceType(), c->deviceId());
switch (c->deviceType()) {
@@ -25,7 +35,7 @@ void add(tensor_t c, tensor_t a, tensor_t b) {
return cpu::add(c->data(), a->data(), b->data(), c->dtype(), c->numel());
#ifdef ENABLE_NVIDIA_API
case LLAISYS_DEVICE_NVIDIA:
- TO_BE_IMPLEMENTED();
+ nvidia::add(c->data(), a->data(), b->data(), c->dtype(), c->numel());
return;
#endif
default:
diff --git a/src/ops/add/op.hpp b/src/ops/add/op.hpp
index 62ef1ac87..7d4cf6e5d 100644
--- a/src/ops/add/op.hpp
+++ b/src/ops/add/op.hpp
@@ -1,7 +1,15 @@
-#pragma once
-
-#include "../../tensor/tensor.hpp"
-
-namespace llaisys::ops {
-void add(tensor_t c, tensor_t a, tensor_t b);
-}
+/**
+ * Add 算子对外接口声明。
+ *
+ * 语义：逐元素加法 c = a + b。c、a、b 为同形状、同 dtype、同设备的张量，
+ * 调用方需保证 c 已分配好内存；本算子只负责写入 c，不负责分配。
+ */
+#pragma once
+
+#include "../../tensor/tensor.hpp"
+
+namespace llaisys::ops {
+
+/// 逐元素加法：c[i] = a[i] + b[i]，c / a / b 需同 shape、同 dtype、同 device，且连续
+void add(tensor_t c, tensor_t a, tensor_t b);
+}
diff --git a/src/ops/argmax/op.cpp b/src/ops/argmax/op.cpp
index 6dc37d426..19e4b93c8 100644
--- a/src/ops/argmax/op.cpp
+++ b/src/ops/argmax/op.cpp
@@ -1,7 +1,85 @@
#include "op.hpp"
+#include "../../core/llaisys_core.hpp"
+#include "../../utils.hpp"
+
+#include
+#ifdef ENABLE_NVIDIA_API
+#include "llaisys/ops_nvidia.h"
+#endif
+
+namespace {
+
+template
+void argmax_impl(int64_t *max_idx, T *max_val, const T *vals, size_t numel) {
+ if (numel == 0) return;
+ size_t idx = 0;
+ T best = vals[0];
+ if constexpr (std::is_same_v || std::is_same_v) {
+ float best_f = llaisys::utils::cast(vals[0]);
+ for (size_t i = 1; i < numel; i++) {
+ float v = llaisys::utils::cast(vals[i]);
+ if (v > best_f) {
+ best_f = v;
+ best = vals[i];
+ idx = i;
+ }
+ }
+ } else {
+ for (size_t i = 1; i < numel; i++) {
+ if (vals[i] > best) {
+ best = vals[i];
+ idx = i;
+ }
+ }
+ }
+ *max_idx = static_cast(idx);
+ *max_val = best;
+}
+
+void argmax_cpu(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t vals_type, size_t numel) {
+ int64_t *out_idx = reinterpret_cast(max_idx);
+ switch (vals_type) {
+ case LLAISYS_DTYPE_F32:
+ argmax_impl(out_idx, reinterpret_cast(max_val), reinterpret_cast(vals), numel);
+ return;
+ case LLAISYS_DTYPE_F16:
+ argmax_impl(out_idx, reinterpret_cast(max_val), reinterpret_cast(vals), numel);
+ return;
+ case LLAISYS_DTYPE_BF16:
+ argmax_impl(out_idx, reinterpret_cast(max_val), reinterpret_cast(vals), numel);
+ return;
+ default:
+ EXCEPTION_UNSUPPORTED_DATATYPE(vals_type);
+ }
+}
+
+} // namespace
+
namespace llaisys::ops {
void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals) {
- TO_BE_IMPLEMENTED();
+ CHECK_SAME_DEVICE(max_idx, max_val, vals);
+ ASSERT(max_idx->dtype() == LLAISYS_DTYPE_I64, "argmax: max_idx must be int64");
+ ASSERT(max_val->dtype() == vals->dtype(), "argmax: max_val dtype must match vals");
+ ASSERT(vals->isContiguous() && max_idx->isContiguous() && max_val->isContiguous(), "argmax: all tensors must be contiguous");
+ ASSERT(vals->ndim() == 1, "argmax: vals must be 1D");
+ ASSERT(max_idx->numel() == 1 && max_val->numel() == 1, "argmax: max_idx and max_val must have one element");
+
+ if (vals->deviceType() == LLAISYS_DEVICE_CPU) {
+ return argmax_cpu(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
+ }
+
+ llaisys::core::context().setDevice(vals->deviceType(), vals->deviceId());
+ switch (vals->deviceType()) {
+ case LLAISYS_DEVICE_CPU:
+ return argmax_cpu(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
+#ifdef ENABLE_NVIDIA_API
+ case LLAISYS_DEVICE_NVIDIA:
+ nvidia::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
+ return;
+#endif
+ default:
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ }
}
} // namespace llaisys::ops
diff --git a/src/ops/argmax/op.hpp b/src/ops/argmax/op.hpp
index 433fdacdb..8bbb3267a 100644
--- a/src/ops/argmax/op.hpp
+++ b/src/ops/argmax/op.hpp
@@ -1,7 +1,7 @@
-#pragma once
-
-#include "../../tensor/tensor.hpp"
-
-namespace llaisys::ops {
-void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
-}
+#pragma once
+
+#include "../../tensor/tensor.hpp"
+
+namespace llaisys::ops {
+void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals);
+}
diff --git a/src/ops/embedding/op.cpp b/src/ops/embedding/op.cpp
index 84b9a5d06..969a99c1b 100644
--- a/src/ops/embedding/op.cpp
+++ b/src/ops/embedding/op.cpp
@@ -1,7 +1,65 @@
#include "op.hpp"
+#include "../../core/llaisys_core.hpp"
+#include "../../utils.hpp"
+
+#include
+#ifdef ENABLE_NVIDIA_API
+#include "llaisys/ops_nvidia.h"
+#endif
+
+namespace {
+
+void embedding_cpu(std::byte *out, const std::byte *weight, const int64_t *index, size_t num_index, size_t embed_dim, size_t vocab_size, size_t elem_size) {
+ size_t row_bytes = embed_dim * elem_size;
+ for (size_t i = 0; i < num_index; i++) {
+ int64_t row_idx = index[i];
+ ASSERT(row_idx >= 0 && static_cast(row_idx) < vocab_size, "embedding: index out of range");
+ const std::byte *src = weight + static_cast(row_idx) * row_bytes;
+ std::memcpy(out + i * row_bytes, src, row_bytes);
+ }
+}
+
+} // namespace
+
namespace llaisys::ops {
void embedding(tensor_t out, tensor_t index, tensor_t weight) {
- TO_BE_IMPLEMENTED();
+ CHECK_SAME_DEVICE(out, index, weight);
+ ASSERT(index->dtype() == LLAISYS_DTYPE_I64, "embedding: index must be int64");
+ ASSERT(out->dtype() == weight->dtype(), "embedding: out dtype must match weight");
+ ASSERT(out->isContiguous() && index->isContiguous() && weight->isContiguous(), "embedding: all tensors must be contiguous");
+ ASSERT(weight->ndim() == 2, "embedding: weight must be 2D");
+ ASSERT(index->ndim() == 1, "embedding: index must be 1D");
+ ASSERT(out->ndim() == 2, "embedding: out must be 2D");
+ size_t num_index = index->numel();
+ size_t vocab_size = weight->shape()[0];
+ size_t embed_dim = weight->shape()[1];
+ ASSERT(out->shape()[0] == num_index && out->shape()[1] == embed_dim, "embedding: out shape must be (index_len, embed_dim)");
+
+ switch (out->dtype()) {
+ case LLAISYS_DTYPE_F32:
+ case LLAISYS_DTYPE_F16:
+ case LLAISYS_DTYPE_BF16:
+ break;
+ default:
+ EXCEPTION_UNSUPPORTED_DATATYPE(out->dtype());
+ }
+
+ if (out->deviceType() == LLAISYS_DEVICE_CPU) {
+ return embedding_cpu(out->data(), weight->data(), reinterpret_cast(index->data()), num_index, embed_dim, vocab_size, out->elementSize());
+ }
+
+ llaisys::core::context().setDevice(out->deviceType(), out->deviceId());
+ switch (out->deviceType()) {
+ case LLAISYS_DEVICE_CPU:
+ return embedding_cpu(out->data(), weight->data(), reinterpret_cast(index->data()), num_index, embed_dim, vocab_size, out->elementSize());
+#ifdef ENABLE_NVIDIA_API
+ case LLAISYS_DEVICE_NVIDIA:
+ nvidia::embedding(out->data(), weight->data(), reinterpret_cast(index->data()), num_index, embed_dim, vocab_size, out->elementSize());
+ return;
+#endif
+ default:
+ EXCEPTION_UNSUPPORTED_DEVICE;
+ }
}
} // namespace llaisys::ops
diff --git a/src/ops/embedding/op.hpp b/src/ops/embedding/op.hpp
index 37216c0cf..f1546e259 100644
--- a/src/ops/embedding/op.hpp
+++ b/src/ops/embedding/op.hpp
@@ -1,7 +1,7 @@
-#pragma once
-
-#include "../../tensor/tensor.hpp"
-
-namespace llaisys::ops {
-void embedding(tensor_t out, tensor_t index, tensor_t weight);
-}
+#pragma once
+
+#include "../../tensor/tensor.hpp"
+
+namespace llaisys::ops {
+void embedding(tensor_t out, tensor_t index, tensor_t weight);
+}
diff --git a/src/ops/linear/op.cpp b/src/ops/linear/op.cpp
index 97d1f8655..dc980b024 100644
--- a/src/ops/linear/op.cpp
+++ b/src/ops/linear/op.cpp
@@ -1,7 +1,201 @@
+// 系统 SIMD 头文件须在项目头文件之前包含，避免项目宏 __C 与系统头文件中的 __C 冲突
+#ifdef __AVX2__
+#include
+#endif
+#include
+
#include "op.hpp"
+#include "../../core/llaisys_core.hpp"
+#include "../../utils.hpp"
+#ifdef ENABLE_NVIDIA_API
+#include "llaisys/ops_nvidia.h"
+#endif
+
+#ifdef _OPENMP
+#include
+#endif
+
+namespace {
+
+// 线性层（矩阵乘法）核心实现，计算公式为 Y = X W^T + b。
+// 参数说明：out 形状为 (B, M)，in 形状为 (B, K)，weight 形状为 (M, K)，bias 形状为 (M,) 或者是空指针。
+// 注意 weight 的形状是 (M, K) 而不是 (K, M)，这表明权重在内存中是以转置后的形态连续存储的。
+// OpenMP：外层 B 维并行，多核同时计算多行输出。
+template
+void linear_impl(T *out, const T *in, const T *weight, const T *bias, size_t B, size_t M, size_t K) {
+ // 外层循环：遍历输入张量的批次大小或序列长度维度 B。
+ // MSVC OpenMP 要求 index 为 signed integral type，故用 ptrdiff_t。
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
+ for (ptrdiff_t i = 0; i < static_cast(B); i++) {
+ // 中层循环：遍历输出特征维度 M。
+ for (size_t j = 0; j < M; j++) {
+
+ // 编译期条件分支：if constexpr 是 C++17 特性，用于在编译阶段静态计算表达式。
+ // std::is_same_v 用于类型萃取，判断当前模板类型 T 是否为 16 位浮点数（bf16_t 或 fp16_t）。
+ // 这种写法确保了在程序运行时，这里没有任何 if/else 的条件跳转指令开销，极大优化了指令执行效率。
+ if constexpr (std::is_same_v || std::is_same_v) {
+
+ // 声明一个 32 位单精度浮点数 sum_f 作为累加器。
+ // 16 位浮点数的尾数位数较少，如果在长循环的矩阵点积中直接用 16 位类型进行累加，
+ // 会产生严重的舍入误差，甚至可能导致数值溢出。使用 float 作为高精度中间变量是深度学习推理的通用做法。
+ float sum_f = 0;
+
+ // 内层循环：执行向量点积计算。遍历特征维度 K。
+ for (size_t k = 0; k < K; k++)
+ // 内存访问模式解析：
+ // in[i * K + k] 是按行主序访问输入的第 i 行第 k 列。
+ // weight[j * K + k] 是按行主序访问权重的第 j 行第 k 列。
+ // 由于计算公式是 X * W^T，正常应该访问权重的第 k 行第 j 列，但因为我们的 weight 传入时就是 (M, K) 形状，
+ // 这使得我们在内层循环 k 递增时，in 和 weight 的内存地址都是线性且连续递增的。
+ // 这种连续的内存访问模式能够最大化 CPU L1/L2 缓存的命中率（Cache Prefetching），是性能优化的关键。
+ sum_f += llaisys::utils::cast(in[i * K + k]) * llaisys::utils::cast(weight[j * K + k]);
+
+ // 如果传入了偏置项指针，将其对应元素转换为 float 后加到累加器上。
+ if (bias) sum_f += llaisys::utils::cast(bias[j]);
+
+ // 将 32 位的高精度累加结果强制向下转换为原始的 16 位类型 T，并写入输出张量。
+ out[i * M + j] = llaisys::utils::cast