diff --git a/.gitignore b/.gitignore index e38cf5747..2937b851f 100644 --- a/.gitignore +++ b/.gitignore @@ -87,4 +87,8 @@ htmlcov/ # Windows Thumbs.db ehthumbs.db -desktop.ini \ No newline at end of file +desktop.ini + +# Model weights (large, downloaded separately) +DeepSeek-R1-Distill-Qwen-1.5B/ +*.safetensors diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..ba5820691 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,200 @@ +# CLAUDE.md + +Guidance for Claude Code when working in this repository. + +## Project + +LLAISYS ("Let's Learn AI SYStem") — an educational C++/Python AI systems project. C++ backend (`src/`) exposes a C API (`include/llaisys/*.h`), wrapped by Python ctypes (`python/llaisys/libllaisys/`) and a friendlier Python layer (`python/llaisys/`). Full assignment structure: `README.md` / `README_ZN.md`. + +Build: `xmake` (compile) → `xmake install` (copies `.so` into `python/llaisys/libllaisys/`) → `pip install -e ./python` (editable install picks up `.py` changes immediately, no reinstall needed). + +## Assignment #3 — Qwen2 inference: DONE + +Goal: implement full Qwen2 (prefill + incremental decode) inference in C++, driven by a thin Python wrapper. No PyTorch allowed in the Python compute path — only for loading/comparing weights and running the HF reference in `test/test_infer.py`. Target model: DeepSeek-R1-Distill-Qwen-1.5B (28 layers, hs=1536, nh=12, nkvh=2, dh=128, di=8960, theta=10000, eps=1e-6, tie_word_embeddings=false). Full task breakdown: `docs/QWEN2_INFERENCE_ZH.md`. + +`test/test_infer.py --model DeepSeek-R1-Distill-Qwen-1.5B --test` passes (`Test passed!`, greedy-decode token-for-token match against HF `transformers`). Two bugs fixed after the initial "implementation done" pass, both by the user with step-by-step guidance, not written for them: +- `python/llaisys/models/qwen2.py`: absolute import `from python.llaisys.libllaisys.llaisys_types import DataType` only worked when run as a script from repo root, not as an installed package — changed to relative import `from ..libllaisys.llaisys_types import DataType`. +- `python/llaisys/models/qwen2.py` `__init__`: `end_token` was hardcoded to `151646`, which is actually `bos_token_id` — the real `eos_token_id` (151643) was already sitting in the same `config` dict as the other meta fields, just needed `config["eos_token_id"]` instead of the literal. +- `python/llaisys/models/qwen2.py` `generate`: was returning only newly-generated tokens; the test expects the HF-`generate()` convention (prompt + generated tokens concatenated) — fixed by appending `next_token` to `generated_tokens` (the list already seeded with `list(inputs)`) each loop iteration instead of a separate now-dead `output_tokens` list, and returning `generated_tokens`. + +## Next: Assignment #4 — CUDA integration + +**Core correctness done.** Full task breakdown: `docs/CUDA_INTEGRATION_ZH.md`. This machine has a real GPU reachable from WSL2 (`nvidia-smi` shows an RTX 5070 Ti Laptop GPU, `nvcc` is CUDA 12.9) so this can be developed and tested locally, not just on camp-provided remote resources. `xmake build llaisys --nv-gpu=y` now succeeds end-to-end, and all 8 non-deprecated ops pass `python test/ops/.py --device nvidia` (`add`, `embedding`, `argmax`, `rope`, `linear`, `swiglu`, `rms_norm`, `self_attention`). Remaining work is the performance-optimization pass below, not correctness. + +### Per-op status (`src/ops/*/nvidia/*.cu`) — all done, wired, and passing `test/ops/*.py --device nvidia` + +- **`add`, `embedding`**: done (elementwise, no reduction needed). +- **`argmax`**: done, back on the original V1 hand-written two-round kernel. Round 1 launches with `grid_size = ceil(numel/block_size)` blocks, each block does a thread-local grid-stride scan + shared-memory tree reduction (writes one (val, idx) candidate per block into a `cudaMalloc`'d intermediate buffer); round 2 launches `<<<1, block_size>>>` over that buffer (`numel` becomes `grid_size`) to produce the final `max_idx`/`max_val`. `argmax_kernel` takes an extra `const int64_t *idx_map` param (`nullptr` in round 1, the round-1 index buffer in round 2) so round 2 can translate "position within the candidate buffer" back to the original vocab index. + - **Detour, since undone**: briefly rewritten around `cudnnReduceTensor` (`CUDNN_REDUCE_TENSOR_MAX` + `CUDNN_REDUCE_TENSOR_FLATTENED_INDICES`, one call gives both `max_val` and the flattened index) to explore a cuDNN-backed path. Worked for F32/F16 but errored `CUDNN_STATUS_NOT_SUPPORTED` on BF16 regardless of 32-bit vs 64-bit index width. Investigating that turned up that `cudnnReduceTensor`/`cudnnSetReduceTensorDescriptor`/`cudnnReduceTensorIndices_t` are all marked `CUDNN_DEPRECATED` in the cuDNN version this project actually builds against, and the non-deprecated Graph API's `ReductionNode` has no index output at all (`X`→`Y` only, see `graph_properties.h`'s `Reduction_attributes`) — so there's currently no non-deprecated cuDNN path that produces argmax's value+index in one call. Reverted back to V1 for all three dtypes rather than keep a deprecated-API dependency for a partial (F32/F16-only) win. +- **`rms_norm`**: thread-local sum-of-squares over the row → shared-memory tree reduction, accumulated/stored in `float` regardless of `T` for precision → `rms = sqrtf(shared[0]/d + eps)` read from `shared[0]` by every thread → normalize+scale+write loop using the same `row*d+i` indexing as the read phase. +- **`rearrange`**: deprecated, will not be implemented — `src/ops/rearrange/nvidia/rearrange_cuda.cu` stays a 0-line stub on purpose. Nothing in the NVIDIA path needs it: the KV-cache copy (the one place CPU code used it as a substitute) uses a flat `std::memcpy` instead, since both sides are provably contiguous. Don't propose implementing this op unless the user explicitly reopens it. +- **`swiglu`**: done, unchanged since Assignment #4's initial pass. +- **`linear`**: done, now cuBLAS-backed — see the "Performance optimization plan" entry below for the full story (superseded the earlier hand-written tiled-GEMM rewrite). +- **`rope`**: done, back on the original V1 hand-written kernel (grid = `(seqlen, nhead)`, one thread block per token/head pair, each thread handles a grid-stride chunk of the `d/2` rotation pairs). + - **Detour, since undone**: explored cuDNN Graph API's `RoPE` node (`CUDNN_BACKEND_OPERATION_ROPE_FWD_DESCRIPTOR`) — but per the official docs (https://docs.nvidia.com/deeplearning/cudnn/latest/operations/RoPE.html), that op only supports **f16/bf16** for the X/Y tensors, not f32 (the `FREQS` angle tensor is always f32 regardless). Would've needed an F32→V1-fallback split like `self_attention`'s d%8 branch, plus a separate kernel to precompute the `FREQS` angle tensor (`phi(i,j) = pos_ids[i] / theta^(2j/d)`) since cuDNN's RoPE node takes precomputed angles, not `theta`/`pos_ids` directly — cuDNN only calls `sincosf` on them internally. Set aside for now (paused, not abandoned outright — the plumbing/design notes are gone from the file since it was reverted to V1, but this summary + the git history covers what was learned). +- **`self_attention`**: **cuDNN acceleration removed (2026-08-06); NVIDIA now dispatches by shape to hand-written flash-attention kernels (prefill and decode, both DONE — see "Performance optimization plan" below), falling back to this V1 kernel for everything else** — one block per (query token `i`, head `h`) pair, causal + GQA-aware (grid-stride dot product + shared-memory tree reduction for `max_score`, in-place `exp`/sum-reduction for `sum_exp`, final weighted-sum-by-`dv` pass). Iluvatar still uses this V1 kernel unconditionally (no flash-attention port done for that platform). + - There used to be a second path here: a **cuDNN `cudnn_frontend` SDPA graph** for the `d`/`dv`-multiple-of-8 case. Decided against continuing down that road and removed it — see the "Performance optimization plan" entry below for the full reasoning (graph-rebuild-per-call overhead was the dominant decode-time cost, the fix for that turned out to need per-bucket padding + explicit masking that was never worked out safely, plus real cross-platform/version fragility: a cuDNN 9.24.0 SDPA runtime crash on this machine's sm_120 GPU, and Iluvatar's cuDNN 7.6.5 not having the Graph API at all). + - `llaisys::device::nvidia::Resource`/`llaisys::device::iluvatar::Resource` no longer carry a `cudnnHandle_t` (only `cublasHandle_t`, still used by `linear`), and `self_attention`'s signature (`op.cpp`, `self_attention_cuda.cuh`/`self_attention_iluvatar.cuh`) dropped the now-unused `llaisys::device::DeviceResource *resource` param. `xmake/nvidia.lua`/`xmake/iluvatar.lua` no longer link `cudnn`/`nvrtc` — nothing in the project uses cuDNN anymore. + - Re-verified after the removal: `xmake build llaisys --nv-gpu=y` (no cuDNN linked) builds clean, all 8 ops' `test/ops/*.py --device nvidia` pass, and `test/test_infer.py --device nvidia --test` still matches HF token-for-token — and end-to-end generation is noticeably faster than the cuDNN-based run (no more per-call graph rebuild). + +### Iluvatar CoreX (天数) support — DONE: build, all 8 ops, and full end-to-end inference all verified on the remote box + +Assignment #4 requires picking two of {Nvidia, Iluvatar, Metax, Moore Threads} (see `README_ZN.md:319`, `docs/CUDA_INTEGRATION_ZH.md:12`). Iluvatar (天数智芯) is the second platform, developed on a remote cloud machine reachable only by generating a diagnostic/build shell script here, having the user run it there, and pasting the output back — there is no direct SSH access from this agent. Everything below is from that remote box (Iluvatar BI-V150, driver/SDK `corex-4.4.0`), not this machine. + +**Key findings from remote probing (all empirically confirmed, not guessed):** +- `/usr/local/corex/bin/nvcc` is **not a real compiler** — `file` shows it's a 6-line bash script that just `cat`s a fake `nvcc --version`-style banner and exits 0 no matter what arguments (including `--list-gpu-arch`, `-v -c foo.cu -o foo.o`) are passed. It never produces real object files. Any diagnosis based on invoking this `nvcc` is worthless — this cost two rounds of wasted probing (a false "CUDA 10.2, bf16 might not exist" alarm) before catching it via `file`/`cat` on the binary itself. +- The **real** compiler is `/usr/local/corex/bin/clang++` (→ `clang-18`, LLVM-based) invoked with **`-x ivcore`** instead of the standard `-x cuda` (using `-x cuda` prints a warning that it's unsupported for "ivcore" targets). Confirmed via real `.o`/executable output (not just exit code, which is meaningless here) that `clang++ -x ivcore -std=c++17 ...` successfully compiles: a bare `__global__` kernel, `cuda_bf16.h`/`__nv_bfloat16`, `cuda_fp16.h`/`__half`, C++17 structured bindings (`auto [a,b] = ...`, which `cudnn_frontend` depends on), and a `cublasSgemmEx` call with `CUDA_R_16BF` compute type that both compiles **and links** against their cuBLAS-compatible lib. +- `cudnn_frontend.h` does **not exist** anywhere on this box, and `cudnn.h` (found at `/usr/local/corex-4.4.0/include/cudnn.h`, no separate `cudnn_version.h`) reports `CUDNN_MAJOR=7, MINOR=6, PATCHLEVEL=5` — cuDNN 7.6.5, which predates the Graph API's existence entirely. This was one of the findings that fed into the later decision to drop cuDNN acceleration from `self_attention` altogether (see the per-op status entry above) — Iluvatar could never have run the Graph API SDPA path at all, so `self_attention` now uses the same cuDNN-free V1 kernel on every platform, no per-platform branching needed. +- Basic CUDA-runtime-API-level source compatibility (thread/block/grid built-ins, `__shared__`/`extern __shared__`, `__syncthreads()`, `dim3`, standard math intrinsics) is real and documented by Iluvatar's own programming-interface guide, not just inferred. One concrete divergence worth remembering for any future warp-level optimization work: **`warpSize` is 64 on this hardware, not 32.** Nothing currently in this codebase uses warp-shuffle intrinsics or otherwise assumes 32, so this doesn't bite today, but any hand-written warp-level reduction added later (e.g. if `self_attention`'s abandoned warp-shuffle optimization stage ever gets revisited) must read `warpSize` rather than hardcode 32. + +**Design decisions made:** +- New device type **`ILUVATAR`** added to `llaisysDeviceType_t` (`include/llaisys.h`, value `2`) and `DeviceType` (`python/llaisys/libllaisys/llaisys_types.py`, `ILUVATAR = 2`, `COUNT` bumped to `3`) — not reusing `NVIDIA`, since the runtime/library surface genuinely differs (no Graph API, different driver, etc.) even though the language/kernel level is source-compatible. +- **Reversed from the original plan: source files are duplicated, not reused.** The original idea (previous version of this note) was to have `llaisys-ops-iluvatar`/`llaisys-device-iluvatar` point `add_files` at the *same* `nvidia/*.cu` paths, compiled by a different toolchain into a different static lib, with no new source directories. The user overrode this: real `src/ops/*/iluvatar/` (one per op, 9 total including the `rearrange` stub) and `src/device/iluvatar/` directories now exist, seeded as verbatim copies of the `nvidia/` sources. Reasoning: having only `cpu`/`nvidia` subfolders per op was judged not to generalize, and ops are expected to eventually need real Iluvatar-specific fixes ("许多算子还是有细微的差别" — several ops have subtle platform differences), which needs its own editable location. Content in the new `iluvatar/` files is currently an unmodified copy — no Iluvatar-specific fixes have been made yet, this is bootstrap state only. + - Consequence of duplicating rather than reusing: the copied files still had `namespace llaisys::ops::cuda`/`llaisys::device::nvidia` from the source they were copied from, which would have caused **duplicate-symbol link errors** if a build ever enabled both `--nv-gpu=y` and `--iluvatar-gpu=y` at once (two static libs each defining the same mangled symbol, e.g. `llaisys::ops::cuda::add`). Fixed by renaming every copied file's namespace to `llaisys::ops::iluvatar` / `llaisys::device::iluvatar` (`sed` over `src/ops/*/iluvatar/*.cu(h)` and `src/device/iluvatar/*.cu(h)` — opening/closing namespace lines only, no other code needed changes). `linear_iluvatar.cu` and `self_attention_iluvatar.cu` (the two ops that reach into the device resource handle) had their `#include "../../../device/nvidia/nvidia_resource.cuh"` and `static_cast` updated to point at `device/iluvatar/iluvatar_resource.cuh` / `llaisys::device::iluvatar::Resource` accordingly. + - Filenames were also renamed for consistency (matches `cpu`'s `add_cpu.cpp`-style convention): `src/ops/*/iluvatar/_cuda.cu(h)` → `_iluvatar.cu(h)`; `src/device/iluvatar/nvidia_resource.cu(h)` → `iluvatar_resource.cu(h)`; `nvidia_runtime_api.cu` → `iluvatar_runtime_api.cu`. All self-include lines (`#include "_iluvatar.cuh"` etc.) and the two cross-file device-resource includes were updated to match — verified no leftover references to the old `_cuda`/`nvidia_resource` names anywhere under `src/*/iluvatar/`. + - `xmake/iluvatar.lua`'s `add_files` uses glob patterns (`../src/ops/*/iluvatar/*.cu`, `../src/device/iluvatar/*.cu`), so none of the renames required touching it. + +**xmake mechanism research (confirmed by reading xmake v2.8.7's own Lua source under `/usr/share/xmake/`, not guessed):** +- The stock `toolchains/cuda/xmake.lua` already declares `set_toolset("cu", "nvcc", "clang")` — clang was always a supported fallback compiler for `.cu` files, not something to bolt on from scratch. +- `modules/core/tools/nvcc.lua`'s `compargv` (the function that actually builds the compile command line) is generic: `return self:program(), table.join("-c", flags, "-o", objectfile, sourcefile)` — no nvcc-specific assumptions baked into the invocation itself, so pointing the `cu` toolset at Iluvatar's `clang++` should drive it correctly through the existing machinery. +- `rules/cuda/gencodes/xmake.lua` (backs `add_cugencodes(...)`) already branches on `target:has_tool("cu", "nvcc")`: nvcc gets `-gencode arch=...,code=...`, anything else gets `--cuda-gpu-arch=sm_XX`. The "unknown architecture: sm_120" warning seen locally comes from this file's hardcoded `known_r_archs`/`known_v_archs` hashsets only going up to `sm_90` (Hopper) — unrelated to nvcc-vs-clang, just a stale architecture table in this xmake version. + +**`xmake/iluvatar.lua` — first real on-box build test passed (`add` op only):** +- `toolchain("iluvatar")`: `set_kind("standalone")`, `set_toolset("cu", "/usr/local/corex/bin/clang++")` (also sets `cc`/`cxx`/`ld`/`sh` to plain system `/usr/bin/gcc`/`/usr/bin/g++` — probably moot since both targets below are `.cu`-only, kept as a harmless fallback), `on_check` via `find_tool("clang++", {paths = "/usr/local/corex/bin"})` modeled on `toolchains/clang/xmake.lua`. **The `/usr/local/corex/bin` path and `/usr/local/corex-4.4.0/lib64` linkdir (below) are carried over from the earlier probing session and have not been re-verified against this specific toolchain config on the actual box** — first on-box test should confirm both still resolve. +- `llaisys-device-iluvatar` / `llaisys-ops-iluvatar` targets mirror `xmake/nvidia.lua`'s two targets with the deltas that were planned: explicit `set_toolchains("iluvatar")` (required — the box's PATH also has the fake `nvcc`, so xmake's auto toolchain detection can't be trusted); `add_cuflags("-fPIC", {force = true})` with no `-Xcompiler=` prefix (that's nvcc-only flag-forwarding syntax, irrelevant once `clang++` is invoked directly) plus a separate `add_cuflags("-x", "ivcore", {force = true})` line; `add_cugencodes("native")` dropped entirely for this first pass (untested whether Iluvatar's device-probe/arch-numbering even works through this path); explicit `add_linkdirs("/usr/local/corex-4.4.0/lib64")` (bare `clang++` doesn't know CUDA lib paths the way nvcc does); links trimmed to `cublas` only (`cudnn`/`nvrtc` were linked at the time this was first written, for `self_attention`'s now-removed cuDNN SDPA path — dropped since nothing in the project uses cuDNN anymore) — add back if link errors show a missing symbol actually needs it. +- `xmake.lua` wiring: `option("iluvatar-gpu")` declared (mirrors `nv-gpu`'s block exactly), `add_defines("ENABLE_ILUVATAR_API")` (**note the casing** — matches the project's `ENABLE_NVIDIA_API` convention; an earlier draft of this line had it as `ENABLE_iluvatar_API`, which would have silently broken every `#ifdef ENABLE_ILUVATAR_API` check written to the project's normal all-caps convention), `includes("xmake/iluvatar.lua")` gated behind `has_config("iluvatar-gpu")`. + +**Still open / not done:** +- **First real on-box build test passed** (2026-08-05): narrowed `llaisys-ops-iluvatar`'s `add_files` to just `src/ops/add/iluvatar/add_iluvatar.cu`, ran `xmake f -c --iluvatar-gpu=y` + `xmake build -v llaisys-ops-iluvatar` on the remote box — `[100%]: build ok`, real non-empty `.o` produced and archived into `libllaisys-ops-iluvatar.a` (verified via `file`, not just exit code — the fake-`nvcc` lesson still applies). Confirms both the `/usr/local/corex/bin` toolchain path and the `/usr/local/corex-4.4.0/lib64` linkdir resolve correctly on the actual box. + - One real bug found and fixed along the way: `set_languages("cxx17")` only propagates to the `cc`/`cxx` tools, not the custom `cu` tool (bare `clang++` invoked via `-x ivcore`) — so the `.cu` compile command was missing `-std=c++17` entirely, and C++17 features used by the copied `nvidia/` sources (`std::byte`, `if constexpr`, nested namespace defs like `namespace llaisys::ops::iluvatar {`) failed with real compiler errors (not warnings) until an explicit `add_cuflags("-std=c++17", {force = true})` was added to both targets in `xmake/iluvatar.lua`. This is a `cu`-toolset-specific gotcha that doesn't affect nvidia's `nvcc`-based toolset (nvcc apparently gets the standard flag some other way xmake handles internally) — worth remembering for any other custom (non-nvcc) `cu` toolchain. + - `xmake` itself also had to be installed on the remote box (wasn't present) via `curl -fsSL https://xmake.io/shget.text | bash` (note: pass no `--branch` flag — the install script's arg parsing chokes on `--branch dev` and treats `--branch` itself as the branch name), plus `export XMAKE_ROOT=y` / `~/.bashrc` entry since the box runs as root and xmake refuses root by default without that override. + - **Full-glob build also passed** (2026-08-05, same session): widened `llaisys-ops-iluvatar` back to all 9 ops and built `llaisys-device-iluvatar` too — both compiled clean, `libllaisys-device-iluvatar.a`/`libllaisys-ops-iluvatar.a` both real non-empty archives. `linear`/`self_attention` (the two ops that reach into the device resource's `cublasHandle_t`) compiled without any Iluvatar-specific changes needed. So **source-level compilation is now fully verified for all 9 ops** on the actual Iluvatar box; nothing left to fix at the `xmake`/toolchain layer for a first pass. (This predated `self_attention`'s cuDNN removal — the op has since dropped `cudnnHandle_t`/the SDPA path entirely, see the per-op status entry above; re-verified against that change on 2026-08-10, see below.) +- **Aggregate deps, dispatch layer, and test-script wiring: all DONE** (2026-08-05). `xmake.lua`'s `llaisys-device`/`llaisys-ops` targets now `add_deps("llaisys-device-iluvatar")`/`add_deps("llaisys-ops-iluvatar")` under `has_config("iluvatar-gpu")`. `src/device/runtime_api.{hpp,cpp}` and `device_resource.{hpp,cpp}` have `#ifdef ENABLE_ILUVATAR_API`-guarded `case LLAISYS_DEVICE_ILUVATAR:` branches mirroring the NVIDIA ones. All 9 ops' `op.cpp` got the matching `#ifdef ENABLE_ILUVATAR_API` case calling into `llaisys::ops::iluvatar::...` (verified every iluvatar `.cuh` signature matches its nvidia counterpart exactly before wiring, so this was a mechanical 1:1 mirror — done directly rather than left for step-by-step guidance, per user's explicit call). `test/test_utils.py`'s `torch_device`/`llaisys_device`/`device_name` and all 10 test scripts' `--device` choices now have an `"iluvatar"` branch/option. + - Verified no regression: rebuilt locally with `--nv-gpu=y` (this machine has no Iluvatar toolchain, so `ENABLE_ILUVATAR_API` is undefined here and those branches compile out as dead code) — `build ok`, then `test/ops/{add,self_attention,linear}.py --device nvidia` all still pass. + - **First full `xmake build llaisys --iluvatar-gpu=y` attempt (2026-08-05) found one more real bug**, this time at the *final shared-library link* step, not per-op compilation (all 9 ops + `llaisys-device-iluvatar` still compiled/archived cleanly, same as before): the link command for `libllaisys.so` failed with `cannot find -lcudadevrt` / `cannot find -lcudart_static`. Root cause traced by reading xmake's own Lua source (`/usr/share/xmake/rules/cuda/env/xmake.lua`'s `cuda.env` rule, auto-attached to any target containing `.cu` files via `languages/cuda/xmake.lua`'s `language("cuda")` declaring `add_rules("cuda")`): its `after_load` unconditionally does `target:add("syslinks", "cudadevrt")`, and separately adds `"cudart_static"` unless the target's `links`/`syslinks` already contain `"cudart"` or `"cudart_static"` — this runs regardless of the `cuda.rdc` value (that value only gates the `-rdc=true` cuflag, not these two syslinks). Confirmed via `find`/`ls` on the remote box that Iluvatar's `corex` SDK only ships `libcudart.so` (real, shared) — no `libcudart_static.a` and no `libcudadevrt` under any name. + - Fix (in `xmake/iluvatar.lua`): added `add_links("cudart")` to both `llaisys-device-iluvatar` and `llaisys-ops-iluvatar` — since this is declared at parse time, it's already present by the time `cuda.env`'s `after_load` runs its check, so `cudart_static` never gets auto-added, and `cudart` resolves to the real `libcudart.so`. For `cudadevrt` — genuinely absent on this platform and genuinely unneeded (both targets set `cuda.rdc = false`, i.e. no relocatable-device-code linking) — there's no config value that suppresses xmake's unconditional syslink add, so `scripts/iluvatar/verify_ops.sh` now creates an empty stub `libcudadevrt.a` (`ar rcs` with zero object files, confirmed locally this produces a valid empty archive) in `/usr/local/corex-4.4.0/lib64` before building, satisfying the linker without providing or needing any real symbols. + - **Re-verified (2026-08-05): `libllaisys.so` now links and installs cleanly** with the stub in place. `test/ops/*.py --device iluvatar` results: **7 of 8 ops pass** (`add`, `argmax`, `embedding`, `rms_norm`, `rope`, `self_attention`, `swiglu`), plus `test_runtime.py --device iluvatar` (correctly detects "1 iluvatar devices"). Only `linear` fails, and only on **bf16** (f32/f16 both pass). + - **`linear`'s bf16 failure root-caused**: `linear_iluvatar.cu`'s BF16 branch calls `cublasSgemmEx(..., CUDA_R_16BF, ...)`, copied verbatim from the NVIDIA implementation. Added a debug `printf` of the returned `cublasStatus_t` right after the call (first attempt used `assert()`, which turned out to be silently compiled out — xmake's release build always defines `-DNDEBUG`, so ``'s `assert()` is a no-op; had to switch to `printf`, which isn't compiled out). Confirmed output: `dtype bf16 15` — `15` is `CUBLAS_STATUS_NOT_SUPPORTED`. Root cause: Iluvatar's `corex` SDK ships `libcublas.so.10.2.3.254`, a CUDA-10.2-era cuBLAS compat build; NVIDIA's own `cublasSgemmEx` only gained `CUDA_R_16BF` support in CUDA 11+, so this old ABI genuinely doesn't support the bf16 data type in that call — it's not a bug in llaisys's code, it's a platform capability gap. Since the call fails, `out`'s device buffer is left with whatever was there before (uninitialized/stale), which is why the result is wrong rather than an obvious crash. + - **Fixed (2026-08-05): switched the bf16 branch from `cublasSgemmEx` to `cublasGemmEx`.** Researched via `scripts/iluvatar/research_bf16_gemm.sh` (dumps the actual `cublas_api.h`/`library_types.h` declarations on the box rather than guessing): confirmed `CUBLAS_VER_MAJOR/MINOR = 10/2` (genuinely CUDA-10.2-era cuBLAS) but the header also declares `cublasComputeType_t` and a `CUBLAS_COMPUTE_32F_FAST_16BF` enum value (CUDA-11+ concepts) — Iluvatar's `corex` header is a patched hybrid, not a stock 10.2 header, so a bf16-capable path was plausible even though `cublasSgemmEx` itself rejects `CUDA_R_16BF`. `cublasGemmEx` (a distinct, more general function from `cublasSgemmEx`) is declared with the *older*-style signature — `cudaDataType computeType`, not `cublasComputeType_t` (that newer enum is only used by an unrelated function, `cublasGemmGroupedBatchedEx`, elsewhere in the same header — confirmed no `#if`/`#else` version-gates either declaration, both are unconditionally visible). First attempt passed `CUBLAS_COMPUTE_32F` (wrong enum type — a C++ compile error, since unscoped enums don't implicitly convert to a different enum type) instead of the correct `CUDA_R_32F`; fixed as a one-line mechanical correction once pointed out. With that fix, `test/ops/linear.py --device iluvatar` now reports `dtype bf16 0` (`CUBLAS_STATUS_SUCCESS`) and **passes for all shapes/dtypes, including the large 512×4096 case** — `linear` is now fully correct on Iluvatar. + - **Net result: all 8 ops now pass `test/ops/*.py --device iluvatar`** (previously 7/8, `linear` was the last holdout). Debug `assert`/`printf` lines added during this investigation are still in `linear_iluvatar.cu` and should be cleaned up before considering this op done. +- **End-to-end inference verified (2026-08-10).** The original remote Iluvatar instance from the sessions above was destroyed at some point (cloud instance lifecycle, not something in this repo) — verification had to restart from a freshly-provisioned instance of the same course image (same corex SDK / Iluvatar-patched PyTorch build, just no code/xmake/model weights on it yet). Two new scripts cover this: `scripts/iluvatar/verify_infer.sh` (repo already cloned on the box, just needs a rebuild + test) and `scripts/iluvatar/bootstrap_and_verify.sh` (fully fresh instance — installs xmake, clones the repo over HTTPS from the Gitee remote, installs the platform-agnostic Python deps like `ml_dtypes`/`transformers` without touching the special torch build, then runs the same build+test sequence). Ran `bootstrap_and_verify.sh` on the new instance: full build (`xmake build llaisys --iluvatar-gpu=y`) clean, all 8 ops pass, and — since no local model weights existed on the fresh instance — `test/test_infer.py --device iluvatar --test` auto-downloaded `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B` from Hugging Face (`load_hf_model`'s `snapshot_download` fallback, reused for the llaisys side too since both share the same `model_path`) and passed, token-for-token matching HF. This confirms the cuDNN-removal change from the note above didn't regress Iluvatar. +- Important prerequisite learned the hard way: the remote box only gets code via `git`, and this repo has three remotes (`origin`=GitHub personal fork, `origincn`=Gitee, `upstream`=read-only course repo) — the Iluvatar box pulls from **`origincn` (Gitee)**, not `origin`. Any local commits meant to be tested remotely must be pushed to `origincn` specifically (`git push origincn main`), or the remote box's `git pull` silently gets stale code with no error. + +**Note on the `argmax`/`rope` cuDNN detours mentioned above**: this machine has (had) **three** cuDNN installs (apt-installed 8.9.2 under `/usr/include/x86_64-linux-gnu/` + `/usr/lib/x86_64-linux-gnu/`; a private 9.18.1 copy bundled under ollama's install dir; and the one the project's `.cu` files actually resolved `#include ` to, bundled with the CUDA 12.9 toolkit under `/usr/local/cuda-12.9/targets/x86_64-linux/` — confirmed via `nvcc -E`'s actual include path, not guessed). Checking API behavior/support against the wrong installed copy gives wrong answers — a real gotcha hit during that exploration, worth remembering if cuDNN ever gets reconsidered for anything in this project. As of 2026-08-06, nothing in the project links or includes cuDNN anymore (see `self_attention`'s per-op status entry above), so this is dormant history rather than active guidance. + +### Recurring bug class: `op.cpp` missing the `nvidia/*.cuh` include + +Every op's `op.cpp` NVIDIA wiring hit the same two mistakes at least once — worth recognizing on sight rather than re-diagnosing each time: +1. **Missing include entirely**: `op.cpp` calls `cuda::foo(...)` but never `#include`s the header declaring `llaisys::ops::cuda::foo` — compiles to `error: 'cuda' has not been declared`. Fix: add `#include "nvidia/foo_cuda.cuh"` (op.cpp lives one directory *above* `nvidia/`, so the include needs that prefix — see `add`/`swiglu`'s `op.cpp` for the pattern that was always correct). +2. **Wrong prefix inside the `.cu` file itself**: the same `nvidia/foo_cuda.cuh` prefix is wrong *inside* `foo_cuda.cu`, since that file already lives in the `nvidia/` directory — it needs the bare `#include "foo_cuda.cuh"`. This bit `self_attention_cuda.cu` once when the `nvidia/` prefix was copy-pasted into the wrong file during the `op.cpp` fix. + +This hit `linear`, `self_attention`, `rms_norm`, `argmax`, and `rope`'s `op.cpp` in turn — all now fixed. + +### Test-script bug fixed along the way (not project code) + +`test/ops/self_attention.py`'s `torch_self_attention` reference (line 18) built its causal-mask tensor with `torch.ones(L, S, dtype=torch.bool)` — no `device=`, so it defaulted to CPU while `attn_bias` (line 16) correctly used `device=query.device`. Under `--device cpu` this never surfaced; under `--device nvidia` it crashed with a device-mismatch `RuntimeError` in `masked_fill_`, unrelated to `self_attention_cuda.cu`'s correctness (already independently verified — see below). Fixed by adding `device=query.device` to the `torch.ones(...)` call. + +### Performance optimization plan (linear + self_attention → flash attention) + +- **`linear`: DONE — now cuBLAS-backed**, superseding the earlier hand-written tiled-GEMM rewrite (see `613f360 Switch linear's NVIDIA backend to cuBLAS, add per-device Resource plumbing`). `src/ops/linear/nvidia/linear_cuda.cu` dispatches to `cublasSgemm` (F32) / `cublasSgemmEx` (BF16/F16) via the per-device `cublasHandle_t` owned by `llaisys::device::nvidia::Resource` (`src/device/nvidia/nvidia_resource.{cu,cuh}`). Bias is applied by a small separate `add_bias_kernel` pass after the GEMM (cuBLAS's GEMM call has no bias operand). The original hand-written tiled-GEMM kernel (`TILE_SIZE=16` static shared-memory tiling, described in old notes below) is still in the file as a commented-out block for reference — not compiled, not called. + - The benchmark note that used to live here (tiled kernel ~28% slower than the naive kernel at `M=1` decode-step shapes) **described the hand-written kernel and no longer applies** — it was a trade-off in that specific implementation, not a property of `linear` itself, and nobody has re-benchmarked `M=1` (decode-step) shapes against the current cuBLAS path. If `linear` optimization work resumes, start with a fresh `M=1` benchmark against cuBLAS rather than assuming the old naive/tiled trade-off still holds — cuBLAS may already dispatch to a different algorithm for tall-skinny shapes. +- **`self_attention`**: **flash attention prefill + decode DONE, wired, and verified (supersedes the "mid-way, not wired in" state described in older notes below).** `src/ops/self_attention/nvidia/self_attention_cuda.cu` (V1) is untouched and stays the permanent fallback for non-128 head-dim / non-first-prefill shapes. `op.cpp`'s NVIDIA branch now really does dispatch on shape: `seqlen>1 && total_len==seqlen && d==dv==128` → `flash_attention_prefill_cuda.cu`; `seqlen==1 && d==dv==128` → `flash_attention_decode_cuda.cu` (one warp per head, scans the whole KV cache); everything else → V1. Two real bugs hit while wiring this up (an ODR weak-symbol collision between same-named `__global__` kernel templates in different `.cu` files, and a cross-chunk shared-memory data race from a missing `__syncthreads()`) are written up in full in `docs/FLASH_ATTENTION_DEBUG_LOG_ZH.md` — read that before touching either file, the lessons there (anonymous-namespace every new kernel file, and always test a shape that forces >1 tiling iteration) apply to any further kernel work in this area. + - Design rationale for the two-level tiling / online-softmax approach used by the prefill kernel is preserved in that file's own header comment (kept up to date as the source of truth). + - **Benchmarked (2026-08-08, `test/benchmark_infer.py --device nvidia`, same prompt/`max_steps=64`, this machine's RTX 5070 Ti Laptop)**: flash attention vs. V1-for-everything — prefill 520.77ms vs 1068.98ms (**~2.05x**), decode 16.49ms/token vs 19.99ms/token (**~1.21x**). The much smaller decode win is expected: the decode kernel only launches `nhead` (12 for this model) blocks of 1 warp (32 threads) each — 384 threads total against this GPU's **46 SMs** (`torch.cuda.get_device_properties(0).multi_processor_count`), so most SMs sit idle during decode's `self_attention` call regardless of how fast the per-warp math is. + - The previously-explored cuDNN path (`cudnn_frontend` SDPA graph for the `d`/`dv`-multiple-of-8 case) was tried and fully reverted (2026-08-06) before this rewrite started: the cuDNN graph was rebuilt from scratch on every call (validate → build_operation_graph → create_execution_plans → build_plans), confirmed as the dominant real decode-time cost (~20ms/call at the real decode shape, ~570ms/token across 28 layers); fixing that properly would have needed bucketing `total_len` plus padding K/V with an explicit mask, on top of real version/platform fragility (a cuDNN 9.24.0 SDPA runtime crash on this machine's sm_120 GPU, and no Graph API at all on Iluvatar's cuDNN 7.6.5). No cuDNN dependency remains anywhere in the project. + +### Decode split-KV (flash-decoding) — DONE (2026-08-10) + +Goal: fix the decode occupancy problem above by splitting the KV direction instead of the (nonexistent, at `seqlen=1`) query-row direction — i.e. "flash-decoding". `src/ops/self_attention/nvidia/flash_attention_decode_splitkv_cuda.cu`. `flash_attention_decode_cuda.cu` is untouched and stays the fallback for short `total_len` (splitting a short KV cache into many tiny pieces isn't worth the extra kernel launch + reduction overhead). + +**Note on how this got finished**: phase 1 (below) had already been fixed by the user by the time this pass started — the three bugs described in the previous version of this note (duplicated shared-mem decl, tile bound checked against `total_len` instead of `split_end`, missing write-back) were gone. Given the imminent submission deadline, the user explicitly asked to set aside the usual hands-off/step-by-step rule for this file specifically and have phase 2 + host orchestration + wiring written directly rather than guided — a one-time exception, not a change to the standing working-style rule for the rest of the project. + +**Design, as implemented:** +- Two-phase structure, since splits run in different blocks/SMs and can't `__syncthreads()` with each other: **phase 1** — one warp per `(head, split_idx)` pair (`dim3 grid(nhead, num_splits)`, `blockDim=WARP_SIZE`) computes a local online-softmax state over just its `[split_start, split_end)` slice of the KV cache and writes `(m, l, acc[4])` to an intermediate buffer instead of the final normalized output. **Phase 2** — `flash_attention_decode_splitkv_phase2_kernel`, one warp per head, serially loops over the (small, single-digit-to-tens) `num_splits` local results and rescales/accumulates by `exp(m_split - m_final)` (every lane redundantly computes the same `final_max` via its own scan rather than doing a cross-warp broadcast — cheap given how small `num_splits` is, and simpler), skipping any split with `m == -INFINITY` (the empty-split case below) to avoid `exp(-inf - -inf) = nan`. Modeled on `argmax_cuda.cu`'s existing round-1/round-2 `cudaMalloc`-intermediate-buffer pattern. +- Buffer layout: `partial_m`/`partial_l` : `[nhead, num_splits]`; `partial_acc` : `[nhead, num_splits, HEAD_DIM]` with the same `item*WARP_SIZE+lane` interleaving the existing decode kernel already uses for `output`. +- Empty-split handling (`split_start >= total_len`, possible whenever `num_splits`/`split_size` don't evenly divide `total_len`): phase 1 writes the neutral triple `(m=-INFINITY, l=0, acc=0)` and returns early *before* the tile loop; phase 2 skips these splits entirely (see above). +- `choose_num_splits(total_len, nhead)`: targets `nhead * num_splits ≈ 64` total blocks for SM occupancy, floored so a split is never smaller than `TILE_K=32` (`kMinSplitSize`), capped at `kMaxNumSplits=32`. +- `op.cpp` dispatch: `seqlen==1 && d==dv==128 && total_len > 256` → `flash_attention_decode_splitkv`; `seqlen==1 && d==dv==128` (i.e. `total_len <= 256`) → falls back to the plain `flash_attention_decode`. The `256` threshold matches `kMinSplitSize`-driven reasoning in the `.cu` file (below `total_len≈256`, `choose_num_splits` would only produce 1 useful split, so splitting isn't worth the extra kernel launch + phase-2 pass) and is empirically confirmed by the microbenchmark below. (One incidental fix needed while wiring this in: the `NVIDIA` switch-case had to be wrapped in `{ }` since the new `constexpr size_t kSplitKVThreshold` inside a `case` label was crossing into `default:`, which GCC rejects as "jump to case label".) + +**Verification (2026-08-10):** +- `xmake build llaisys --nv-gpu=y` clean, `test/ops/self_attention.py --device nvidia` passes, `test/test_infer.py --device nvidia --test` still matches HF token-for-token. +- Standalone nvcc harness (scratchpad, not committed — same pattern as the "Testing approach" section below) comparing `flash_attention_decode_splitkv`'s output directly against the already-verified `flash_attention_decode` bit-for-bit (within fp tolerance): F32/F16/BF16 × `total_len ∈ {1, 31, 32, 255, 256, 257, 1000, 4000}` (covers the `TILE_K=32` and threshold=256 boundaries) plus a non-GQA (`nhead==nkvhead`) case — all passed, max abs diff ≤ ~4e-6 (bf16) / ~1e-7 (f32/f16). +- Kernel-level microbenchmark (cudaEvent timing, bf16, real-model shape `nhead=12, nkvhead=2, d=dv=128`, 200 iters/point after 20 warmup): `flash_attention_decode` (V1) vs `flash_attention_decode_splitkv` latency — + + | total_len | decode | splitkv | speedup | + |---|---|---|---| + | 128 | 0.058 ms | 0.084 ms | 0.69x (splitkv slower — confirms the 256 threshold) | + | 256 | 0.110 ms | 0.092 ms | 1.20x | + | 512 | 0.213 ms | 0.096 ms | 2.23x | + | 1000 | 0.443 ms | 0.141 ms | 3.14x | + | 2000 | 0.920 ms | 0.235 ms | 3.91x | + | 4000 | 2.362 ms | 0.468 ms | 5.05x | + | 8000 | 5.706 ms | 0.829 ms | 6.88x | + | 16000 | 11.463 ms | 1.567 ms | 7.31x | + + Confirms the expected shape: V1's single warp scans `total_len` serially (linear growth), splitkv's block count scales with `total_len` so SM occupancy stays high (sub-linear growth) — and confirms `256` is a reasonable cutover point (below it, splitkv's extra kernel-launch/phase-2 overhead isn't paid back yet). +- End-to-end A/B on the real model (`test/benchmark_infer.py --device nvidia`, DeepSeek-R1-Distill-Qwen-1.5B, prompt engineered to avoid early EOS — a multi-step arithmetic word problem — `max_steps=400`, `total_len` grows 67→467 over the run so both the plain-decode and splitkv regimes are actually exercised): split-KV enabled 17.41 ms/token avg (57.45 tok/s, 7.43s total) vs disabled (threshold temporarily set to a huge value, rebuilt, re-measured, then reverted to 256) 22.42 ms/token avg (44.60 tok/s, 9.49s total) — **~1.29x end-to-end decode speedup**. Lower than the kernel-level numbers because this run blends the pre-256 and post-256 regimes; the kernel-level table above is the cleaner number for isolating the algorithm's own effect. + +Same hands-off/step-by-step guidance rule applies to any further hand-written-kernel work in this project going forward — see "Working style for this project" below. The exception made for this file was scoped to finishing this one feature under a submission deadline, not a standing change. + +### Testing approach + +The full project now builds (`xmake build llaisys --nv-gpu=y`) and `python test/ops/.py --device nvidia` is the primary end-to-end check — no more workaround needed for finished ops. The standalone-nvcc-harness pattern used earlier (compile one op's `.cu` file plus a small scratchpad `main()` directly via `nvcc`, comparing against a from-scratch CPU reference, bypassing the full build) is still useful for iterating on the optimization work above without a full rebuild each time — see `self_attention`'s and `rms_norm`'s scratchpad `main.cu` for the pattern (random-input cases spanning edge shapes like `d`/`total_len`/`dv` not a multiple of `block_size`, f32 + bf16, real-model shapes). + +### What's implemented — status as of end of the Infer/Python-glue pass + +All three pieces below are now implemented, compiled, and verified working end-to-end (manual `generate()` smoke test produced coherent text for a real prompt against the real DeepSeek-R1-Distill-Qwen-1.5B weights). Same hands-off/step-by-step-guidance rule still applies to any further changes in these files — don't rewrite user's logic for them, review/point-at-line-numbers instead. + +- `src/llaisys/models/qwen2.cc`: `LlaisysQwen2Model` struct + `Create`/`Destroy`/`Weights`/`Infer` all implemented by the user (embedding → per-layer loop [rms_norm, qkv proj (2D out, matches `linear`'s shape assert), view to 3D, rope on q/k only, kv-cache write via `slice(cur_len,total_len)->load(...)`, kv-cache read via `slice(0,total_len)`, self_attention, attn_o proj, residual, mlp rms_norm, gate/up proj, swiglu, down proj, residual] → final rms_norm → out_embed proj → slice last token → argmax → `cur_len = total_len` → return). Two bugs found via self-review and fixed by the user: (1) `q`/`k`/`v` tensors were created 3D (`{ntoken, nh, dh}`) and passed straight into `ops::linear`, which asserts a 2D `out` shape matching `weight`'s dout — fixed by creating them 2D (`{ntoken, nh*dh}`) and keeping the existing `->view(...)` to reshape to 3D before RoPE; (2) `argmax`'s `max_val` tensor was created with `LLAISYS_DTYPE_I64` instead of `meta.dtype` — `argmax` asserts `max_val`'s dtype matches the input values' dtype, only `max_idx` must be I64. +- `python/llaisys/libllaisys/model.py` (singular, not `models.py`): ctypes bindings for `LlaisysQwen2Meta`/`LlaisysQwen2Weights`/the four `llaisysQwen2Model*` functions, wired into `__init__.py` via `load_model(lib)`. Done. +- `python/llaisys/models/qwen2.py`: both TODOs done. + - `__init__`: builds `self.meta` (stored on `self`, needed by `generate()` for `end_token`), calls `Create`/`Weights`, then loads safetensors weights into the pre-allocated tensors via two lookup tables — `GLOBAL_MAP` (full key name → attribute on `weights.contents`, for `in_embed`/`out_embed`/`out_norm_w`) and `LAYER_MAP` (key with `"model.layers.{i}."` prefix stripped → attribute, indexed by the parsed-out layer number), then `LIB_LLAISYS.tensorLoad(handle, ptr)` per weight. Requires `import ml_dtypes` before `safetensors.safe_open(..., framework="numpy", ...)` — numpy has no native bf16 dtype, and `ml_dtypes` registers the extension dtype safetensors needs on import (installed via `pip install ml_dtypes`; without the import, even with the package installed, `get_tensor` raises `TypeError: data type 'bfloat16' not understood`). + - `generate`: prefill and decode share one loop — `current_input` is the full prompt on the first iteration only, then always `[next_token]` after; `generated_tokens` (seeded with `list(inputs)`) has `next_token` appended each iteration and is what gets returned (matches HF `generate()`'s convention of prompt+generated concatenated); breaks on `next_token == self.meta.end_token`. +- Reference (non-compiled, not imported by anything) implementations the user can diff against once stuck: `docs/qwen2_infer_reference.cc`, `docs/qwen2_python_bindings_reference.py`, `docs/qwen2_model_reference.py`. + +### Still open +- dtype/contiguity sanity-check TODO left as a comment in `qwen2.py`'s weight-loading loop was never actually implemented (low priority — real run already proves the bf16 raw-memcpy path works for this checkpoint). + +### Fixed (2026-08-05): segfault in `qwen2.cc`'s `Infer` on NVIDIA — full end-to-end inference was never actually re-verified after the cuBLAS/cuDNN backend swaps + +Discovered while building a performance-analysis benchmark script (`test/benchmark_infer.py`, see below) and trying to profile a real `generate()` call on NVIDIA — `test/test_infer.py --device nvidia` (and any other full-model call) segfaulted immediately on the very first `Infer` call. Bisected via `git stash` (confirmed present with the pristine, pre-session `qwen2.py`, so unrelated to this session's Iluvatar work) and a debug-mode (`xmake f --mode=debug`) `gdb` backtrace, which pointed at the exact line: `src/llaisys/models/qwen2.cc:251`, `return *reinterpret_cast(max_idx->data());`. `max_idx` is created via `Tensor::create({1}, LLAISYS_DTYPE_I64, device, device_id)` — on NVIDIA, `->data()` is a **device (GPU) pointer**, and dereferencing it directly from host code is an illegal memory access. Silently "worked" on CPU only because CPU device pointers and host pointers are the same address space there, masking the bug — this is why `--device cpu` (the only device end-to-end inference had actually been tested on since the cuBLAS/cuDNN switch) never caught it. + +Fix (mirrors the existing device-to-host readback pattern in `Tensor::debug()`, `src/tensor/tensor.cpp:149-164`): branch on `device == LLAISYS_DEVICE_CPU`; on GPU, copy the single int64 back to a host-side local via `llaisys::core::context().runtime().api()->memcpy_sync(&host_max_idx, max_idx->data(), sizeof(int64_t), LLAISYS_MEMCPY_D2H)` before returning it. (`core::` needs the `llaisys::` prefix here since — unlike `tensor.cpp` — `qwen2.cc`'s code isn't inside `namespace llaisys { ... }`, it only pulls in `Tensor`/`tensor_t` via per-name `using` declarations.) Verified: `test/test_infer.py --model DeepSeek-R1-Distill-Qwen-1.5B --device nvidia --max_steps 32 --test` now passes, token-for-token match against HF, in release mode. + +Side note for future performance work: that same run showed HF taking 5.59s vs llaisys's 27.28s for the same 32 tokens — llaisys is currently ~5x slower end-to-end. Not investigated yet, but a concrete starting point once full-model profiling resumes. + +### Performance-analysis tooling added (2026-08-05) + +- `test/benchmark_infer.py`: new script, separate from `test_infer.py` (which only does a wall-clock-timed correctness comparison against HF). Loads only the llaisys model, drives `generate()` with a `step_context` callback (see below) to time the prefill step separately from each decode step, and reports prefill latency, decode steps/avg-latency/tokens-per-sec, and total wall time. Optional `--nvtx` flag wraps each step in an `torch.cuda.nvtx` range (labeled `prefill_N`/`decode_N`) for use under `nsys`. +- `python/llaisys/models/qwen2.py`'s `generate()` grew an optional `step_context=None` parameter: if provided, it's called as `step_context(step, is_prefill)` and must return a context manager that wraps that step's `Infer` call (default: `contextlib.nullcontext()`, a complete no-op — existing callers like `test_infer.py` that don't pass it are unaffected). This is the hook `benchmark_infer.py` uses for both timing and NVTX ranges without `qwen2.py` itself needing to depend on `torch` or any profiling library (still holds to the "no PyTorch in the compute path" rule). +- Profiling-tool findings from this machine (NVIDIA, WSL2): `nsys` works out of the box for CUDA API-level tracing. (At the time this was written, `self_attention` still used cuDNN's SDPA graph, and profiling it under `nsys` reproducibly made the GQA test case's graph build fail — `self_attention` no longer depends on cuDNN, see the per-op status entry above, so that finding no longer applies.) `ncu` (Nsight Compute, for per-kernel deep-dive metrics) requires elevated GPU performance-counter permissions (`ERR_NVGPUCTRPERM`) — common on WSL2, needs an interactive `sudo` the agent can't supply. +- Not yet done: getting `nsys stats --report cuda_gpu_kern_sum` to actually return kernel-level timing data (one attempt reported "does not contain CUDA kernel data" — cuBLAS's kernels may launch via the driver API in a way that needs different trace flags); adding NVTX instrumentation inside the C++ ops themselves for per-op (not just per-step) resolution in a full-model trace; investigating whether Iluvatar's `corex` SDK ships any profiling tool of its own. + +### Op signatures already implemented (Assignment #2, reusable as-is) + +All under `src/ops//op.hpp`, namespace `llaisys::ops`: + +- `embedding(out[ntoken,hs], index[ntoken] i64, weight[voc,hs])` +- `rms_norm(out[ntoken,hs], in[ntoken,hs], weight[hs], eps)` +- `linear(out[ntoken,dout], in[ntoken,din], weight[dout,din], bias[dout]_or_nullptr)` — bias arg is `tensor_t`, pass `nullptr` for no bias (`attn_o_w` and all MLP linears have no bias; QKV projections do). +- `rope(out[seqlen,nhead,d], in[seqlen,nhead,d], pos_ids[seqlen] i64, theta)` +- `self_attention(attn_val[seqlen,nhead,dv], q[seqlen,nhead,d], k[total_len,nkvhead,d], v[total_len,nkvhead,dv], scale)` — causal, GQA-aware (`nhead` must be a multiple of `nkvhead`); internally `causal_offset = total_len - seqlen`. +- `swiglu(out[ntoken,di], gate[ntoken,di], up[ntoken,di])` +- `add(c, a, b)` — same shape, contiguous only +- `argmax(max_idx[1] i64, max_val[1], vals[N] 1D)` — `vals` must be 1D and contiguous, so slice logits down to just the last token's row before calling. + +**`ops::rearrange(out, in)` is deprecated — a permanent `TO_BE_IMPLEMENTED()` stub, not just unfinished** (an Assignment #2 leftover) — don't call it, and don't propose implementing it. For copying newly-computed K/V into the KV cache, both sides are provably fully-contiguous (`Tensor::isContiguous()` treats a size-1 leading dim as contiguous regardless of its stride there), so a flat `std::memcpy` of `numel()*elementSize()` bytes is a correct, simpler substitute — see the reference file. + +### KV cache design + +`LlaisysQwen2Model::k_cache`/`v_cache` are `std::vector`, one entry per layer (`k_cache[l]` = layer `l`'s cache — each layer has independent K/V projection weights, so caches can't be shared across layers). Each is shaped `[maxseq, nkvh, dh]`. `model->cur_len` tracks how many positions are actually filled so far (independent of `maxseq`, which is just the preallocated capacity, persists across `Infer` calls on the same model instance). + +Per `Infer` call: `total_len = cur_len + ntoken`; slice `[cur_len, total_len)` on dim 0 to write this step's new K/V, slice `[0, total_len)` to read the full history into `self_attention`. `meta.maxseq` should be capped well below the config's `max_position_embeddings` (e.g. 4096) in the Python wrapper — using the raw config value would allocate a multi-GB KV cache; this machine only has ~5.5GB free RAM. + +### Local test model + +`DeepSeek-R1-Distill-Qwen-1.5B/` (untracked) — `model.safetensors` is now the real ~3.5GB weights (the earlier git-lfs-pointer problem was resolved). `test/test_infer.py --model --test` is the acceptance test (greedy-decode token-for-token match against HF `transformers`); passes. + +## Working style for this project + +The user is doing these assignments to learn. **Do not generate finished implementations for the parts they're actively working through** — this now means the Assignment #4 CUDA work (`xmake/nvidia.lua`, `src/device/nvidia/nvidia_runtime_api.cu`, `src/ops/*/nvidia/*.cu`). Prefer step-by-step guidance: explain what a compiler error means, describe in words what the next piece of logic needs to do, answer conceptual questions directly, and let them write the actual code. Mechanical/environment-only fixes (wrong import path, wrong constant sourced from a config dict that's already being read for other fields) are fine to just fix and explain, rather than treating as "logic they're working through" — judgment call based on whether it's a design/algorithm decision vs. a one-line correction. The `docs/qwen2_*_reference.*` files exist for them to check against once stuck on Assignment #3 — don't paste their contents proactively; no CUDA reference files exist yet for Assignment #4. diff --git a/README_ZN.md b/README_ZN.md index 4d2026ff5..24187cf86 100644 --- a/README_ZN.md +++ b/README_ZN.md @@ -327,7 +327,7 @@ python test/test_infer.py --model [dir_path/to/model] --test 每个 **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`` 宏保护： - +SM(Streaming Multiprocessor,流多处理器) ```c++ #ifdef ENABLE_NVIDIA_API namespace nvidia { diff --git a/REPORT.md b/REPORT.md new file mode 100644 index 000000000..1eaec4adf --- /dev/null +++ b/REPORT.md @@ -0,0 +1,115 @@ +# LLAISYS Assignment #4 报告 + +CUDA 集成 + 双平台适配（NVIDIA、天数智芯 Iluvatar CoreX），以及性能优化。 + +## 评分环境说明 + +开发/验证硬件（见第 1 节）与官方评分环境（NV A100 / 天数 TG150-200）不同，未实测过。已知差异与应对： +- **NVIDIA**：`xmake/nvidia.lua` 原来的 `-gencode` 只用 `native`（构建时探测本机 GPU），若评分构建环境探测不到 GPU 会静默不生成任何机器码；已改为 `add_cugencodes("native", "sm_80")`，保证 A100 一定有对应机器码。 +- **Iluvatar**：`xmake/iluvatar.lua` 中 `/usr/local/corex-4.4.0/lib64` 链接路径写死了 SDK 版本号，若 TG150/200 机器 corex 版本不同，可能需要按实际版本调整该路径（未做防御性修改，无法在本机验证）。 + +## 1. 环境 + +**NVIDIA（本地开发机，WSL2）** +- GPU：RTX 5070 Ti Laptop GPU（`sm_120`，Blackwell） +- CUDA：12.9 + +**天数智芯 Iluvatar CoreX（远程云平台）** +- 硬件：Iluvatar BI-V150，SDK `corex-4.4.0` +- 编译器：`/usr/local/corex/bin/clang++`（用 `-x ivcore` 而非标准 `-x cuda`） +- cuBLAS：`libcublas.so.10.2.3.254`（CUDA 10.2 时代兼容库） + +项目不依赖 cuDNN。 + +## 2. 复现步骤 + +**NVIDIA：** +```bash +xmake f --nv-gpu=y && xmake && xmake install +pip install -e ./python + +python test/ops/.py --device nvidia +python test/test_infer.py --model --test --device nvidia +``` + +**Iluvatar（远程机器）：** +```bash +export XMAKE_ROOT=y # 容器内是 root，xmake 默认拒绝 root 运行 +xmake f --iluvatar-gpu=y && xmake && xmake install +pip install -e ./python + +python test/ops/.py --device iluvatar +``` +远程机器需先安装 xmake：`curl -fsSL https://xmake.io/shget.text | bash`（不要带 `--branch` 参数）。 + +## 3. 逐算子正确性 + +| 算子 | NVIDIA | Iluvatar | +| --- | --- | --- | +| add | ✅ | ✅ | +| embedding | ✅ | ✅ | +| argmax | ✅ | ✅ | +| rope | ✅ | ✅ | +| linear | ✅（cuBLAS） | ✅（见下方 bf16 说明） | +| swiglu | ✅ | ✅ | +| rms_norm | ✅ | ✅ | +| self_attention | ✅（V1 kernel + NVIDIA 上的 flash attention，见第 5 节） | ✅（V1 kernel） | +| rearrange | 不实现（作业 #2 遗留的废弃算子） | 同左 | + +**Iluvatar `linear` 的 bf16 修复**：Iluvatar 的 cuBLAS 兼容库对 `cublasSgemmEx` 的 `CUDA_R_16BF` 返回 `CUBLAS_STATUS_NOT_SUPPORTED`。改用更通用的 `cublasGemmEx`（`computeType` 传旧式的 `CUDA_R_32F` 而非新枚举 `CUBLAS_COMPUTE_32F`）后数值完全正确，所有 shape 通过。 + +## 4. 完整端到端推理（`test_infer.py --test`） + +| 平台 | 状态 | +| --- | --- | +| CPU | ✅ | +| NVIDIA | ✅ | +| Iluvatar | ✅（远程机器上跑通，见下方说明） | + +**修复过的 bug**：`qwen2.cc` 里 `argmax` 结果 `max_idx` 在 GPU 上是设备指针，早期代码直接在主机端解引用导致段错误（只在 CPU 设备下"凑巧"能跑，因为 CPU 的显存指针和主机指针是同一地址空间）。参照 `Tensor::debug()` 的做法，改为 `memcpy_sync(..., LLAISYS_MEMCPY_D2H)` 读回主机端后再使用。修复后 32 个 token 逐个匹配 HF 参考实现。 + +**Iluvatar 端到端验证**：原来的远程实例中途被销毁，重新开了一台同镜像的新实例后补跑：`test/test_infer.py --device iluvatar --test`（模型权重从 HuggingFace 自动下载）8 个算子 + 完整推理全部通过，token 输出逐个匹配 HF 参考实现。 + +## 5. 性能优化：`self_attention` + +### 5.1 Flash Attention 手写 kernel（NVIDIA，prefill + decode） + +在 V1 kernel 基础上，为 `d=dv=128`（DeepSeek-R1-Distill-Qwen-1.5B 的真实 head_dim）单独写了两条 tiled + online-softmax 的 flash attention kernel，其余 shape 仍 fallback 到 V1： + +- **Prefill**：两级 tiling + online softmax，一次处理整段 prompt。 +- **Decode**：一个 warp 处理一个 head，扫全部 KV cache。 + +`test/benchmark_infer.py` 实测（同一 prompt，`max_steps=64`）： + +| 阶段 | V1（全程） | Flash Attention | 加速比 | +| --- | --- | --- | --- | +| Prefill | 1068.98 ms | 520.77 ms | 2.05x | +| Decode | 19.99 ms/token | 16.49 ms/token | 1.21x | + +Decode 加速比明显小于 prefill：decode kernel 只发射 `nhead`（12）个 block、每个 1 个 warp，对这块 GPU 的 46 个 SM 来说远未打满，大部分 SM 在 decode 阶段闲置。 + +### 5.2 Decode Split-KV（flash-decoding），解决 SM 占用不足 + +针对 5.1 里发现的 decode 占用率问题，把 KV 方向切成多段并行处理（而非 query 行方向，因为 decode 时 `seqlen=1` 没有行可切）： + +- **Phase 1**：一个 warp 处理一个 `(head, split)`，在自己负责的 `[split_start, split_end)` 区间内做局部 online softmax，把 `(m, l, acc)` 写入中间 buffer。 +- **Phase 2**：一个 warp 处理一个 head，按 `exp(m_split - m_final)` 重新缩放并合并所有 split 的局部结果。 +- `num_splits` 按 `nhead * num_splits ≈ 64`（目标 block 数覆盖常见 GPU 的 SM 数量级）动态选取，且不低于 `TILE_K=32` 一个 tile 的量。 +- `op.cpp` 按 `total_len > 256` 切换到 split-KV，否则用原 decode kernel——阈值选取见下方分析。 + +**Kernel 级 microbenchmark**（bf16，真实模型 shape `nhead=12, nkvhead=2, d=dv=128`）： + +| total_len | 原 decode kernel | Split-KV | 加速比 | +| --- | --- | --- | --- | +| 128 | 0.058 ms | 0.084 ms | 0.69x（更慢） | +| 256 | 0.110 ms | 0.092 ms | 1.20x | +| 512 | 0.213 ms | 0.096 ms | 2.23x | +| 1000 | 0.443 ms | 0.141 ms | 3.14x | +| 4000 | 2.362 ms | 0.468 ms | 5.05x | +| 16000 | 11.463 ms | 1.567 ms | 7.31x | + +`total_len` 越长优势越大：原 kernel 单 warp 串行扫描、耗时线性增长；split-KV 的 block 数随 `total_len` 一起涨，SM 占用率保持住，增长明显更平缓。`total_len=128` 时 split-KV 反而更慢——多一次 kernel launch、多一趟中间结果的显存读写、外加当前实现每次调用都现分配/释放三个中间 buffer 的固定开销，在计算量本身很小时盖不过收益，这正是选 256 作为切换阈值的依据。 + +端到端 A/B（真实模型，`total_len` 从 67 涨到 467 的一次生成，两种模式各跑一遍）：split-KV 开启 17.41 ms/token（57.45 tok/s）vs 关闭 22.42 ms/token（44.60 tok/s），**约 1.29x**。低于 kernel 级数字是因为这次测量本身跨越了阈值前后两种状态，被平均拉低。 + +Iluvatar 侧未移植上述优化，仍为 V1 kernel（正确性已验证，见第 3、4 节）。 diff --git a/docs/CORE_PACKAGE_ZH.md b/docs/CORE_PACKAGE_ZH.md new file mode 100644 index 000000000..89fdd83a4 --- /dev/null +++ b/docs/CORE_PACKAGE_ZH.md @@ -0,0 +1,197 @@ +# `src/core/` —— LLAISYS 核心运行时模块 + +> 本文单独介绍 `src/core/` 这一个软件包，是 `docs/PROJECT_OVERVIEW_ZH.md` 第 5 节内容的展开版。已对照当前仓库源码（`src/core/**`、`include/llaisys/runtime.h`、`xmake.lua`）核对。 + +## 1. 这个包是做什么的 + +`src/core` 是整个 LLAISYS 里**唯一负责"设备资源管理"的模块**：它不知道张量（Tensor）长什么样，也不知道任何算子（Op）的计算逻辑，只回答四个问题： + +1. 当前线程正在用哪个设备？（`Context`） +2. 这个设备的显存 / 流 / API 函数表在哪？（`Runtime`） +3. 一块分配出来的内存归谁管、什么时候释放？（`Storage`） +4. 内存到底怎么分配、怎么释放？（`MemoryAllocator`） + +上层的 `src/tensor`（张量）和 `src/ops`（算子）都构建在这四个概念之上：`Tensor` 内部持有一个 `core::storage_t`（即 `std::shared_ptr`），算子在派发到具体设备实现前也要通过 `core::context()` 确认/切换当前设备。 + +## 2. 目录结构 + +```text +src/core/ +├── core.hpp # 前向声明 + 全局入口 context() +├── llaisys_core.hpp # 汇总头文件，外部只需 #include 这一个 +├── context/ +│ ├── context.hpp / .cpp # Context：线程局部单例，管理所有设备的 Runtime +├── runtime/ +│ ├── runtime.hpp / .cpp # Runtime：单个设备实例的资源持有者 +├── storage/ +│ ├── storage.hpp / .cpp # Storage：一块设备/主机内存的 RAII 包装 +└── allocator/ + ├── allocator.hpp # MemoryAllocator：分配策略的抽象接口 + └── naive_allocator.hpp/.cpp # NaiveAllocator：唯一实现，直接透传给设备 API +``` + +对应的 xmake 目标是 `llaisys-core`（`xmake.lua:52-66`），静态库，依赖 `llaisys-utils` 和 `llaisys-device`，被 `llaisys-tensor` 依赖： + +```text +llaisys-utils ──┐ +llaisys-device ──┼──▶ llaisys-core ──▶ llaisys-tensor ──▶ llaisys-ops ──▶ llaisys（共享库） +``` + +## 3. 四个核心类的关系 + +```text + thread_local +Context ───────────────────────▶ 持有每种设备类型 × 每个 device_id 的 Runtime* + │ setDevice(type, id) （_runtime_map，惰性创建，当前激活的是 _current_runtime） + ▼ +Runtime ── 持有一个设备实例的资源： + ├── const LlaisysRuntimeAPI *_api （设备无关函数表，来自 src/device） + ├── llaisysStream_t _stream （create_stream 得到） + └── MemoryAllocator *_allocator （策略对象，当前是 NaiveAllocator） + │ allocateDeviceStorage(size) / allocateHostStorage(size) + ▼ +Storage ── RAII 包装一块内存： + ├── std::byte *_memory + ├── size_t _size + ├── Runtime &_runtime （记住是哪个 Runtime 分配的） + └── bool _is_host + 析构时自动调用 _runtime.freeStorage(this) 归还内存 +``` + +外部（`src/tensor`）看到的只是 `core::storage_t = std::shared_ptr`，多个 `Tensor` 视图（`view`/`permute`/`slice`）可以共享同一个 `Storage`，引用计数归零时才真正释放。 + +## 4. 逐个组件详解 + +### 4.1 `Context`（`context/context.hpp` / `.cpp`） + +线程局部单例，通过函数内 `static thread_local` 实现： + +```cpp +Context &context() { + thread_local Context thread_context; + return thread_context; +} +``` + +- **构造**：遍历所有 `llaisysDeviceType_t`（NVIDIA 等排在前面，`LLAISYS_DEVICE_CPU` 特意放最后作为兜底），对每种设备调用 `get_device_count()` 探测数量，为每个 `device_id` 预留一个 `Runtime*` 槽位。第一个被发现的可用设备会立刻 `new Runtime` 并激活为 `_current_runtime`。由于 CPU 后端 `get_device_count()` 恒为 1，而 NVIDIA 在未实现时返回 0，**默认总是落到 CPU**。 +- **`setDevice(device_type, device_id)`**：若与当前激活的 Runtime 不匹配，则 `_deactivate()` 旧的、惰性 `new Runtime(...)`（如果该槽位还没创建过）、再 `_activate()` 新的。 +- **`runtime()`**：返回当前激活的 `Runtime&`；如果从未 `setDevice` 过且构造时也没有可用设备，会 `ASSERT` 失败——所以规则是**用之前必须确保有一个激活的 Runtime**（构造函数已经保证至少 CPU 可用）。 +- 显式 `delete` 拷贝/移动构造，保证每个线程只有一份状态、且不会被意外复制。 + +### 4.2 `Runtime`（`runtime/runtime.hpp` / `.cpp`） + +代表"一个设备实例"的资源持有者，构造函数私有，只能由 `Context`（`friend`）创建： + +```cpp +Runtime(llaisysDeviceType_t device_type, int device_id) + : _api(llaisys::device::getRuntimeAPI(device_type)), + _stream(_api->create_stream()), + _allocator(new allocators::NaiveAllocator(_api)) {} +``` + +对外方法： + +| 方法 | 作用 | +| --- | --- | +| `deviceType()` / `deviceId()` / `isActive()` | 查询自身状态 | +| `api()` | 拿到底层 `LlaisysRuntimeAPI` 函数表指针，供算子层直接调用（如 `memcpy_sync`） | +| `allocateDeviceStorage(size)` | 走 `_allocator->allocate(size)`，包成 `Storage`，`is_host=false` | +| `allocateHostStorage(size)` | 走 `_api->malloc_host(size)`，包成 `Storage`，`is_host=true` | +| `freeStorage(Storage*)` | 根据 `storage->isHost()` 决定走 `_api->free_host()` 还是 `_allocator->release()`，由 `Storage` 析构时自动触发 | +| `stream()` / `synchronize()` | 暴露当前流，或阻塞等待其上所有操作完成 | + +析构时 `delete _allocator` 并 `_api->destroy_stream(_stream)`；如果析构时 `_is_active` 仍是 `false`，会打印一条警告（当前实现里只是 `std::cerr`，不算致命错误）。 + +### 4.3 `Storage`（`storage/storage.hpp` / `.cpp`） + +一块内存的 RAII 包装，**构造函数私有，只有 `Runtime` 是 `friend`**——这保证了"内存块的生命周期必须由分配它的 `Runtime` 来管理"这一不变式，外部代码不可能绕过 `Runtime` 直接构造出一个 `Storage`。 + +```cpp +class Storage { + std::byte *_memory; size_t _size; Runtime &_runtime; bool _is_host; + ~Storage() { _runtime.freeStorage(this); } // 析构自动归还内存 +public: + std::byte *memory() const; + size_t size() const; + llaisysDeviceType_t deviceType() const; // is_host 时恒为 CPU + int deviceId() const; // is_host 时恒为 0 + bool isHost() const; +}; +``` + +`Tensor` 通过 `core::storage_t`（`std::shared_ptr`）持有它；`view`/`permute`/`slice` 产生的多个 `Tensor` 视图可以共享同一个 `Storage`，最后一个 `shared_ptr` 析构时才真正释放内存——这是整个项目里内存自动管理的根基。 + +### 4.4 `MemoryAllocator` / `NaiveAllocator`（`allocator/`） + +`MemoryAllocator` 是分配策略的抽象接口，只有两个纯虚函数： + +```cpp +class MemoryAllocator { +protected: + const LlaisysRuntimeAPI *_api; +public: + virtual std::byte *allocate(size_t size) = 0; + virtual void release(std::byte *memory) = 0; +}; +``` + +当前唯一实现 `NaiveAllocator` 直接透传给 `_api->malloc_device()` / `_api->free_device()`，**不做池化、不做复用**——这是刻意简化的教学版本。如果要做显存池、按大小分级复用等性能优化（作业范围之外的扩展方向），`MemoryAllocator` 这个接口就是天然的切入点：只需新写一个子类，在 `Runtime` 构造函数里换成新的实现即可，`Runtime`/`Storage`/`Context` 都不需要改动。 + +### 4.5 `LlaisysRuntimeAPI`：设备无关的函数表（跨到 `src/device`） + +`core` 包本身不实现任何设备相关的系统调用，而是通过 `include/llaisys/runtime.h` 定义的一张 12 个函数指针组成的 C 结构体来间接调用： + +```c +struct LlaisysRuntimeAPI { + get_device_count_api get_device_count; set_device_api set_device; + device_synchronize_api device_synchronize; + create_stream_api create_stream; destroy_stream_api destroy_stream; + stream_synchronize_api stream_synchronize; + malloc_device_api malloc_device; free_device_api free_device; + malloc_host_api malloc_host; free_host_api free_host; + memcpy_sync_api memcpy_sync; memcpy_async_api memcpy_async; +}; +``` + +`Runtime` 构造时通过 `llaisys::device::getRuntimeAPI(device_type)`（`src/device/runtime_api.cpp`）拿到这张表： + +- `LLAISYS_DEVICE_CPU` → CPU 实现（`src/device/cpu/`），全部基于 `std::malloc/std::free/std::memcpy`，"设备内存"和"主机内存"其实是同一块内存； +- `LLAISYS_DEVICE_NVIDIA` → 若编译时定义了 `ENABLE_NVIDIA_API` 则调用 NVIDIA 实现，否则退回 `getUnsupportedRuntimeAPI()`（所有函数都 `throw std::runtime_error`，用于给出明确报错而非链接失败）。 + +这种"胖接口 + 函数表分发"的设计，使得 `core` 里的 `Context`/`Runtime`/`Storage`/`Allocator` **完全不需要关心具体是哪种设备**，新增一种设备后端只需要在 `src/device` 下补一张函数表，`core` 包不用改一行代码。 + +## 5. 典型调用时序 + +以"创建一个 CPU 张量"为例（发生在 `src/tensor/tensor.cpp` 的 `Tensor::create` 里），串起 `core` 包的四个类： + +```text +Tensor::create(shape, dtype, LLAISYS_DEVICE_CPU, 0) + │ + ▼ +core::context() // 拿到当前线程的 Context（首次调用触发构造，默认已激活 CPU Runtime） + │ + ▼ +context().setDevice(LLAISYS_DEVICE_CPU, 0) // 若已经是 CPU，直接跳过 + │ + ▼ +context().runtime() // 拿到当前激活的 Runtime& + │ + ▼ +runtime.allocateDeviceStorage(bytes) // Runtime -> NaiveAllocator::allocate -> api->malloc_device + │ + ▼ +new Storage(ptr, size, runtime, /*is_host=*/false) // Runtime 是 friend，可以调用私有构造函数 + │ + ▼ +返回 shared_ptr 给 Tensor 持有 +``` + +当 `shared_ptr` 引用计数归零时，`~Storage()` 自动调用 `runtime.freeStorage(this)`，再根据 `is_host` 决定走分配器释放还是 `free_host`——**整条链路上没有任何一处需要手动 `delete` 内存**。 + +## 6. 使用 `core` 包时需要记住的几条规则 + +- **必须先有激活的 Runtime 才能 `context().runtime()`**：`Context` 构造时已保证至少 CPU 可用，正常使用不需要手动处理，但如果要切到非默认设备，记得先 `setDevice()`。 +- **`Storage`/`Runtime`/`Context` 都不可拷贝、不可移动**，全部通过指针/引用/`shared_ptr` 传递，这是为了保证"一块内存只有一个归属者"的不变式。 +- **`Storage` 的构造函数是私有的**，只能通过 `Runtime::allocateDeviceStorage` / `allocateHostStorage` 获得，不要试图绕开 `Runtime` 直接构造。 +- **`MemoryAllocator` 是唯一为将来扩展预留的接口**：如果要实现显存池等优化，只需要新增一个 `MemoryAllocator` 子类并在 `Runtime` 构造函数里替换 `NaiveAllocator`。 +- **`core` 包完全不知道 `Tensor` 的存在**：它只提供"内存在哪、怎么分配、什么时候释放"这三件事，任何张量形状/步长/dtype 相关的逻辑都在 `src/tensor` 层，不应该也不需要下沉到这里。 diff --git a/docs/CUDA_INTEGRATION_ZH.md b/docs/CUDA_INTEGRATION_ZH.md new file mode 100644 index 000000000..58f2f0fff --- /dev/null +++ b/docs/CUDA_INTEGRATION_ZH.md @@ -0,0 +1,135 @@ +# Assignment #4:CUDA 集成 —— 任务说明 + +> 本文梳理 Assignment #4(`README.md` "Integrate CUDA into LLAISYS" 一节)要完成的工作,对照当前仓库里已有的骨架代码整理成一份可执行的任务清单,写法上跟 [[QWEN2_INFERENCE_ZH]]（`docs/QWEN2_INFERENCE_ZH.md`）保持一致。Assignment #3 已经全部完成并通过验收测试(`test/test_infer.py --test` 输出 `Test passed!`),这是下一阶段的工作。 + +## 0. 本机环境核对 + +这台机器(WSL2)实际上已经具备做 CUDA 开发的条件,不用等训练营分配远程资源就能先跑通本地部分: + +- `nvidia-smi` 能看到一块 `NVIDIA GeForce RTX 5070 Ti Laptop GPU`。 +- `nvcc --version` 是 CUDA 12.9。 + +也就是说 Nvidia 这一个平台可以完全在本地开发调试。README 要求"从 Nvidia / Iluvatar / Metax / Moore Threads 里选两个平台",第二个平台需要训练营另外批的账号/资源,跟 Nvidia 这部分互不影响,可以先把 Nvidia 走完再申请另一个。 + +## 1. 目标 + +给 LLAISYS 加上 CUDA 后端支持,让 Runtime API、所有算子(Assignment #2)、Qwen2 模型推理(Assignment #3)都能在 `LLAISYS_DEVICE_NVIDIA` 上跑,验收标准是: + +```bash +python test/test_runtime.py --device nvidia +python test/test_infer.py --model <本地模型目录> --test --device nvidia +``` + +第二条要求 `--device nvidia` 跑出来的 token 序列跟 HuggingFace CPU/GPU 推理完全一致(逻辑上和 CPU 版本的验收标准相同,只是换了设备)。 + +## 2. 现状盘点:骨架已经搭到什么程度 + +跟 Assignment #3 开始时"C 头文件定义好、其余全空"的情况不太一样,这次框架已经把 CUDA 分支的**接口骨架**都占好位了,到处都是 `TO_BE_IMPLEMENTED()`,你要做的是把这些占位填成真正的 CUDA 实现。 + +### 已经搭好、不用你新建的部分 + +| 文件/位置 | 现状 | +|---|---| +| `include/llaisys.h` | `LLAISYS_DEVICE_NVIDIA = 1` 已定义 | +| `xmake.lua` | `option("nv-gpu")` 开关、`ENABLE_NVIDIA_API` 宏、`includes("xmake/nvidia.lua")`(条件包含)都已经写好 | +| `src/device/runtime_api.hpp` / `.cpp` | `getRuntimeAPI(device_type)` 的 switch 已经有 `LLAISYS_DEVICE_NVIDIA` 分支,`#ifdef ENABLE_NVIDIA_API` 时会调 `llaisys::device::nvidia::getRuntimeAPI()` | +| `src/device/nvidia/nvidia_runtime_api.cu` | 12 个 Runtime API 函数(`getDeviceCount`/`setDevice`/`deviceSynchronize`/`createStream`/`destroyStream`/`streamSynchronize`/`mallocDevice`/`freeDevice`/`mallocHost`/`freeHost`/`memcpySync`/`memcpyAsync`)**签名都在,函数体全是 `TO_BE_IMPLEMENTED()`** | +| `src/device/nvidia/nvidia_resource.cuh` / `.cu` | `Resource` 类骨架已有(继承 `DeviceResource`),构造函数已经调好基类,目前不需要改 | +| `src/ops/*/op.cpp`(全部 9 个算子:`add`/`argmax`/`embedding`/`linear`/`rearrange`/`rms_norm`/`rope`/`self_attention`/`swiglu`) | 每个的 device switch 里已经有 `#ifdef ENABLE_NVIDIA_API case LLAISYS_DEVICE_NVIDIA: TO_BE_IMPLEMENTED();` 分支占位 | + +### 完全是空的、需要你从头写的部分 + +| 文件 | 现状 | +|---|---| +| `xmake/nvidia.lua` | **不存在**。`xmake.lua` 第 15 行 `includes("xmake/nvidia.lua")` 会直接找不到文件报错——这是你要做的第一件事 | +| `src/ops/*/nvidia/*.cu`(9 个算子各自的 nvidia 子目录) | **不存在**,对照 `src/ops/*/cpu/*.cpp` 的组织方式,每个算子都要补一个 nvidia 版本 | +| `src/llaisys/models/qwen2.cc` 里对多设备的支持 | 目前的 `Create` 只按 `device`/`device_ids` 存了下来,但内部分配张量、KV Cache 时要确认走的是设备无关的路径(理论上应该是,因为都是通过 `Tensor`/`tensor_t` 的构造函数走 `core::context()` 分配,不需要在模型层写 if-cpu-else-nvidia) | + +## 3. 要做的工作,按 README 给的顺序拆解 + +### 3.1 先把编译打通:`xmake/nvidia.lua` + +参考 `xmake/cpu.lua` 的写法(`llaisys-device-cpu` 和 `llaisys-ops-cpu` 两个 target),照着建两个对应的 CUDA target: + +- `llaisys-device-nvidia`:编译 `src/device/nvidia/*.cu`,注意 `.cu` 文件要用 CUDA 规则编译(xmake 里一般是 `add_rules("cuda")` 或直接靠文件后缀识别,具体语法查 xmake 官方 CUDA 支持文档),还需要设置 CUDA 架构(`add_cugencodes` 或类似 API,对应这块 RTX 5070 Ti 的 compute capability)。 +- `llaisys-ops-nvidia`:编译 `src/ops/*/nvidia/*.cu`(通配符路径参照 `cpu.lua` 里 `"../src/ops/*/cpu/*.cpp"` 的写法)。 + +然后回到根 `xmake.lua`,在 `llaisys-device` 和 `llaisys-ops` 两个 target 里,仿照现有 `add_deps("llaisys-device-cpu")` / `add_deps("llaisys-ops-cpu")`,在 `has_config("nv-gpu")` 时额外 `add_deps("llaisys-device-nvidia")` / `add_deps("llaisys-ops-nvidia")`。 + +打通后应该能跑: + +```bash +xmake f --nv-gpu=y -cv +xmake +xmake install +``` + +编译不报错(哪怕运行时全是 `TO_BE_IMPLEMENTED()` 抛异常)就算这一步完成。 + +### 3.2 Runtime API:`src/device/nvidia/nvidia_runtime_api.cu` + +对照 `src/device/cpu/cpu_runtime_api.cpp` 那份已经写好的 CPU 实现,每个函数找 CUDA Runtime API 里的对应物: + +| 函数 | 对应的 CUDA API(大致) | +|---|---| +| `getDeviceCount` | `cudaGetDeviceCount` | +| `setDevice` | `cudaSetDevice` | +| `deviceSynchronize` | `cudaDeviceSynchronize` | +| `createStream` / `destroyStream` | `cudaStreamCreate` / `cudaStreamDestroy` | +| `streamSynchronize` | `cudaStreamSynchronize` | +| `mallocDevice` / `freeDevice` | `cudaMalloc` / `cudaFree` | +| `mallocHost` / `freeHost` | `cudaMallocHost` / `cudaFreeHost`(pinned memory,不是普通 `malloc`) | +| `memcpySync` | `cudaMemcpy`,注意 `llaisysMemcpyKind_t` 要映射到 `cudaMemcpyKind`(H2D/D2H/D2D 等) | +| `memcpyAsync` | `cudaMemcpyAsync`,多一个 `stream` 参数 | + +写完后跑: + +```bash +python test/test_runtime.py --device nvidia +``` + +这个测试逻辑很直接(见 `test/test_runtime.py`):分配两块 device 内存,做 H2D → D2D → D2H 三次拷贝,最后用 `torch.testing.assert_close` 比较,只验证内存管理和拷贝对不对,不涉及计算。 + +### 3.3 CUDA 算子:`src/ops/*/nvidia/*.cu` + +9 个算子(`add`/`argmax`/`embedding`/`linear`/`rearrange`/`rms_norm`/`rope`/`self_attention`/`swiglu`)逐个补齐。`rearrange` 的 CPU 版本本身也还是 `TO_BE_IMPLEMENTED()`(见 [[QWEN2_INFERENCE_ZH]] 里的说明,Assignment #3 靠 `memcpy` 绕过了它),可以放到最后再做,优先级最低。 + +每个算子的套路是一致的,可以参考同目录下 `cpu/*.cpp` 的实现思路搬到 CUDA kernel 上: + +1. 在 `src/ops//nvidia/` 下新建 `.cuh`(声明)+ `.cu`(kernel 实现),函数签名对照 `src/ops//cpu/_cpu.hpp` 抄一份。 +2. 在 `src/ops//op.cpp` 里把 `#ifdef ENABLE_NVIDIA_API case LLAISYS_DEVICE_NVIDIA: TO_BE_IMPLEMENTED();` 换成实际调用(参照同一个文件里 `case LLAISYS_DEVICE_CPU` 那一行怎么调 `cpu::xxx(...)`)。 +3. `src/ops//CMakeLists`-等价物这里是靠 `xmake/nvidia.lua` 里 `"../src/ops/*/nvidia/*.cu"` 的通配符自动纳入,不需要逐个算子改 xmake 配置。 + +建议顺序:先做 `add`(最简单,纯逐元素操作,用来验证整条编译+调用链路通不通),再做 `embedding`/`rms_norm`/`swiglu`/`argmax` 这几个逐元素或简单归约的算子,最后做 `linear`/`rope`/`self_attention` 这几个涉及矩阵乘法或者复杂索引的。`linear`(矩阵乘法)可以考虑直接用 cuBLAS(`cublasSgemm`/`cublasGemmEx` 等,注意这个模型权重是 BF16),不用手写 kernel。 + +每实现完一个算子,可以用 Assignment #2 现成的算子测试脚本加 `--device nvidia` 跑(`test/ops/` 目录下,具体参照 `test/ops/*.py` 里已有的 CPU 用例怎么写,应该有 `--device` 参数)。 + +### 3.4 把 Qwen2 模型接到 CUDA + +理论上如果 3.1~3.3 做完了,`src/llaisys/models/qwen2.cc` 不需要额外改动——你在 Assignment #3 里写的 `Infer` 全部是通过 `ops::xxx(...)` 调用算子、通过 `Tensor` 构造函数分配张量,只要这些底层调用是设备无关的(靠 `tensor->deviceType()` 在 op.cpp 内部分发),模型层的 C++ 代码本身不用感知 CPU/NVIDIA 的区别。 + +需要确认(不代表一定要改,先去看代码是否已经这样写的): +- `LlaisysQwen2ModelCreate` 里分配权重张量、KV Cache 张量时用的 `device`/`device_ids` 参数,有没有正确传给 `Tensor` 的构造(而不是被忽略、默认写死成 CPU)。 +- Python 侧 `python/llaisys/models/qwen2.py` 的 `__init__` 接的 `device: DeviceType` 参数有没有真的传下去(现在的实现已经支持传参,回顾 `LIB_LLAISYS.llaisysQwen2ModelCreate(ctypes.byref(self.meta), device.value, device_ids, 1)` 这一行)。 + +验收: + +```bash +python test/test_infer.py --model <本地模型目录> --test --device nvidia +``` + +## 4. 调试建议 + +- 跟 Assignment #3 一样,先用最小规模验证:单个算子(比如 `add`)先跑通,再逐步加算子。 +- CUDA kernel 写错很容易表现为"编译通过但结果不对"或者直接 core dump,建议开发时先在小 shape、固定输入下用 `cuda-memcheck` / `compute-sanitizer` 排查越界访问,比对着 CPU 版本结果肉眼找 bug 快得多。 +- BF16 在 CUDA 上要用 `__nv_bfloat16` / ``,跟 CPU 端用 `ml_dtypes`/内存直接 memcpy 的方式不一样,注意类型转换。 +- `self_attention` 是最复杂的一个,可以先在 CPU 上把结果存下来(用 `tensor.debug()`),CUDA 版本跑出来后逐值比对。 + +## 5. 验收 + +```bash +python test/test_runtime.py --device nvidia +python test/test_infer.py --model <本地模型目录> --test --device nvidia +``` + +两条都通过、`Test passed!` 打印出来即算完成。之后按 README 说明 commit + push,CI 里 Assignment #4 对应步骤应该跑绿(注意 CI runner 大概率没有 GPU,这部分很可能是训练营额外配置的自跑或人工验收环节,具体以 README "Assignment Submission Requirements" 一节的最新说明为准)。 diff --git a/docs/FLASH_ATTENTION_DEBUG_LOG_ZH.md b/docs/FLASH_ATTENTION_DEBUG_LOG_ZH.md new file mode 100644 index 000000000..d094a4617 --- /dev/null +++ b/docs/FLASH_ATTENTION_DEBUG_LOG_ZH.md @@ -0,0 +1,133 @@ +# `self_attention` NVIDIA Flash Attention 开发笔记 + +> 个人开发过程记录，不是提交评分用的文档（评分文档见 `REPORT.md`）。记录 `src/ops/self_attention/nvidia/flash_attention_cuda.cu`（prefill）+ `flash_attention_decode_cuda.cu`（decode）从接线到跑通期间踩到的两个真实 bug，方便以后遇到同类问题时回来查。 + +## 背景 + +`self_attention` 原来在 NVIDIA 上只有 V1 手写 kernel（一个 block 处理一个 `(query token, head)`，两遍扫描）。这次重写目标是给 prefill/decode 各写一个专门优化过的 flash attention kernel（K/V 分块 + online softmax），`op.cpp` 按 shape 分流： + +```cpp +if (seqlen > 1 && total_len == seqlen && d == 128 && dv == 128) { + flash_attention(...); // prefill：K/V 分块 + online softmax，TILE_Q=8 行/block +} else if (seqlen == 1 && d == 128 && dv == 128) { + flash_attention_decode(...); // decode：单 query 行，一个 warp 处理一个 head +} else { + self_attention(...); // 其余情况（非 128 的 head_dim、非首次 prefill）落回 V1 +} +``` + +写完主体逻辑、刚接上 dispatch 时，跑 `test/ops/self_attention.py --device nvidia`（自带的两组小 shape，`hd=4`/`hd=8`，逻辑上根本不会走 flash 分支）就间歇性崩溃——这是第一个 bug；解决之后用专门覆盖 `hd=128` flash 路径的对拍脚本测更大的 shape，又炸出第二个 bug。两个 bug 完全独立，分开记。 + +## Bug 1：ODR（重复定义）导致的间歇性 `illegal memory access` + +### 现象 + +`test/ops/self_attention.py --device nvidia` 连续跑几次，大概 30%-50% 概率随机报： + +``` +torch.AcceleratorError: CUDA error: an illegal memory access was encountered +CUDA error at src/device/nvidia/nvidia_runtime_api.cu:69: an illegal memory access was encountered +``` + +诡异的地方：这个测试脚本只用 `hd=4`/`hd=8` 两组 shape，`op.cpp` 的分流条件要求 `d==128 && dv==128` 才会调用 flash kernel——按代码逻辑，flash kernel **根本不应该被调用**。而且用 `compute-sanitizer --tool memcheck` 跑同一个命令反而 0 errors、稳定通过；`git stash` 掉 `op.cpp` 的 dispatch 改动（纯 V1）跑 15 次全过。说明问题跟"是否真的调用了 flash kernel"无关，是别的东西在起作用。 + +### 排查 + +`git stash`/`pop` 反复横跳（配合多次 `xmake build -r` 强制全量重编）之后，怀疑是不是链接层面的问题，直接对比两个 `.cu` 文件里同名符号的完整签名： + +```cpp +// self_attention_cuda.cu（V1） +template +__global__ void self_attention_kernel(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) + +// flash_attention_cuda.cu（当时的写法） +template +__global__ void self_attention_kernel(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) +``` + +**两个不同 `.cu` 文件，函数模板同名、参数类型完全一样，函数体完全不同。** 两者都在文件全局作用域（没有 `namespace { ... }`、没有 `static`），都是外部链接。C++ 里函数模板的每个实例化默认是弱符号（weak/linkonce），多个 `.o` 里出现同一个弱符号时，链接器只保留一份——具体保留哪一份跟目标文件在静态库里的排列顺序、优化级别等因素有关，是未定义行为，这正好解释了"随机"这个现象：**V1 的调用点有一定概率被链接器悄悄换成了 flash kernel 的机器码**，而 V1 的 launcher 只按 `total_len*sizeof(float)` 分配了很小的 shared memory，flash kernel 却需要 `Bc*(d+dv)*sizeof(float)`（`hd=4/8` 时按 V1 的公式算出来的 shared memory 远小于 flash kernel 实际会访问的范围）——一旦被换过去就是越界访问。 + +`launch_self_attention`（host 端 launcher）也是同样的重名情况，一并确认。 + +### 修复 + +把 `flash_attention_cuda.cu` 里的 `self_attention_kernel`/`launch_self_attention` 改名成 `flash_attention_kernel`/`launch_flash_attention`，跟 V1 不再同名。（`flash_attention_decode_cuda.cu` 从一开始就把所有内部符号包在匿名 `namespace { ... }` 里，天然拿到内部链接，不会有这个问题——这是更稳妥的写法，以后新增 `.cu` 文件时应该照这个来，不要依赖"记得手动改名"。） + +### 验证 + +`xmake build -r`（强制全量重编，排除任何缓存 `.o` 的影响）之后，`test/ops/self_attention.py --device nvidia` 连续跑 10 次，10/10 通过（改之前大概 30%-50% 崩溃率，个别情况下 python 进程直接 core dump 或者 hang 住）。 + +### 教训 + +- **C++ 里 `__global__`/普通函数模板默认外部链接**，不同 `.cu` 文件里出现同名同签名的模板函数，就算内容完全不同，也不会在编译期报错——只有链接期弱符号合并的未定义行为，behavior 取决于构建细节，非常难复现和定位。新加一个跟已有 kernel 结构相似的 `.cu` 文件（复制粘贴改一份是常见操作）时，命名要么加前缀避免撞名，要么直接包进匿名 `namespace`。 +- **"错误跟代码逻辑对不上"是重要信号**：guard 条件明明排除了 flash 路径，却还是间歇性触发本该只有 flash 路径才会踩到的越界访问——遇到这种"逻辑上不可能，但现象上确实发生"的情况，第一反应不该是继续在业务逻辑里找 bug，而是往"链接/构建层面是不是有问题"这个方向想。`compute-sanitizer` 跑几次全过、纯手跑反而间歇失败，这种"加了诊断工具后 bug 消失"的模式也是一个提示（诊断工具往往会改变时序或强制同步，容易让竞态/未定义行为暂时不出现）。 + +## Bug 2：跨 chunk 的 shared memory data race + +### 现象 + +修完 Bug 1 之后，用专门覆盖 flash 路径（`hd=128`，触发 `op.cpp` 的分流条件）的对拍脚本测：`qlen ∈ {2, 8, 9, 17, 33}`，GQA(`nh=12,nkvh=2`) 和非 GQA(`nh=4,nkvh=4`) 各一遍，三种精度都测。`qlen=33` 在 f32 精度下断言失败，`qlen<=17` 全过。 + +单独写脚本把 llaisys 输出和 PyTorch 参考实现逐 query 行比较误差（`(got - ref).abs().amax(dim=(1,2))`）： + +``` +row 0: max_abs_err=0.000000e+00 +row 1: max_abs_err=1.192093e-07 # float32 舍入误差量级，正常 +... +row 31: max_abs_err=1.192093e-07 +row 32: max_abs_err=2.323037e-02 # 差两个数量级，逻辑错误 +``` + +只有第 32 行（`qlen=33` 里最后一行）误差异常。 + +### 定位 + +kernel 行方向 `TILE_Q=8` 行一组、列方向 `Bc=32` 个 key 一组分块。`qlen=33` 时第 32 行单独落在最后一个 tile（`tile_start=32`，这个 tile 只有第 32 行存在）。这个 tile 的 `tile_max_limit = min(32+8-1, 33-1) = 32`，K/V 分块循环 `for (j=0; j<=tile_max_limit; j+=Bc)` 要跑**两轮**（`j=0` 覆盖 key 0-31，`j=32` 只覆盖 key 32）。`qlen<=32` 的所有 tile，`tile_max_limit` 都小于 `Bc=32`，分块循环只跑一轮——也就是说**在这次调试之前，"多轮分块累加"这个 flash attention 真正的核心逻辑，从来没被实际测到过**，小 shape 测试全部只经历单轮分块。 + +检查 kernel 主循环，每轮分块结束（读完 `K_chunk`/`V_chunk` 算完这一轮 softmax）之后直接回到循环顶部做下一轮的协作搬运，中间没有 `__syncthreads()`： + +```cpp +for (int j = 0; j <= tile_max_limit; j += Bc) { + // 协作搬运 K/V 进 shared memory + ... + __syncthreads(); // 搬运完 → 读之前，有同步 + + if (i < seqlen) { + ... 用 K_chunk/V_chunk 算 score/softmax/累加 ... + } + // 直接回到 for 循环顶部，没有同步！ +} +``` + +`qlen=33` 那个 tile 里 8 个 warp 只有 `warp_id=0`（第 32 行）满足 `i 本文描述的是当前仓库快照（已核对 `src/`、`include/`、`xmake.lua` 源码）。该仓库本身是一套待完成的课程骨架，很多函数故意保留为 `TO_BE_IMPLEMENTED()`，不应把所有已声明接口理解为已经可用。 + +## 2. 整体架构 + +一次典型调用会经过以下链路： + +```text +用户 / Python 测试 + │ + ▼ +Python 友好接口（python/llaisys/*.py） + │ + ▼ ctypes +C ABI 绑定（python/llaisys/libllaisys + include/llaisys） + │ + ▼ +C++ 边界层（src/llaisys/*.cc）── unwrap LlaisysTensor -> tensor_t，调用 C++ 实现，再包装返回 + │ + ├── 张量（src/tensor）── TensorMeta + Storage(shared_ptr) + offset + ├── 算子（src/ops）── 参数校验 -> dtype 分派 -> CPU/NVIDIA kernel + └── 核心运行时（src/core）── Context（线程局部单例） -> Runtime -> MemoryAllocator + │ + ▼ +设备后端（src/device/cpu 完整实现；nvidia 为预留骨架，受 ENABLE_NVIDIA_API 宏保护） +``` + +这种结构有两个主要教学价值：一是展示 Python 框架接口如何跨语言进入原生后端；二是把与设备无关的张量/算子逻辑和 CPU、GPU 等设备实现分离。 + +## 3. 目录说明 + +| 路径 | 作用 | +| --- | --- | +| `include/` | 对外公开的 C API：`llaisys.h`（设备/数据类型枚举）、`llaisys/{runtime,tensor,ops}.h`、`llaisys/models/qwen2.h`。 | +| `src/llaisys/` | C ABI 到内部 C++ 对象的适配边界（`tensor.cc`、`ops.cc`、`runtime.cc`）。 | +| `src/core/` | `context/`、`runtime/`、`storage/`、`allocator/` —— 设备无关的运行时基础设施。 | +| `src/device/` | 设备资源与 `LlaisysRuntimeAPI` 函数表；`cpu/` 是完整实现，`nvidia/` 为预留骨架。 | +| `src/tensor/` | `Tensor`/`TensorMeta` —— 形状、步长、存储引用、视图变换和数据搬运。 | +| `src/ops/` | 每个算子一个子目录：`add`、`argmax`、`embedding`、`linear`、`rms_norm`、`rope`、`self_attention`、`swiglu`、`rearrange`，各自再分 `cpu/`（未来还有 `nvidia/`）。 | +| `src/utils/`、`src/utils.hpp` | `check.hpp`（断言/校验宏）、`types.hpp`/`types.cpp`（dtype 大小、fp16/bf16 转换）。 | +| `python/llaisys/libllaisys/` | 共享库加载和底层 `ctypes` 函数签名，一一对应 `include/llaisys/*.h`。 | +| `python/llaisys/` | `RuntimeAPI`、`Tensor`、`Ops` 等 Python 接口。 | +| `python/llaisys/models/` | 模型层，目前 `qwen2.py` 提供待实现的 Qwen2 骨架（构造函数和 `generate` 均为 TODO）。 | +| `test/` | 运行时、张量、算子和端到端推理测试，以 PyTorch 为参考。 | +| `xmake.lua`、`xmake/` | C++ 构建配置：`llaisys-utils`/`llaisys-device`/`llaisys-core`/`llaisys-tensor`/`llaisys-ops` 静态库 + `llaisys` 共享库；`xmake/cpu.lua` 定义 CPU 子目标。 | +| `.github/workflows/build.yaml` | Windows、Ubuntu 上的构建和分阶段作业测试。 | + +## 4. C++ 构建体系与静态库依赖图 + +`xmake.lua` 把 C++ 代码拆成多个静态库目标，最终一起链接进 `llaisys` 共享库，依赖关系如下： + +```text +llaisys-utils (src/utils/*.cpp：dtype 大小、fp16/bf16 转换) + ▲ + │ +llaisys-device-cpu (src/device/cpu/*.cpp) ──┐ + ▲ │ + │ ▼ +llaisys-device (src/device/*.cpp) ← 依赖 llaisys-utils + llaisys-device-cpu + ▲ + │ +llaisys-core (src/core/*/*.cpp) ← 依赖 llaisys-utils + llaisys-device + ▲ + │ +llaisys-tensor (src/tensor/*.cpp) ← 依赖 llaisys-core + ▲ + │ +llaisys-ops-cpu (src/ops/*/cpu/*.cpp) ← 依赖 llaisys-tensor + ▲ + │ +llaisys-ops (src/ops/*/*.cpp) ← 依赖 llaisys-ops-cpu + │ + ▼ +llaisys（共享库，src/llaisys/*.cc）← 依赖以上全部 + │ + ▼ xmake install 后自动 os.cp +python/llaisys/libllaisys/*.so（或 Windows 下 *.dll） +``` + +几个值得注意的构建细节： + +- 所有目标统一 `set_languages("cxx17")` 并开启 `set_warnings("all", "error")`（警告即错误），非 Windows 平台加 `-fPIC -Wno-unknown-pragmas`。 +- `nv-gpu` 是一个 xmake `option`，默认 `false`；打开后会 `add_defines("ENABLE_NVIDIA_API")` 并 `includes("xmake/nvidia.lua")`——但该文件目前**不存在于仓库**，所以 `--nv-gpu=y` 现在还编译不过，是留给作业 #4 的扩展点。 +- `xmake install` 的 `after_install` 钩子会把编译好的 `lib/*.so`（Linux）拷贝到 `python/llaisys/libllaisys/`，这是 Python 侧 `ctypes.CDLL` 能加载到最新库的关键一步——改完 C++ 代码后必须重新 `xmake && xmake install` 才会生效。 + +## 5. 核心运行时（`src/core/`）详解 + +### 5.1 Context：线程局部单例 + +`Context` 通过函数内 `static thread_local` 实现每线程唯一实例： + +```cpp +Context &context() { + thread_local Context thread_context; + return thread_context; +} +``` + +构造函数会遍历所有 `llaisysDeviceType_t`（NVIDIA 排在前面，CPU 特意放最后作为兜底），对每种设备调用 `llaisysGetRuntimeAPI(device_type)->get_device_count()` 探测可用设备数，并为每个设备号预留一个 `Runtime*` 槽位（此时先不创建，除非它是第一个被找到的可用设备，那个会被立即创建并激活为 `_current_runtime`）。由于 CPU 后端的 `get_device_count()` 恒为 `1`，NVIDIA 分支在没有实现 GPU Runtime API 时通过 `getUnsupportedRuntimeAPI()` 返回 `get_device_count() == 0`，所以**默认总是落到 CPU**。 + +`setDevice(device_type, device_id)` 会检查当前激活的 Runtime 是否已经匹配，若不匹配则 `_deactivate()` 旧 Runtime、惰性 `new Runtime(...)`（如果该槽位还没创建过）、再 `_activate()` 新 Runtime。`Context` 禁止拷贝和移动（`= delete`），保证每个线程只有一份状态。 + +### 5.2 Runtime：单设备资源管理器 + +`Runtime` 持有一个设备的 `LlaisysRuntimeAPI` 函数表指针、一条 `llaisysStream_t`、一个 `MemoryAllocator*`。构造时通过 `llaisys::device::getRuntimeAPI(device_type)` 拿到函数表，再 `_api->create_stream()` 建流、`new allocators::NaiveAllocator(_api)` 建分配器。 + +`allocateDeviceStorage(size)` / `allocateHostStorage(size)` 分别调用分配器的 `allocate()` 或运行时 API 的 `malloc_host()`，把裸指针包进一个新建的 `Storage`（`is_host` 标志区分二者）。`freeStorage(Storage*)` 则根据 `storage->isHost()` 决定走 `_api->free_host()` 还是分配器的 `release()`——这个判断逻辑是 `Storage` 析构时自动触发的（见 5.3）。 + +### 5.3 Storage：跨张量共享的内存所有权 + +```cpp +class Storage { + std::byte *_memory; size_t _size; Runtime &_runtime; bool _is_host; + ~Storage() { _runtime.freeStorage(this); } +}; +``` + +`Storage` 的构造函数是私有的，只有 `Runtime` 是 `friend` 能创建它，这保证了"内存块的生命周期必须由分配它的 Runtime 来管理"这一不变式。`Tensor` 通过 `core::storage_t`（即 `std::shared_ptr`）持有它，多个 `Tensor`（例如 `view`/`permute`/`slice` 产生的视图）可以共享同一个 `Storage`，最后一个 `shared_ptr` 析构时才真正释放内存。 + +### 5.4 MemoryAllocator：抽象分配策略 + +`MemoryAllocator` 是一个只有两个纯虚函数的接口（`allocate`/`release`），当前唯一实现 `NaiveAllocator` 直接透传给 `_api->malloc_device()` / `_api->free_device()`，不做池化或复用——这是刻意简化的教学版本，如果要做性能优化（作业之外的扩展方向），这里是天然的切入点。 + +### 5.5 LlaisysRuntimeAPI：设备无关的函数表 + +`include/llaisys/runtime.h` 定义了一张 12 个函数指针组成的 C 结构体（设备数/切换设备/同步、流的创建销毁同步、设备内存与主机内存的分配释放、同步/异步拷贝）。`src/device/runtime_api.cpp` 里的 `getRuntimeAPI(device_type)` 按 `device_type` 分发： + +- `LLAISYS_DEVICE_CPU` → `llaisys::device::cpu::getRuntimeAPI()`（`src/device/cpu/cpu_runtime_api.cpp`，全部基于 `std::malloc/std::free/std::memcpy` 实现，"设备内存"和"主机内存"其实是同一块主机内存，`memcpySync`/`Async` 不区分 `kind` 直接 `memcpy`——这让 CPU 后端非常适合先验证跨设备抽象是否设计正确，再去接入真正异构的 NVIDIA 后端）； +- `LLAISYS_DEVICE_NVIDIA` → 若定义了 `ENABLE_NVIDIA_API` 则调用 `nvidia::getRuntimeAPI()`（目前仓库中**没有对应源文件**，只在头文件里声明了函数原型），否则退回 `getUnsupportedRuntimeAPI()`——这是一张所有函数都直接 `throw std::runtime_error` 的"空实现"函数表，用于在未启用 NVIDIA 支持时给出明确报错而不是链接失败。 + +## 6. 张量实现（`src/tensor/tensor.hpp` / `tensor.cpp`） + +### 6.1 数据结构 + +```cpp +struct TensorMeta { + llaisysDataType_t dtype; + std::vector shape; + std::vector strides; +}; + +class Tensor { + TensorMeta _meta; + core::storage_t _storage; // shared_ptr + size_t _offset; // 相对 Storage 起点的字节偏移 +}; +``` + +`Tensor` 的构造函数是私有的，只能通过静态工厂 `Tensor::create(shape, dtype, device_type, device)` 或内部的 `new Tensor(meta, storage, offset)` 创建，返回类型统一是 `tensor_t = std::shared_ptr`。 + +### 6.2 `Tensor::create` 的设备放置逻辑 + +```cpp +if (device_type == LLAISYS_DEVICE_CPU + && core::context().runtime().deviceType() != LLAISYS_DEVICE_CPU) { + // 当前激活的是非 CPU 设备，但要建的张量指定为 CPU —— + // 分配"主机存储"（pinned/host memory），不改变当前激活设备 + auto storage = core::context().runtime().allocateHostStorage(bytes); +} else { + // 目标设备就是当前设备，或本来就要切设备：先 setDevice 再分配设备存储 + core::context().setDevice(device_type, device); + auto storage = core::context().runtime().allocateDeviceStorage(bytes); +} +``` + +新建张量默认按行主序（C-contiguous）计算 `strides`：从最后一维往前累乘 `shape`，这也是 `numel()`（`std::accumulate` 对 `shape` 做乘积）和 `elementSize()`（`utils::dsize(dtype)`）的基础。 + +### 6.3 `debug()` 与模板化打印 + +`info()` 输出形状/步长/dtype 的字符串摘要；`debug()` 先做 `device_synchronize()`，若张量在设备上则先 `memcpy_sync(..., D2H)` 拷到一块临时 CPU 张量，再调用按 dtype 分派的 `debug_print`（`print_data` 模板函数按 stride 递归打印每一维，`fp16_t`/`bf16_t` 会先 `utils::cast` 再打印）。这个函数是调试作业实现时对拍 PyTorch 张量数值的主要工具。 + +### 6.4 已实现 vs 待实现的方法（按源码逐一核对） + +| 方法 | 状态 | 说明 | +| --- | --- | --- | +| `create` / `data` / `ndim` / `shape` / `strides` / `dtype` / `deviceType` / `deviceId` / `numel` / `elementSize` | ✅ 已实现 | 元信息查询与工厂函数，作业 #1 之前就已提供。 | +| `info()` / `debug()` | ✅ 已实现 | 调试打印，覆盖全部 dtype。 | +| `isContiguous()`（`tensor.cpp:166`） | ❌ `TO_BE_IMPLEMENTED()` | 任务 1.2：需要根据 `shape`/`strides` 判断是否行主序连续。 | +| `permute(order)`（`tensor.cpp:171`） | ❌ `TO_BE_IMPLEMENTED()` | 任务 1.4：只调整 `shape`/`strides` 顺序，共享 `_storage`。 | +| `view(shape)`（`tensor.cpp:176`） | ❌ `TO_BE_IMPLEMENTED()` | 任务 1.3：合并/拆分维度且不搬数据；若原张量非连续（如转置后的视图）则应报错。 | +| `slice(dim, start, end)`（`tensor.cpp:181`） | ❌ `TO_BE_IMPLEMENTED()` | 任务 1.5：调整对应维度的 `shape` 与 `_offset`，共享 `_storage`。 | +| `load(src)`（`tensor.cpp:186`） | ❌ `TO_BE_IMPLEMENTED()` | 任务 1.1：需要从当前设备 Runtime 拿 API 做 host→device 的 `memcpy`。 | +| `contiguous()` / `reshape(shape)` / `to(device_type, device)`（`tensor.cpp:190/195/200`） | ❌ `TO_BE_IMPLEMENTED()` | 进阶挑战：把非连续张量整理为连续、通用 reshape（必要时拷贝）、跨设备搬运。 | + +`view`/`permute`/`slice` 三者应遵循同一原则：**只描述新的形状/步长/偏移，绝不复制底层数据**，因为它们返回的新 `Tensor` 与原张量共享同一个 `_storage`（`std::shared_ptr` 引用计数 +1）。 + +## 7. C ABI 边界层（`src/llaisys/*.cc`） + +这一层是 C++ 世界和外部世界（Python `ctypes` / 未来的其他语言绑定）之间唯一允许出现 `extern "C"` 的地方。核心技巧是一个包装结构体： + +```cpp +// src/llaisys/llaisys_tensor.hpp +typedef struct LlaisysTensor { + llaisys::tensor_t tensor; // shared_ptr +} LlaisysTensor; +``` + +`llaisysTensor_t` 在 C API 里只是一个不透明指针（`struct LlaisysTensor *`）。每个导出函数的模式高度一致： + +```cpp +llaisysTensor_t tensorView(llaisysTensor_t tensor, size_t *shape, size_t ndim) { + std::vector shape_vec(shape, shape + ndim); // 1. C 数组 -> std::vector + return new LlaisysTensor{tensor->tensor->view(shape_vec)}; // 2. 调 C++ 实现 3. 包装成新的不透明指针 +} +``` + +`tensorDestroy` 对应 `delete tensor`——注意这里 `delete` 的只是 `LlaisysTensor` 包装体，真正的张量内存由 `shared_ptr` 引用计数决定何时释放。`src/llaisys/ops.cc` 里的每个 `llaisysXxx` 函数同样只做"解包 `->tensor` → 调 `llaisys::ops::xxx` → （若有返回值）重新包装"，不包含任何业务逻辑；`src/llaisys/runtime.cc` 则直接转发 `llaisysGetRuntimeAPI` / `llaisysSetContextRuntime` 到 `core::context()`。这种"零逻辑边界层"是刻意设计：出问题时可以放心假设 bug 在更下层的 C++ 实现里。 + +## 8. 算子层设计（`src/ops/`） + +每个算子固定包含：`op.hpp`（声明，参数类型都是 `tensor_t`）、`op.cpp`（设备无关的公共层：参数校验 + 按 `deviceType()` 分派）、`cpu/xxx_cpu.hpp`+`.cpp`（CPU 实现，未来还会有 `nvidia/` 子目录）。 + +### 8.1 校验宏（`src/utils/check.hpp`） + +- `CHECK_SAME_DEVICE(a, b, c...)`：所有张量的 `deviceType()`/`deviceId()` 必须一致，否则抛 `EXCEPTION_DEVICE_MISMATCH`。 +- `CHECK_SAME_SHAPE(...)` / `CHECK_SAME_DTYPE(...)`：基于同一个 `CHECK_SAME(ERR, FIRST, ...)` 宏模板，用花括号初始化列表遍历比较。 +- `ASSERT(condition, message)` / `CHECK_ARGUMENT(condition, message)`：前者用于内部不变式（抛 `runtime_error`），后者用于外部输入校验（抛 `invalid_argument`），都会打印文件名、行号、函数名。 +- `TO_BE_IMPLEMENTED()`：统一的"未实现"占位，打印位置信息后 `throw std::runtime_error("Unimplemented function")`——这也是为什么运行到未实现算子时测试会直接崩溃报错，而不是静默返回错误结果。 +- `EXCEPTION_UNSUPPORTED_DATATYPE(dtype)` / `EXCEPTION_UNSUPPORTED_DEVICE`：dtype `switch`/设备 `switch` 的 `default` 分支统一使用。 + +### 8.2 唯一的完整范例：`add` + +```cpp +// src/ops/add/op.cpp —— 设备无关公共层 +void add(tensor_t c, tensor_t a, tensor_t b) { + CHECK_SAME_DEVICE(c, a, b); + 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(), "..."); + + if (c->deviceType() == LLAISYS_DEVICE_CPU) { + return cpu::add(c->data(), a->data(), b->data(), c->dtype(), c->numel()); + } + llaisys::core::context().setDevice(c->deviceType(), c->deviceId()); + switch (c->deviceType()) { + case LLAISYS_DEVICE_CPU: return cpu::add(...); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: TO_BE_IMPLEMENTED(); return; // 唯一的 stub，只因 NVIDIA 后端未接入 +#endif + default: EXCEPTION_UNSUPPORTED_DEVICE; + } +} +``` + +```cpp +// src/ops/add/cpu/add_cpu.cpp —— CPU kernel，按 dtype 分派到模板函数 +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] = utils::cast(utils::cast(a[i]) + utils::cast(b[i])); // 升到 float 再算 + } else { + c[i] = a[i] + b[i]; + } + } +} +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), ...); + case LLAISYS_DTYPE_BF16: return add_(reinterpret_cast(c), ...); + case LLAISYS_DTYPE_F16: return add_(reinterpret_cast(c), ...); + default: EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +``` + +`add` 是仓库里唯一端到端跑通的算子，其余算子都应该照此范式实现：**公共层只做校验和分派，具体计算下沉到 `cpu::xxx`，浮点 kernel 内部对 fp16/bf16 都先 `cast` 提升精度再计算，写回时再 `cast` 转回原类型**——这是因为 `fp16_t`/`bf16_t` 只是包了一个 `uint16_t` 的自定义结构体，不支持原生算术运算符，必须借助 `utils::cast<>` 与 IEEE754 转换函数往返。 + +### 8.3 其余算子当前状态 + +| 算子 | 头文件签名 | 状态 | 数学契约（详见 README 作业 #2） | +| --- | --- | --- | --- | +| `add` | `add(c, a, b)` | ✅ CPU 完整实现（仅 NVIDIA 分支未实现） | `c = a + b`，逐元素 | +| `argmax` | `argmax(max_idx, max_val, vals)` | ❌ stub | 1D 张量求最大值及下标 | +| `embedding` | `embedding(out, index, weight)` | ❌ stub | 按 `index`(int64) 从 `weight` 表查行 | +| `linear` | `linear(out, in, weight, bias)` | ❌ stub | `Y = X Wᵀ + b`，`bias` 可为空 | +| `rms_norm` | `rms_norm(out, in, weight, eps)` | ❌ stub | 按最后一维做 RMSNorm | +| `rope` | `rope(out, in, pos_ids, theta)` | ❌ stub | 对 Q/K 做旋转位置编码 | +| `self_attention` | `self_attention(attn_val, q, k, v, scale)` | ❌ stub | 因果 softmax 注意力，需自行拼接 KV cache | +| `swiglu` | `swiglu(out, gate, up)` | ❌ stub | `out = up ⊙ silu(gate)` | +| `rearrange` | `rearrange(out, in)` | ❌ stub | 把非连续/不同布局的 `in` 数据整理进连续的 `out` | + +课程要求核心浮点算子至少覆盖 **Float32、Float16、BFloat16** 三种 dtype。 + +## 9. 数据类型系统（`llaisys.h` + `src/utils/`） + +`llaisysDataType_t` 是一个覆盖布尔、8~64 位有符号/无符号整数、byte、fp8/fp16/bf16/fp32/fp64、以及 complex16/32/64/128 的枚举（`include/llaisys.h`）。`src/utils/types.hpp` 提供： + +- `dsize(dtype)`：返回每种 dtype 的字节数（`switch`，未知类型抛异常）； +- `dtype_to_str(dtype)`：用于报错信息里可读的类型名； +- `fp16_t` / `bf16_t`：都只是 `struct { uint16_t _v; }` 的薄包装，**不重载任何运算符**，因此不能直接 `+`/`*`； +- `utils::cast(val)`：一个基于 `if constexpr` 的模板分派器，覆盖了"同类型直通"、"float↔fp16"、"float↔bf16"、"任意类型→fp16/bf16 先转 float 再转"、"fp16/bf16→任意类型先转 float 再 static_cast"等所有组合，底层依赖 `src/utils/types.cpp` 中实现的 `_f16_to_f32`/`_f32_to_f16`/`_bf16_to_f32`/`_f32_to_bf16` 四个真正做位运算的转换函数。 + +写算子 kernel 时的标准模式就是：**用 `switch(dtype)` 把 `std::byte*` reinterpret 成具体类型指针，模板函数内部对非 float 的窄浮点类型统一 `cast` 升精度计算，最后 `cast` 写回**，`add` 已完整示范了这一模式。 + +## 10. Python 接口与模型层 + +Python 包分为两层：`libllaisys` 精确描述 C 类型和函数签名（与 `include/llaisys/*.h` 一一对应）；外层 `RuntimeAPI`、`Tensor`、`Ops` 提供更自然的 Python 调用方式。测试会在 PyTorch 与 LLAISYS 之间搬运相同数据，并比较形状、步长、dtype 和数值结果。 + +`python/llaisys/models/qwen2.py` 中的 `Qwen2` 类目前只有两个方法的骨架：`__init__` 遍历 safetensors 文件枚举权重名但没有实际加载逻辑，`generate` 直接 `return []`。按作业 #3 要求，模型的前向计算必须用 C/C++ 在 LLAISYS 后端实现，Python 侧只能做胶水代码（不允许借助 PyTorch 等框架实现推理逻辑本身），并且需要实现 KV Cache 否则推理速度会难以接受。 + +## 11. 构建与运行 + +前置条件：Xmake、支持 C++17 的编译器、Python 3.9 或更高版本。Python 包声明依赖 PyTorch、Transformers 和 Accelerate。 + +```bash +# 构建 C++ 共享库 +xmake + +# 将共享库安装/复制到 Python 包目录（改完 C++ 代码后必须重新执行这一步） +xmake install + +# 安装 Python 包及依赖 +python -m pip install ./python +``` + +可以按课程顺序执行测试： + +```bash +python test/test_runtime.py --device cpu +python test/test_tensor.py +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 +``` + +端到端推理测试需要本地模型目录，或者允许脚本从 Hugging Face 下载约 1.5B 参数的模型： + +```bash +python test/test_infer.py --model /path/to/DeepSeek-R1-Distill-Qwen-1.5B +``` + +## 12. 当前实现状态（逐文件核对） + +- ✅ 项目分层、公开 C API、CPU Runtime、内存分配（`NaiveAllocator`）、`Context`/`Runtime`/`Storage` 核心框架均已完成且可编译。 +- ✅ `add` 算子端到端完整（公共校验 + CPU kernel + fp32/fp16/bf16），是其余算子的实现模板；其 `op.cpp` 中唯一的 `TO_BE_IMPLEMENTED()` 只存在于 `ENABLE_NVIDIA_API` 分支。 +- ✅ 张量创建、元信息查询（`shape`/`strides`/`dtype`/`numel`/...）、`debug()`/`info()` 调试打印已经存在。 +- ❌ `src/tensor/tensor.cpp` 中 `isContiguous`、`view`、`permute`、`slice`、`load`、`contiguous`、`reshape`、`to` 共 8 个方法仍是 `TO_BE_IMPLEMENTED()`（课程任务 1.1–1.5 及进阶挑战）。 +- ❌ 除 `add` 外，`argmax`/`embedding`/`linear`/`rms_norm`/`rope`/`self_attention`/`swiglu`/`rearrange` 的 `op.cpp` 入口全部是 `TO_BE_IMPLEMENTED()`。 +- ❌ `python/llaisys/models/qwen2.py` 的权重加载与 `generate` 均为 TODO 占位。 +- ⏳ NVIDIA 设备代码：`include`/`src/device/runtime_api.hpp` 中已经预留 `nvidia::getRuntimeAPI()` 声明和 `ENABLE_NVIDIA_API` 宏开关，但 `xmake/nvidia.lua`、`src/device/nvidia/*` 均不存在，`--nv-gpu=y` 目前无法编译通过。 + +因此，当前仓库更准确的定位是"可编译的教学起点"，而不是已经完成的推理引擎。预置的 `.so` 文件也不能替代从当前源码重新构建和测试，因为它可能与工作区源码状态不同。 + +## 13. 推荐阅读与实现顺序 + +1. 先阅读 `include/llaisys.h` 和 `include/llaisys/runtime.h`，理解设备、dtype 和函数表结构。 +2. 沿 `python/llaisys/runtime.py` → `src/llaisys/runtime.cc` → `src/core/context/context.cpp` → `src/device/cpu/cpu_runtime_api.cpp` 跟踪一次跨语言调用，理解 `Context`（线程局部单例）→ `Runtime`（单设备资源）→ `LlaisysRuntimeAPI`（函数表）三层关系。 +3. 阅读 `src/tensor/tensor.hpp`/`.cpp`，重点理解 `TensorMeta`（dtype/shape/strides）+ `Storage`（共享内存所有权）+ `offset`（视图偏移）三件套，再实现并通过张量测试（`isContiguous`/`view`/`permute`/`slice`/`load`）。 +4. 以 `src/ops/add/`（`op.cpp` 校验分派 + `cpu/add_cpu.cpp` 模板 kernel）为模板逐个实现 CPU 算子，注意 fp16/bf16 要借助 `utils::cast` 升精度计算，并用对应 PyTorch 测试验证。 +5. 最后实现 Qwen2 权重映射（`python/llaisys/models/qwen2.py` + 对应 C/C++ 后端）、前向传播、KV cache 和 token 生成，再运行端到端推理测试。 +6. 若进入作业 #4（CUDA 集成），参考 `src/device/cpu/cpu_runtime_api.cpp` 的写法实现 `src/device/nvidia/`，新增 `xmake/nvidia.lua`，并在每个算子目录下新增 `nvidia/` 子目录接入 CUDA kernel。 + +## 14. 项目特点与边界 + +LLAISYS 的优势是规模小、分层清楚、测试直接对齐 PyTorch，适合理解 AI 框架最核心的机制：跨语言 ABI 边界、线程局部的设备上下文管理、张量的 stride/offset/共享存储模型、算子的"校验-分派-kernel"范式，以及窄浮点类型的手工转换。它当前强调正确性和教学可读性，而非生产级性能、完整算子覆盖、自动求导、分布式训练或成熟 GPU 优化（`NaiveAllocator` 不做内存池化就是典型例子）。学习时应把重点放在接口边界、内存布局、设备派发和 Transformer 推理数据流上。 diff --git a/docs/QWEN2_INFERENCE_ZH.md b/docs/QWEN2_INFERENCE_ZH.md new file mode 100644 index 000000000..b0c816e83 --- /dev/null +++ b/docs/QWEN2_INFERENCE_ZH.md @@ -0,0 +1,142 @@ +# Assignment #3:Qwen2 大模型推理 —— 任务说明 + +> 本文梳理 Assignment #3(`README.md` "Large Language Model Inference" 一节)要完成的工作,对照当前仓库里已有的骨架代码(`include/llaisys/models/qwen2.h`、`python/llaisys/models/qwen2.py`、`test/test_infer.py`)整理成一份可执行的任务清单。目标模型是 [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B),它的注意力机制细节见 [[SELF_ATTENTION_ZH.md]] 同目录下的分析文档。 + +## 1. 目标 + +用 C/C++ 在 LLAISYS 后端实现 Qwen2 模型的**完整推理**(prefill + 增量解码),Python 侧只做权重加载和薄封装,**不允许**用 PyTorch 之类的框架在 Python 里代劳计算逻辑。验收标准是 `test/test_infer.py --test` 通过——用 `top_k=1`(贪心采样)时,LLAISYS 生成的 token 序列要和 HuggingFace `transformers` 跑出来的完全一致。 + +## 2. 现状盘点:哪些是已经搭好的骨架,哪些是空的 + +### 已经定义好、不用你操心的部分 + +**C 结构体和函数原型**(`include/llaisys/models/qwen2.h`),已经给出: + +```c +struct LlaisysQwen2Meta { + llaisysDataType_t dtype; + size_t nlayer, hs, nh, nkvh, dh, di, maxseq, voc; + float epsilon, theta; + int64_t end_token; +}; + +struct LlaisysQwen2Weights { + llaisysTensor_t in_embed; + llaisysTensor_t out_embed; + llaisysTensor_t out_norm_w; + llaisysTensor_t *attn_norm_w; // 每层一个,数组长度 = nlayer + llaisysTensor_t *attn_q_w; + llaisysTensor_t *attn_q_b; + llaisysTensor_t *attn_k_w; + llaisysTensor_t *attn_k_b; + llaisysTensor_t *attn_v_w; + llaisysTensor_t *attn_v_b; + llaisysTensor_t *attn_o_w; + llaisysTensor_t *mlp_norm_w; + llaisysTensor_t *mlp_gate_w; + llaisysTensor_t *mlp_up_w; + llaisysTensor_t *mlp_down_w; +}; + +struct LlaisysQwen2Model *llaisysQwen2ModelCreate(const LlaisysQwen2Meta *meta, llaisysDeviceType_t device, int *device_ids, int ndevice); +void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model *model); +struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model *model); +int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model *model, int64_t *token_ids, size_t ntoken); +``` + +注意 `attn_q_w/attn_q_b`、`attn_k_w/attn_k_b`、`attn_v_w/attn_v_b` 都有 bias,但 `attn_o_w` 没有——这正好对应 Qwen2 架构里 QKV 投影带 bias、输出投影不带 bias 的设计(细节见 `SELF_ATTENTION_ZH.md` 里 §5 的说明)。 + +- **测试脚本** `test/test_infer.py`:已经写好了完整的对比逻辑——用 HuggingFace `transformers` 跑一遍 `model.generate(...)`,再用你实现的 `llaisys.models.Qwen2` 跑一遍,`--test` 模式下强制 `top_k=1, top_p=1.0, temperature=1.0`(即贪心解码),两边 token 序列必须完全一致(`assert llaisys_tokens == tokens`)。 +- **算子层**:`self_attention`、`rope`、`rms_norm`、`swiglu`、`linear`、`embedding`、`add`、`argmax` 都已经在 Assignment #2 里实现完了(见 `docs/SELF_ATTENTION_ZH.md`),可以直接复用,不用重新实现计算逻辑。 + +### 完全是空的、需要你从头写的部分 + +| 文件 | 现状 | +|---|---| +| `src/llaisys/models/`(或类似目录) | **不存在**,`llaisysQwen2Model*` 四个 C API 函数完全没有实现 | +| `python/llaisys/libllaisys/models.py`(或类似文件) | **不存在**,没有任何 ctypes 包装能调用上面这四个函数 | +| `python/llaisys/models/qwen2.py` | 只有函数签名和两个 `# TODO` 注释,构造函数和 `generate` 都是空的 | +| `xmake.lua` | `src/llaisys/*.cc` 目前只 glob 了 `src/llaisys/` 一层目录(`ops.cc`/`runtime.cc`/`tensor.cc`),如果新文件放在子目录(比如 `src/llaisys/models/qwen2.cc`)需要检查 glob 规则要不要改成 `src/llaisys/**.cc` 或显式加路径 | + +## 3. 要做的工作,按调用链从下往上拆解 + +### 3.1 C++ 后端:实现模型本体和四个导出函数 + +新建类似 `src/llaisys/models/qwen2.cc`(或直接 `src/llaisys/models.cc`,取决于你想不想拆子目录),内部至少需要: + +1. **`LlaisysQwen2Model` 的具体定义**(头文件里只是前向声明 `struct LlaisysQwen2Model;`,具体成员由你定义),通常要持有: + - `LlaisysQwen2Meta` 的一份拷贝(层数、head 数、维度这些超参数)。 + - `LlaisysQwen2Weights`,以及背后实际的 `tensor_t`(权重张量本身,在 `Create` 时按 `meta` 里的形状分配好,等 Python 侧再通过 `tensorLoad` 把 safetensors 里的数据拷进来)。 + - **KV Cache**:每层一份 K、V 的缓存张量,形状类似 `[maxseq, nkvh, dh]`,以及一个记录"当前已经缓存了多少 token"的计数器(对应 `self_attention` 里的 `total_len`)。 + - 当前设备类型/编号(`device`、`device_ids`、`ndevice`,这次作业阶段只用管 CPU,`ndevice` 大概率恒为 1,多卡是 Assignment #4 CUDA 阶段以后再考虑的事)。 + +2. **`llaisysQwen2ModelCreate`**:按 `meta` 分配好所有权重张量和 KV Cache 张量,返回一个 `LlaisysQwen2Model*`。 + +3. **`llaisysQwen2ModelWeights`**:返回内部持有的 `LlaisysQwen2Weights*`,好让 Python 侧能拿到每个 `llaisysTensor_t` 逐个调用 `tensorLoad` 把 safetensors 数据灌进去。 + +4. **`llaisysQwen2ModelInfer`**:核心推理函数,输入 `token_ids`(长度 `ntoken`),跑一次前向,返回**下一个 token 的 id**(`int64_t`)。内部要做的事情,對每一层(`nlayer` 层)循环: + - `embedding`:token id 查表得到输入向量(首次调用时对 `ntoken` 个 token 一起做;后续增量解码时 `ntoken` 通常是 1)。 + - `rms_norm`(pre-attn) → `linear` 算 Q/K/V(注意 QKV 带 bias) → `rope` 对 Q、K 做旋转位置编码 → 把新算出的 K/V 写入 KV Cache 对应位置 → `self_attention`(这里要传 `seqlen`=本次新增 token 数,`total_len`=写入后的缓存总长度,`nkvhead`=`meta.nkvh`,跟你在 Assignment #2 里实现的接口完全对应)→ `linear` 算输出投影 → 残差相加(`add`)。 + - `rms_norm`(pre-mlp) → `linear` 算 gate/up → `swiglu` → `linear` 算 down → 残差相加。 + - 最后一层结束后:`rms_norm`(`out_norm_w`)→ `linear`(`out_embed`)算 logits → `argmax` 取概率最大的 token id 返回。 + + `end_token`(`meta.end_token`,即 eos token id)由 Python 侧生成循环里判断是否停止,C 层只管返回下一个 token。 + +5. **`llaisysQwen2ModelDestroy`**:释放上面分配的所有资源。 + +6. 在 `xmake.lua` 里把新增的 `.cc` 文件纳入 `llaisys` target 的编译范围(检查现有 `add_files("src/llaisys/*.cc")` 这一行的 glob 模式)。 + +### 3.2 Python ctypes 包装层 + +新建 `python/llaisys/libllaisys/models.py`(参考同目录下 `ops.py`/`tensor.py` 的写法),需要: + +- 用 `ctypes.Structure` 把 `LlaisysQwen2Meta` 和 `LlaisysQwen2Weights` 各自定义一份对应的 Python 结构体(字段名、类型、顺序要跟 C 头文件严格一致,数组字段用 `POINTER(llaisysTensor_t)`)。 +- 给 `llaisysQwen2ModelCreate` / `llaisysQwen2ModelDestroy` / `llaisysQwen2ModelWeights` / `llaisysQwen2ModelInfer` 四个函数分别设置 `argtypes`/`restype`(照抄 `load_ops(lib)` 的模式,写一个 `load_models(lib)` 之类的函数)。 +- 在 `python/llaisys/libllaisys/__init__.py` 里调用这个新的 `load_xxx(LIB_LLAISYS)`(参照现有的 `load_ops(LIB_LLAISYS)` 那一行)。 + +### 3.3 Python 模型类:`python/llaisys/models/qwen2.py` + +**构造函数 `__init__`**(第一个 TODO): +1. 读 `model_path` 目录下的 `config.json`,填出一份 `LlaisysQwen2Meta`(`nlayer`←`num_hidden_layers`,`hs`←`hidden_size`,`nh`←`num_attention_heads`,`nkvh`←`num_key_value_heads`,`dh`←`hidden_size/num_attention_heads`,`di`←`intermediate_size`,`maxseq`←`max_position_embeddings`,`voc`←`vocab_size`,`epsilon`←`rms_norm_eps`,`theta`←`rope_theta`,`end_token`←`eos_token_id`)。DeepSeek-R1-Distill-Qwen-1.5B 具体数值见 `SELF_ATTENTION_ZH.md`(28 层、12 head、2 kv head、head_dim=128、theta=10000 等)。 +2. 调 `llaisysQwen2ModelCreate` 拿到底层模型指针,再调 `llaisysQwen2ModelWeights` 拿到权重张量的句柄。 +3. 遍历 `model_path` 下所有 `*.safetensors` 文件(现有代码已经在做这一步),把每个 tensor 名字映射到 `LlaisysQwen2Weights` 里对应的字段,再用 `tensorLoad` 把数据拷进去。命名映射(标准 HF Qwen2 checkpoint 命名): + + | safetensors 里的名字 | 对应字段 | + |---|---| + | `model.embed_tokens.weight` | `in_embed` | + | `lm_head.weight` | `out_embed`(注意 `tie_word_embeddings=false`,这两个不是同一份权重) | + | `model.norm.weight` | `out_norm_w` | + | `model.layers.{i}.input_layernorm.weight` | `attn_norm_w[i]` | + | `model.layers.{i}.self_attn.q_proj.weight` / `.bias` | `attn_q_w[i]` / `attn_q_b[i]` | + | `model.layers.{i}.self_attn.k_proj.weight` / `.bias` | `attn_k_w[i]` / `attn_k_b[i]` | + | `model.layers.{i}.self_attn.v_proj.weight` / `.bias` | `attn_v_w[i]` / `attn_v_b[i]` | + | `model.layers.{i}.self_attn.o_proj.weight` | `attn_o_w[i]`(无 bias) | + | `model.layers.{i}.post_attention_layernorm.weight` | `mlp_norm_w[i]` | + | `model.layers.{i}.mlp.gate_proj.weight` | `mlp_gate_w[i]` | + | `model.layers.{i}.mlp.up_proj.weight` | `mlp_up_w[i]` | + | `model.layers.{i}.mlp.down_proj.weight` | `mlp_down_w[i]` | + +**`generate` 函数**(第二个 TODO): +1. `inputs`(prompt 对应的 token id 列表)整体喂给 `llaisysQwen2ModelInfer` 做一次 **prefill**(`ntoken = len(inputs)`),拿到第一个新 token。 +2. 循环:每次只把**上一步生成的那一个 token** 喂给 `llaisysQwen2ModelInfer`(`ntoken = 1`,增量解码,依赖 KV Cache 记住之前的上下文),拿到下一个 token,直到生成满 `max_new_tokens` 个,或者遇到 `end_token` 提前停止。 +3. `--test` 模式下 `top_k=1/top_p=1.0/temperature=1.0` 等价于贪心采样,C 层的 `argmax` 已经是贪心的,所以这个模式下 Python 侧基本不需要额外采样逻辑;如果要支持非 `--test` 模式下的 `top_k`/`top_p`/`temperature` 采样,那部分逻辑需要在 Python 侧或者额外的 C 接口里处理(具体实现方式取决于你想把随机采样放在哪一层)。 + +## 4. KV Cache:为什么必须做 + +`README.md` 原文强调"不实现 KV Cache 模型会慢到不可用",原因和你在 Assignment #2 里实现 `self_attention` 时看到的参数是同一件事: + +- 没有 KV Cache:每生成一个新 token,都要把**从头到当前位置的所有 token** 重新过一遍 embedding → attention → mlp 全部 28 层,复杂度随生成长度呈平方增长。 +- 有 KV Cache:每层的 K、V 一旦算出来就存住,新 token 只需要算它自己的 Q/K/V,新的 K/V 追加进缓存,`self_attention` 用**全部缓存的 K/V**(`total_len` 个)去和**这一个新 token 的 Q**(`seqlen=1`)做注意力——这正是你已经实现好的 `causal_offset = total_len - seqlen` 逻辑在真实推理场景里的用法:prefill 阶段 `causal_offset=0`,增量解码阶段 `causal_offset` 等于已经生成的历史长度。 + +## 5. 调试建议 + +- 用 `tensor.debug()`(Python 侧,底层调 `tensorDebug`)在每一步(embedding 之后、每层 attention/mlp 之后、最终 logits 之前)打印张量数据,和 HuggingFace 模型在同样位置的中间结果对比,定位第一个数值开始偏离的地方,比逐字排查整个前向过程高效得多。 +- 建议先只跑 1~2 层、固定输入(比如全 0 或固定 token id),确认单层结果对得上,再扩展到全部层数,最后再接上 KV Cache 和多步生成。 + +## 6. 验收 + +```bash +python test/test_infer.py --model <本地模型目录> --test +``` + +`--test` 模式下用贪心解码(`top_k=1`),要求 `llaisys` 生成的 `tokens` 和 HuggingFace `transformers` 生成的 `tokens` 完全一致(`assert llaisys_tokens == tokens`),打印 `Test passed!` 即算通过。之后按 README 里的说明 commit + push,CI 里 `Assignment-3` 那一步(`python test/test_infer.py --test`)应该跑绿——但注意 CI 环境要能拿到模型权重(`load_hf_model` 在没传 `--model` 时会用 `huggingface_hub.snapshot_download` 自动下载),如果 CI 里没联网或没配好缓存,这一步可能需要额外配置。 diff --git a/docs/SELF_ATTENTION_ZH.md b/docs/SELF_ATTENTION_ZH.md new file mode 100644 index 000000000..1e798b13a --- /dev/null +++ b/docs/SELF_ATTENTION_ZH.md @@ -0,0 +1,152 @@ +# `self_attention` 算子实现过程记录 + +> 本文记录 `src/ops/self_attention/` 从空实现到完整实现(V1 → V4)的完整迭代过程,包含每一版的设计动机、核心代码、遇到的问题和排查过程。已对照当前仓库源码(`src/ops/self_attention/**`)核对。 + +## 1. 目标与整体接口 + +实现的是标准的因果(causal)、支持 GQA 的 scaled dot-product self-attention: + +``` +attn_val = softmax(scale * q @ k^T + causal_mask) @ v +``` + +最终的张量形状约定(`op.cpp` 从 tensor shape 里推导,传给 CPU kernel): + +| 张量 | 形状 | 含义 | +|---|---|---| +| `q` | `[seqlen, nhead, d]` | 本次新增的 query | +| `k`, `v` | `[total_len, nkvhead, d]` / `[total_len, nkvhead, dv]` | 含 KV Cache 历史在内的完整 key/value | +| `attn_val` | `[seqlen, nhead, dv]` | 输出 | + +其中 `nhead` 必须是 `nkvhead` 的整数倍(GQA),`causal_offset = total_len - seqlen` 表示 KV Cache 里历史 token 的长度,query 位置 `i` 能看到的 key 范围是 `j <= i + causal_offset`。 + +最终 CPU kernel 接口(`cpu/self_attention.hpp`): + +```cpp +namespace llaisys::ops::cpu { +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); +} +``` + +## 2. 迭代路线:V1 → V4 + +没有一次写全,而是刻意拆成四个版本,每一版只加一个新概念,便于逐步验证: + +| 版本 | 新增内容 | 循环方式 | 是否过测试 | +|---|---|---|---| +| V1 | 最朴素实现 | 单遍扫描,直接 `exp(score)` | 部分(仅 `nh==nkvh` 且数值范围温和时) | +| V2 | + 因果掩码 | 单遍扫描,`j` 范围收窄到 `[0, i+causal_offset]` | `nh==nkvh` 用例全过 | +| V3 | + GQA | 单遍扫描,`k`/`v` 按 `kvh = h/group` 取 | 全部用例过 | +| V4 | + 数值稳定 softmax | 两遍扫描(先求 max,再减 max 求 exp) | 全部用例过,且对极端数值更鲁棒 | + +### V1:最简单版本 + +- **设计取舍**:不支持 GQA(假设 `nhead == nkvhead`)、不做因果掩码(每个 `i` 都看全部 `total_len` 个 key)、不做 softmax 数值稳定处理(直接 `exp(score)`,不减 max)。明确"这一版不需要过测试"。 +- **核心逻辑**(单遍扫描,打分、`exp`、累加一次做完): + +```cpp +float sum_exp = 0.0f; +std::vector acc(dv, 0.0f); +for (size_t j = 0; j < total_len; j++) { + float score = 0.0f; + for (size_t dim = 0; dim < d; dim++) { + score += cast(q[...]) * cast(k[...]); + } + score *= scale; + float e = std::exp(score); + sum_exp += e; + for (size_t t = 0; t < dv; t++) { + acc[t] += e * cast(v[...]); + } +} +for (size_t t = 0; t < dv; t++) { + attn_val[...] = cast(acc[t] / sum_exp); +} +``` + +- **`acc` 的含义**:不是 QKᵀ 矩阵,而是 softmax 权重对 V 的加权和(输出的分子部分),长度是 `dv` 而不是 `total_len`。 + +### V2:加因果掩码 + +- 只改一处:`j` 的循环范围从 `[0, total_len)` 收窄成 `[0, i + causal_offset]`(闭区间),其余打分/exp/累加/归一化逻辑不变。 +- `causal_offset = total_len - seqlen`,含义是 KV Cache 里历史 token 的长度。 +- 验证方式:第一组测试形状 `(qlen=2, kvlen=2, nh=1, nkvh=1, hd=4)` 里,`i=1` 时两种实现(有/无掩码)结果本来就一样(因为 `total_len` 刚好等于看得到的范围),但 `i=0` 时不一样——这正好是验证掩码是否生效的天然信号。 + +### V3:加 GQA(Grouped-Query Attention) + +- `op.cpp` 的校验从 `nhead == nkvhead` 放宽为 `nhead % nkvhead == 0`。 +- kernel 里新增 `group = nhead / nkvhead`,每个 query head `h` 对应的 kv head 是 `kvh = h / group`。 +- 把原来所有 `k`/`v` 下标里的 `nhead`/`h` 换成 `nkvhead`/`kvh`(`q`、`attn_val` 的下标不用换,它们本来就按 `nhead` 排列)。 + +```cpp +size_t kvh = h / group; +... +score += cast(q[i*nhead*d + h*d + dim]) * cast(k[j*nkvhead*d + kvh*d + dim]); +... +acc[t] += e * cast(v[j*nkvhead*dv + kvh*dv + t]); +``` + +### V4:数值稳定 softmax + +- **动机**:直接 `exp(score)` 在 `score` 较大时会溢出成 `inf`,进而导致 `sum_exp` 变成 `inf`、最终输出变成 `NaN`。数学上 softmax 对所有输入减去任意常数结果不变(`exp(c)` 项在分子分母间约掉),所以可以安全地统一减去 `max_score`,让指数的输入恒 `<= 0`,分子的每一项被压缩到 `(0, 1]`,`sum_exp` 也因此有下界(至少为 1),彻底避免 `inf`/`NaN`。 +- **代价**:单遍扫描不再够用。第二步 `exp(score - max_score)` 依赖一个要扫完整个 `j` 范围才能确定的全局量 `max_score`,所以必须拆成两遍;而两遍之间要把每个 `score[j]` 重新用上,于是引入 `std::vector scores` 把第一遍算出的打分存下来,避免第二遍重新做一次 `O(d)` 的点积(空间换时间)。 + +```cpp +// 第一遍:算出所有 score[j] 并记录最大值 max_score。 +std::vector scores(limit + 1); +float max_score = -std::numeric_limits::infinity(); +for (size_t j = 0; j <= limit; j++) { + float score = 0.0f; + for (size_t dim = 0; dim < d; dim++) { + score += cast(q[...]) * cast(k[...]); + } + score *= scale; + scores[j] = score; + max_score = std::max(max_score, score); +} + +float sum_exp = 0.0f; +std::vector acc(dv, 0.0f); +// 第二遍:用 exp(score - max_score) 做数值稳定的 softmax 加权求和。 +for (size_t j = 0; j <= limit; j++) { + float e = std::exp(scores[j] - max_score); + sum_exp += e; + for (size_t t = 0; t < dv; t++) { + acc[t] += e * cast(v[...]); + } +} +``` + +## 3. 过程中排查过的问题 + +实现过程中不是一路顺风,记录几个有代表性的坑: + +1. **全角标点导致语法错误**:注释里混入了全角逗号 `、` 和被意外换行拆断的 `//acc` → `/` + `acc[t]...`,编译器把残余的 `/` 当除号解析,报 `expected primary-expression before '/' token`。这类问题的特征是报错行号和实际"看起来有问题"的代码对不上,需要用 `cat -A` 或 `sed -n` 看原始字节才能定位。 + +2. **`-Werror=unused-variable` / `-Werror=sign-compare`**:项目开了 `-Werror`,每次先加一个后面才会用到的变量(比如 V2 的 `causal_offset`、V3 的 `group`、V4 的 `max_score`)都需要先 `(void)var;` 占位,否则编译不过;循环变量声明成 `int` 和 `size_t` 比较也会报错,统一改成 `size_t`。 + +3. **`.so` 没同步导致的假报错**:改完代码跑 Python 测试,报错信息(`Shapes mismatch`)和当前源码对不上号。排查后发现是两层缓存没刷新: + - `xmake build` 编译产物在 `build/`,需要额外 `xmake install` 才会同步到 `python/llaisys/libllaisys/libllaisys.so`; + - 更隐蔽的是,`test/ops/self_attention.py` 用的 `sys.path.insert` 加的是 `test/` 目录而不是 `python/` 目录,实际 `import llaisys` 落到了之前 `pip install ./python/` 装到 conda `site-packages` 里的旧副本,而不是仓库本地这份。用 `gdb -batch -ex "catch throw" ... --args python ...` 打断点看堆栈,才在帧里看到 `.so` 的真实加载路径,定位到问题。 + - 修复方式:改完代码后固定跑 `xmake build && xmake install && pip install ./python/`,确保三层(编译产物 / xmake install 目标 / pip 包)都同步。 + +4. **半成品 TODO 被当成代码**:自己实现两遍扫描版本时,把我给的 TODO 提示文字(`max_score = std::max(max_score, scores[j]); // 这里是伪代码...`)直接留在了代码里,`j` 没有对应的循环声明,编译报 `'j' was not declared in this scope`。提醒:TODO 注释只是提示,真正要写一个完整的 `for` 循环替换掉它。 + +## 4. 最终测试结果 + +``` +qlen=2 kvlen=2 nh=1 nkvh=1 hd=4 f32 / f16 / bf16 ✓ +qlen=5 kvlen=11 nh=4 nkvh=2 hd=8 f32 / f16 / bf16 ✓ +Test passed! +``` + +覆盖了 `nh==nkvh`(无分组)和 `nh=4,nkvh=2`(2 倍分组)两种 GQA 配置,三种精度(`f32`/`f16`/`bf16`)全部通过 `test/ops/self_attention.py`。 + +## 5. 涉及的文件 + +- `src/ops/self_attention/op.hpp` / `op.cpp`:对外接口,做形状校验、从 tensor shape 推导出 `seqlen/total_len/nhead/nkvhead/d/dv`,按 device 分发。 +- `src/ops/self_attention/cpu/self_attention.hpp` / `self_attention.cpp`:CPU kernel,按 dtype(`f32`/`bf16`/`f16`)分发到同一个模板函数 `self_attention_`。 +- `test/ops/self_attention.py`:用 PyTorch 的 `scaled_dot_product_attention` 等价实现做基准,对比 `llaisys.Ops.self_attention` 的输出。 diff --git a/docs/qwen2_infer_reference.cc b/docs/qwen2_infer_reference.cc new file mode 100644 index 000000000..358d52007 --- /dev/null +++ b/docs/qwen2_infer_reference.cc @@ -0,0 +1,129 @@ +// Reference implementation of llaisysQwen2ModelInfer, kept here (outside +// src/, so xmake never compiles it) for the user to check their own +// from-scratch implementation against. Not part of the build. +// +// See src/llaisys/models/qwen2.cc for the live file the user is writing. + +#include "llaisys/models/qwen2.h" // 声明 +#include "../../src/ops/add/op.hpp" +#include "../../src/ops/argmax/op.hpp" +#include "../../src/ops/embedding/op.hpp" +#include "../../src/ops/linear/op.hpp" +#include "../../src/ops/rms_norm/op.hpp" +#include "../../src/ops/rope/op.hpp" +#include "../../src/ops/self_attention/op.hpp" +#include "../../src/ops/swiglu/op.hpp" +#include "../../src/tensor/tensor.hpp" +#include "../../src/llaisys/llaisys_tensor.hpp" + +#include +#include + +int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken) { + using llaisys::Tensor; + using llaisys::tensor_t; + namespace ops = llaisys::ops; + + const auto &meta = model->meta; + auto device = model->device; + int dev_id = model->device_ids.empty() ? 0 : model->device_ids[0]; + auto dtype = meta.dtype; + + size_t cur_len = model->cur_len; + size_t total_len = cur_len + ntoken; + + // token ids -> device tensor + tensor_t idx = Tensor::create({ntoken}, LLAISYS_DTYPE_I64, device, dev_id); + idx->load(token_ids); + + // position ids for the newly-fed tokens: [cur_len, cur_len+1, ...] + std::vector pos_host(ntoken); + for (size_t i = 0; i < ntoken; i++) { + pos_host[i] = static_cast(cur_len + i); + } + tensor_t pos = Tensor::create({ntoken}, LLAISYS_DTYPE_I64, device, dev_id); + pos->load(pos_host.data()); + + tensor_t x = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::embedding(x, idx, model->weights.in_embed->tensor); + + float scale = 1.0f / std::sqrt(static_cast(meta.dh)); + + for (size_t l = 0; l < meta.nlayer; l++) { + tensor_t normed = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::rms_norm(normed, x, model->attn_norm_w[l]->tensor, meta.epsilon); + + tensor_t q = Tensor::create({ntoken, meta.nh * meta.dh}, dtype, device, dev_id); + ops::linear(q, normed, model->attn_q_w[l]->tensor, model->attn_q_b[l]->tensor); + tensor_t k = Tensor::create({ntoken, meta.nkvh * meta.dh}, dtype, device, dev_id); + ops::linear(k, normed, model->attn_k_w[l]->tensor, model->attn_k_b[l]->tensor); + tensor_t v = Tensor::create({ntoken, meta.nkvh * meta.dh}, dtype, device, dev_id); + ops::linear(v, normed, model->attn_v_w[l]->tensor, model->attn_v_b[l]->tensor); + + tensor_t q3 = q->view({ntoken, meta.nh, meta.dh}); + tensor_t k3 = k->view({ntoken, meta.nkvh, meta.dh}); + tensor_t v3 = v->view({ntoken, meta.nkvh, meta.dh}); + + tensor_t q_rope = Tensor::create({ntoken, meta.nh, meta.dh}, dtype, device, dev_id); + ops::rope(q_rope, q3, pos, meta.theta); + tensor_t k_rope = Tensor::create({ntoken, meta.nkvh, meta.dh}, dtype, device, dev_id); + ops::rope(k_rope, k3, pos, meta.theta); + + // Both sides are always fully-contiguous [ntoken, nkvh, dh] buffers + // (a slice of the leading dim of a contiguous cache tensor stays + // contiguous), so a flat memcpy stands in for a rearrange copy here + // (ops::rearrange itself is still a TO_BE_IMPLEMENTED stub). + tensor_t k_cache_slice = model->k_cache[l]->slice(0, cur_len, total_len); + tensor_t v_cache_slice = model->v_cache[l]->slice(0, cur_len, total_len); + std::memcpy(k_cache_slice->data(), k_rope->data(), k_rope->numel() * k_rope->elementSize()); + std::memcpy(v_cache_slice->data(), v3->data(), v3->numel() * v3->elementSize()); + + tensor_t k_all = model->k_cache[l]->slice(0, 0, total_len); + tensor_t v_all = model->v_cache[l]->slice(0, 0, total_len); + + tensor_t attn_out = Tensor::create({ntoken, meta.nh, meta.dh}, dtype, device, dev_id); + ops::self_attention(attn_out, q_rope, k_all, v_all, scale); + + tensor_t attn_out2 = attn_out->view({ntoken, meta.nh * meta.dh}); + tensor_t o = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::linear(o, attn_out2, model->attn_o_w[l]->tensor, nullptr); + + tensor_t x_attn = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::add(x_attn, x, o); + x = x_attn; + + tensor_t normed2 = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::rms_norm(normed2, x, model->mlp_norm_w[l]->tensor, meta.epsilon); + + tensor_t gate = Tensor::create({ntoken, meta.di}, dtype, device, dev_id); + ops::linear(gate, normed2, model->mlp_gate_w[l]->tensor, nullptr); + tensor_t up = Tensor::create({ntoken, meta.di}, dtype, device, dev_id); + ops::linear(up, normed2, model->mlp_up_w[l]->tensor, nullptr); + tensor_t swiglu_out = Tensor::create({ntoken, meta.di}, dtype, device, dev_id); + ops::swiglu(swiglu_out, gate, up); + tensor_t down = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::linear(down, swiglu_out, model->mlp_down_w[l]->tensor, nullptr); + + tensor_t x_mlp = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::add(x_mlp, x, down); + x = x_mlp; + } + + tensor_t final_normed = Tensor::create({ntoken, meta.hs}, dtype, device, dev_id); + ops::rms_norm(final_normed, x, model->weights.out_norm_w->tensor, meta.epsilon); + + tensor_t logits = Tensor::create({ntoken, meta.voc}, dtype, device, dev_id); + ops::linear(logits, final_normed, model->weights.out_embed->tensor, nullptr); + + tensor_t last_logits = logits->slice(0, ntoken - 1, ntoken)->view({meta.voc}); + + tensor_t max_idx = Tensor::create({1}, LLAISYS_DTYPE_I64, device, dev_id); + tensor_t max_val = Tensor::create({1}, dtype, device, dev_id); + ops::argmax(max_idx, max_val, last_logits); + + int64_t next_token = 0; + std::memcpy(&next_token, max_idx->data(), sizeof(int64_t)); + + model->cur_len = total_len; + return next_token; +} diff --git a/docs/qwen2_model_reference.py b/docs/qwen2_model_reference.py new file mode 100644 index 000000000..e1fe8116b --- /dev/null +++ b/docs/qwen2_model_reference.py @@ -0,0 +1,143 @@ +from typing import Sequence +from ..libllaisys import LIB_LLAISYS +from ..libllaisys import DeviceType +from ..libllaisys import DataType +from ..libllaisys import LlaisysQwen2Meta + +from pathlib import Path +import ctypes +import json +import re +import safetensors + +# Cap on how many tokens (prompt + generated) a single model instance can +# hold in its KV cache. Real Qwen2 configs advertise max_position_embeddings +# in the tens of thousands, which would allocate a KV cache far larger than +# needed for a short prompt + a few hundred generated tokens. +_MAX_SEQ_LEN = 4096 + +_HF_DTYPE_TO_LLAISYS = { + "bfloat16": DataType.BF16, + "float16": DataType.F16, + "float32": DataType.F32, +} + +# safetensors name suffix (after "model.layers.{i}.") -> LlaisysQwen2Weights field name +_LAYER_FIELD_MAP = { + "input_layernorm.weight": "attn_norm_w", + "self_attn.q_proj.weight": "attn_q_w", + "self_attn.q_proj.bias": "attn_q_b", + "self_attn.k_proj.weight": "attn_k_w", + "self_attn.k_proj.bias": "attn_k_b", + "self_attn.v_proj.weight": "attn_v_w", + "self_attn.v_proj.bias": "attn_v_b", + "self_attn.o_proj.weight": "attn_o_w", + "post_attention_layernorm.weight": "mlp_norm_w", + "mlp.gate_proj.weight": "mlp_gate_w", + "mlp.up_proj.weight": "mlp_up_w", + "mlp.down_proj.weight": "mlp_down_w", +} + +_LAYER_NAME_RE = re.compile(r"model\.layers\.(\d+)\.(.+)") + + +class Qwen2: + + def __init__(self, model_path, device: DeviceType = DeviceType.CPU): + model_path = Path(model_path) + + with open(model_path / "config.json", "r") as f: + config = json.load(f) + + hidden_size = config["hidden_size"] + num_attention_heads = config["num_attention_heads"] + head_dim = config.get("head_dim", hidden_size // num_attention_heads) + dtype = _HF_DTYPE_TO_LLAISYS.get(config.get("torch_dtype", "float32"), DataType.F32) + + eos_token_id = config["eos_token_id"] + self._end_token = eos_token_id[0] if isinstance(eos_token_id, list) else eos_token_id + + self._meta = LlaisysQwen2Meta( + dtype=dtype, + nlayer=config["num_hidden_layers"], + hs=hidden_size, + nh=num_attention_heads, + nkvh=config["num_key_value_heads"], + dh=head_dim, + di=config["intermediate_size"], + maxseq=min(config.get("max_position_embeddings", _MAX_SEQ_LEN), _MAX_SEQ_LEN), + voc=config["vocab_size"], + epsilon=config["rms_norm_eps"], + theta=config["rope_theta"], + end_token=self._end_token, + ) + + device_ids = (ctypes.c_int * 1)(0) + self._model = LIB_LLAISYS.llaisysQwen2ModelCreate( + ctypes.byref(self._meta), ctypes.c_int(device), device_ids, ctypes.c_int(1) + ) + weights = LIB_LLAISYS.llaisysQwen2ModelWeights(self._model).contents + + def load(handle, tensor): + tensor = tensor.contiguous() + LIB_LLAISYS.tensorLoad(handle, ctypes.c_void_p(tensor.data_ptr())) + + for file in sorted(model_path.glob("*.safetensors")): + data_ = safetensors.safe_open(file, framework="pt", device="cpu") + for name_ in data_.keys(): + tensor = data_.get_tensor(name_) + + if name_ == "model.embed_tokens.weight": + load(weights.in_embed, tensor) + continue + if name_ == "lm_head.weight": + load(weights.out_embed, tensor) + continue + if name_ == "model.norm.weight": + load(weights.out_norm_w, tensor) + continue + + m = _LAYER_NAME_RE.match(name_) + if not m: + continue + layer_idx, suffix = int(m.group(1)), m.group(2) + field = _LAYER_FIELD_MAP.get(suffix) + if field is None: + continue + load(getattr(weights, field)[layer_idx], tensor) + + def __del__(self): + if getattr(self, "_model", None): + LIB_LLAISYS.llaisysQwen2ModelDestroy(self._model) + self._model = None + + def generate( + self, + inputs: Sequence[int], + max_new_tokens: int = None, + top_k: int = 1, + top_p: float = 0.8, + temperature: float = 0.8, + ): + if max_new_tokens is None: + max_new_tokens = 128 + + prompt_tokens = list(inputs) + result_tokens = list(prompt_tokens) + + token_array = (ctypes.c_int64 * len(prompt_tokens))(*prompt_tokens) + next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( + self._model, token_array, ctypes.c_size_t(len(prompt_tokens)) + ) + result_tokens.append(next_token) + + steps = 1 + while next_token != self._end_token and steps < max_new_tokens: + token_array = (ctypes.c_int64 * 1)(next_token) + next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( + self._model, token_array, ctypes.c_size_t(1) + ) + result_tokens.append(next_token) + steps += 1 + + return result_tokens diff --git a/docs/qwen2_python_bindings_reference.py b/docs/qwen2_python_bindings_reference.py new file mode 100644 index 000000000..3c2b0a78f --- /dev/null +++ b/docs/qwen2_python_bindings_reference.py @@ -0,0 +1,69 @@ +import ctypes +from ctypes import POINTER, Structure, c_float, c_int64, c_size_t, c_int, c_void_p + +from .tensor import llaisysTensor_t +from .llaisys_types import llaisysDataType_t, llaisysDeviceType_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), + ("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)), + ] + + +# Opaque handle to the C++-side LlaisysQwen2Model. +llaisysQwen2Model_t = c_void_p + + +def load_models(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.llaisysQwen2ModelWeights.argtypes = [llaisysQwen2Model_t] + lib.llaisysQwen2ModelWeights.restype = POINTER(LlaisysQwen2Weights) + + lib.llaisysQwen2ModelInfer.argtypes = [ + llaisysQwen2Model_t, + POINTER(c_int64), + c_size_t, + ] + lib.llaisysQwen2ModelInfer.restype = c_int64 diff --git a/include/llaisys.h b/include/llaisys.h index 73ca7eead..6eab75b62 100644 --- a/include/llaisys.h +++ b/include/llaisys.h @@ -22,8 +22,9 @@ // Device Types typedef enum { LLAISYS_DEVICE_CPU = 0, - //// TODO: Add more device types here. Numbers need to be consecutive. + // 新增设备类型编号需连续 LLAISYS_DEVICE_NVIDIA = 1, + LLAISYS_DEVICE_ILUVATAR = 2, LLAISYS_DEVICE_TYPE_COUNT } llaisysDeviceType_t; diff --git a/include/llaisys/tensor.h b/include/llaisys/tensor.h index 76f13fbc3..72660307c 100644 --- a/include/llaisys/tensor.h +++ b/include/llaisys/tensor.h @@ -63,6 +63,11 @@ __C { size_t dim, size_t start, size_t end); + + __export llaisysTensor_t tensorReshape( + llaisysTensor_t tensor, + size_t * shape, + size_t ndim); } #endif // LLAISYS_TENSOR_H diff --git a/python/llaisys/libllaisys/__init__.py b/python/llaisys/libllaisys/__init__.py index f536fb527..417fb9f19 100644 --- a/python/llaisys/libllaisys/__init__.py +++ b/python/llaisys/libllaisys/__init__.py @@ -12,8 +12,8 @@ from .tensor import llaisysTensor_t from .tensor import load_tensor from .ops import load_ops - - +from .model import load_model +from .model import LlaisysQwen2Meta,LlaisysQwen2Weights,LlaisysQwen2Model def load_shared_library(): lib_dir = Path(__file__).parent @@ -38,7 +38,7 @@ def load_shared_library(): load_runtime(LIB_LLAISYS) load_tensor(LIB_LLAISYS) load_ops(LIB_LLAISYS) - +load_model(LIB_LLAISYS) __all__ = [ "LIB_LLAISYS", @@ -52,4 +52,7 @@ def load_shared_library(): "llaisysMemcpyKind_t", "MemcpyKind", "llaisysStream_t", + "LlaisysQwen2Meta", + "LlaisysQwen2Weights", + "LlaisysQwen2Model" ] diff --git a/python/llaisys/libllaisys/llaisys_types.py b/python/llaisys/libllaisys/llaisys_types.py index c5a0b4679..84c761b73 100644 --- a/python/llaisys/libllaisys/llaisys_types.py +++ b/python/llaisys/libllaisys/llaisys_types.py @@ -6,7 +6,8 @@ class DeviceType(IntEnum): CPU = 0 NVIDIA = 1 - COUNT = 2 + ILUVATAR = 2 + COUNT = 3 llaisysDeviceType_t = ctypes.c_int diff --git a/python/llaisys/libllaisys/model.py b/python/llaisys/libllaisys/model.py new file mode 100644 index 000000000..e031d00cd --- /dev/null +++ b/python/llaisys/libllaisys/model.py @@ -0,0 +1,53 @@ +import ctypes +from ctypes import c_size_t, c_int, c_int64,POINTER +from .llaisys_types import llaisysDataType_t,llaisysDeviceType_t +from .tensor import llaisysTensor_t +class LlaisysQwen2Meta(ctypes.Structure): + _fields_ = [ + ("dtype", llaisysDataType_t), + ("nlayer", ctypes.c_size_t), + ("hs", ctypes.c_size_t), + ("nh", ctypes.c_size_t), + ("nkvh", ctypes.c_size_t), + ("dh", ctypes.c_size_t), + ("di", ctypes.c_size_t), + ("maxseq", ctypes.c_size_t), + ("voc", ctypes.c_size_t), + ("epsilon", ctypes.c_float), + ("theta", ctypes.c_float), + ("end_token", ctypes.c_int64), + ] + +LlaisysQwen2Model = ctypes.c_void_p + +class LlaisysQwen2Weights(ctypes.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)), + ] +def load_model(lib): + lib.llaisysQwen2ModelCreate.argtypes = [POINTER(LlaisysQwen2Meta), llaisysDeviceType_t, POINTER(c_int), c_int] + lib.llaisysQwen2ModelCreate.restype = LlaisysQwen2Model + + lib.llaisysQwen2ModelDestroy.argtypes = [LlaisysQwen2Model] + lib.llaisysQwen2ModelDestroy.restype = None + + lib.llaisysQwen2ModelWeights.argtypes = [LlaisysQwen2Model] + lib.llaisysQwen2ModelWeights.restype = POINTER(LlaisysQwen2Weights) + + lib.llaisysQwen2ModelInfer.argtypes = [LlaisysQwen2Model, POINTER(c_int64) , c_size_t ] + lib.llaisysQwen2ModelInfer.restype = c_int64 + diff --git a/python/llaisys/libllaisys/tensor.py b/python/llaisys/libllaisys/tensor.py index b58057883..21649339f 100644 --- a/python/llaisys/libllaisys/tensor.py +++ b/python/llaisys/libllaisys/tensor.py @@ -15,60 +15,45 @@ def load_tensor(lib): ] 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 @@ -76,3 +61,6 @@ def load_tensor(lib): c_size_t, # end : exclusive ] lib.tensorSlice.restype = llaisysTensor_t + + lib.tensorReshape.argtypes = [llaisysTensor_t, POINTER(c_size_t), c_size_t] + lib.tensorReshape.restype = llaisysTensor_t diff --git a/python/llaisys/models/qwen2.py b/python/llaisys/models/qwen2.py index 0d07b0b21..d6d4078c1 100644 --- a/python/llaisys/models/qwen2.py +++ b/python/llaisys/models/qwen2.py @@ -1,23 +1,91 @@ from typing import Sequence + +import contextlib +import ctypes +from ..libllaisys.llaisys_types import DataType from ..libllaisys import LIB_LLAISYS from ..libllaisys import DeviceType - +import json from pathlib import Path +import ml_dtypes # noqa: F401 -- import registers numpy's bfloat16 dtype, needed before safetensors reads BF16 weights import safetensors +from ..libllaisys import LlaisysQwen2Meta, LlaisysQwen2Weights, LlaisysQwen2Model class Qwen2: def __init__(self, model_path, device: DeviceType = DeviceType.CPU): - # TODO: Implement model constructor - model_path = Path(model_path) + config_path = model_path / "config.json" + with open(config_path, "r") as f: + config = json.load(f) + self.meta = LlaisysQwen2Meta( + dtype=DataType.BF16, + nlayer=config["num_hidden_layers"], + hs=config["hidden_size"], + nh=config["num_attention_heads"], + nkvh=config["num_key_value_heads"], + dh=config["hidden_size"] // config["num_attention_heads"], + di=config["intermediate_size"], + maxseq=2048, # 上限，别用 config 里的 max_position_embeddings（太大，内存装不下） + voc=config["vocab_size"], + epsilon=config["rms_norm_eps"], + theta=config["rope_theta"], + end_token=config["eos_token_id"], + ) + + device_ids = (ctypes.c_int * 1)(0) # 单卡 + self.model = LIB_LLAISYS.llaisysQwen2ModelCreate( + ctypes.byref(self.meta), device.value, device_ids, 1 + ) + self.weights = LIB_LLAISYS.llaisysQwen2ModelWeights(self.model) + + w = self.weights.contents + + # 完整 safetensors 名 -> w 上的属性名 + GLOBAL_MAP = { + "model.embed_tokens.weight": "in_embed", + "lm_head.weight": "out_embed", + "model.norm.weight": "out_norm_w", + } + + # 去掉 "model.layers.{i}." 前缀后的 key -> w 上的属性名（数组，按 layer 索引） + LAYER_MAP = { + "input_layernorm.weight": "attn_norm_w", + "self_attn.q_proj.weight": "attn_q_w", + "self_attn.q_proj.bias": "attn_q_b", + "self_attn.k_proj.weight": "attn_k_w", + "self_attn.k_proj.bias": "attn_k_b", + "self_attn.v_proj.weight": "attn_v_w", + "self_attn.v_proj.bias": "attn_v_b", + "self_attn.o_proj.weight": "attn_o_w", + # Attention 后、MLP 前的 RMSNorm + "post_attention_layernorm.weight": "mlp_norm_w", + # MLP + "mlp.gate_proj.weight": "mlp_gate_w", + "mlp.up_proj.weight": "mlp_up_w", + "mlp.down_proj.weight": "mlp_down_w", + } 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 + arr = data_.get_tensor(name_) + # TODO: 加载前没校验 arr.dtype/contiguity 是否匹配 C++ 侧分配的 meta.dtype， + # 目前是直接信任这次真实权重跑通了的 raw memcpy。 + ptr = arr.ctypes.data_as(ctypes.c_void_p) + + if name_.startswith("model.layers."): + _ ,_, remain = name_.partition("model.layers.") + layer_idx,_,key = remain.partition(".") + layer_idx = int(layer_idx) + attr = LAYER_MAP[key] + handle = getattr(w, attr)[layer_idx] + else: + attr = GLOBAL_MAP[name_] + handle = getattr(w, attr) + + LIB_LLAISYS.tensorLoad(handle, ptr) def generate( self, @@ -26,8 +94,29 @@ def generate( top_k: int = 1, top_p: float = 0.8, temperature: float = 0.8, + step_context=None, ): + if max_new_tokens is None: + max_new_tokens = 1024 + if len(inputs) == 0: + raise ValueError("inputs cannot be empty") + generated_tokens = list(inputs) - # TODO: Implement generate function + # prefill/decode 共用一个循环：第一轮 current_input 是整个 prompt（prefill）， + # 之后每轮只喂上一步生成的那个 token（decode）。 + current_input = generated_tokens + for step in range(max_new_tokens): + token_array = (ctypes.c_int64 * len(current_input))(*current_input) + # 让调用方（如 benchmark 脚本）包一层 context 分别计时 prefill/decode， + # 而不需要 generate() 本身依赖任何 profiling 库。 + ctx = step_context(step, step == 0) if step_context else contextlib.nullcontext() + with ctx: + next_token = LIB_LLAISYS.llaisysQwen2ModelInfer( + self.model, token_array, len(current_input) + ) + generated_tokens.append(next_token) + if next_token == self.meta.end_token: + break + current_input = [next_token] - return [] + return generated_tokens diff --git a/python/llaisys/ops.py b/python/llaisys/ops.py index ed0180bc8..7274c1294 100644 --- a/python/llaisys/ops.py +++ b/python/llaisys/ops.py @@ -19,9 +19,12 @@ 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: Tensor = 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 diff --git a/python/llaisys/tensor.py b/python/llaisys/tensor.py index 1466d851e..4deca8361 100644 --- a/python/llaisys/tensor.py +++ b/python/llaisys/tensor.py @@ -95,3 +95,9 @@ def slice(self, dim: int, start: int, end: int): self._tensor, c_size_t(dim), c_size_t(start), c_size_t(end) ) ) + + def reshape(self, *shape: int) -> llaisysTensor_t: + _shape = (c_size_t * len(shape))(*shape) + return Tensor( + tensor=LIB_LLAISYS.tensorReshape(self._tensor, _shape, c_size_t(len(shape))) + ) diff --git a/python/setup.cfg b/python/setup.cfg index b35fc65f7..6e753b6a2 100644 --- a/python/setup.cfg +++ b/python/setup.cfg @@ -13,6 +13,8 @@ install_requires = torch>=2.4.0 transformers accelerate + safetensors + ml_dtypes [options.package_data] llaisys = diff --git a/scripts/iluvatar/bootstrap_and_verify.sh b/scripts/iluvatar/bootstrap_and_verify.sh new file mode 100755 index 000000000..87f0a0cce --- /dev/null +++ b/scripts/iluvatar/bootstrap_and_verify.sh @@ -0,0 +1,156 @@ +#!/usr/bin/env bash +# 全新天数（Iluvatar）实例的一键 clone + build + 验证脚本。 +# 前提：跟上次用的是同一个课程镜像/模板（corex SDK、天数版 PyTorch 已经预装好， +# 只是这次是全新开的机器，代码、模型权重、xmake 都没有）。 +# +# 用法（在新实例的任意目录下）： +# curl -fsSL https://gitee.com/tang-jinchi/llaisys/raw/main/scripts/iluvatar/bootstrap_and_verify.sh -o bootstrap_and_verify.sh +# bash bootstrap_and_verify.sh +# 或者先手动 clone 仓库、cd 进去，再 bash scripts/iluvatar/bootstrap_and_verify.sh +# +# 把完整输出贴回来。 + +set -u +REPO_URL="https://gitee.com/tang-jinchi/llaisys.git" +REPO_DIR="$HOME/llaisys" + +echo "====================================================" +echo "[0] 环境自检：确认这是预期的天数镜像" +echo "====================================================" +echo "-- clang++ (真正的编译器，不是那个假 nvcc) --" +/usr/local/corex/bin/clang++ --version 2>&1 +echo +echo "-- corex SDK 版本目录 --" +ls -d /usr/local/corex-*/ 2>&1 +echo +echo "-- 天数版 PyTorch（应该已经预装，脚本不会碰它）--" +python3 -c "import torch; print('torch', torch.__version__); print('torch.cuda.is_available():', torch.cuda.is_available())" 2>&1 + +echo +echo "====================================================" +echo "[1] xmake（这次是全新机器，大概率没装）" +echo "====================================================" +export XMAKE_ROOT=y +if command -v xmake >/dev/null 2>&1; then + echo "xmake 已存在: $(xmake --version | head -1)" +else + echo "安装 xmake（注意不带 --branch 参数，装脚本对这个 flag 解析有 bug）..." + curl -fsSL https://xmake.io/shget.text | bash + export PATH="$HOME/.local/bin:$PATH" + echo "'export XMAKE_ROOT=y' 和 PATH 已在这次 shell 里生效；" \ + "如果重开一个 shell，记得先手动 export 一遍或者把它们写进 ~/.bashrc。" +fi +xmake --version 2>&1 + +echo +echo "====================================================" +echo "[2] clone / 更新仓库到 $REPO_DIR" +echo "====================================================" +if [ -d "$REPO_DIR/.git" ]; then + echo "仓库已存在，git pull" + (cd "$REPO_DIR" && git pull 2>&1) +else + git clone "$REPO_URL" "$REPO_DIR" 2>&1 +fi +cd "$REPO_DIR" || { echo ">>> cd 失败，终止"; exit 1; } +git log --oneline -3 +echo "(期望最新 commit 是 fc46265 或更新)" + +echo +echo "====================================================" +echo "[3] Python 依赖检查（只装跟平台无关的库，不碰 torch）" +echo "====================================================" +for pkg_import in "transformers:transformers" "huggingface_hub:huggingface_hub" "safetensors:safetensors" "ml_dtypes:ml_dtypes"; do + mod="${pkg_import%%:*}" + pipname="${pkg_import##*:}" + if python3 -c "import ${mod}" >/dev/null 2>&1; then + echo "${mod}: 已安装" + else + echo "${mod}: 缺失，尝试 pip install ${pipname}" + pip install "${pipname}" 2>&1 + fi +done + +echo +echo "====================================================" +echo "[4] 确保 libcudadevrt 桩库存在" +echo "====================================================" +STUB_DIR="/usr/local/corex-4.4.0/lib64" +STUB_LIB="${STUB_DIR}/libcudadevrt.a" +if [ -f "$STUB_LIB" ]; then + echo "已存在: $STUB_LIB" +else + echo "创建空桩库: $STUB_LIB" + ar rcs "$STUB_LIB" +fi + +echo +echo "====================================================" +echo "[5] xmake 配置 + 完整构建 (iluvatar-gpu=y)" +echo "====================================================" +xmake f -c --iluvatar-gpu=y 2>&1 +echo +xmake build -v llaisys 2>&1 +BUILD_STATUS=$? +echo "--- build exit code: $BUILD_STATUS ---" +if [ $BUILD_STATUS -ne 0 ]; then + echo ">>> 构建失败，后面步骤大概率也会失败，但继续跑方便一次性看到所有问题" +fi + +echo +echo "====================================================" +echo "[6] xmake install + pip 可编辑安装" +echo "====================================================" +xmake install 2>&1 +pip install -e ./python 2>&1 + +echo +echo "====================================================" +echo "[7] 8 个算子回归" +echo "====================================================" +for op in add argmax embedding linear rms_norm rope self_attention swiglu; do + echo "---- test/ops/${op}.py --device iluvatar ----" + python3 test/ops/${op}.py --device iluvatar 2>&1 + echo "---- exit code: $? ----" +done + +echo +echo "====================================================" +echo "[8] 查找本地模型权重（没有的话让脚本自己从 HuggingFace 下载）" +echo "====================================================" +MODEL_DIR="" +for candidate in \ + "$REPO_DIR/DeepSeek-R1-Distill-Qwen-1.5B" \ + "$HOME/DeepSeek-R1-Distill-Qwen-1.5B" \ + /root/DeepSeek-R1-Distill-Qwen-1.5B \ + /data/DeepSeek-R1-Distill-Qwen-1.5B \ + /models/DeepSeek-R1-Distill-Qwen-1.5B +do + if [ -f "$candidate/config.json" ] && [ -f "$candidate/model.safetensors" ]; then + MODEL_DIR="$candidate" + break + fi +done +if [ -n "$MODEL_DIR" ]; then + echo "找到本地模型目录: $MODEL_DIR" +else + echo "没找到本地权重。test_infer.py 在不传 --model 时会自动从 HuggingFace" \ + "(deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B, 约 3.5GB) 下载到本地 HF 缓存并复用同一份路径，" \ + "前提是这台机器能访问外网——接下来这步会自动尝试。" +fi + +echo +echo "====================================================" +echo "[9] 完整端到端推理 test_infer.py --device iluvatar --test" +echo "====================================================" +if [ -n "$MODEL_DIR" ]; then + python3 test/test_infer.py --model "$MODEL_DIR" --device iluvatar --test 2>&1 +else + python3 test/test_infer.py --device iluvatar --test 2>&1 +fi +echo "--- exit code: $? ---" + +echo +echo "====================================================" +echo "全部结束，请把以上完整输出贴回去" +echo "====================================================" diff --git a/scripts/iluvatar/research_bf16_gemm.sh b/scripts/iluvatar/research_bf16_gemm.sh new file mode 100755 index 000000000..26fa61c52 --- /dev/null +++ b/scripts/iluvatar/research_bf16_gemm.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# 研究用：为什么 linear 的 bf16 在 Iluvatar 上会失败。 +# 已经确认 cublasSgemmEx(..., CUDA_R_16BF, ...) 返回 status 15 +# (CUBLAS_STATUS_NOT_SUPPORTED)。这一步不改代码，只是挖一下： +# 1. 这台机器的 cublas 头文件里到底有没有 bf16 相关的声明/函数 +# （尤其是更通用的 cublasGemmEx，它和 cublasSgemmEx 是两个不同的函数）。 +# 2. cublasGemmEx 的签名长什么样（老版本 cuBLAS 用 cudaDataType 作为 +# computeType，新版本引入了专门的 cublasComputeType_t 枚举——两者 +# 不兼容，得先看清楚这台机器头文件里实际是哪种，再决定要不要写测试代码调用）。 +# +# 用法: cd 到 llaisys 仓库根目录， +# bash scripts/iluvatar/research_bf16_gemm.sh +# 把完整输出贴回来。 + +set -u + +echo "====================================================" +echo "[0] 确认 cublas 头文件实际路径和版本" +echo "====================================================" +CUBLAS_H=$(find /usr/local/corex* -iname "cublas_v2.h" 2>/dev/null | head -1) +CUBLAS_API_H=$(find /usr/local/corex* -iname "cublas_api.h" 2>/dev/null | head -1) +echo "cublas_v2.h: $CUBLAS_H" +echo "cublas_api.h: $CUBLAS_API_H" + +echo +echo "====================================================" +echo "[1] cublasGemmEx 的声明（更通用的函数，和 cublasSgemmEx 不是一个）" +echo "====================================================" +grep -n -A 20 "cublasGemmEx" "$CUBLAS_API_H" 2>/dev/null | head -60 + +echo +echo "====================================================" +echo "[2] cublasComputeType_t 是否存在（新枚举，还是老的直接用 cudaDataType）" +echo "====================================================" +grep -n "cublasComputeType_t" "$CUBLAS_API_H" 2>/dev/null | head -20 +echo "---" +grep -rn "CUBLAS_COMPUTE_" /usr/local/corex*/include/*.h 2>/dev/null | head -20 + +echo +echo "====================================================" +echo "[3] CUDA_R_16BF 这个数据类型枚举是否定义、在哪个头文件" +echo "====================================================" +grep -rn "CUDA_R_16BF" /usr/local/corex*/include/*.h 2>/dev/null + +echo +echo "====================================================" +echo "[4] 头文件里 bf16/bfloat16 相关的所有提及（找找有没有专门的 bf16 GEMM 接口）" +echo "====================================================" +grep -rln "bfloat16\|BF16" /usr/local/corex*/include/*.h 2>/dev/null + +echo +echo "====================================================" +echo "[5] cublasSgemmEx 本身的声明（对照确认我们调用的这个函数的真实签名）" +echo "====================================================" +grep -n -A 20 "cublasSgemmEx" "$CUBLAS_API_H" 2>/dev/null | head -30 + +echo +echo "====================================================" +echo "全部输出结束，请把上面内容贴回去" +echo "====================================================" diff --git a/scripts/iluvatar/verify_add_narrow.sh b/scripts/iluvatar/verify_add_narrow.sh new file mode 100755 index 000000000..cd3720401 --- /dev/null +++ b/scripts/iluvatar/verify_add_narrow.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# 在天数智芯远程机器上运行，验证 iluvatar 工具链是否真的能编译 llaisys 代码。 +# 把这个脚本拷到远程机器上（比如 scp 或直接粘贴内容），cd 到 llaisys 仓库根目录后执行： +# bash verify_iluvatar.sh +# 然后把全部输出贴回来。 + +set -u # 不用 -e，因为我们要收集失败信息而不是中途退出 + +REPO_DIR="$(pwd)" +echo "====================================================" +echo "[0] 基本信息" +echo "====================================================" +echo "repo dir: $REPO_DIR" +if [ ! -d "$REPO_DIR/src/ops/add/iluvatar" ]; then + echo "!! 当前目录看起来不是 llaisys 仓库根目录（找不到 src/ops/add/iluvatar），请 cd 到正确目录后重跑" + exit 1 +fi + +echo +echo "====================================================" +echo "[1] 工具链 / SDK 路径检查" +echo "====================================================" +for p in /usr/local/corex/bin/clang++ /usr/local/corex/bin/nvcc /usr/local/corex-4.4.0/lib64 /usr/local/corex-4.4.0/include/cudnn.h; do + if [ -e "$p" ]; then + echo "OK $p" + else + echo "MISSING $p" + fi +done + +echo +echo "clang++ 真实身份 (file):" +file /usr/local/corex/bin/clang++ 2>&1 +echo +echo "clang++ --version:" +/usr/local/corex/bin/clang++ --version 2>&1 + +echo +echo "cudnn 版本 (来自 cudnn.h):" +grep -E "CUDNN_MAJOR|CUDNN_MINOR|CUDNN_PATCHLEVEL" /usr/local/corex-4.4.0/include/cudnn.h 2>&1 + +echo +echo "libcublas / libcudnn 是否在 lib64 里:" +ls /usr/local/corex-4.4.0/lib64/ 2>&1 | grep -iE "cublas|cudnn" + +echo +echo "GPU 是否可见 (torch):" +python3 -c "import torch; print('cuda available:', torch.cuda.is_available()); print('device name:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A')" 2>&1 + +echo +echo "====================================================" +echo "[2] 独立编译测试：直接用 clang++ 编译 add_iluvatar.cu（不经过 xmake）" +echo "====================================================" +TESTOUT="/tmp/add_iluvatar_test.o" +rm -f "$TESTOUT" + +CMD=(/usr/local/corex/bin/clang++ + -x ivcore + -std=c++17 + -fPIC + -c "$REPO_DIR/src/ops/add/iluvatar/add_iluvatar.cu" + -I "$REPO_DIR/src" + -I "$REPO_DIR/include" + -o "$TESTOUT") + +echo "执行命令:" +printf ' %q' "${CMD[@]}" +echo +echo "---- 输出 ----" +"${CMD[@]}" +STATUS=$? +echo "---- exit code: $STATUS ----" + +if [ -f "$TESTOUT" ]; then + echo "产物存在，大小/类型：" + ls -la "$TESTOUT" + file "$TESTOUT" + if [ -s "$TESTOUT" ]; then + echo ">>> 看起来是真实的非空目标文件" + else + echo ">>> 警告：产物是空文件！" + fi +else + echo ">>> 没有产物文件，编译失败" +fi + +echo +echo "====================================================" +echo "[3] xmake 真实构建测试：只编 add 这一个 iluvatar 算子" +echo "====================================================" +echo "临时把 xmake/iluvatar.lua 里 llaisys-ops-iluvatar 的 add_files 收窄成只有 add_iluvatar.cu" + +cp xmake/iluvatar.lua /tmp/iluvatar.lua.bak + +python3 - "$REPO_DIR/xmake/iluvatar.lua" <<'PYEOF' +import sys, re +path = sys.argv[1] +with open(path) as f: + content = f.read() + +target_re = re.compile(r'(target\("llaisys-ops-iluvatar"\).*?)add_files\("\.\./src/ops/\*/iluvatar/\*\.cu"\)', re.S) +new_content, n = target_re.subn( + r'\1add_files("../src/ops/add/iluvatar/add_iluvatar.cu")', + content +) +if n != 1: + print("!! 没能找到预期的 add_files 行，脚本假设与当前 iluvatar.lua 内容不匹配，跳过第 3 步", file=sys.stderr) + sys.exit(1) + +with open(path, "w") as f: + f.write(new_content) +print("已临时收窄 add_files") +PYEOF + +if [ $? -eq 0 ]; then + echo + echo "--- xmake f -c --iluvatar-gpu=y ---" + xmake f -c --iluvatar-gpu=y 2>&1 + echo + echo "--- xmake build -v llaisys-ops-iluvatar ---" + xmake build -v llaisys-ops-iluvatar 2>&1 + XBUILD_STATUS=$? + echo "--- exit code: $XBUILD_STATUS ---" + + echo + echo "查找生成的 .o 产物：" + find build -iname "*add_iluvatar*" 2>&1 + find build -iname "*add_iluvatar*" -exec file {} \; 2>&1 +else + echo "跳过 xmake 构建测试（收窄 add_files 失败）" +fi + +echo +echo "====================================================" +echo "[4] 恢复 xmake/iluvatar.lua" +echo "====================================================" +cp /tmp/iluvatar.lua.bak xmake/iluvatar.lua +git -C "$REPO_DIR" diff --stat xmake/iluvatar.lua +echo "已恢复原始 xmake/iluvatar.lua（如果 git diff 还有差异，说明恢复有问题，检查 /tmp/iluvatar.lua.bak）" + +echo +echo "====================================================" +echo "全部测试结束，请把以上完整输出贴回去" +echo "====================================================" diff --git a/scripts/iluvatar/verify_full.sh b/scripts/iluvatar/verify_full.sh new file mode 100755 index 000000000..7de319f6e --- /dev/null +++ b/scripts/iluvatar/verify_full.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# 第二阶段验证：add 算子已经在 xmake 下编译通过（含 -std=c++17 修复）。 +# 这一步不再收窄 add_files，直接用 xmake/iluvatar.lua 里已有的完整 glob +# （../src/ops/*/iluvatar/*.cu 和 ../src/device/iluvatar/*.cu）， +# 编译全部 9 个算子 + device-iluvatar，看哪些算子会因为平台差异报错。 +# +# 用法: cd 到 llaisys 仓库根目录（先 git pull 到带 -std=c++17 修复的最新代码）， +# bash verify_iluvatar_full.sh +# 把完整输出贴回来。 + +set -u +REPO_DIR="$(pwd)" + +echo "====================================================" +echo "[0] 确认代码是最新的（应该能看到 -std=c++17 那次提交）" +echo "====================================================" +git log --oneline -5 +echo +grep -n "std=c++17" xmake/iluvatar.lua + +echo +echo "====================================================" +echo "[1] xmake 配置 (iluvatar-gpu=y)" +echo "====================================================" +xmake f -c --iluvatar-gpu=y 2>&1 + +echo +echo "====================================================" +echo "[2] 构建 llaisys-device-iluvatar（全部文件，不收窄）" +echo "====================================================" +xmake build -v llaisys-device-iluvatar 2>&1 +echo "--- exit code: $? ---" + +echo +echo "====================================================" +echo "[3] 构建 llaisys-ops-iluvatar（全部 9 个算子，不收窄）" +echo "====================================================" +xmake build -v llaisys-ops-iluvatar 2>&1 +echo "--- exit code: $? ---" + +echo +echo "====================================================" +echo "[4] 产物检查" +echo "====================================================" +find build -iname "*iluvatar*.o" 2>&1 +echo +find build -iname "libllaisys-*iluvatar*.a" -exec file {} \; 2>&1 + +echo +echo "====================================================" +echo "全部测试结束，请把以上完整输出贴回去" +echo "====================================================" diff --git a/scripts/iluvatar/verify_infer.sh b/scripts/iluvatar/verify_infer.sh new file mode 100755 index 000000000..f7e356d51 --- /dev/null +++ b/scripts/iluvatar/verify_infer.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash +# 端到端推理验证：确认自 self_attention 去掉 cuDNN 加速之后（commit 171bb0b/a72386d）, +# 完整模型推理在天数机器上没有回归——这一步之前只做过 8 个算子的单测，没跑过完整的 +# test_infer.py --device iluvatar --test。 +# +# 用法: cd 到 llaisys 仓库根目录， +# bash scripts/iluvatar/verify_infer.sh +# 把完整输出贴回来。 + +set -u +REPO_DIR="$(pwd)" + +echo "====================================================" +echo "[0] 拉最新代码并确认版本" +echo "====================================================" +git pull 2>&1 +git log --oneline -3 +echo "(期望最新 commit 是 04f5738 或更新)" + +echo +echo "====================================================" +echo "[1] 确保 libcudadevrt 桩库存在" +echo "====================================================" +STUB_DIR="/usr/local/corex-4.4.0/lib64" +STUB_LIB="${STUB_DIR}/libcudadevrt.a" +if [ -f "$STUB_LIB" ]; then + echo "已存在: $STUB_LIB" +else + echo "创建空桩库: $STUB_LIB" + ar rcs "$STUB_LIB" +fi + +echo +echo "====================================================" +echo "[2] xmake 配置 + 完整构建 (iluvatar-gpu=y)" +echo "====================================================" +xmake f -c --iluvatar-gpu=y 2>&1 +echo +xmake build -v llaisys 2>&1 +BUILD_STATUS=$? +echo "--- build exit code: $BUILD_STATUS ---" +if [ $BUILD_STATUS -ne 0 ]; then + echo ">>> 构建失败，后面步骤大概率也会失败，但继续跑方便一次性看到所有问题" +fi + +echo +echo "====================================================" +echo "[3] xmake install + pip 可编辑安装" +echo "====================================================" +xmake install 2>&1 +pip install -e ./python 2>&1 + +echo +echo "====================================================" +echo "[4] 8 个算子回归（确认 cuDNN 移除 + 注释清理没有引入问题）" +echo "====================================================" +for op in add argmax embedding linear rms_norm rope self_attention swiglu; do + echo "---- test/ops/${op}.py --device iluvatar ----" + python3 test/ops/${op}.py --device iluvatar 2>&1 + echo "---- exit code: $? ----" +done + +echo +echo "====================================================" +echo "[5] 查找本地模型权重（DeepSeek-R1-Distill-Qwen-1.5B）" +echo "====================================================" +MODEL_DIR="" +for candidate in \ + "$REPO_DIR/DeepSeek-R1-Distill-Qwen-1.5B" \ + "$HOME/DeepSeek-R1-Distill-Qwen-1.5B" \ + /root/DeepSeek-R1-Distill-Qwen-1.5B \ + /data/DeepSeek-R1-Distill-Qwen-1.5B \ + /data/*/DeepSeek-R1-Distill-Qwen-1.5B \ + /models/DeepSeek-R1-Distill-Qwen-1.5B +do + if [ -f "$candidate/config.json" ] && [ -f "$candidate/model.safetensors" ]; then + MODEL_DIR="$candidate" + break + fi +done +if [ -z "$MODEL_DIR" ]; then + echo "没在常见路径下找到，扩大范围用 find 搜（可能会慢一点）..." + FOUND=$(find / -maxdepth 6 -iname "config.json" -path "*DeepSeek*" 2>/dev/null | head -1) + if [ -n "$FOUND" ]; then + MODEL_DIR=$(dirname "$FOUND") + fi +fi + +if [ -n "$MODEL_DIR" ]; then + echo "找到模型目录: $MODEL_DIR" + ls -la "$MODEL_DIR" +else + echo ">>> 没找到 DeepSeek-R1-Distill-Qwen-1.5B 的权重目录。" + echo ">>> 需要包含 config.json / tokenizer 相关文件 / model.safetensors（约 3.5GB）的目录，" + echo ">>> 放到这台机器上（比如 repo 根目录下），然后重跑这个脚本，或手动执行：" + echo ">>> python3 test/test_infer.py --model <目录路径> --device iluvatar --test" +fi + +echo +echo "====================================================" +echo "[6] 完整端到端推理 test_infer.py --device iluvatar --test" +echo "====================================================" +if [ -n "$MODEL_DIR" ]; then + python3 test/test_infer.py --model "$MODEL_DIR" --device iluvatar --test 2>&1 + echo "--- exit code: $? ---" +else + echo "跳过（第 [5] 步没找到模型权重）" +fi + +echo +echo "====================================================" +echo "全部结束，请把以上完整输出贴回去" +echo "====================================================" diff --git a/scripts/iluvatar/verify_ops.sh b/scripts/iluvatar/verify_ops.sh new file mode 100755 index 000000000..cbfc7f41b --- /dev/null +++ b/scripts/iluvatar/verify_ops.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# 第三阶段验证：派发层(op.cpp / runtime_api.cpp / device_resource.cpp)、 +# xmake 聚合 deps、测试脚本 --device 选项都已经接好了。 +# 这一步做完整的 `xmake build llaisys --iluvatar-gpu=y`（不再是单独编某个 target）， +# 然后 `xmake install` 把 .so 拷进 python 包，再逐个跑 8 个算子的 +# `test/ops/*.py --device iluvatar`，最后跑 test_runtime.py。 +# +# 用法: cd 到 llaisys 仓库根目录（先 git pull 到最新代码）， +# bash scripts/iluvatar/verify_ops.sh +# 把完整输出贴回来。 + +set -u +REPO_DIR="$(pwd)" + +echo "====================================================" +echo "[0] 确认代码版本" +echo "====================================================" +git log --oneline -3 + +echo +echo "====================================================" +echo "[1] 确保 libcudadevrt 桩库存在" +echo "====================================================" +# 天数智芯的 corex SDK 没有 libcudadevrt.a（用于 relocatable device code 链接）， +# 但我们两个 target 都设了 cuda.rdc=false，实际功能上用不到这个库—— +# xmake 内置的 cuda 规则却无条件在链接命令里加 -lcudadevrt，不管 rdc 是否为 false。 +# 建一个空的桩静态库满足链接器即可，不会被真正用到任何符号。 +STUB_DIR="/usr/local/corex-4.4.0/lib64" +STUB_LIB="${STUB_DIR}/libcudadevrt.a" +if [ -f "$STUB_LIB" ]; then + echo "已存在: $STUB_LIB" + file "$STUB_LIB" +else + echo "创建空桩库: $STUB_LIB" + ar rcs "$STUB_LIB" + ls -la "$STUB_LIB" + file "$STUB_LIB" +fi + +echo +echo "====================================================" +echo "[2] xmake 配置 + 完整构建 (iluvatar-gpu=y)" +echo "====================================================" +xmake f -c --iluvatar-gpu=y 2>&1 +echo +xmake build -v llaisys 2>&1 +BUILD_STATUS=$? +echo "--- build exit code: $BUILD_STATUS ---" + +if [ $BUILD_STATUS -ne 0 ]; then + echo ">>> 构建失败，后面的测试大概率也会失败，但还是继续跑，方便一次性看到所有问题" +fi + +echo +echo "====================================================" +echo "[3] xmake install（把 .so 拷进 python 包）+ pip 可编辑安装" +echo "====================================================" +xmake install 2>&1 +echo +python3 -c "import llaisys" 2>/dev/null +if [ $? -ne 0 ]; then + echo "llaisys 包还没装，执行 pip install -e ./python" + pip install -e ./python 2>&1 +else + echo "llaisys 包已可以 import，跳过 pip install" +fi + +echo +echo "====================================================" +echo "[4] 逐个跑 8 个算子的 --device iluvatar 测试" +echo "====================================================" +for op in add argmax embedding linear rms_norm rope self_attention swiglu; do + echo "---- test/ops/${op}.py --device iluvatar ----" + python3 test/ops/${op}.py --device iluvatar 2>&1 + echo "---- exit code: $? ----" + echo +done + +echo +echo "====================================================" +echo "[5] test_runtime.py --device iluvatar" +echo "====================================================" +python3 test/test_runtime.py --device iluvatar 2>&1 +echo "--- exit code: $? ---" + +echo +echo "====================================================" +echo "全部测试结束，请把以上完整输出贴回去" +echo "====================================================" diff --git a/scripts/print_model.py b/scripts/print_model.py new file mode 100644 index 000000000..045098d06 --- /dev/null +++ b/scripts/print_model.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""打印当前 Qwen2/DeepSeek 模型的结构和基本信息。 + +默认使用 meta device 仅构建模型结构，不读取或分配模型权重。 +如果需要打印实际加载后的模型，可使用 ``--load-weights``。 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from transformers import AutoConfig, AutoModelForCausalLM + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_MODEL_PATH = PROJECT_ROOT / "DeepSeek-R1-Distill-Qwen-1.5B" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="打印 Hugging Face 模型结构") + parser.add_argument( + "--model", + type=Path, + default=DEFAULT_MODEL_PATH, + help=f"本地模型目录（默认：{DEFAULT_MODEL_PATH}）", + ) + parser.add_argument( + "--load-weights", + action="store_true", + help="加载真实权重；默认只在 meta device 上构建并打印结构", + ) + parser.add_argument( + "--device", + default="cpu", + help="加载真实权重时使用的设备，例如 cpu 或 cuda:0", + ) + return parser.parse_args() + + +def is_git_lfs_pointer(path: Path) -> bool: + if not path.is_file() or path.stat().st_size > 1024: + return False + return path.read_bytes().startswith(b"version https://git-lfs.github.com/spec/") + + +def main() -> None: + args = parse_args() + model_path = args.model.expanduser().resolve() + config_path = model_path / "config.json" + + if not config_path.is_file(): + raise SystemExit(f"找不到模型配置文件：{config_path}") + + config = AutoConfig.from_pretrained(model_path, local_files_only=True) + + if args.load_weights: + weight_files = list(model_path.glob("*.safetensors")) + lfs_pointers = [path for path in weight_files if is_git_lfs_pointer(path)] + if not weight_files: + raise SystemExit(f"目录中没有 Safetensors 权重：{model_path}") + if lfs_pointers: + names = ", ".join(path.name for path in lfs_pointers) + raise SystemExit( + f"权重仍是 Git LFS 指针（{names}），请先执行 `git lfs pull`。" + ) + + dtype = getattr(config, "torch_dtype", None) or torch.float32 + model = AutoModelForCausalLM.from_pretrained( + model_path, + local_files_only=True, + torch_dtype=dtype, + ).to(args.device) + mode = f"真实权重，设备={args.device}，dtype={dtype}" + else: + # meta device 只保留参数的形状和类型，不为数十亿参数分配实际内存。 + with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config) + mode = "仅模型结构（meta device，未加载权重）" + + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + + print(f"模型目录：{model_path}") + print(f"模型类型：{config.model_type}") + print(f"架构：{', '.join(config.architectures or [])}") + print(f"打印模式：{mode}") + print(f"参数量：{parameter_count:,}（约 {parameter_count / 1e9:.3f}B）") + print("\n=== 模型结构 ===") + print(model) + + +if __name__ == "__main__": + main() diff --git a/src/core/runtime/runtime.cpp b/src/core/runtime/runtime.cpp index 7f03a8622..e65f815f2 100644 --- a/src/core/runtime/runtime.cpp +++ b/src/core/runtime/runtime.cpp @@ -7,8 +7,10 @@ namespace llaisys::core { Runtime::Runtime(llaisysDeviceType_t device_type, int device_id) : _device_type(device_type), _device_id(device_id), _is_active(false) { _api = llaisys::device::getRuntimeAPI(_device_type); + _resource = llaisys::device::getDeviceResource(_device_type, _device_id); _stream = _api->create_stream(); _allocator = new allocators::NaiveAllocator(_api); + } Runtime::~Runtime() { @@ -19,6 +21,8 @@ Runtime::~Runtime() { _allocator = nullptr; _api->destroy_stream(_stream); _api = nullptr; + delete _resource; + _resource = nullptr; } void Runtime::_activate() { @@ -70,4 +74,7 @@ void Runtime::synchronize() const { _api->stream_synchronize(_stream); } +llaisys::device::DeviceResource* Runtime::resource() const{ + return _resource; +} } // namespace llaisys::core diff --git a/src/core/runtime/runtime.hpp b/src/core/runtime/runtime.hpp index 43235824e..ee4c05dbb 100644 --- a/src/core/runtime/runtime.hpp +++ b/src/core/runtime/runtime.hpp @@ -3,7 +3,7 @@ #include "../../device/runtime_api.hpp" #include "../allocator/allocator.hpp" - +#include "../../device/device_resource.hpp" namespace llaisys::core { class Runtime { private: @@ -15,6 +15,7 @@ class Runtime { void _activate(); void _deactivate(); llaisysStream_t _stream; + llaisys::device::DeviceResource* _resource; Runtime(llaisysDeviceType_t device_type, int device_id); public: @@ -43,5 +44,6 @@ class Runtime { llaisysStream_t stream() const; void synchronize() const; + llaisys::device::DeviceResource *resource() const; }; } // namespace llaisys::core diff --git a/src/device/cpu/cpu_resource.cpp b/src/device/cpu/cpu_resource.cpp index 4fb28bd06..c16668025 100644 --- a/src/device/cpu/cpu_resource.cpp +++ b/src/device/cpu/cpu_resource.cpp @@ -2,4 +2,8 @@ namespace llaisys::device::cpu { Resource::Resource() : llaisys::device::DeviceResource(LLAISYS_DEVICE_CPU, 0) {} + +DeviceResource *getDeviceResource(){ + return new Resource(); +} } // namespace llaisys::device::cpu diff --git a/src/device/device_resource.cpp b/src/device/device_resource.cpp new file mode 100644 index 000000000..894898e8c --- /dev/null +++ b/src/device/device_resource.cpp @@ -0,0 +1,28 @@ +#include "device_resource.hpp" + +namespace llaisys::device { + +DeviceResource * getDeviceResource(llaisysDeviceType_t device_type, int device_id){ + switch (device_type) { + case LLAISYS_DEVICE_CPU: + return cpu::getDeviceResource(); + case LLAISYS_DEVICE_NVIDIA: +#ifdef ENABLE_NVIDIA_API + return nvidia::getDeviceResource(device_id); +#else + EXCEPTION_UNSUPPORTED_DEVICE; + return nullptr; +#endif + case LLAISYS_DEVICE_ILUVATAR: +#ifdef ENABLE_ILUVATAR_API + return iluvatar::getDeviceResource(device_id); +#else + EXCEPTION_UNSUPPORTED_DEVICE; + return nullptr; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + return nullptr; + } +} +} diff --git a/src/device/device_resource.hpp b/src/device/device_resource.hpp index e9062e510..03f902ad4 100644 --- a/src/device/device_resource.hpp +++ b/src/device/device_resource.hpp @@ -14,9 +14,20 @@ class DeviceResource { : _device_type(device_type), _device_id(device_id) { } - ~DeviceResource() = default; + virtual ~DeviceResource() = default; llaisysDeviceType_t getDeviceType() const { return _device_type; } int getDeviceId() const { return _device_id; }; }; + +DeviceResource *getDeviceResource(llaisysDeviceType_t device_type, int device_id); +namespace cpu { DeviceResource *getDeviceResource(); } +#ifdef ENABLE_NVIDIA_API +namespace nvidia { DeviceResource *getDeviceResource(int device_id); } +#endif +#ifdef ENABLE_ILUVATAR_API +namespace iluvatar { DeviceResource *getDeviceResource(int device_id); } +#endif + + } // namespace llaisys::device diff --git a/src/device/iluvatar/iluvatar_resource.cu b/src/device/iluvatar/iluvatar_resource.cu new file mode 100644 index 000000000..8194795ce --- /dev/null +++ b/src/device/iluvatar/iluvatar_resource.cu @@ -0,0 +1,18 @@ +#include "iluvatar_resource.cuh" + +namespace llaisys::device::iluvatar { + +Resource::Resource(int device_id) : llaisys::device::DeviceResource(LLAISYS_DEVICE_NVIDIA, device_id) { + cudaSetDevice(device_id); + cublasCreate(&_cublas_handle); +} + +Resource::~Resource() { + cublasDestroy(_cublas_handle); +} + +DeviceResource *getDeviceResource(int device_id){ + return new Resource(device_id); +} + +} // namespace llaisys::device::iluvatar diff --git a/src/device/iluvatar/iluvatar_resource.cuh b/src/device/iluvatar/iluvatar_resource.cuh new file mode 100644 index 000000000..7510a37dd --- /dev/null +++ b/src/device/iluvatar/iluvatar_resource.cuh @@ -0,0 +1,18 @@ +#pragma once +#include +#include "../device_resource.hpp" + +namespace llaisys::device::iluvatar { +class Resource : public llaisys::device::DeviceResource { +private: + cublasHandle_t _cublas_handle; +public: + Resource(int device_id); + ~Resource(); + + + cublasHandle_t cublasHandle() const { + return _cublas_handle; + } +}; +} // namespace llaisys::device::iluvatar diff --git a/src/device/iluvatar/iluvatar_runtime_api.cu b/src/device/iluvatar/iluvatar_runtime_api.cu new file mode 100644 index 000000000..b441157d6 --- /dev/null +++ b/src/device/iluvatar/iluvatar_runtime_api.cu @@ -0,0 +1,109 @@ +#include "../runtime_api.hpp" + +#include +#include +#define CHECK_CUDA(call) \ + do { \ + cudaError_t error_ = (call); \ + if (error_ != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error at %s:%d: %s\n", \ + __FILE__, __LINE__, cudaGetErrorString(error_)); \ + std::abort(); \ + } \ + } while (0) +namespace llaisys::device::iluvatar { + +namespace runtime_api { + +static cudaMemcpyKind convertMemcpyKind(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: + std::abort(); + } +} + + +int getDeviceCount() { + int count = 0; + CHECK_CUDA(cudaGetDeviceCount(&count)); + return count; +} + +void setDevice(int device) { + CHECK_CUDA(cudaSetDevice(device)); +} + +void deviceSynchronize() { + CHECK_CUDA(cudaDeviceSynchronize()); +} + +llaisysStream_t createStream() { + cudaStream_t stream = nullptr; + CHECK_CUDA(cudaStreamCreate(&stream)); + return reinterpret_cast(stream); +} + +void destroyStream(llaisysStream_t stream) { + CHECK_CUDA(cudaStreamDestroy(reinterpret_cast(stream))); +} + +void streamSynchronize(llaisysStream_t stream) { + CHECK_CUDA(cudaStreamSynchronize(reinterpret_cast (stream))); +} + +void *mallocDevice(size_t size) { + void * ptr =nullptr; + CHECK_CUDA(cudaMalloc(&ptr, size)); + return ptr; +} + +void freeDevice(void *ptr) { + CHECK_CUDA(cudaFree(ptr)); +} + +void *mallocHost(size_t size) { + void * ptr =nullptr; + CHECK_CUDA(cudaMallocHost(&ptr, size)); + return ptr; +} + +void freeHost(void *ptr) { + CHECK_CUDA(cudaFreeHost(ptr)); +} + +void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) { + CHECK_CUDA(cudaMemcpy(dst, src, size, convertMemcpyKind(kind))); +} + +void memcpyAsync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind, llaisysStream_t stream) { + CHECK_CUDA(cudaMemcpyAsync(dst, src, size, convertMemcpyKind(kind), reinterpret_cast(stream))); +} + +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::iluvatar diff --git a/src/device/nvidia/nvidia_resource.cu b/src/device/nvidia/nvidia_resource.cu index 2e63647e5..ee5f0aaaa 100644 --- a/src/device/nvidia/nvidia_resource.cu +++ b/src/device/nvidia/nvidia_resource.cu @@ -2,6 +2,17 @@ namespace llaisys::device::nvidia { -Resource::Resource(int device_id) : llaisys::device::DeviceResource(LLAISYS_DEVICE_NVIDIA, device_id) {} +Resource::Resource(int device_id) : llaisys::device::DeviceResource(LLAISYS_DEVICE_NVIDIA, device_id) { + cudaSetDevice(device_id); + cublasCreate(&_cublas_handle); +} + +Resource::~Resource() { + cublasDestroy(_cublas_handle); +} + +DeviceResource *getDeviceResource(int device_id){ + return new Resource(device_id); +} } // namespace llaisys::device::nvidia diff --git a/src/device/nvidia/nvidia_resource.cuh b/src/device/nvidia/nvidia_resource.cuh index a3002170b..3253c4731 100644 --- a/src/device/nvidia/nvidia_resource.cuh +++ b/src/device/nvidia/nvidia_resource.cuh @@ -1,11 +1,18 @@ #pragma once - +#include #include "../device_resource.hpp" namespace llaisys::device::nvidia { class Resource : public llaisys::device::DeviceResource { +private: + cublasHandle_t _cublas_handle; public: Resource(int device_id); ~Resource(); + + + cublasHandle_t cublasHandle() const { + return _cublas_handle; + } }; } // namespace llaisys::device::nvidia diff --git a/src/device/nvidia/nvidia_runtime_api.cu b/src/device/nvidia/nvidia_runtime_api.cu index cab928261..10517c617 100644 --- a/src/device/nvidia/nvidia_runtime_api.cu +++ b/src/device/nvidia/nvidia_runtime_api.cu @@ -2,55 +2,89 @@ #include #include - +#define CHECK_CUDA(call) \ + do { \ + cudaError_t error_ = (call); \ + if (error_ != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error at %s:%d: %s\n", \ + __FILE__, __LINE__, cudaGetErrorString(error_)); \ + std::abort(); \ + } \ + } while (0) namespace llaisys::device::nvidia { namespace runtime_api { + +static cudaMemcpyKind convertMemcpyKind(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: + std::abort(); + } +} + + int getDeviceCount() { - TO_BE_IMPLEMENTED(); + int count = 0; + CHECK_CUDA(cudaGetDeviceCount(&count)); + return count; } -void setDevice(int) { - TO_BE_IMPLEMENTED(); +void setDevice(int device) { + CHECK_CUDA(cudaSetDevice(device)); } void deviceSynchronize() { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaDeviceSynchronize()); } llaisysStream_t createStream() { - TO_BE_IMPLEMENTED(); + cudaStream_t stream = nullptr; + CHECK_CUDA(cudaStreamCreate(&stream)); + return reinterpret_cast(stream); } void destroyStream(llaisysStream_t stream) { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaStreamDestroy(reinterpret_cast(stream))); } + void streamSynchronize(llaisysStream_t stream) { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaStreamSynchronize(reinterpret_cast (stream))); } void *mallocDevice(size_t size) { - TO_BE_IMPLEMENTED(); + void * ptr =nullptr; + CHECK_CUDA(cudaMalloc(&ptr, size)); + return ptr; } void freeDevice(void *ptr) { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaFree(ptr)); } void *mallocHost(size_t size) { - TO_BE_IMPLEMENTED(); + void * ptr =nullptr; + CHECK_CUDA(cudaMallocHost(&ptr, size)); + return ptr; } void freeHost(void *ptr) { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaFreeHost(ptr)); } void memcpySync(void *dst, const void *src, size_t size, llaisysMemcpyKind_t kind) { - TO_BE_IMPLEMENTED(); + CHECK_CUDA(cudaMemcpy(dst, src, size, convertMemcpyKind(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) { + CHECK_CUDA(cudaMemcpyAsync(dst, src, size, convertMemcpyKind(kind), reinterpret_cast(stream))); } static const LlaisysRuntimeAPI RUNTIME_API = { diff --git a/src/device/runtime_api.cpp b/src/device/runtime_api.cpp index 2de3eca02..1a8dfc6be 100644 --- a/src/device/runtime_api.cpp +++ b/src/device/runtime_api.cpp @@ -80,6 +80,12 @@ const LlaisysRuntimeAPI *getRuntimeAPI(llaisysDeviceType_t device_type) { return llaisys::device::nvidia::getRuntimeAPI(); #else return getUnsupportedRuntimeAPI(); +#endif + case LLAISYS_DEVICE_ILUVATAR: +#ifdef ENABLE_ILUVATAR_API + return llaisys::device::iluvatar::getRuntimeAPI(); +#else + return getUnsupportedRuntimeAPI(); #endif default: EXCEPTION_UNSUPPORTED_DEVICE; diff --git a/src/device/runtime_api.hpp b/src/device/runtime_api.hpp index e6b9f80d6..12ebdc40f 100644 --- a/src/device/runtime_api.hpp +++ b/src/device/runtime_api.hpp @@ -17,4 +17,10 @@ namespace nvidia { const LlaisysRuntimeAPI *getRuntimeAPI(); } #endif + +#ifdef ENABLE_ILUVATAR_API +namespace iluvatar { +const LlaisysRuntimeAPI *getRuntimeAPI(); +} +#endif } // namespace llaisys::device diff --git a/src/llaisys/models/qwen2.cc b/src/llaisys/models/qwen2.cc new file mode 100644 index 000000000..cd785da20 --- /dev/null +++ b/src/llaisys/models/qwen2.cc @@ -0,0 +1,248 @@ +#include "llaisys/models/qwen2.h" +#include "../../ops/add/op.hpp" +#include "../../ops/argmax/op.hpp" +#include "../../ops/embedding/op.hpp" +#include "../../ops/linear/op.hpp" +#include "../../ops/rearrange/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 "../../tensor/tensor.hpp" +#include "../llaisys_tensor.hpp" // LlaisysTensor{tensor_t} 包装 +#include +// 头文件里 LlaisysQwen2Model 只是前向声明，"真身"定义在这里（不透明指针模式）。 +struct LlaisysQwen2Model { + LlaisysQwen2Meta meta; + LlaisysQwen2Weights weights{}; + std::vector attn_norm_w; + std::vector attn_q_w, attn_q_b; + std::vector attn_k_w, attn_k_b; + std::vector attn_v_w, attn_v_b; + std::vector attn_o_w; + std::vector mlp_norm_w; + std::vector mlp_gate_w, mlp_up_w, mlp_down_w; + + std::vector k_cache, v_cache; + size_t cur_len = 0; // 已经缓存了多少个 token + + llaisysDeviceType_t device; + std::vector device_ids; +}; + +__C { + struct LlaisysQwen2Model *llaisysQwen2ModelCreate(const LlaisysQwen2Meta *meta, llaisysDeviceType_t device, int *device_ids, int ndevice) { + LlaisysQwen2Model *model = new LlaisysQwen2Model(); + model->meta = *meta; + model->device = device; + model->device_ids.assign(device_ids, device_ids + ndevice); + + int dev_id = (ndevice > 0 && device_ids) ? device_ids[0] : 0; + size_t nlayer = meta->nlayer; + + // ---- single-tensor weights ---- + model->weights.in_embed = new LlaisysTensor{ + llaisys::Tensor::create({meta->voc, meta->hs}, meta->dtype, device, dev_id)}; + model->weights.out_embed = new LlaisysTensor{ + llaisys::Tensor::create({meta->voc, meta->hs}, meta->dtype, device, dev_id)}; + model->weights.out_norm_w = new LlaisysTensor{ + llaisys::Tensor::create({meta->hs}, meta->dtype, device, dev_id)}; + + // ---- attn_norm_w ---- + model->attn_norm_w.resize(nlayer); + for (size_t i = 0; i < nlayer; i++) { + model->attn_norm_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->hs}, meta->dtype, device, dev_id)}; + } + model->weights.attn_norm_w = model->attn_norm_w.data(); + + // ---- attention q/k/v/o ---- + model->attn_q_w.resize(nlayer); + model->attn_q_b.resize(nlayer); + model->attn_k_w.resize(nlayer); + model->attn_k_b.resize(nlayer); + model->attn_v_w.resize(nlayer); + model->attn_v_b.resize(nlayer); + model->attn_o_w.resize(nlayer); + for (size_t i = 0; i < nlayer; i++) { + model->attn_q_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nh * meta->dh, meta->hs}, meta->dtype, device, dev_id)}; + model->attn_q_b[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nh * meta->dh}, meta->dtype, device, dev_id)}; + + model->attn_k_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nkvh * meta->dh, meta->hs}, meta->dtype, device, dev_id)}; + model->attn_k_b[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nkvh * meta->dh}, meta->dtype, device, dev_id)}; + + model->attn_v_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nkvh * meta->dh, meta->hs}, meta->dtype, device, dev_id)}; + model->attn_v_b[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->nkvh * meta->dh}, meta->dtype, device, dev_id)}; + + model->attn_o_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->hs, meta->nh * meta->dh}, meta->dtype, device, dev_id)}; + } + model->weights.attn_q_w = model->attn_q_w.data(); + model->weights.attn_q_b = model->attn_q_b.data(); + model->weights.attn_k_w = model->attn_k_w.data(); + model->weights.attn_k_b = model->attn_k_b.data(); + model->weights.attn_v_w = model->attn_v_w.data(); + model->weights.attn_v_b = model->attn_v_b.data(); + model->weights.attn_o_w = model->attn_o_w.data(); + + // ---- mlp norm/gate/up/down ---- + model->mlp_norm_w.resize(nlayer); + model->mlp_gate_w.resize(nlayer); + model->mlp_up_w.resize(nlayer); + model->mlp_down_w.resize(nlayer); + for (size_t i = 0; i < nlayer; i++) { + model->mlp_norm_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->hs}, meta->dtype, device, dev_id)}; + model->mlp_gate_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->di, meta->hs}, meta->dtype, device, dev_id)}; + model->mlp_up_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->di, meta->hs}, meta->dtype, device, dev_id)}; + model->mlp_down_w[i] = new LlaisysTensor{ + llaisys::Tensor::create({meta->hs, meta->di}, meta->dtype, device, dev_id)}; + } + model->weights.mlp_norm_w = model->mlp_norm_w.data(); + model->weights.mlp_gate_w = model->mlp_gate_w.data(); + model->weights.mlp_up_w = model->mlp_up_w.data(); + model->weights.mlp_down_w = model->mlp_down_w.data(); + + // ---- KV cache (internal only, not exposed to Python) ---- + model->k_cache.resize(nlayer); + model->v_cache.resize(nlayer); + for (size_t i = 0; i < nlayer; i++) { + model->k_cache[i] = llaisys::Tensor::create({meta->maxseq, meta->nkvh, meta->dh}, meta->dtype, device, dev_id); + model->v_cache[i] = llaisys::Tensor::create({meta->maxseq, meta->nkvh, meta->dh}, meta->dtype, device, dev_id); + } + + return model; + } + + void llaisysQwen2ModelDestroy(struct LlaisysQwen2Model * model) { + delete model->weights.in_embed; + delete model->weights.out_embed; + delete model->weights.out_norm_w; + for (size_t i = 0; i < model->meta.nlayer; i++) { + delete model->attn_norm_w[i]; + delete model->attn_q_w[i]; + delete model->attn_q_b[i]; + delete model->attn_k_w[i]; + delete model->attn_k_b[i]; + delete model->attn_v_w[i]; + delete model->attn_v_b[i]; + delete model->attn_o_w[i]; + delete model->mlp_norm_w[i]; + delete model->mlp_gate_w[i]; + delete model->mlp_up_w[i]; + delete model->mlp_down_w[i]; + } + delete model; + } + + struct LlaisysQwen2Weights *llaisysQwen2ModelWeights(struct LlaisysQwen2Model * model) { + return &model->weights; + } + + int64_t llaisysQwen2ModelInfer(struct LlaisysQwen2Model * model, int64_t * token_ids, size_t ntoken) { + using llaisys::Tensor; + using llaisys::tensor_t; + namespace ops = llaisys::ops; + + const auto &meta = model->meta; + auto device = model->device; + auto device_id = (model->device_ids.size() > 0) ? model->device_ids[0] : 0; + size_t cur_len = model->cur_len; // 已缓存的 token 数 + size_t total_len = cur_len + ntoken; // 这次调用结束后 cache 里的 token 数 + tensor_t input_ids = Tensor::create({ntoken}, LLAISYS_DTYPE_I64, device, device_id); + input_ids->load(token_ids); + std::vector pos_ids_vec(ntoken); + for (size_t i = 0; i < ntoken; i++) { + pos_ids_vec[i] = cur_len + i; + } + tensor_t pos_ids = Tensor::create({ntoken}, LLAISYS_DTYPE_I64, device, device_id); + pos_ids->load(pos_ids_vec.data()); + + tensor_t input_embeds = Tensor::create({ntoken, meta.hs}, meta.dtype, device, device_id); + ops::embedding(input_embeds, input_ids, model->weights.in_embed->tensor); + + float scale = 1.0f / std::sqrt(static_cast(meta.dh)); + for(size_t layer_idx = 0; layer_idx < meta.nlayer; layer_idx++) { + tensor_t normed = Tensor::create({ntoken, meta.hs}, meta.dtype, device, device_id); + ops::rms_norm(normed, input_embeds, model->weights.attn_norm_w[layer_idx]->tensor, meta.epsilon); + + // Q,K,V projections + tensor_t q = Tensor::create({ntoken, meta.nh*meta.dh}, meta.dtype, device, device_id); + tensor_t k = Tensor::create({ntoken, meta.nkvh*meta.dh}, meta.dtype, device, device_id); + tensor_t v = Tensor::create({ntoken, meta.nkvh*meta.dh}, meta.dtype, device, device_id); + ops::linear(q, normed, model->weights.attn_q_w[layer_idx]->tensor, model->weights.attn_q_b[layer_idx]->tensor); + ops::linear(k, normed, model->weights.attn_k_w[layer_idx]->tensor, model->weights.attn_k_b[layer_idx]->tensor); + ops::linear(v, normed, model->weights.attn_v_w[layer_idx]->tensor, model->weights.attn_v_b[layer_idx]->tensor); + + tensor_t q_reshaped = q->view({ntoken, meta.nh, meta.dh}); + tensor_t k_reshaped = k->view({ntoken, meta.nkvh, meta.dh}); + tensor_t v_reshaped = v->view({ntoken, meta.nkvh, meta.dh}); + + tensor_t q_rope = Tensor::create({ntoken, meta.nh, meta.dh}, meta.dtype, device, device_id); + tensor_t k_rope = Tensor::create({ntoken, meta.nkvh, meta.dh}, meta.dtype, device, device_id); + ops::rope(q_rope, q_reshaped, pos_ids, meta.theta); + ops::rope(k_rope, k_reshaped, pos_ids, meta.theta); + + // 写入本次新算出的 [cur_len, total_len) 段，再读出 [0, total_len) 完整历史 + model->k_cache[layer_idx]->slice(0,cur_len,total_len)->load(k_rope->data()); + model->v_cache[layer_idx]->slice(0,cur_len,total_len)->load(v_reshaped->data()); + tensor_t k_all = model->k_cache[layer_idx]->slice(0, 0, total_len); + tensor_t v_all = model->v_cache[layer_idx]->slice(0, 0, total_len); + + tensor_t attn_val = Tensor::create({ntoken, meta.nh, meta.dh}, meta.dtype, device, device_id); + ops::self_attention(attn_val, q_rope, k_all, v_all, scale); + + tensor_t attn_val_flat = attn_val->view({ntoken, meta.nh * meta.dh}); + tensor_t attn_out = Tensor::create({ntoken,meta.hs},meta.dtype,device,device_id); + ops::linear(attn_out,attn_val_flat,model->weights.attn_o_w[layer_idx]->tensor,nullptr); + tensor_t residual = Tensor::create({ntoken,meta.hs},meta.dtype,device,device_id); + ops::add(residual,input_embeds, attn_out); + + // mlp + tensor_t post_attention_normed = Tensor::create({ntoken,meta.hs},meta.dtype,device,device_id); + ops::rms_norm(post_attention_normed, residual, model->weights.mlp_norm_w[layer_idx]->tensor, meta.epsilon); + tensor_t after_swiglu = Tensor::create({ntoken,meta.di},meta.dtype,device,device_id); + tensor_t gate = Tensor::create({ntoken,meta.di},meta.dtype,device,device_id); + tensor_t up = Tensor::create({ntoken,meta.di},meta.dtype,device,device_id); + ops::linear(gate, post_attention_normed, model->weights.mlp_gate_w[layer_idx]->tensor, nullptr); + ops::linear(up, post_attention_normed, model->weights.mlp_up_w[layer_idx]->tensor, nullptr); + ops::swiglu(after_swiglu,gate,up); + tensor_t after_down = Tensor::create({ntoken,meta.hs},meta.dtype,device,device_id); + ops::linear(after_down,after_swiglu,model->weights.mlp_down_w[layer_idx]->tensor, nullptr); + ops::add(input_embeds,after_down,residual); + } + tensor_t after_final_normed = Tensor::create({ntoken, meta.hs}, meta.dtype, device, device_id); + ops::rms_norm(after_final_normed, input_embeds, model->weights.out_norm_w->tensor, meta.epsilon); + tensor_t outputs_prob = Tensor::create({ntoken,meta.voc}, meta.dtype, device, device_id); + ops::linear(outputs_prob,after_final_normed, model->weights.out_embed->tensor, nullptr); + + tensor_t last_logits = outputs_prob->slice(0, ntoken - 1, ntoken); + tensor_t last_logits_1d = last_logits->view({meta.voc}); + tensor_t max_idx = Tensor::create({1}, LLAISYS_DTYPE_I64, device, device_id); + tensor_t max_val = Tensor::create({1}, meta.dtype, device, device_id); + ops::argmax(max_idx,max_val,last_logits_1d); + + model->cur_len = total_len; + if (device == LLAISYS_DEVICE_CPU) { + return *reinterpret_cast(max_idx->data()); + } else { + int64_t host_max_idx = 0; + llaisys::core::context().runtime().api()->memcpy_sync( + &host_max_idx, + max_idx->data(), + sizeof(int64_t), + LLAISYS_MEMCPY_D2H + ); + return host_max_idx; + } + } +} + diff --git a/src/llaisys/ops.cc b/src/llaisys/ops.cc index c99fbc32f..ca8de527d 100644 --- a/src/llaisys/ops.cc +++ b/src/llaisys/ops.cc @@ -23,7 +23,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); diff --git a/src/llaisys/tensor.cc b/src/llaisys/tensor.cc index 5e6e50124..82024ffed 100644 --- a/src/llaisys/tensor.cc +++ b/src/llaisys/tensor.cc @@ -93,4 +93,12 @@ __C { size_t end) { return new LlaisysTensor{tensor->tensor->slice(dim, start, end)}; } + + llaisysTensor_t tensorReshape( + llaisysTensor_t tensor, + size_t * shape, + size_t ndim) { + std::vector shape_vec(shape, shape + ndim); + return new LlaisysTensor{tensor->tensor->reshape(shape_vec)}; + } } diff --git a/src/ops/add/iluvatar/add_iluvatar.cu b/src/ops/add/iluvatar/add_iluvatar.cu new file mode 100644 index 000000000..8c9f3b09e --- /dev/null +++ b/src/ops/add/iluvatar/add_iluvatar.cu @@ -0,0 +1,35 @@ +#include "../../../utils.hpp" +#include "add_iluvatar.cuh" +#include +#include +#include + +template +__global__ void add_kernel(T *c, const T *a, const T *b, size_t numel) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if(idx < numel){ + c[idx] = a[idx] + b[idx]; + } +} + +namespace llaisys::ops::iluvatar { +void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t numel) { + constexpr int block_size = 256; + int grid_size = static_cast((numel + block_size - 1) / block_size); + switch (type) { + case LLAISYS_DTYPE_F32: + add_kernel<<>>(reinterpret_cast(c), reinterpret_cast(a), reinterpret_cast(b), numel); + return; + case LLAISYS_DTYPE_BF16: + add_kernel<<>>(reinterpret_cast<__nv_bfloat16 *>(c), reinterpret_cast(a), + reinterpret_cast(b), numel); + return; + case LLAISYS_DTYPE_F16: + add_kernel<<>>(reinterpret_cast<__half *>(c), reinterpret_cast(a), + reinterpret_cast(b), numel); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} \ No newline at end of file diff --git a/src/ops/add/iluvatar/add_iluvatar.cuh b/src/ops/add/iluvatar/add_iluvatar.cuh new file mode 100644 index 000000000..a90d4e703 --- /dev/null +++ b/src/ops/add/iluvatar/add_iluvatar.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + +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 diff --git a/src/ops/add/nvidia/add_cuda.cu b/src/ops/add/nvidia/add_cuda.cu new file mode 100644 index 000000000..1412ce605 --- /dev/null +++ b/src/ops/add/nvidia/add_cuda.cu @@ -0,0 +1,35 @@ +#include "../../../utils.hpp" +#include "add_cuda.cuh" +#include +#include +#include + +template +__global__ void add_kernel(T *c, const T *a, const T *b, size_t numel) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if(idx < numel){ + c[idx] = a[idx] + b[idx]; + } +} + +namespace llaisys::ops::cuda { +void add(std::byte *c, const std::byte *a, const std::byte *b, llaisysDataType_t type, size_t numel) { + constexpr int block_size = 256; + int grid_size = static_cast((numel + block_size - 1) / block_size); + switch (type) { + case LLAISYS_DTYPE_F32: + add_kernel<<>>(reinterpret_cast(c), reinterpret_cast(a), reinterpret_cast(b), numel); + return; + case LLAISYS_DTYPE_BF16: + add_kernel<<>>(reinterpret_cast<__nv_bfloat16 *>(c), reinterpret_cast(a), + reinterpret_cast(b), numel); + return; + case LLAISYS_DTYPE_F16: + add_kernel<<>>(reinterpret_cast<__half *>(c), reinterpret_cast(a), + reinterpret_cast(b), numel); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} \ No newline at end of file diff --git a/src/ops/add/nvidia/add_cuda.cuh b/src/ops/add/nvidia/add_cuda.cuh new file mode 100644 index 000000000..3eb9126d8 --- /dev/null +++ b/src/ops/add/nvidia/add_cuda.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +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 diff --git a/src/ops/add/op.cpp b/src/ops/add/op.cpp index a057330d7..86a01c5c4 100644 --- a/src/ops/add/op.cpp +++ b/src/ops/add/op.cpp @@ -2,7 +2,8 @@ #include "../../core/llaisys_core.hpp" #include "../../utils.hpp" - +#include "nvidia/add_cuda.cuh" +#include "iluvatar/add_iluvatar.cuh" #include "cpu/add_cpu.hpp" namespace llaisys::ops { @@ -25,8 +26,11 @@ 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(); - return; + return cuda::add(c->data(), a->data(), b->data(), c->dtype(), c->numel()); +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + return iluvatar::add(c->data(), a->data(), b->data(), c->dtype(), c->numel()); #endif default: EXCEPTION_UNSUPPORTED_DEVICE; diff --git a/src/ops/argmax/cpu/argmax_cpu.cpp b/src/ops/argmax/cpu/argmax_cpu.cpp new file mode 100644 index 000000000..cc0c111ae --- /dev/null +++ b/src/ops/argmax/cpu/argmax_cpu.cpp @@ -0,0 +1,42 @@ +#include "../../../utils.hpp" +#include "argmax_cpu.hpp" +#include + +template +void argmax_(int64_t *max_idx, T *max_val, const T *vals, size_t numel) { + for (size_t i = 0; i < numel; i++) { + if constexpr (std::is_same_v || std::is_same_v){ + float val = llaisys::utils::cast(vals[i]); + if (i == 0 || val > llaisys::utils::cast(*max_val)) { + *max_val = llaisys::utils::cast(val); + *max_idx = static_cast(i); + } + }else{ + if (i == 0 || vals[i] > *max_val) { + *max_val = vals[i]; + *max_idx = static_cast(i); + } + } + + } +} + +namespace llaisys::ops::cpu { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel) { + + switch (type) { + case LLAISYS_DTYPE_F32: + return argmax_(reinterpret_cast(max_idx), reinterpret_cast(max_val), + reinterpret_cast(vals), numel); + case LLAISYS_DTYPE_BF16: + return argmax_(reinterpret_cast(max_idx), reinterpret_cast(max_val), + reinterpret_cast(vals), numel); + case LLAISYS_DTYPE_F16: + return argmax_(reinterpret_cast(max_idx), reinterpret_cast(max_val), + reinterpret_cast(vals), numel); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } + +} +} \ No newline at end of file diff --git a/src/ops/argmax/cpu/argmax_cpu.hpp b/src/ops/argmax/cpu/argmax_cpu.hpp new file mode 100644 index 000000000..37fab43c0 --- /dev/null +++ b/src/ops/argmax/cpu/argmax_cpu.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "llaisys.h" + +#include + +namespace llaisys::ops::cpu { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t size); +} + diff --git a/src/ops/argmax/iluvatar/argmax_iluvatar.cu b/src/ops/argmax/iluvatar/argmax_iluvatar.cu new file mode 100644 index 000000000..de96704ce --- /dev/null +++ b/src/ops/argmax/iluvatar/argmax_iluvatar.cu @@ -0,0 +1,107 @@ +#include "../../../utils.hpp" +#include "argmax_iluvatar.cuh" +#include +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t error = (call); \ + if (error != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(error)); \ + std::abort(); \ + } \ + } while (0) + +// 两轮 kernel（手写，cuDNN 的归约节点没有能一次给出 value+index 的非弃用 API，见 CLAUDE.md）： +// 第一轮多 block 各自规约出局部候选，第二轮单 block 对候选再规约一次得到全局结果， +// idx_map 非空时把"候选缓冲区里的位置"翻译回原始下标。 +template +__global__ void argmax_kernel(int64_t *block_idx, T *block_val, const T *vals, size_t numel, const int64_t *idx_map) { + + // 线程局部扫描 + T best_val = -INFINITY;; + int64_t best_idx; + size_t tid = threadIdx.x; + size_t idx = static_cast(blockIdx.x) * blockDim.x + tid; + + for (size_t i = idx; i < numel; i += blockDim.x * gridDim.x) { + if(vals[i] > best_val){ + best_val = vals[i]; + if(idx_map == nullptr){ + best_idx = i; + }else{ + best_idx = idx_map[i]; + } + } + } + // 块内规约 + extern __shared__ unsigned char shared_mem[]; + T *shared_val = reinterpret_cast(shared_mem); + int64_t *shared_idx = reinterpret_cast(shared_val + blockDim.x); + + T local_max = -INFINITY; + // 加载数据 + shared_val[tid] = idx < numel ? best_val : local_max; + shared_idx[tid] = idx < numel ? static_cast(best_idx) : 0.0; + __syncthreads(); + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride && shared_val[tid] < shared_val[tid + stride]) { + shared_val[tid] = shared_val[tid + stride]; + shared_idx[tid] = shared_idx[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + block_idx[blockIdx.x] = shared_idx[0]; + block_val[blockIdx.x] = shared_val[0]; + } +} + +template +void launch_argmax(int64_t *max_idx, T *max_val, const T *vals, size_t numel) { + constexpr int block_size = 256; + int grid_size = static_cast((numel + block_size - 1) / block_size); + size_t shared_bytes = block_size * sizeof(T) + block_size * sizeof(int64_t); + T *val_ptr = nullptr; + CHECK_CUDA(cudaMalloc(reinterpret_cast(&val_ptr), grid_size * sizeof(T))); + int64_t *idx_ptr = nullptr; + CHECK_CUDA(cudaMalloc(reinterpret_cast(&idx_ptr), grid_size * sizeof(int64_t))); + argmax_kernel<<>>(idx_ptr,val_ptr, vals, numel,nullptr); + argmax_kernel<<<1, block_size, shared_bytes>>>(max_idx, max_val, val_ptr, grid_size,idx_ptr); + CHECK_CUDA(cudaFree(val_ptr)); + CHECK_CUDA(cudaFree(idx_ptr)); +} + +namespace llaisys::ops::iluvatar { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel) { + switch (type) { + case LLAISYS_DTYPE_F32: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast(max_val), + reinterpret_cast(vals), + numel); + + case LLAISYS_DTYPE_BF16: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast<__nv_bfloat16 *>(max_val), + reinterpret_cast(vals), + numel); + + case LLAISYS_DTYPE_F16: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast<__half *>(max_val), + reinterpret_cast(vals), + numel); + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::iluvatar diff --git a/src/ops/argmax/iluvatar/argmax_iluvatar.cuh b/src/ops/argmax/iluvatar/argmax_iluvatar.cuh new file mode 100644 index 000000000..81f70b8ba --- /dev/null +++ b/src/ops/argmax/iluvatar/argmax_iluvatar.cuh @@ -0,0 +1,9 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t size); +} \ No newline at end of file diff --git a/src/ops/argmax/nvidia/argmax_cuda.cu b/src/ops/argmax/nvidia/argmax_cuda.cu new file mode 100644 index 000000000..7006d120f --- /dev/null +++ b/src/ops/argmax/nvidia/argmax_cuda.cu @@ -0,0 +1,107 @@ +#include "../../../utils.hpp" +#include "argmax_cuda.cuh" +#include +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t error = (call); \ + if (error != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(error)); \ + std::abort(); \ + } \ + } while (0) + +// 两轮 kernel（手写，cuDNN 的归约节点没有能一次给出 value+index 的非弃用 API，见 CLAUDE.md）： +// 第一轮多 block 各自规约出局部候选，第二轮单 block 对候选再规约一次得到全局结果， +// idx_map 非空时把"候选缓冲区里的位置"翻译回原始下标。 +template +__global__ void argmax_kernel(int64_t *block_idx, T *block_val, const T *vals, size_t numel, const int64_t *idx_map) { + + // 线程局部扫描 + T best_val = -INFINITY;; + int64_t best_idx; + size_t tid = threadIdx.x; + size_t idx = static_cast(blockIdx.x) * blockDim.x + tid; + + for (size_t i = idx; i < numel; i += blockDim.x * gridDim.x) { + if(vals[i] > best_val){ + best_val = vals[i]; + if(idx_map == nullptr){ + best_idx = i; + }else{ + best_idx = idx_map[i]; + } + } + } + // 块内规约 + extern __shared__ unsigned char shared_mem[]; + T *shared_val = reinterpret_cast(shared_mem); + int64_t *shared_idx = reinterpret_cast(shared_val + blockDim.x); + + T local_max = -INFINITY; + // 加载数据 + shared_val[tid] = idx < numel ? best_val : local_max; + shared_idx[tid] = idx < numel ? static_cast(best_idx) : 0.0; + __syncthreads(); + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride && shared_val[tid] < shared_val[tid + stride]) { + shared_val[tid] = shared_val[tid + stride]; + shared_idx[tid] = shared_idx[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + block_idx[blockIdx.x] = shared_idx[0]; + block_val[blockIdx.x] = shared_val[0]; + } +} + +template +void launch_argmax(int64_t *max_idx, T *max_val, const T *vals, size_t numel) { + constexpr int block_size = 256; + int grid_size = static_cast((numel + block_size - 1) / block_size); + size_t shared_bytes = block_size * sizeof(T) + block_size * sizeof(int64_t); + T *val_ptr = nullptr; + CHECK_CUDA(cudaMalloc(reinterpret_cast(&val_ptr), grid_size * sizeof(T))); + int64_t *idx_ptr = nullptr; + CHECK_CUDA(cudaMalloc(reinterpret_cast(&idx_ptr), grid_size * sizeof(int64_t))); + argmax_kernel<<>>(idx_ptr,val_ptr, vals, numel,nullptr); + argmax_kernel<<<1, block_size, shared_bytes>>>(max_idx, max_val, val_ptr, grid_size,idx_ptr); + CHECK_CUDA(cudaFree(val_ptr)); + CHECK_CUDA(cudaFree(idx_ptr)); +} + +namespace llaisys::ops::cuda { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel) { + switch (type) { + case LLAISYS_DTYPE_F32: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast(max_val), + reinterpret_cast(vals), + numel); + + case LLAISYS_DTYPE_BF16: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast<__nv_bfloat16 *>(max_val), + reinterpret_cast(vals), + numel); + + case LLAISYS_DTYPE_F16: + return launch_argmax( + reinterpret_cast(max_idx), + reinterpret_cast<__half *>(max_val), + reinterpret_cast(vals), + numel); + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::cuda diff --git a/src/ops/argmax/nvidia/argmax_cuda.cuh b/src/ops/argmax/nvidia/argmax_cuda.cuh new file mode 100644 index 000000000..d41f1f1d7 --- /dev/null +++ b/src/ops/argmax/nvidia/argmax_cuda.cuh @@ -0,0 +1,9 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { +void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t size); +} \ No newline at end of file diff --git a/src/ops/argmax/op.cpp b/src/ops/argmax/op.cpp index 6dc37d426..6f3a6a68b 100644 --- a/src/ops/argmax/op.cpp +++ b/src/ops/argmax/op.cpp @@ -1,7 +1,44 @@ #include "op.hpp" +#include "nvidia/argmax_cuda.cuh" +#include "iluvatar/argmax_iluvatar.cuh" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/argmax_cpu.hpp" namespace llaisys::ops { void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(max_idx, max_val, vals); + CHECK_SAME_DTYPE(max_val->dtype(), vals->dtype()); + ASSERT(max_idx->dtype() == LLAISYS_DTYPE_I64, "Argmax: max_idx must be int64."); + + // 当前只支持一维输入 + ASSERT(vals->ndim() == 1, "Argmax: vals must be 1D."); + CHECK_SAME_SHAPE(max_idx->shape(), max_val->shape(), std::vector{1}); + + ASSERT(max_idx->isContiguous() && max_val->isContiguous() && vals->isContiguous(), + "Argmax: all tensors must be contiguous."); + if (max_idx->deviceType() == LLAISYS_DEVICE_CPU) { + return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel()); + } + + llaisys::core::context().setDevice(max_idx->deviceType(), max_idx->deviceId()); + + switch (max_idx->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel()); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + cuda::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel()); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + iluvatar::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel()); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } -} // namespace llaisys::ops +} \ No newline at end of file diff --git a/src/ops/embedding/cpu/embedding.cpp b/src/ops/embedding/cpu/embedding.cpp new file mode 100644 index 000000000..563d7c66c --- /dev/null +++ b/src/ops/embedding/cpu/embedding.cpp @@ -0,0 +1,18 @@ +#include "embedding.hpp" +#include "../../../utils.hpp" +#include +namespace llaisys::ops::cpu { + + +void embedding(std::byte *out, const std::byte *index, const std::byte + *weight, llaisysDataType_t type, size_t num_indices, size_t embd_dim) { + size_t esz = llaisys::utils::dsize(type); + const int64_t *index_ptr = reinterpret_cast(index); + + for (size_t i = 0; i < num_indices; i++) { + std::memcpy(out + i * embd_dim * esz, + weight + index_ptr[i] * embd_dim * esz, + embd_dim * esz); + } +} +} \ No newline at end of file diff --git a/src/ops/embedding/cpu/embedding.hpp b/src/ops/embedding/cpu/embedding.hpp new file mode 100644 index 000000000..6d0433c84 --- /dev/null +++ b/src/ops/embedding/cpu/embedding.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "llaisys.h" + +#include + +namespace llaisys::ops::cpu { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t type, size_t num_indices, size_t embd_dim); +} diff --git a/src/ops/embedding/iluvatar/embedding_iluvatar.cu b/src/ops/embedding/iluvatar/embedding_iluvatar.cu new file mode 100644 index 000000000..8f3ad1d85 --- /dev/null +++ b/src/ops/embedding/iluvatar/embedding_iluvatar.cu @@ -0,0 +1,41 @@ +#include "embedding_iluvatar.cuh" +#include +#include "../../../utils.hpp" +#include +#include + +#define CHECK_CUDA(call) do { cudaError_t error = (call); if (error != cudaSuccess) { std::fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(error)); std::abort(); } } while (0) + + +__global__ void embedding_kernel(unsigned char *out, const int64_t *index, const unsigned char *weight, + size_t num_indices, size_t rowbytes){ + size_t offset = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + size_t total_bytes = num_indices * rowbytes; + if (offset < total_bytes){ + size_t row = offset / rowbytes; + size_t col = offset % rowbytes; + size_t weight_off = index[row] * rowbytes + col; + out[offset] =weight[weight_off]; + } + } + +namespace llaisys::ops::iluvatar { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t type, size_t num_indices, size_t embd_dim){ + size_t esz = llaisys::utils::dsize(type); + size_t row_bytes = embd_dim * esz; + size_t total_bytes = num_indices * row_bytes; + + constexpr int block_size = 256; + int grid_size = (total_bytes + block_size -1)/block_size; + + embedding_kernel<<>>( + reinterpret_cast(out), + reinterpret_cast(index), + reinterpret_cast(weight), + num_indices, + row_bytes + ); + CHECK_CUDA(cudaGetLastError()); + } + } \ No newline at end of file diff --git a/src/ops/embedding/iluvatar/embedding_iluvatar.cuh b/src/ops/embedding/iluvatar/embedding_iluvatar.cuh new file mode 100644 index 000000000..0d4070b09 --- /dev/null +++ b/src/ops/embedding/iluvatar/embedding_iluvatar.cuh @@ -0,0 +1,10 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t type, size_t num_indices, size_t embd_dim); +} \ No newline at end of file diff --git a/src/ops/embedding/nvidia/embedding_cuda.cu b/src/ops/embedding/nvidia/embedding_cuda.cu new file mode 100644 index 000000000..9e96e9d93 --- /dev/null +++ b/src/ops/embedding/nvidia/embedding_cuda.cu @@ -0,0 +1,41 @@ +#include "embedding_cuda.cuh" +#include +#include "../../../utils.hpp" +#include +#include + +#define CHECK_CUDA(call) do { cudaError_t error = (call); if (error != cudaSuccess) { std::fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(error)); std::abort(); } } while (0) + + +__global__ void embedding_kernel(unsigned char *out, const int64_t *index, const unsigned char *weight, + size_t num_indices, size_t rowbytes){ + size_t offset = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + size_t total_bytes = num_indices * rowbytes; + if (offset < total_bytes){ + size_t row = offset / rowbytes; + size_t col = offset % rowbytes; + size_t weight_off = index[row] * rowbytes + col; + out[offset] =weight[weight_off]; + } + } + +namespace llaisys::ops::cuda { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t type, size_t num_indices, size_t embd_dim){ + size_t esz = llaisys::utils::dsize(type); + size_t row_bytes = embd_dim * esz; + size_t total_bytes = num_indices * row_bytes; + + constexpr int block_size = 256; + int grid_size = (total_bytes + block_size -1)/block_size; + + embedding_kernel<<>>( + reinterpret_cast(out), + reinterpret_cast(index), + reinterpret_cast(weight), + num_indices, + row_bytes + ); + CHECK_CUDA(cudaGetLastError()); + } + } \ No newline at end of file diff --git a/src/ops/embedding/nvidia/embedding_cuda.cuh b/src/ops/embedding/nvidia/embedding_cuda.cuh new file mode 100644 index 000000000..33f4289d0 --- /dev/null +++ b/src/ops/embedding/nvidia/embedding_cuda.cuh @@ -0,0 +1,10 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { +void embedding(std::byte *out, const std::byte *index, const std::byte *weight, + llaisysDataType_t type, size_t num_indices, size_t embd_dim); +} \ No newline at end of file diff --git a/src/ops/embedding/op.cpp b/src/ops/embedding/op.cpp index 84b9a5d06..19a06438f 100644 --- a/src/ops/embedding/op.cpp +++ b/src/ops/embedding/op.cpp @@ -1,7 +1,41 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/embedding.hpp" +#include "nvidia/embedding_cuda.cuh" +#include "iluvatar/embedding_iluvatar.cuh" 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(weight->ndim() == 2, "Embedding: weight must be 2D."); + if (out->deviceType() == LLAISYS_DEVICE_CPU) { + return cpu::embedding(out->data(), index->data(), weight->data(), + out->dtype(), index->numel(), weight->shape()[1]); + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::embedding(out->data(), index->data(), weight->data(), + out->dtype(), index->numel(), weight->shape()[1]); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + cuda::embedding(out->data(), index->data(), weight->data(), + out->dtype(), index->numel(), weight->shape()[1]); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + iluvatar::embedding(out->data(), index->data(), weight->data(), + out->dtype(), index->numel(), weight->shape()[1]); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } -} // namespace llaisys::ops +} \ No newline at end of file diff --git a/src/ops/linear/cpu/linear_cpu.cpp b/src/ops/linear/cpu/linear_cpu.cpp new file mode 100644 index 000000000..cf7c953c4 --- /dev/null +++ b/src/ops/linear/cpu/linear_cpu.cpp @@ -0,0 +1,57 @@ +#include "linear_cpu.hpp" +#include "../../../utils.hpp" +// in :[M,K], weight :[N,K], bias :[N], out :[M,N] +// M: in's first dimension, N: weight's first dimension, K: in's second dimension +// out = in * weight^T + bias + +template +void linear_(T *out, const T *in, const T *weight, const T *bias, size_t M, size_t N, size_t K) { + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + if constexpr (std::is_same_v || std::is_same_v) { + float acc = 0.0f; + for (size_t k = 0; k < K; k++) { + acc += llaisys::utils::cast(in[m * K + k]) * llaisys::utils::cast(weight[n * K + k]); + } + out[m * N + n] = llaisys::utils::cast(acc); + } else { + T acc = 0; + for (size_t k = 0; k < K; k++) { + acc += in[m * K + k] * weight[n * K + k]; + } + out[m * N + n] = acc; + } + } + } + if (bias != nullptr) { + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + if constexpr (std::is_same_v || std::is_same_v) { + out[m * N + n] = llaisys::utils::cast(llaisys::utils::cast(out[m * N + n]) + llaisys::utils::cast(bias[n])); + } else { + out[m * N + n] += bias[n];; + } + + } + } + } +} + +namespace llaisys::ops::cpu { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, + const std::byte *bias, llaisysDataType_t type, size_t M, size_t N, size_t K) { + switch (type) { + case LLAISYS_DTYPE_F32: + return linear_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), reinterpret_cast(bias), M, N, K); + case LLAISYS_DTYPE_BF16: + return linear_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), reinterpret_cast(bias), M, N, K); + case LLAISYS_DTYPE_F16: + return linear_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), reinterpret_cast(bias), M, N, K); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::cpu \ No newline at end of file diff --git a/src/ops/linear/cpu/linear_cpu.hpp b/src/ops/linear/cpu/linear_cpu.hpp new file mode 100644 index 000000000..ba8de5cd1 --- /dev/null +++ b/src/ops/linear/cpu/linear_cpu.hpp @@ -0,0 +1,11 @@ +#pragma once +#include "llaisys.h" + +#include +// in :[M,K], weight :[N,K], bias :[N], out :[M,N] +// M: in's first dimension, N: weight's first dimension, K: in's second dimension +// out = in * weight^T + bias +namespace llaisys::ops::cpu { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t type, size_t M, size_t N, size_t K); +} diff --git a/src/ops/linear/iluvatar/linear_iluvatar.cu b/src/ops/linear/iluvatar/linear_iluvatar.cu new file mode 100644 index 000000000..8db84a77c --- /dev/null +++ b/src/ops/linear/iluvatar/linear_iluvatar.cu @@ -0,0 +1,150 @@ +#include "../../../utils.hpp" +#include "linear_iluvatar.cuh" +#include +#include +#include +#include "../../../device/iluvatar/iluvatar_resource.cuh" +#define CEIL(a, b) (((a) + (b) - 1) / (b)) +// F32/BF16/F16 都走 cublas；原来练手写的 tiled-GEMM kernel 见 git 历史。 + +template +__global__ void add_bias_kernel(T *out, const T *bias, size_t M, size_t N) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < M * N) out[idx] = out[idx] + bias[idx % N]; +} + + +// in : [M, K], weight : [N, K], bias : [N], out : [M, N]；out = in * weight^T + bias +namespace llaisys::ops::iluvatar { + +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t type, size_t M, size_t N, size_t K,llaisys::device::DeviceResource *resource) { + auto handle = static_cast(resource)->cublasHandle(); + switch (type) { + case LLAISYS_DTYPE_F32: + { + const float alpha = 1.0f; + const float beta = 0.0f; + + const auto *in_f32 = reinterpret_cast(in); + const auto *weight_f32 = reinterpret_cast(weight); + const auto *bias_f32 = reinterpret_cast(bias); + auto *out_f32 = reinterpret_cast(out); + + cublasStatus_t status = cublasSgemm( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + static_cast(N), + static_cast(M), + static_cast(K), + &alpha, + weight_f32, + static_cast(K), + in_f32, + static_cast(K), + &beta, + out_f32, + static_cast(N) + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + add_bias_kernel<<>>(out_f32, bias_f32 , M, N); + + } + return; + } + case LLAISYS_DTYPE_BF16: + { + // TODO: bias 可以广播进 out + beta=1.0f 融合掉，省下面单独的 add_bias_kernel + const float alpha = 1.0f; + const float beta = 0.0f; + + const int m = static_cast(M); + const int n = static_cast(N); + const int k = static_cast(K); + + + const auto *in_bf16 = reinterpret_cast(in); + const auto *weight_bf16 = reinterpret_cast(weight); + const auto *bias_bf16 = reinterpret_cast(bias); + auto *out_bf16 = reinterpret_cast<__nv_bfloat16 *>(out); + + cublasStatus_t status = cublasGemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + n, + m, + k, + &alpha, + weight_bf16, + CUDA_R_16BF, + k, + in_bf16, + CUDA_R_16BF, + k, + &beta, + out_bf16, + CUDA_R_16BF, + n, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + + add_bias_kernel<<>>(out_bf16, bias_bf16 , M, N); + + } + return; + } + + case LLAISYS_DTYPE_F16: + { + // TODO: bias 可以广播进 out + beta=1.0f 融合掉，省下面单独的 add_bias_kernel + const float alpha = 1.0f; + const float beta = 0.0f; + + const auto *in_f16 = reinterpret_cast(in); + const auto *weight_f16 = reinterpret_cast(weight); + const auto *bias_f16 = reinterpret_cast(bias); + auto *out_f16 = reinterpret_cast<__half *>(out); + + cublasStatus_t status = cublasSgemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + static_cast(N), + static_cast(M), + static_cast(K), + &alpha, + weight_f16, + CUDA_R_16F, + static_cast(K), + in_f16, + CUDA_R_16F, + static_cast(K), + &beta, + out_f16, + CUDA_R_16F, + static_cast(N) + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + + add_bias_kernel<<>>(out_f16, bias_f16 , M, N); + + } + return; + } + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::iluvatar diff --git a/src/ops/linear/iluvatar/linear_iluvatar.cuh b/src/ops/linear/iluvatar/linear_iluvatar.cuh new file mode 100644 index 000000000..8f5501087 --- /dev/null +++ b/src/ops/linear/iluvatar/linear_iluvatar.cuh @@ -0,0 +1,11 @@ +#pragma once +#include "llaisys.h" +#include "../../../device/device_resource.hpp" +#include +// in :[M,K], weight :[N,K], bias :[N], out :[M,N] +// M: in's first dimension, N: weight's first dimension, K: in's second dimension +// out = in * weight^T + bias +namespace llaisys::ops::iluvatar { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t type, size_t M, size_t N, size_t K,llaisys::device::DeviceResource *resource); +} diff --git a/src/ops/linear/nvidia/linear_cuda.cu b/src/ops/linear/nvidia/linear_cuda.cu new file mode 100644 index 000000000..4fe2d7ef5 --- /dev/null +++ b/src/ops/linear/nvidia/linear_cuda.cu @@ -0,0 +1,141 @@ +#include "../../../utils.hpp" +#include "linear_cuda.cuh" +#include +#include +#include +#include "../../../device/nvidia/nvidia_resource.cuh" +#define CEIL(a, b) (((a) + (b) - 1) / (b)) +// F32/BF16/F16 都走 cublas；原来练手写的 tiled-GEMM kernel 见 git 历史（613f360 之前）。 + +template +__global__ void add_bias_kernel(T *out, const T *bias, size_t M, size_t N) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < M * N) out[idx] = out[idx] + bias[idx % N]; +} + + +// in : [M, K], weight : [N, K], bias : [N], out : [M, N]；out = in * weight^T + bias +namespace llaisys::ops::cuda { + +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t type, size_t M, size_t N, size_t K,llaisys::device::DeviceResource *resource) { + auto handle = static_cast(resource)->cublasHandle(); + switch (type) { + case LLAISYS_DTYPE_F32: + { + const float alpha = 1.0f; + const float beta = 0.0f; + + const auto *in_f32 = reinterpret_cast(in); + const auto *weight_f32 = reinterpret_cast(weight); + const auto *bias_f32 = reinterpret_cast(bias); + auto *out_f32 = reinterpret_cast(out); + + cublasStatus_t status = cublasSgemm( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + static_cast(N), + static_cast(M), + static_cast(K), + &alpha, + weight_f32, + static_cast(K), + in_f32, + static_cast(K), + &beta, + out_f32, + static_cast(N) + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + add_bias_kernel<<>>(out_f32, bias_f32 , M, N); + + } + return; + } + case LLAISYS_DTYPE_BF16: + { + const float alpha = 1.0f; + const float beta = 0.0f; + + const auto *in_bf16 = reinterpret_cast(in); + const auto *weight_bf16 = reinterpret_cast(weight); + const auto *bias_bf16 = reinterpret_cast(bias); + auto *out_bf16 = reinterpret_cast<__nv_bfloat16 *>(out); + + cublasStatus_t status = cublasSgemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + static_cast(N), + static_cast(M), + static_cast(K), + &alpha, + weight_bf16, + CUDA_R_16BF, + static_cast(K), + in_bf16, + CUDA_R_16BF, + static_cast(K), + &beta, + out_bf16, + CUDA_R_16BF, + static_cast(N) + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + + add_bias_kernel<<>>(out_bf16, bias_bf16 , M, N); + + } + return; + } + + case LLAISYS_DTYPE_F16: + { + const float alpha = 1.0f; + const float beta = 0.0f; + + const auto *in_f16 = reinterpret_cast(in); + const auto *weight_f16 = reinterpret_cast(weight); + const auto *bias_f16 = reinterpret_cast(bias); + auto *out_f16 = reinterpret_cast<__half *>(out); + + cublasStatus_t status = cublasSgemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + static_cast(N), + static_cast(M), + static_cast(K), + &alpha, + weight_f16, + CUDA_R_16F, + static_cast(K), + in_f16, + CUDA_R_16F, + static_cast(K), + &beta, + out_f16, + CUDA_R_16F, + static_cast(N) + ); + if(bias != nullptr){ + constexpr int block_size = 256; + int grid_size = static_cast(CEIL(M * N,block_size)); + + add_bias_kernel<<>>(out_f16, bias_f16 , M, N); + + } + return; + } + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/linear/nvidia/linear_cuda.cuh b/src/ops/linear/nvidia/linear_cuda.cuh new file mode 100644 index 000000000..f93e8ba07 --- /dev/null +++ b/src/ops/linear/nvidia/linear_cuda.cuh @@ -0,0 +1,11 @@ +#pragma once +#include "llaisys.h" +#include "../../../device/device_resource.hpp" +#include +// in :[M,K], weight :[N,K], bias :[N], out :[M,N] +// M: in's first dimension, N: weight's first dimension, K: in's second dimension +// out = in * weight^T + bias +namespace llaisys::ops::cuda { +void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias, + llaisysDataType_t type, size_t M, size_t N, size_t K,llaisys::device::DeviceResource *resource); +} diff --git a/src/ops/linear/op.cpp b/src/ops/linear/op.cpp index 97d1f8655..3ffc2be71 100644 --- a/src/ops/linear/op.cpp +++ b/src/ops/linear/op.cpp @@ -1,7 +1,52 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" +#include "nvidia/linear_cuda.cuh" +#include "iluvatar/linear_iluvatar.cuh" +#include "cpu/linear_cpu.hpp" namespace llaisys::ops { void linear(tensor_t out, tensor_t in, tensor_t weight, tensor_t bias) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(out, in, weight); + CHECK_SAME_DTYPE(out->dtype(), in->dtype(), weight->dtype()); + ASSERT(in->ndim() == 2, "Linear: in must be 2D."); + ASSERT(weight->ndim() == 2, "Linear: weight must be 2D."); + ASSERT(in->shape()[1] == weight->shape()[1], "Linear: in's second dimension must match weight's second dimension."); + ASSERT(out->shape()[0] == in->shape()[0], "Linear: out's first dimension must match in's first dimension."); + ASSERT(out->shape()[1] == weight->shape()[0], "Linear: out's second dimension must match weight's first dimension."); + + std::byte *bias_ptr = nullptr; + if (bias) { + CHECK_SAME_DEVICE(out, bias); + CHECK_SAME_DTYPE(out->dtype(), bias->dtype()); + ASSERT(bias->ndim() == 1 && bias->shape()[0] == weight->shape()[0], + "Linear: bias must be 1D with size equal to weight's first dimension."); + bias_ptr = bias->data(); + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::linear(out->data(), in->data(), weight->data(), bias_ptr, + out->dtype(), in->shape()[0], weight->shape()[0], in->shape()[1]); + #ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: { + auto resource = llaisys::core::context().runtime().resource(); + cuda::linear(out->data(), in->data(), weight->data(), bias_ptr, + out->dtype(), in->shape()[0], weight->shape()[0], in->shape()[1],resource); + return; + } + #endif + #ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: { + auto resource = llaisys::core::context().runtime().resource(); + iluvatar::linear(out->data(), in->data(), weight->data(), bias_ptr, + out->dtype(), in->shape()[0], weight->shape()[0], in->shape()[1],resource); + return; + } + #endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/rearrange/iluvatar/rearrange_iluvatar.cu b/src/ops/rearrange/iluvatar/rearrange_iluvatar.cu new file mode 100644 index 000000000..e69de29bb diff --git a/src/ops/rearrange/iluvatar/rearrange_iluvatar.cuh b/src/ops/rearrange/iluvatar/rearrange_iluvatar.cuh new file mode 100644 index 000000000..cf1e32350 --- /dev/null +++ b/src/ops/rearrange/iluvatar/rearrange_iluvatar.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + + + +} \ No newline at end of file diff --git a/src/ops/rearrange/nvidia/rearrange_cuda.cu b/src/ops/rearrange/nvidia/rearrange_cuda.cu new file mode 100644 index 000000000..e69de29bb diff --git a/src/ops/rearrange/nvidia/rearrange_cuda.cuh b/src/ops/rearrange/nvidia/rearrange_cuda.cuh new file mode 100644 index 000000000..65023fca7 --- /dev/null +++ b/src/ops/rearrange/nvidia/rearrange_cuda.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + + + +} \ No newline at end of file diff --git a/src/ops/rearrange/op.cpp b/src/ops/rearrange/op.cpp index 017a6ae59..68c534eae 100644 --- a/src/ops/rearrange/op.cpp +++ b/src/ops/rearrange/op.cpp @@ -1,7 +1,36 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + namespace llaisys::ops { void rearrange(tensor_t out, tensor_t in) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(out, in); + + // 弃用：永久 stub，不会实现（KV cache 拷贝等场景改用 memcpy，见 CLAUDE.md）。 + if (out->deviceType() == LLAISYS_DEVICE_CPU) { + TO_BE_IMPLEMENTED(); + return; + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + TO_BE_IMPLEMENTED(); + return; +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + TO_BE_IMPLEMENTED(); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + TO_BE_IMPLEMENTED(); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/rms_norm/cpu/rms_norm.cpp b/src/ops/rms_norm/cpu/rms_norm.cpp new file mode 100644 index 000000000..2ec863b93 --- /dev/null +++ b/src/ops/rms_norm/cpu/rms_norm.cpp @@ -0,0 +1,39 @@ +#include "rms_norm.hpp" +#include "../../../utils.hpp" +#include +// out/in : [rows, d]，weight : [d] +// rms_m = sqrt(mean(in[m,:]^2) + eps)，out[m,k] = weight[k] * in[m,k] / rms_m +template +void rms_norm_(T *out, const T *in, const T *weight, float eps, size_t rows, size_t d) { + for (size_t m = 0; m < rows; m++) { + float mean_square = 0.0f; + for (size_t k = 0; k < d; k++) { + mean_square += llaisys::utils::cast(in[m * d + k]) * llaisys::utils::cast(in[m * d + k]); + } + mean_square /= static_cast(d); + float denom = std::sqrt(mean_square + eps); + for (size_t k = 0; k < d; k++) { + out[m * d + k] = llaisys::utils::cast(llaisys::utils::cast(weight[k]) * llaisys::utils::cast(in[m * d + k]) / denom); + } + + } +} + +namespace llaisys::ops::cpu { +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + return rms_norm_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), eps, rows, d); + case LLAISYS_DTYPE_BF16: + return rms_norm_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), eps, rows, d); + case LLAISYS_DTYPE_F16: + return rms_norm_(reinterpret_cast(out), reinterpret_cast(in), + reinterpret_cast(weight), eps, rows, d); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::cpu \ No newline at end of file diff --git a/src/ops/rms_norm/cpu/rms_norm.hpp b/src/ops/rms_norm/cpu/rms_norm.hpp new file mode 100644 index 000000000..14c6d069c --- /dev/null +++ b/src/ops/rms_norm/cpu/rms_norm.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "llaisys.h" + +#include +// out[i][j] = weight[j] * in[i][j] / sqrt(mean_j(in[i][j]^2) + eps) +// in :[rows, d], weight :[d], out :[rows, d] +namespace llaisys::ops::cpu { +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d); +} diff --git a/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cu b/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cu new file mode 100644 index 000000000..2457a6cbb --- /dev/null +++ b/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cu @@ -0,0 +1,86 @@ +#include "../../../utils.hpp" +#include "rms_norm_iluvatar.cuh" +#include +#include +#include + +// rows: token 数；d: hidden_size。一个 block 处理一行，先规约出 sum-of-squares 再归一化。 +template +__global__ void rms_norm_kernel(T *out, const T *in, const T *weight, float eps, size_t rows, size_t d) { + size_t row = blockIdx.x; + size_t tid = threadIdx.x; + + float temp = 0.0f; + for (size_t i = tid; i < d; i+=blockDim.x) + { + size_t idx = row*d + i; + temp += (float)in[idx] * (float)in[idx]; + } + + + extern __shared__ float shared[]; + shared[tid] = temp; + __syncthreads(); + for(size_t stride = blockDim.x/2; stride>0;stride/=2){ + if (tid < stride){ + shared[tid] += shared[tid + stride]; + } + __syncthreads(); + } + float rms = sqrtf(shared[0] / static_cast(d) + eps); + for(size_t i = tid; i < d; i+=blockDim.x){ + out[row * d + i] = (float)weight[i] * (float)in[row * d + i] / rms; + } + + return; +} + +template +void launch_rms_norm(T *out, const T *in, const T *weight, float eps, size_t rows, size_t d) { + constexpr int block_size = 256; + size_t grid_size = rows; + size_t shared_bytes = block_size * sizeof(float); + rms_norm_kernel<<>>(out, in, weight,eps,rows,d); +} + +namespace llaisys::ops::iluvatar { + +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_rms_norm( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + case LLAISYS_DTYPE_BF16: + launch_rms_norm( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + case LLAISYS_DTYPE_F16: + launch_rms_norm( + reinterpret_cast<__half *>(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::iluvatar \ No newline at end of file diff --git a/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cuh b/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cuh new file mode 100644 index 000000000..5d360948d --- /dev/null +++ b/src/ops/rms_norm/iluvatar/rms_norm_iluvatar.cuh @@ -0,0 +1,12 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d); + +} \ No newline at end of file diff --git a/src/ops/rms_norm/nvidia/rms_norm_cuda.cu b/src/ops/rms_norm/nvidia/rms_norm_cuda.cu new file mode 100644 index 000000000..b07bf2284 --- /dev/null +++ b/src/ops/rms_norm/nvidia/rms_norm_cuda.cu @@ -0,0 +1,86 @@ +#include "../../../utils.hpp" +#include "rms_norm_cuda.cuh" +#include +#include +#include + +// rows: token 数；d: hidden_size。一个 block 处理一行，先规约出 sum-of-squares 再归一化。 +template +__global__ void rms_norm_kernel(T *out, const T *in, const T *weight, float eps, size_t rows, size_t d) { + size_t row = blockIdx.x; + size_t tid = threadIdx.x; + + float temp = 0.0f; + for (size_t i = tid; i < d; i+=blockDim.x) + { + size_t idx = row*d + i; + temp += (float)in[idx] * (float)in[idx]; + } + + + extern __shared__ float shared[]; + shared[tid] = temp; + __syncthreads(); + for(size_t stride = blockDim.x/2; stride>0;stride/=2){ + if (tid < stride){ + shared[tid] += shared[tid + stride]; + } + __syncthreads(); + } + float rms = sqrtf(shared[0] / static_cast(d) + eps); + for(size_t i = tid; i < d; i+=blockDim.x){ + out[row * d + i] = (float)weight[i] * (float)in[row * d + i] / rms; + } + + return; +} + +template +void launch_rms_norm(T *out, const T *in, const T *weight, float eps, size_t rows, size_t d) { + constexpr int block_size = 256; + size_t grid_size = rows; + size_t shared_bytes = block_size * sizeof(float); + rms_norm_kernel<<>>(out, in, weight,eps,rows,d); +} + +namespace llaisys::ops::cuda { + +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_rms_norm( + reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + case LLAISYS_DTYPE_BF16: + launch_rms_norm( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + case LLAISYS_DTYPE_F16: + launch_rms_norm( + reinterpret_cast<__half *>(out), + reinterpret_cast(in), + reinterpret_cast(weight), + eps, + rows, + d); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda \ No newline at end of file diff --git a/src/ops/rms_norm/nvidia/rms_norm_cuda.cuh b/src/ops/rms_norm/nvidia/rms_norm_cuda.cuh new file mode 100644 index 000000000..522e162b3 --- /dev/null +++ b/src/ops/rms_norm/nvidia/rms_norm_cuda.cuh @@ -0,0 +1,12 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +void rms_norm(std::byte *out, const std::byte *in, const std::byte *weight, + float eps, llaisysDataType_t type, size_t rows, size_t d); + +} \ No newline at end of file diff --git a/src/ops/rms_norm/op.cpp b/src/ops/rms_norm/op.cpp index 529553d9d..206d104b0 100644 --- a/src/ops/rms_norm/op.cpp +++ b/src/ops/rms_norm/op.cpp @@ -1,7 +1,44 @@ #include "op.hpp" - +#include "cpu/rms_norm.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" +#include "nvidia/rms_norm_cuda.cuh" +#include "iluvatar/rms_norm_iluvatar.cuh" +// y[i] = w[i] * x[i] / sqrt(mean(x^2) + eps)；in/out : [rows, d], weight : [d] namespace llaisys::ops { void rms_norm(tensor_t out, tensor_t in, tensor_t weight, float eps) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(out, in, weight); + CHECK_SAME_DTYPE(out->dtype(), in->dtype(), weight->dtype()); + ASSERT(out->isContiguous() && in->isContiguous() && weight->isContiguous(), "RMSNorm: all tensors must be contiguous."); + ASSERT(in->ndim() == 2, "RMSNorm: in must be 2D."); + ASSERT(weight->ndim() == 1, "RMSNorm: weight must be 1D."); + ASSERT(in->shape()[1] == weight->shape()[0], "RMSNorm: in's second dimension must match weight's first dimension."); + ASSERT(out->shape()[0] == in->shape()[0], "RMSNorm: out's first dimension must match in's first dimension."); + ASSERT(out->shape()[1] == in->shape()[1], "RMSNorm: out's second dimension must match in's second dimension."); + + if (out->deviceType() == LLAISYS_DEVICE_CPU) { + llaisys::ops::cpu::rms_norm(out->data(), in->data(), weight->data(), eps, out->dtype(), in->shape()[0], in->shape()[1]); + return; + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + llaisys::ops::cpu::rms_norm(out->data(), in->data(), weight->data(), eps, out->dtype(), in->shape()[0], in->shape()[1]); + return; +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + llaisys::ops::cuda::rms_norm(out->data(), in->data(), weight->data(), eps, out->dtype(), in->shape()[0], in->shape()[1]); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + llaisys::ops::iluvatar::rms_norm(out->data(), in->data(), weight->data(), eps, out->dtype(), in->shape()[0], in->shape()[1]); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/rope/cpu/rope_cpu.cpp b/src/ops/rope/cpu/rope_cpu.cpp new file mode 100644 index 000000000..2210e8497 --- /dev/null +++ b/src/ops/rope/cpu/rope_cpu.cpp @@ -0,0 +1,39 @@ +#include "rope_cpu.hpp" +#include "../../../utils.hpp" +#include +// a' = a*cos(phi) - b*sin(phi)，b' = b*cos(phi) + a*sin(phi)，(a,b) = (in[j], in[j+d/2]) +// phi = pos_ids[i] / theta^(2j/d), j = 0 .. d/2-1 +template +void rope_(T *out, const T *in, const int64_t *pos_ids, float theta, size_t seqlen, size_t nhead, size_t d) { + for(size_t i = 0; i < seqlen; i++) { + for(size_t h = 0; h < nhead; h++) { + size_t base = i * nhead * d + h * d; + for(size_t j = 0; j < d / 2; j++) { + float phi = pos_ids[i] / std::pow(theta, 2.0f * j / d); + float cos_phi = std::cos(phi); + float sin_phi = std::sin(phi); + float a = llaisys::utils::cast(in[base + j]); + float b = llaisys::utils::cast(in[base + j + d/2]); + out[base + j] = llaisys::utils::cast(a * cos_phi - b * sin_phi); + out[base + j + d/2] = llaisys::utils::cast(b * cos_phi + a * sin_phi); + } + } + } +} + +namespace llaisys::ops::cpu { +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta) { + const int64_t *pos = reinterpret_cast(pos_ids); + switch (type) { + case LLAISYS_DTYPE_F32: + return rope_(reinterpret_cast(out), reinterpret_cast(in), pos, theta, seqlen, nhead, d); + case LLAISYS_DTYPE_BF16: + return rope_(reinterpret_cast(out), reinterpret_cast(in), pos, theta, seqlen, nhead, d); + case LLAISYS_DTYPE_F16: + return rope_(reinterpret_cast(out), reinterpret_cast(in), pos, theta, seqlen, nhead, d); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} diff --git a/src/ops/rope/cpu/rope_cpu.hpp b/src/ops/rope/cpu/rope_cpu.hpp new file mode 100644 index 000000000..7d1e773f5 --- /dev/null +++ b/src/ops/rope/cpu/rope_cpu.hpp @@ -0,0 +1,10 @@ +#pragma once + +#include "llaisys.h" + +#include +// in/out :[seqlen, nhead, d], pos_ids :[seqlen] +namespace llaisys::ops::cpu { +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta); +} \ No newline at end of file diff --git a/src/ops/rope/iluvatar/rope_iluvatar.cu b/src/ops/rope/iluvatar/rope_iluvatar.cu new file mode 100644 index 000000000..ed6b0502c --- /dev/null +++ b/src/ops/rope/iluvatar/rope_iluvatar.cu @@ -0,0 +1,71 @@ +#include "../../../utils.hpp" +#include "rope_iluvatar.cuh" +#include +#include +#include + +// 手写 kernel（cuDNN RoPE Graph 节点只支持 f16/bf16 不支持 f32，探索后放弃，见 CLAUDE.md）。 +template +__global__ void rope_kernel(T *out, const T *in, const int64_t *pos_ids, float theta, + size_t nhead, size_t d) { + size_t i = blockIdx.x; // token 位置 + size_t h = blockIdx.y; // head 编号 + size_t tid = threadIdx.x; + size_t base = i * nhead * d + h * d; + size_t stride = blockDim.x; + for(size_t j = tid; j< d/2;j+=stride){ + float phi = pos_ids[i] / pow(theta, 2.0f * j/d); + float cos_phi = cos(phi); + float sin_phi = sin(phi); + float a = (float)in[base + j]; + float b = (float)in[base + j + d/2]; + out[base + j] = (float) (a * cos_phi - b * sin_phi); + out[base + j + d/2] = (float) (b * cos_phi + a * sin_phi); + } +} + +template +void launch_rope(T *out, const T *in, const int64_t *pos_ids, float theta, + size_t seqlen, size_t nhead, size_t d) { + constexpr int block_size = 256; + dim3 grid(static_cast(seqlen), static_cast(nhead)); + + rope_kernel<<>>(out, in, pos_ids, theta, nhead, d); +} + +namespace llaisys::ops::iluvatar { + +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta) { + const int64_t *pos = reinterpret_cast(pos_ids); + switch (type) { + case LLAISYS_DTYPE_F32: + launch_rope( + reinterpret_cast(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + case LLAISYS_DTYPE_BF16: + launch_rope( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + case LLAISYS_DTYPE_F16: + launch_rope( + reinterpret_cast<__half *>(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::iluvatar diff --git a/src/ops/rope/iluvatar/rope_iluvatar.cuh b/src/ops/rope/iluvatar/rope_iluvatar.cuh new file mode 100644 index 000000000..ee317a567 --- /dev/null +++ b/src/ops/rope/iluvatar/rope_iluvatar.cuh @@ -0,0 +1,12 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta); + +} \ No newline at end of file diff --git a/src/ops/rope/nvidia/rope_cuda.cu b/src/ops/rope/nvidia/rope_cuda.cu new file mode 100644 index 000000000..2db0e4c58 --- /dev/null +++ b/src/ops/rope/nvidia/rope_cuda.cu @@ -0,0 +1,71 @@ +#include "../../../utils.hpp" +#include "rope_cuda.cuh" +#include +#include +#include + +// 手写 kernel（cuDNN RoPE Graph 节点只支持 f16/bf16 不支持 f32，探索后放弃，见 CLAUDE.md）。 +template +__global__ void rope_kernel(T *out, const T *in, const int64_t *pos_ids, float theta, + size_t nhead, size_t d) { + size_t i = blockIdx.x; // token 位置 + size_t h = blockIdx.y; // head 编号 + size_t tid = threadIdx.x; + size_t base = i * nhead * d + h * d; + size_t stride = blockDim.x; + for(size_t j = tid; j< d/2;j+=stride){ + float phi = pos_ids[i] / pow(theta, 2.0f * j/d); + float cos_phi = cos(phi); + float sin_phi = sin(phi); + float a = (float)in[base + j]; + float b = (float)in[base + j + d/2]; + out[base + j] = (float) (a * cos_phi - b * sin_phi); + out[base + j + d/2] = (float) (b * cos_phi + a * sin_phi); + } +} + +template +void launch_rope(T *out, const T *in, const int64_t *pos_ids, float theta, + size_t seqlen, size_t nhead, size_t d) { + constexpr int block_size = 256; + dim3 grid(static_cast(seqlen), static_cast(nhead)); + + rope_kernel<<>>(out, in, pos_ids, theta, nhead, d); +} + +namespace llaisys::ops::cuda { + +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta) { + const int64_t *pos = reinterpret_cast(pos_ids); + switch (type) { + case LLAISYS_DTYPE_F32: + launch_rope( + reinterpret_cast(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + case LLAISYS_DTYPE_BF16: + launch_rope( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + case LLAISYS_DTYPE_F16: + launch_rope( + reinterpret_cast<__half *>(out), + reinterpret_cast(in), + pos, theta, seqlen, nhead, d + ); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/rope/nvidia/rope_cuda.cuh b/src/ops/rope/nvidia/rope_cuda.cuh new file mode 100644 index 000000000..c48d14ecc --- /dev/null +++ b/src/ops/rope/nvidia/rope_cuda.cuh @@ -0,0 +1,12 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +void rope(std::byte *out, const std::byte *in, const std::byte *pos_ids, + llaisysDataType_t type, size_t seqlen, size_t nhead, size_t d, float theta); + +} \ No newline at end of file diff --git a/src/ops/rope/op.cpp b/src/ops/rope/op.cpp index d60dbe64e..d3b42c0bf 100644 --- a/src/ops/rope/op.cpp +++ b/src/ops/rope/op.cpp @@ -1,7 +1,43 @@ #include "op.hpp" +#include "nvidia/rope_cuda.cuh" +#include "iluvatar/rope_iluvatar.cuh" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" +#include "cpu/rope_cpu.hpp" +// a' = a*cos(phi) - b*sin(phi)，b' = b*cos(phi) + a*sin(phi) +// out/in : [seqlen, nhead_or_nkvhead, d]；pos_ids : [seqlen,] int64；theta：频率向量基值。 namespace llaisys::ops { void rope(tensor_t out, tensor_t in, tensor_t pos_ids, float theta) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(out, in, pos_ids); + CHECK_SAME_DTYPE(out->dtype(), in->dtype()); + ASSERT(in->ndim() == 3, "RoPE: in must be 3D."); + ASSERT(pos_ids->ndim() == 1, "RoPE: pos_ids must be 1D."); + ASSERT(out->shape()[0] == in->shape()[0], "RoPE: out's first dimension must match in's first dimension."); + ASSERT(out->shape()[1] == in->shape()[1], "RoPE: out's second dimension must match in's second dimension."); + + // TODO: 没校验 pos_ids 的 dtype 恒为 int64 + if (out->deviceType() == LLAISYS_DEVICE_CPU) { + return cpu::rope(out->data(), in->data(), pos_ids->data(), out->dtype(), in->shape()[0], in->shape()[1], in->shape()[2], theta); + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::rope(out->data(), in->data(), pos_ids->data(), out->dtype(), in->shape()[0], in->shape()[1], in->shape()[2], theta); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + cuda::rope(out->data(), in->data(), pos_ids->data(), out->dtype(), in->shape()[0], in->shape()[1], in->shape()[2], theta); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + iluvatar::rope(out->data(), in->data(), pos_ids->data(), out->dtype(), in->shape()[0], in->shape()[1], in->shape()[2], theta); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/self_attention/cpu/self_attention.cpp b/src/ops/self_attention/cpu/self_attention.cpp new file mode 100644 index 000000000..bbffc3d04 --- /dev/null +++ b/src/ops/self_attention/cpu/self_attention.cpp @@ -0,0 +1,79 @@ +#include "self_attention.hpp" +#include "../../../utils.hpp" +#include +#include +#include + +// 两遍扫描做数值稳定的 softmax：第一遍算 score[j] 和 max_score，第二遍用 +// exp(score[j]-max_score)（输入恒 <=0，不会溢出）做加权求和。 +// q : [seqlen, nhead, d ] +// k : [total_len, nkvhead, d ] +// qk^T : [seqlen, nhead, total_len] +// softmax(scale * qk^T) : [seqlen, nhead, total_len] +// v : [total_len, nkvhead, dv] +// attn_val : [seqlen, nhead, dv] +// attn_val = softmax(scale * q * k^T) * v +template +void self_attention_(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + size_t causal_offset = total_len - seqlen; // 位置 i 能看到的 key 是 j = 0..(i + causal_offset) + size_t group = nhead / nkvhead; // 每 group 个 query head 共享 1 个 kv head + + for (size_t i = 0; i < seqlen; i++) { + for (size_t h = 0; h < nhead; h++) { + size_t kvh = h / group; + size_t limit = i + causal_offset; // 因果掩码:j 的范围是 [0, limit](闭区间) + + std::vector scores(limit + 1); + float max_score = -std::numeric_limits::infinity(); + + for (size_t j = 0; j <= limit; j++) { + float score = 0.0f; + for (size_t dim = 0; dim < d; dim++) { + score += llaisys::utils::cast(q[i * nhead * d + h * d + dim]) * llaisys::utils::cast(k[j * nkvhead * d + kvh * d + dim]); + } + score *= scale; + scores[j] = score; + max_score = std::max(max_score, score); + } + + float sum_exp = 0.0f; + std::vector acc(dv, 0.0f); + for (size_t j = 0; j <= limit; j++) { + float e = std::exp(scores[j] - max_score); + sum_exp += e; + for (size_t t = 0; t < dv; t++) { + acc[t] += e * llaisys::utils::cast(v[j * nkvhead * dv + kvh * dv + t]); + } + } + for (size_t t = 0; t < dv; t++) { + attn_val[i * nhead * dv + h * dv + t] = llaisys::utils::cast(acc[t] / sum_exp); + } + } + } +} + +namespace llaisys::ops::cpu { +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + switch (type) { + case LLAISYS_DTYPE_F32: + return self_attention_(reinterpret_cast(attn_val), reinterpret_cast(q), + reinterpret_cast(k), reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + case LLAISYS_DTYPE_BF16: + return self_attention_(reinterpret_cast(attn_val), reinterpret_cast(q), + reinterpret_cast(k), reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + case LLAISYS_DTYPE_F16: + return self_attention_(reinterpret_cast(attn_val), reinterpret_cast(q), + reinterpret_cast(k), reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/self_attention/cpu/self_attention.hpp b/src/ops/self_attention/cpu/self_attention.hpp new file mode 100644 index 000000000..5bd36ff9a --- /dev/null +++ b/src/ops/self_attention/cpu/self_attention.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "llaisys.h" + +#include + +// GQA: nhead 必须是 nkvhead 的整数倍,kvh(h) = h / (nhead/nkvhead)。 +// causal: query 位置 i 只能看到 key 位置 j <= i + (total_len - seqlen)。 +// +// attn_val : [seqlen, nhead, dv] +// q : [seqlen, nhead, d ] +// k : [total_len, nkvhead, d ] +// v : [total_len, nkvhead, dv] +namespace llaisys::ops::cpu { +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); +} diff --git a/src/ops/self_attention/iluvatar/self_attention_iluvatar.cu b/src/ops/self_attention/iluvatar/self_attention_iluvatar.cu new file mode 100644 index 000000000..5bf412fcc --- /dev/null +++ b/src/ops/self_attention/iluvatar/self_attention_iluvatar.cu @@ -0,0 +1,129 @@ +#include "../../../utils.hpp" +#include "self_attention_iluvatar.cuh" +#include +#include +#include +#include + +// 通用兜底 kernel（任意 shape/causal_offset）：一个 block 处理一个 (query token i, head h)， +// 三遍扫描（打分+max / softmax 归一 / 加权求和 v），scores 用动态 shared memory 存一整行。 +template +__global__ void self_attention_kernel(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + size_t i = blockIdx.x; // query token, [0, seqlen) + size_t h = blockIdx.y; // query head, [0, nhead) + size_t tid = threadIdx.x; + + size_t group = nhead / nkvhead; // 每 group 个 query head 共享 1 个 kv head + size_t kvh = h / group; // 该 query head 对应的 kv head + size_t causal_offset = total_len - seqlen; + size_t limit = i + causal_offset; // 因果掩码: j 的范围是 [0, limit]（闭区间） + + extern __shared__ float scores[]; // 大小 total_len，只用到 [0, limit] 这一段 + __shared__ float sdata[256]; + // 第一遍打分：每个线程按 grid-stride 负责一部分 j（j = tid, tid+stride, ... <= limit）。 + float max_score = -INFINITY; + for (size_t j = tid; j <= limit; j += blockDim.x) { + float score = 0.0f; + + for (size_t dim = 0; dim < d; dim++) { + score += float(q[i * nhead * d + h * d + dim]) * float(k[j * nkvhead * d + kvh * d + dim]); + } + score *= scale; + scores[j] = score; + max_score = std::fmaxf(max_score, score); + } + sdata[tid] = max_score; + __syncthreads(); + + // block 内做 max 规约得到 max_score。 + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride) { + sdata[tid] = max(sdata[tid], sdata[tid + stride]); + } + __syncthreads(); + } + max_score = sdata[0]; + + __shared__ float sumexp[256]; + sumexp[tid] = 0.0; + // 第二遍:用 exp(score - max_score) 做数值稳定的 softmax 加权求和。 + for (size_t j = tid; j <= limit; j += blockDim.x) { + scores[j] = std::exp(scores[j] - max_score); + sumexp[tid] += scores[j]; + } + __syncthreads(); + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride) { + sumexp[tid] = sumexp[tid] + sumexp[tid + stride]; + } + __syncthreads(); + } + float sum_exp = sumexp[0]; + + // 第三遍：按 dv 维度分工，加权求和 v 得到最终输出。 + for (size_t t = tid; t < dv; t += blockDim.x) { + float acc = 0.0; + for (size_t j = 0; j <= limit; j += 1) { + acc += (float)(scores[j]) * (float)(v[j * nkvhead * dv + kvh * dv + t]); + } + attn_val[i * nhead * dv + h * dv + t] = acc / sum_exp; + } +} + +template +void launch_self_attention(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + constexpr int block_size = 256; + dim3 grid(static_cast(seqlen), static_cast(nhead)); + size_t shared_bytes = total_len * sizeof(float); + + self_attention_kernel<<>>( + attn_val, q, k, v, seqlen, total_len, nhead, nkvhead, d, dv, scale); +} + +namespace llaisys::ops::iluvatar { +// q : [seqlen, nhead, d ] +// k : [total_len, nkvhead, d ] +// qk^T : [seqlen, nhead, total_len] +// softmax(scale * qk^T) : [seqlen, nhead, total_len] +// v : [total_len, nkvhead, dv] +// attn_val : [seqlen, nhead, dv] +// attn_val = softmax(scale * q * k^T) * v +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_self_attention( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_BF16: + launch_self_attention( + reinterpret_cast<__nv_bfloat16 *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_F16: + launch_self_attention( + reinterpret_cast<__half *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::iluvatar diff --git a/src/ops/self_attention/iluvatar/self_attention_iluvatar.cuh b/src/ops/self_attention/iluvatar/self_attention_iluvatar.cuh new file mode 100644 index 000000000..29eb866ac --- /dev/null +++ b/src/ops/self_attention/iluvatar/self_attention_iluvatar.cuh @@ -0,0 +1,14 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); + +} \ No newline at end of file diff --git a/src/ops/self_attention/nvidia/flash_attention_cuda.cuh b/src/ops/self_attention/nvidia/flash_attention_cuda.cuh new file mode 100644 index 000000000..e8188c2d1 --- /dev/null +++ b/src/ops/self_attention/nvidia/flash_attention_cuda.cuh @@ -0,0 +1,27 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +void flash_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); + +// seqlen=1 的 decode 专用入口，对完整 KV cache 执行 attention。 +void flash_attention_decode(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); + +// seqlen=1 且 total_len 较大时的 decode 入口：把 KV 方向切成多段并行规约（flash-decoding）， +// 缓解 flash_attention_decode 只用 nhead 个 warp、SM 打不满的问题。 +void flash_attention_decode_splitkv(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); + +} diff --git a/src/ops/self_attention/nvidia/flash_attention_decode_cuda.cu b/src/ops/self_attention/nvidia/flash_attention_decode_cuda.cu new file mode 100644 index 000000000..e45da9e2d --- /dev/null +++ b/src/ops/self_attention/nvidia/flash_attention_decode_cuda.cu @@ -0,0 +1,145 @@ +#include "../../../utils.hpp" +#include "flash_attention_cuda.cuh" + +#include +#include +#include + +namespace { + +constexpr int WARP_SIZE = 32; +constexpr int HEAD_DIM = 128; +constexpr int TILE_K = 32; + +// Warp 规约后将 lane 0 的结果广播给所有 lane，便于后续统一更新 online softmax 状态。 +__device__ float warp_max(float value) { + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + value = max(value, __shfl_down_sync(0xffffffff, value, offset)); + } + return __shfl_sync(0xffffffff, value, 0); +} + +__device__ float warp_sum(float value) { + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return __shfl_sync(0xffffffff, value, 0); +} + +template +__global__ void flash_attention_decode_kernel(T *output, const T *q, const T *k, const T *v, + size_t total_len, size_t nhead, size_t nkvhead, + float scale) { + // Decode 时 seqlen=1：一个 block 处理一个 query head，一个 warp 处理该 head 的全部 KV cache。 + const int lane = threadIdx.x; + const size_t head = blockIdx.x; + const size_t kv_head = head / (nhead / nkvhead); // GQA + + // K/V 各占 TILE_K * HEAD_DIM 个 float，每次只保留当前 KV tile。 + extern __shared__ float shared[]; + float *k_tile = shared; + float *v_tile = shared + TILE_K * HEAD_DIM; + + float max_score = -INFINITY; // 已处理 tile 的最大 attention score + float sum_exp = 0.0f; // 以 max_score 为基准的 softmax 分母 + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; // 每个 lane 负责 4 个输出维度 + + for (size_t tile_start = 0; tile_start < total_len; tile_start += TILE_K) { + const size_t tile_size = min(static_cast(TILE_K), total_len - tile_start); + + // 32 个 lane 协作将当前 K/V tile 转为 float 后搬入 shared memory。 + for (size_t flat = lane; flat < tile_size * HEAD_DIM; flat += WARP_SIZE) { + const size_t row = flat / HEAD_DIM; + const size_t dim = flat % HEAD_DIM; + k_tile[flat] = static_cast( + k[(tile_start + row) * nkvhead * HEAD_DIM + kv_head * HEAD_DIM + dim]); + v_tile[flat] = static_cast( + v[(tile_start + row) * nkvhead * HEAD_DIM + kv_head * HEAD_DIM + dim]); + } + __syncwarp(); + + float score = -INFINITY; + if (lane < tile_size) { + score = 0.0f; + for (int dim = 0; dim < HEAD_DIM; ++dim) { + score += static_cast(q[head * HEAD_DIM + dim]) * k_tile[lane * HEAD_DIM + dim]; + } + score *= scale; + } + + // 合并当前 tile 与历史 tile 的 online softmax；max 变化时需重缩放历史累加值。 + const float tile_max = warp_max(score); + const float new_max = max(max_score, tile_max); + const float old_scale = expf(max_score - new_max); + const float probability = lane < tile_size ? expf(score - new_max) : 0.0f; + + sum_exp = sum_exp * old_scale + warp_sum(probability); + for (int item = 0; item < 4; ++item) { + acc[item] *= old_scale; + } + + // 每个 lane 持有一个 key 的权重，通过 shuffle 广播后完成 P*V 累加。 + for (int source = 0; source < TILE_K; ++source) { + const float weight = __shfl_sync(0xffffffff, probability, source); + if (source < tile_size) { + for (int item = 0; item < 4; ++item) { + acc[item] += weight * v_tile[source * HEAD_DIM + item * WARP_SIZE + lane]; + } + } + } + max_score = new_max; + __syncwarp(); + } + + for (int item = 0; item < 4; ++item) { + output[head * HEAD_DIM + item * WARP_SIZE + lane] = static_cast(acc[item] / sum_exp); + } +} + +template +void launch_flash_attention_decode(T *output, const T *q, const T *k, const T *v, + size_t total_len, size_t nhead, size_t nkvhead, float scale) { + const size_t shared_bytes = TILE_K * HEAD_DIM * 2 * sizeof(float); + flash_attention_decode_kernel<<(nhead), WARP_SIZE, shared_bytes>>>( + output, q, k, v, total_len, nhead, nkvhead, scale); +} + +} // namespace + +namespace llaisys::ops::cuda { + +void flash_attention_decode(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + ASSERT(seqlen == 1 && d == HEAD_DIM && dv == HEAD_DIM, + "flash_attention_decode requires seqlen=1 and d=dv=128"); + + switch (type) { + case LLAISYS_DTYPE_F32: + launch_flash_attention_decode(reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + case LLAISYS_DTYPE_BF16: + launch_flash_attention_decode(reinterpret_cast<__nv_bfloat16 *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + case LLAISYS_DTYPE_F16: + launch_flash_attention_decode(reinterpret_cast<__half *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/self_attention/nvidia/flash_attention_decode_splitkv_cuda.cu b/src/ops/self_attention/nvidia/flash_attention_decode_splitkv_cuda.cu new file mode 100644 index 000000000..e34e1cc73 --- /dev/null +++ b/src/ops/self_attention/nvidia/flash_attention_decode_splitkv_cuda.cu @@ -0,0 +1,266 @@ +#include "../../../utils.hpp" +#include "flash_attention_cuda.cuh" + +#include +#include +#include +#include + +// Split-KV / flash-decoding，phase 1（局部计算）。 +// 跟 flash_attention_decode_cuda.cu 的区别：那边一个 warp 扫完整个 [0, total_len)； +// 这里把 [0, total_len) 切成 num_splits 段，一个 warp 只扫自己负责的 [split_start, split_end)， +// 算出局部的 (m, l, acc) 三元组写进中间 buffer，不做最终的 acc/l 归一化——归一化和跨 split +// 合并留给 phase 2（另一个 kernel，负责把同一个 head 的 num_splits 份局部结果 rescale 后加起来）。 +// +// grid 组织：blockIdx.x = head，blockIdx.y = split_idx（对应 dim3 grid(nhead, num_splits)）。 +// +// 中间 buffer 布局（建议，可按你的习惯调整，只要 phase 2 读的时候对得上）： +// partial_m : [nhead, num_splits] 局部 running max +// partial_l : [nhead, num_splits] 局部 sum_exp +// partial_acc: [nhead, num_splits, HEAD_DIM] 局部加权和（未除以 l），dim 排布沿用 +// 原 kernel 的 item*WARP_SIZE+lane 交错方式，phase 2 按同样方式读 + +namespace { + +constexpr int WARP_SIZE = 32; +constexpr int HEAD_DIM = 128; +constexpr int TILE_K = 32; + +// 跟 flash_attention_decode_cuda.cu 里的同名函数完全一样，这里保留一份是因为两个 .cu +// 文件都在各自的匿名 namespace 里（内部链接），不会有 ODR 重名问题。 +__device__ float warp_max(float value) { + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + value = max(value, __shfl_down_sync(0xffffffff, value, offset)); + } + return __shfl_sync(0xffffffff, value, 0); +} + +__device__ float warp_sum(float value) { + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return __shfl_sync(0xffffffff, value, 0); +} + +// 作用是返回一部分kernel 的值 +template +__global__ void flash_attention_decode_splitkv_phase1_kernel( + float *partial_m, float *partial_l, float *partial_acc, + const T *q, const T *k, const T *v, + size_t total_len, size_t nhead, size_t nkvhead, float scale, + size_t split_size, size_t num_splits) { + const int lane = threadIdx.x; + const size_t head = blockIdx.x; + const size_t split_idx = blockIdx.y; + const size_t kv_head = head / (nhead / nkvhead); + + // TODO: 算出这个 split 负责的 KV 范围。 + size_t split_start = split_idx * split_size; + size_t split_end = min(split_start + split_size, total_len); + if (split_start >= total_len) { + partial_m[head * num_splits + split_idx] = -INFINITY; + partial_l[head * num_splits + split_idx] = 0; + for (size_t t = 0; t < 4; t++) { + partial_acc[(head * num_splits + split_idx) * HEAD_DIM + t * WARP_SIZE + lane] = 0; + } + return; + } + extern __shared__ float shared[]; + float *k_tile = shared; + float *v_tile = shared + TILE_K * HEAD_DIM; + + float max_score = -INFINITY; + float sum_exp = 0.0f; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (size_t tile_start = split_start; tile_start < split_end; tile_start += TILE_K) { + const size_t tile_size = min(static_cast(TILE_K), split_end - tile_start); + + // 32 个 lane 协作将当前 K/V tile 转为 float 后搬入 shared memory。 + for (size_t flat = lane; flat < tile_size * HEAD_DIM; flat += WARP_SIZE) { + const size_t row = flat / HEAD_DIM; + const size_t dim = flat % HEAD_DIM; + k_tile[flat] = static_cast( + k[(tile_start + row) * nkvhead * HEAD_DIM + kv_head * HEAD_DIM + dim]); + v_tile[flat] = static_cast( + v[(tile_start + row) * nkvhead * HEAD_DIM + kv_head * HEAD_DIM + dim]); + } + __syncwarp(); + + float score = -INFINITY; + if (lane < tile_size) { + score = 0.0f; + for (int dim = 0; dim < HEAD_DIM; ++dim) { + score += static_cast(q[head * HEAD_DIM + dim]) * k_tile[lane * HEAD_DIM + dim]; + } + score *= scale; + } + + // 合并当前 tile 与历史 tile 的 online softmax；max 变化时需重缩放历史累加值。 + const float tile_max = warp_max(score); + const float new_max = max(max_score, tile_max); + const float old_scale = expf(max_score - new_max); + const float probability = lane < tile_size ? expf(score - new_max) : 0.0f; + + sum_exp = sum_exp * old_scale + warp_sum(probability); + for (int item = 0; item < 4; ++item) { + acc[item] *= old_scale; + } + + // 每个 lane 持有一个 key 的权重，通过 shuffle 广播后完成 P*V 累加。 + for (int source = 0; source < TILE_K; ++source) { + const float weight = __shfl_sync(0xffffffff, probability, source); + if (source < tile_size) { + for (int item = 0; item < 4; ++item) { + acc[item] += weight * v_tile[source * HEAD_DIM + item * WARP_SIZE + lane]; + } + } + } + max_score = new_max; + __syncwarp(); + } + + partial_m[head * num_splits + split_idx] = max_score; + partial_l[head * num_splits + split_idx] = sum_exp; + for (size_t item = 0; item < 4; item++) { + partial_acc[(head * num_splits + split_idx) * HEAD_DIM + item * WARP_SIZE + lane] = acc[item]; + } +} + +template +void launch_flash_attention_decode_splitkv_phase1( + float *partial_m, float *partial_l, float *partial_acc, + const T *q, const T *k, const T *v, + size_t total_len, size_t nhead, size_t nkvhead, float scale, + size_t split_size, size_t num_splits) { + const size_t shared_bytes = TILE_K * HEAD_DIM * 2 * sizeof(float); + dim3 grid(static_cast(nhead), static_cast(num_splits)); + flash_attention_decode_splitkv_phase1_kernel<<>>( + partial_m, partial_l, partial_acc, q, k, v, + total_len, nhead, nkvhead, scale, split_size, num_splits); +} + +// Phase 2：跨 split 归约。一个 block（一个 warp）处理一个 head，串行扫过该 head 的 +// num_splits 份局部结果——num_splits 很小（个位数到几十），串行开销可忽略，且避免了 +// 再开一次跨 block 的规约/同步。每个 lane 各自重复扫一遍算出同样的 final_max，是拿 +// 冗余计算换掉一次 broadcast，逻辑更简单。 +// rescale 公式与 tile 内 online softmax 完全一致：sum/acc 各自乘上 exp(m_split - m_final) +// 再累加；m 为 -INFINITY 的空 split（tile 循环里提前 return 那种）直接跳过，避免 exp(-inf-(-inf))=nan。 +template +__global__ void flash_attention_decode_splitkv_phase2_kernel( + T *output, const float *partial_m, const float *partial_l, const float *partial_acc, + size_t num_splits) { + const int lane = threadIdx.x; + const size_t head = blockIdx.x; + + float final_max = -INFINITY; + for (size_t s = 0; s < num_splits; ++s) { + final_max = max(final_max, partial_m[head * num_splits + s]); + } + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + float final_sum = 0.0f; + for (size_t s = 0; s < num_splits; ++s) { + const float m_s = partial_m[head * num_splits + s]; + if (m_s == -INFINITY) { + continue; + } + const float rescale = expf(m_s - final_max); + final_sum += partial_l[head * num_splits + s] * rescale; + for (int item = 0; item < 4; ++item) { + acc[item] += partial_acc[(head * num_splits + s) * HEAD_DIM + item * WARP_SIZE + lane] * rescale; + } + } + + for (int item = 0; item < 4; ++item) { + output[head * HEAD_DIM + item * WARP_SIZE + lane] = static_cast(acc[item] / final_sum); + } +} + +template +void launch_flash_attention_decode_splitkv_phase2( + T *output, const float *partial_m, const float *partial_l, const float *partial_acc, + size_t nhead, size_t num_splits) { + flash_attention_decode_splitkv_phase2_kernel<<(nhead), WARP_SIZE>>>( + output, partial_m, partial_l, partial_acc, num_splits); +} + +// 根据 total_len 选 num_splits：目标是让 nhead*num_splits 个 block 大致打满常见 GPU 的 +// SM 数量级（kTargetBlocks），同时不能切得比 kMinSplitSize 还碎——太碎的话 phase1 里 +// 大量 tail split 只有一两个 tile,收益被 kernel launch/phase2 归约开销吃掉。 +size_t choose_num_splits(size_t total_len, size_t nhead) { + constexpr size_t kTargetBlocks = 64; + constexpr size_t kMinSplitSize = TILE_K; + constexpr size_t kMaxNumSplits = 32; + + size_t num_splits = (kTargetBlocks + nhead - 1) / nhead; + num_splits = std::max(1, std::min(num_splits, kMaxNumSplits)); + + size_t max_useful_splits = std::max(1, total_len / kMinSplitSize); + num_splits = std::min(num_splits, max_useful_splits); + return num_splits; +} + +template +void launch_flash_attention_decode_splitkv( + T *output, const T *q, const T *k, const T *v, + size_t total_len, size_t nhead, size_t nkvhead, float scale) { + const size_t num_splits = choose_num_splits(total_len, nhead); + const size_t split_size = (total_len + num_splits - 1) / num_splits; + + float *partial_m = nullptr; + float *partial_l = nullptr; + float *partial_acc = nullptr; + cudaMalloc(reinterpret_cast(&partial_m), nhead * num_splits * sizeof(float)); + cudaMalloc(reinterpret_cast(&partial_l), nhead * num_splits * sizeof(float)); + cudaMalloc(reinterpret_cast(&partial_acc), nhead * num_splits * HEAD_DIM * sizeof(float)); + + launch_flash_attention_decode_splitkv_phase1( + partial_m, partial_l, partial_acc, q, k, v, + total_len, nhead, nkvhead, scale, split_size, num_splits); + launch_flash_attention_decode_splitkv_phase2( + output, partial_m, partial_l, partial_acc, nhead, num_splits); + + cudaFree(partial_m); + cudaFree(partial_l); + cudaFree(partial_acc); +} + +} // namespace + +namespace llaisys::ops::cuda { + +void flash_attention_decode_splitkv(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + ASSERT(seqlen == 1 && d == HEAD_DIM && dv == HEAD_DIM, + "flash_attention_decode_splitkv requires seqlen=1 and d=dv=128"); + + switch (type) { + case LLAISYS_DTYPE_F32: + launch_flash_attention_decode_splitkv(reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + case LLAISYS_DTYPE_BF16: + launch_flash_attention_decode_splitkv(reinterpret_cast<__nv_bfloat16 *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + case LLAISYS_DTYPE_F16: + launch_flash_attention_decode_splitkv(reinterpret_cast<__half *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + total_len, nhead, nkvhead, scale); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/self_attention/nvidia/flash_attention_prefill_cuda.cu b/src/ops/self_attention/nvidia/flash_attention_prefill_cuda.cu new file mode 100644 index 000000000..69175e2e4 --- /dev/null +++ b/src/ops/self_attention/nvidia/flash_attention_prefill_cuda.cu @@ -0,0 +1,177 @@ +#include "../../../utils.hpp" +#include "flash_attention_cuda.cuh" +#include +#include +#include +#include +#include +#define CEIL(a, b) (((a) + (b) - 1) / (b)) +// q/k/v/attn_val 布局同 V1：q[i*nhead*d+h*d+dim] / k,v 用 total_len 而非 seqlen 做行数。 +// attn_val = causal_softmax(scale * q * k^T) * v；GQA：kvh = h / (nhead/nkvhead)。 +// 内部累加统一用 float，写回时转回 T。 +// +// 只服务 prefill（op.cpp 按 seqlen>1 && total_len==seqlen && d==dv==128 分流到这里； +// decode 走 flash_attention_decode_cuda.cu，其余 shape 落回 V1 self_attention_cuda.cu）。 +// total_len==seqlen 是因为 qwen2.cc 一次性处理整个 prompt（非 chunked prefill）， +// 所以 causal_offset 恒为 0，limit 直接等于 i；assert(total_len==seqlen) 防止以后 +// 调用方式变了却没人发现（release 下是空操作，仅 debug 生效）。 +// +// Tiling：行方向 TILE_Q=8 行/block，一个 warp 管一行；列方向 K/V 按 Bc=32(=warpSize) +// 一组分块进 shared memory 给 tile 内所有行共享，chunk 大小=warp size 省掉一层跨步 +// 循环。Bc×(d+dv)×4B=32KB（d=dv=128），在 48KB 默认上限内。acc[4] 寄存器数组硬编码 +// dv=128（32 lane×4）。 + +// __shfl_down_sync 规约后只有 lane 0 有完整结果，末尾用 __shfl_sync 广播给全 warp， +// 因为调用方（chunk_max/sum_exp）需要全 warp 一致的值参与后续 rescale。 +__device__ float warp_reduce_max(float value) { + for (int offset = 16; offset > 0; offset /= 2) { + float other = __shfl_down_sync(0xffffffff, value, offset); + value = max(value, other); + } + return __shfl_sync(0xffffffff, value, 0); +} + +__device__ float warp_reduce_sum(float value) { + for (int offset = 16; offset > 0; offset /= 2) { + float other = __shfl_down_sync(0xffffffff, value, offset); + value = value + other; + } + return __shfl_sync(0xffffffff, value, 0); +} + +constexpr int block_size = 256; +constexpr int MAX_WARPS = block_size / 32; +const int TILE_Q = 8; +constexpr int Bc = 32; +template +__global__ void flash_attention_kernel(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + + size_t tile_start = blockIdx.x * TILE_Q; + int lane = threadIdx.x & 31; + int warp_id = threadIdx.x >> 5; // = row_in_tile + size_t i = tile_start + warp_id; // 可能 ≥ seqlen（tile 不满） + size_t h = blockIdx.y; + size_t group = nhead / nkvhead; + size_t kvh = h / group; + assert(total_len == seqlen); // 只服务 prefill，见文件头 + size_t limit = i; + // 协作搬运+__syncthreads() 要求 block 内循环次数一致，不能用各 warp 自己的 limit， + // 要用 tile 内最大值。 + size_t tile_max_limit = min(tile_start + TILE_Q - 1, seqlen - 1); + + extern __shared__ float smem[]; + float *K_chunk = smem; + float *V_chunk = smem + Bc * d; + float m = -INFINITY; + float l = 0.0f; + float acc[4] = {0.0}; + for (int j = 0; j <= tile_max_limit; j += Bc) { + // 协作搬运 K/V 进 shared memory，tile 内所有行共享。 + for (size_t flat = threadIdx.x; flat < Bc * d; flat += blockDim.x) { + size_t dim = flat % d; + size_t row_in_chunk = flat / d; + if (j + row_in_chunk <= tile_max_limit) { + K_chunk[flat] = k[(j + row_in_chunk) * nkvhead * d + kvh * d + dim]; + V_chunk[flat] = v[(j + row_in_chunk) * nkvhead * dv + kvh * dv + dim]; + } + } + __syncthreads(); + + // i 可能 ≥ seqlen（tile 不满）：协作搬运不受影响，但打分/softmax/写回只对 + // 真实存在的行有意义，越界读写 q/attn_val 不安全，所以整段包进判断里。i 在 + // 同一 warp 内一致，不会打破 warp_reduce_* 要求全 32 lane 参与 shuffle 的前提。 + if (i < seqlen) { + float score = 0.0; + for (size_t dim = 0; dim < d; dim++) { + score += float(q[i * nhead * d + h * d + dim]) * K_chunk[lane * d + dim]; + } + score *= scale; + if (j + lane > limit) { + score = -INFINITY; + } + float chunk_max = warp_reduce_max(score); + if (chunk_max > m) { + float suofang = exp(m - chunk_max); + l *= suofang; + for (int t = 0; t < 4; t++) { + acc[t] *= suofang; + } + m = chunk_max; + } + float p = exp(score - m); + float sum_exp = warp_reduce_sum(p); + l += sum_exp; + for (int src = 0; src < 32; src++) { + float p_src = __shfl_sync(0xffffffff, p, src); + for (int t = 0; t < 4; t++) { + acc[t] += p_src * V_chunk[src * dv + t * 32 + lane]; + } + } + } + // 必须等所有线程读完 K_chunk/V_chunk 才能进入下一轮覆盖它们——否则 + // i>=seqlen 的 warp（直接跳过上面的 if）会抢先覆盖，跟还在读的 warp 形成 + // data race（只有 tile 跨多个 chunk 时才触发）。 + __syncthreads(); + } + if (i < seqlen) { + for (size_t t = 0; t < 4; t++) { + attn_val[i * nhead * dv + h * dv + t * 32 + lane] = acc[t] / l; + } + } +} + +// 命名故意跟 V1 的 self_attention_kernel/launch_self_attention 不同：两个 .cu 里 +// 出现同名同签名的函数模板会被链接器当弱符号合并，运行时随机选中哪一份（见 +// docs/FLASH_ATTENTION_DEBUG_LOG_ZH.md Bug 1）。 +template +void launch_flash_attention(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + unsigned int gridx = CEIL(seqlen, TILE_Q); + + dim3 grid(static_cast(gridx), static_cast(nhead)); + size_t shared_bytes = Bc * (d + dv) * sizeof(float); + + flash_attention_kernel<<>>( + attn_val, q, k, v, seqlen, total_len, nhead, nkvhead, d, dv, scale); +} + +namespace llaisys::ops::cuda { + +void flash_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_flash_attention( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_BF16: + launch_flash_attention( + reinterpret_cast<__nv_bfloat16 *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_F16: + launch_flash_attention( + reinterpret_cast<__half *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/self_attention/nvidia/self_attention_cuda.cu b/src/ops/self_attention/nvidia/self_attention_cuda.cu new file mode 100644 index 000000000..ecbe38b13 --- /dev/null +++ b/src/ops/self_attention/nvidia/self_attention_cuda.cu @@ -0,0 +1,129 @@ +#include "../../../utils.hpp" +#include "self_attention_cuda.cuh" +#include +#include +#include +#include + +// 通用兜底 kernel（任意 shape/causal_offset）：一个 block 处理一个 (query token i, head h)， +// 三遍扫描（打分+max / softmax 归一 / 加权求和 v），scores 用动态 shared memory 存一整行。 +template +__global__ void self_attention_kernel(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + size_t i = blockIdx.x; // query token, [0, seqlen) + size_t h = blockIdx.y; // query head, [0, nhead) + size_t tid = threadIdx.x; + + size_t group = nhead / nkvhead; // 每 group 个 query head 共享 1 个 kv head + size_t kvh = h / group; // 该 query head 对应的 kv head + size_t causal_offset = total_len - seqlen; + size_t limit = i + causal_offset; // 因果掩码: j 的范围是 [0, limit]（闭区间） + + extern __shared__ float scores[]; // 大小 total_len，只用到 [0, limit] 这一段 + __shared__ float sdata[256]; + // 第一遍打分：每个线程按 grid-stride 负责一部分 j（j = tid, tid+stride, ... <= limit）。 + float max_score = -INFINITY; + for (size_t j = tid; j <= limit; j += blockDim.x) { + float score = 0.0f; + + for (size_t dim = 0; dim < d; dim++) { + score += float(q[i * nhead * d + h * d + dim]) * float(k[j * nkvhead * d + kvh * d + dim]); + } + score *= scale; + scores[j] = score; + max_score = std::fmaxf(max_score, score); + } + sdata[tid] = max_score; + __syncthreads(); + + // block 内做 max 规约得到 max_score。 + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride) { + sdata[tid] = max(sdata[tid], sdata[tid + stride]); + } + __syncthreads(); + } + max_score = sdata[0]; + + __shared__ float sumexp[256]; + sumexp[tid] = 0.0; + // 第二遍:用 exp(score - max_score) 做数值稳定的 softmax 加权求和。 + for (size_t j = tid; j <= limit; j += blockDim.x) { + scores[j] = std::exp(scores[j] - max_score); + sumexp[tid] += scores[j]; + } + __syncthreads(); + for (size_t stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (tid < stride) { + sumexp[tid] = sumexp[tid] + sumexp[tid + stride]; + } + __syncthreads(); + } + float sum_exp = sumexp[0]; + + // 第三遍：按 dv 维度分工，加权求和 v 得到最终输出。 + for (size_t t = tid; t < dv; t += blockDim.x) { + float acc = 0.0; + for (size_t j = 0; j <= limit; j += 1) { + acc += (float)(scores[j]) * (float)(v[j * nkvhead * dv + kvh * dv + t]); + } + attn_val[i * nhead * dv + h * dv + t] = acc / sum_exp; + } +} + +template +void launch_self_attention(T *attn_val, const T *q, const T *k, const T *v, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, + size_t d, size_t dv, float scale) { + constexpr int block_size = 256; + dim3 grid(static_cast(seqlen), static_cast(nhead)); + size_t shared_bytes = total_len * sizeof(float); + + self_attention_kernel<<>>( + attn_val, q, k, v, seqlen, total_len, nhead, nkvhead, d, dv, scale); +} + +namespace llaisys::ops::cuda { +// q : [seqlen, nhead, d ] +// k : [total_len, nkvhead, d ] +// qk^T : [seqlen, nhead, total_len] +// softmax(scale * qk^T) : [seqlen, nhead, total_len] +// v : [total_len, nkvhead, dv] +// attn_val : [seqlen, nhead, dv] +// attn_val = softmax(scale * q * k^T) * v +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_self_attention( + reinterpret_cast(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_BF16: + launch_self_attention( + reinterpret_cast<__nv_bfloat16 *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + case LLAISYS_DTYPE_F16: + launch_self_attention( + reinterpret_cast<__half *>(attn_val), + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} // namespace llaisys::ops::cuda diff --git a/src/ops/self_attention/nvidia/self_attention_cuda.cuh b/src/ops/self_attention/nvidia/self_attention_cuda.cuh new file mode 100644 index 000000000..f42287a28 --- /dev/null +++ b/src/ops/self_attention/nvidia/self_attention_cuda.cuh @@ -0,0 +1,14 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +void self_attention(std::byte *attn_val, const std::byte *q, const std::byte *k, const std::byte *v, + llaisysDataType_t type, + size_t seqlen, size_t total_len, size_t nhead, size_t nkvhead, size_t d, size_t dv, + float scale); + +} \ No newline at end of file diff --git a/src/ops/self_attention/op.cpp b/src/ops/self_attention/op.cpp index 43d620142..0447963b9 100644 --- a/src/ops/self_attention/op.cpp +++ b/src/ops/self_attention/op.cpp @@ -1,7 +1,84 @@ #include "op.hpp" +#include "./cpu/self_attention.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" +#include "nvidia/self_attention_cuda.cuh" +#include "nvidia/flash_attention_cuda.cuh" +#include "iluvatar/self_attention_iluvatar.cuh" +// GQA: 要求 nhead 是 nkvhead 的整数倍。 +// attn_val, q : [seqlen, nhead, d ] / [seqlen, nhead, dv] +// k, v : [total_len, nkvhead, d ] / [total_len, nkvhead, dv] namespace llaisys::ops { void self_attention(tensor_t attn_val, tensor_t q, tensor_t k, tensor_t v, float scale) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(attn_val, q, k, v); + CHECK_SAME_DTYPE(attn_val->dtype(), q->dtype(), k->dtype(), v->dtype()); + ASSERT(q->ndim() == 3 && k->ndim() == 3 && v->ndim() == 3 && attn_val->ndim() == 3, + "self_attention: q/k/v/attn_val must all be 3D."); + ASSERT(q->shape()[1] % k->shape()[1] == 0, + "self_attention (V3): nhead must be a multiple of nkvhead (GQA)."); + ASSERT(k->shape()[0] == v->shape()[0] && k->shape()[1] == v->shape()[1], + "self_attention: k and v must share the same [total_len, nkvhead]."); + ASSERT(attn_val->shape()[0] == q->shape()[0] && attn_val->shape()[1] == q->shape()[1], + "self_attention: attn_val's [seqlen, nhead] must match q's."); + ASSERT(attn_val->shape()[2] == v->shape()[2], + "self_attention: attn_val's dv must match v's dv."); + + size_t seqlen = q->shape()[0]; + size_t total_len = k->shape()[0]; + size_t nhead = q->shape()[1]; + size_t nkvhead = k->shape()[1]; + size_t d = q->shape()[2]; + size_t dv = v->shape()[2]; + + if (attn_val->deviceType() == LLAISYS_DEVICE_CPU) { + llaisys::ops::cpu::self_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; + } + + llaisys::core::context().setDevice(attn_val->deviceType(), attn_val->deviceId()); + + switch (attn_val->deviceType()) { + case LLAISYS_DEVICE_CPU: + llaisys::ops::cpu::self_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: { + // Prefill 和 decode 使用针对 d=dv=128 的 Flash Attention，其他 shape 保留通用实现兜底。 + // decode 再按 total_len 细分：total_len 较短时 nhead 个 warp 就能扫完、直接用 + // flash_attention_decode；total_len 较长时单 warp 串行扫整个 KV cache 太慢， + // SM 也远没打满（decode 只发射 nhead 个 block），改用按 KV 方向切分的 + // flash_attention_decode_splitkv（flash-decoding）。256 这个阈值对应 + // flash_attention_decode_splitkv_cuda.cu 里 choose_num_splits 的下限 + // （kMinSplitSize=TILE_K=32）——total_len 小于它时 splitkv 顶多切出 1 个 split， + // 多付出的 phase2 kernel launch/中间 buffer 开销就白费了。 + constexpr size_t kSplitKVThreshold = 256; + if (seqlen > 1 && total_len == seqlen && d == 128 && dv == 128) { + llaisys::ops::cuda::flash_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + } else if (seqlen == 1 && d == 128 && dv == 128 && total_len > kSplitKVThreshold) { + llaisys::ops::cuda::flash_attention_decode_splitkv(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + } else if (seqlen == 1 && d == 128 && dv == 128) { + llaisys::ops::cuda::flash_attention_decode(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + } else { + llaisys::ops::cuda::self_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + } + return; + } +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + llaisys::ops::iluvatar::self_attention(attn_val->data(), q->data(), k->data(), v->data(), + attn_val->dtype(), seqlen, total_len, nhead, nkvhead, d, dv, scale); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/ops/swiglu/cpu/swiglu_cpu.cpp b/src/ops/swiglu/cpu/swiglu_cpu.cpp new file mode 100644 index 000000000..758f2e4f9 --- /dev/null +++ b/src/ops/swiglu/cpu/swiglu_cpu.cpp @@ -0,0 +1,34 @@ +#include "swiglu_cpu.hpp" +#include "../../../utils.hpp" +#include + +template +void swiglu_(T *out, const T *gate, const T *up, size_t rows, size_t d) { + for (size_t i = 0; i < rows; i++) { + for (size_t j = 0; j < d; j++) { + float gate_val = llaisys::utils::cast(gate[i * d + j]); + gate_val = gate_val / (1 + std::exp(-gate_val)); // silu(gate) + float up_val = llaisys::utils::cast(up[i * d + j]); + float swiglu_val = gate_val * up_val; + out[i * d + j] = llaisys::utils::cast(swiglu_val); + } + } +} + +namespace llaisys::ops::cpu { +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + return swiglu_(reinterpret_cast(out), reinterpret_cast(gate), + reinterpret_cast(up), rows, d); + case LLAISYS_DTYPE_BF16: + return swiglu_(reinterpret_cast(out), reinterpret_cast(gate), + reinterpret_cast(up), rows, d); + case LLAISYS_DTYPE_F16: + return swiglu_(reinterpret_cast(out), reinterpret_cast(gate), + reinterpret_cast(up), rows, d); + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} +} // namespace llaisys::ops::cpu diff --git a/src/ops/swiglu/cpu/swiglu_cpu.hpp b/src/ops/swiglu/cpu/swiglu_cpu.hpp new file mode 100644 index 000000000..1948e0081 --- /dev/null +++ b/src/ops/swiglu/cpu/swiglu_cpu.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "llaisys.h" + +#include +namespace llaisys::ops::cpu { +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d); +} diff --git a/src/ops/swiglu/iluvatar/swiglu_iluvatar.cu b/src/ops/swiglu/iluvatar/swiglu_iluvatar.cu new file mode 100644 index 000000000..1a4ffe72a --- /dev/null +++ b/src/ops/swiglu/iluvatar/swiglu_iluvatar.cu @@ -0,0 +1,71 @@ +#include "../../../utils.hpp" +#include "swiglu_iluvatar.cuh" +#include +#include +#include + +// swiglu(gate, up) = silu(gate) * up，silu(x) = x/(1+exp(-x))。 +// out/gate/up 都是 [rows,d] 且 contiguous（op.cpp 已断言），按 numel=rows*d 摊平成一维处理， +// 逐元素无跨元素依赖，不需要归约/shared memory。 + +template +__global__ void swiglu_kernel(T *out, const T *gate, const T *up, size_t numel) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + float gate_val = (float)gate[idx]; + gate_val = gate_val / (1 + std::exp(-gate_val)); // silu(gate) + float up_val = (float)up[idx]; + float swiglu_val = gate_val * up_val; + out[idx] = swiglu_val; + } +} + +template +void launch_swiglu(T *out, const T *gate, const T *up,size_t rows,size_t d) { + constexpr int block_size = 256; + int numel = rows * d; + int grid_size = static_cast((numel + block_size - 1) / block_size); + + swiglu_kernel<<>>(out, gate, up, numel); +} + +namespace llaisys::ops::iluvatar { + +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_swiglu( + reinterpret_cast(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + case LLAISYS_DTYPE_BF16: + launch_swiglu( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + case LLAISYS_DTYPE_F16: + launch_swiglu( + reinterpret_cast<__half *>(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} \ No newline at end of file diff --git a/src/ops/swiglu/iluvatar/swiglu_iluvatar.cuh b/src/ops/swiglu/iluvatar/swiglu_iluvatar.cuh new file mode 100644 index 000000000..0cd3065da --- /dev/null +++ b/src/ops/swiglu/iluvatar/swiglu_iluvatar.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::iluvatar { + +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d); + +} \ No newline at end of file diff --git a/src/ops/swiglu/nvidia/swiglu_cuda.cu b/src/ops/swiglu/nvidia/swiglu_cuda.cu new file mode 100644 index 000000000..3fa8e4072 --- /dev/null +++ b/src/ops/swiglu/nvidia/swiglu_cuda.cu @@ -0,0 +1,71 @@ +#include "../../../utils.hpp" +#include "swiglu_cuda.cuh" +#include +#include +#include + +// swiglu(gate, up) = silu(gate) * up，silu(x) = x/(1+exp(-x))。 +// out/gate/up 都是 [rows,d] 且 contiguous（op.cpp 已断言），按 numel=rows*d 摊平成一维处理， +// 逐元素无跨元素依赖，不需要归约/shared memory。 + +template +__global__ void swiglu_kernel(T *out, const T *gate, const T *up, size_t numel) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + float gate_val = (float)gate[idx]; + gate_val = gate_val / (1 + std::exp(-gate_val)); // silu(gate) + float up_val = (float)up[idx]; + float swiglu_val = gate_val * up_val; + out[idx] = swiglu_val; + } +} + +template +void launch_swiglu(T *out, const T *gate, const T *up,size_t rows,size_t d) { + constexpr int block_size = 256; + int numel = rows * d; + int grid_size = static_cast((numel + block_size - 1) / block_size); + + swiglu_kernel<<>>(out, gate, up, numel); +} + +namespace llaisys::ops::cuda { + +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d) { + switch (type) { + case LLAISYS_DTYPE_F32: + launch_swiglu( + reinterpret_cast(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + case LLAISYS_DTYPE_BF16: + launch_swiglu( + reinterpret_cast<__nv_bfloat16 *>(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + case LLAISYS_DTYPE_F16: + launch_swiglu( + reinterpret_cast<__half *>(out), + reinterpret_cast(gate), + reinterpret_cast(up), + rows, + d + ); + return; + + default: + EXCEPTION_UNSUPPORTED_DATATYPE(type); + } +} + +} \ No newline at end of file diff --git a/src/ops/swiglu/nvidia/swiglu_cuda.cuh b/src/ops/swiglu/nvidia/swiglu_cuda.cuh new file mode 100644 index 000000000..daa42fbf1 --- /dev/null +++ b/src/ops/swiglu/nvidia/swiglu_cuda.cuh @@ -0,0 +1,11 @@ +#pragma once + +#include "llaisys.h" + +#include + +namespace llaisys::ops::cuda { + +void swiglu(std::byte *out, const std::byte *gate, const std::byte *up, llaisysDataType_t type, size_t rows, size_t d); + +} \ No newline at end of file diff --git a/src/ops/swiglu/op.cpp b/src/ops/swiglu/op.cpp index 47edbcc97..01f62f57c 100644 --- a/src/ops/swiglu/op.cpp +++ b/src/ops/swiglu/op.cpp @@ -1,7 +1,43 @@ #include "op.hpp" +#include "../../core/llaisys_core.hpp" +#include "../../utils.hpp" + +#include "cpu/swiglu_cpu.hpp" +#include "nvidia/swiglu_cuda.cuh" +#include "iluvatar/swiglu_iluvatar.cuh" + namespace llaisys::ops { void swiglu(tensor_t out, tensor_t gate, tensor_t up) { - TO_BE_IMPLEMENTED(); + CHECK_SAME_DEVICE(out, gate, up); + ASSERT(out->dtype() == gate->dtype() && out->dtype() == up->dtype(), "Swiglu: all tensors must have the same dtype."); + ASSERT(gate->ndim() == 2 && up->ndim() == 2, "Swiglu: gate and up must be 2D."); + ASSERT(out->ndim() == 2, "Swiglu: out must be 2D."); + ASSERT(gate->shape()[0] == up->shape()[0] && gate->shape()[1] == up->shape()[1], "Swiglu: gate and up must have the same shape."); + ASSERT(out->shape()[0] == gate->shape()[0] && out->shape()[1] == gate->shape()[1], "Swiglu: out must have the same shape as gate and up."); + ASSERT(out->isContiguous() && gate->isContiguous() && up->isContiguous(), "Swiglu: all tensors must be contiguous."); + + if (out->deviceType() == LLAISYS_DEVICE_CPU) { + return cpu::swiglu(out->data(), gate->data(), up->data(), out->dtype(), out->shape()[0], out->shape()[1]); + } + + llaisys::core::context().setDevice(out->deviceType(), out->deviceId()); + + switch (out->deviceType()) { + case LLAISYS_DEVICE_CPU: + return cpu::swiglu(out->data(), gate->data(), up->data(), out->dtype(), out->shape()[0], out->shape()[1]); +#ifdef ENABLE_NVIDIA_API + case LLAISYS_DEVICE_NVIDIA: + cuda::swiglu(out->data(), gate->data(), up->data(), out->dtype(), out->shape()[0], out->shape()[1]); + return; +#endif +#ifdef ENABLE_ILUVATAR_API + case LLAISYS_DEVICE_ILUVATAR: + iluvatar::swiglu(out->data(), gate->data(), up->data(), out->dtype(), out->shape()[0], out->shape()[1]); + return; +#endif + default: + EXCEPTION_UNSUPPORTED_DEVICE; + } } } // namespace llaisys::ops diff --git a/src/tensor/tensor.cpp b/src/tensor/tensor.cpp index 2f594bb65..4e585e7ad 100644 --- a/src/tensor/tensor.cpp +++ b/src/tensor/tensor.cpp @@ -164,42 +164,136 @@ void Tensor::debug() const { } bool Tensor::isContiguous() const { - TO_BE_IMPLEMENTED(); + ptrdiff_t expected_stride = 1; + for (size_t i = _meta.shape.size(); i-- > 0;) { + if (_meta.strides[i] != expected_stride && _meta.shape[i] != 1) { + return false; + } + expected_stride *= static_cast(_meta.shape[i]); + } return true; } tensor_t Tensor::permute(const std::vector &order) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + ASSERT(order.size() == ndim(), "permute order size must match tensor ndim"); + std::vector new_shape(order.size()); + std::vector new_strides(order.size()); + + for(size_t i = 0; i < order.size(); ++i) { + new_shape[i] = _meta.shape[order[i]]; + new_strides[i] = _meta.strides[order[i]]; + } + TensorMeta new_meta{_meta.dtype, new_shape, new_strides}; + return std::shared_ptr(new Tensor(new_meta, _storage, _offset)); } tensor_t Tensor::view(const std::vector &shape) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + ASSERT(isContiguous(), "Tensor must be contiguous to view"); + size_t new_numel = std::accumulate(shape.begin(), shape.end(), size_t(1), std::multiplies()); + ASSERT(numel() == new_numel, "Tensor numel does not match view shape"); + std::vector strides(shape.size()); + ptrdiff_t stride = 1; + for (size_t i = shape.size(); i-- > 0;) { + strides[i] = stride; + stride *= static_cast(shape[i]); + } + Tensor new_tensor = *this; + new_tensor._meta.shape = shape; + new_tensor._meta.strides = strides; + return std::shared_ptr(new Tensor(new_tensor._meta, new_tensor._storage, new_tensor._offset)); } tensor_t Tensor::slice(size_t dim, size_t start, size_t end) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + ASSERT(dim < ndim(), "invalid dim"); + ASSERT(start <= end && end <= shape()[dim], "invalid range"); + + TensorMeta new_meta = _meta; + new_meta.shape[dim] = end - start; + + size_t new_offset = _offset + start * static_cast(_meta.strides[dim]) * elementSize(); + + return std::shared_ptr(new Tensor(new_meta, _storage, new_offset)); } void Tensor::load(const void *src_) { - TO_BE_IMPLEMENTED(); + core::context().setDevice(this->deviceType(), this->deviceId()); + core::context().runtime().api()->memcpy_sync( + this->data(), + src_, + this->numel() * this->elementSize(), + LLAISYS_MEMCPY_H2D); } tensor_t Tensor::contiguous() const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + if(isContiguous()) { + return std::shared_ptr(new Tensor(_meta, _storage, _offset)); + } + ASSERT(deviceType() == LLAISYS_DEVICE_CPU, "contiguous is only supported for CPU tensors"); + auto new_tensor = create(shape(), dtype(), deviceType(), deviceId()); + size_t esz = elementSize(); + size_t n = ndim(); + std::vector indices(n, 0); + std::byte *dst = new_tensor->data(); + for(size_t i = 0; i < numel(); ++i) { + size_t src_offset = _offset; + for (size_t d = 0; d < n; d++) { + src_offset += indices[d] + * static_cast(_meta.strides[d]) + * esz; + } + std::memcpy(dst + i * esz, _storage->memory() + src_offset, esz); + for (size_t d = n; d-- > 0;) { + if (++indices[d] < _meta.shape[d]) break; + indices[d] = 0; + } + } + return new_tensor; } tensor_t Tensor::reshape(const std::vector &shape) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + size_t new_numel = std::accumulate(shape.begin(), shape.end(), size_t(1), std::multiplies()); + ASSERT(numel() == new_numel, "Tensor numel does not match reshape shape"); + if (isContiguous()) { + return view(shape); + } + return contiguous()->view(shape); } tensor_t Tensor::to(llaisysDeviceType_t device_type, int device) const { - TO_BE_IMPLEMENTED(); - return std::shared_ptr(new Tensor(_meta, _storage)); + if (device < 0) { + device = 0; + } + if (device_type == this->deviceType() && device == this->deviceId()) { + return std::shared_ptr(new Tensor(_meta, _storage, _offset)); + } + + ASSERT(isContiguous(), "Tensor must be contiguous to change device"); + auto new_tensor = create(shape(), dtype(), device_type, device); + + llaisysMemcpyKind_t kind; + if (this->deviceType() == LLAISYS_DEVICE_CPU && device_type == LLAISYS_DEVICE_CPU) { + kind = LLAISYS_MEMCPY_H2H; + } else if (this->deviceType() == LLAISYS_DEVICE_CPU) { + kind = LLAISYS_MEMCPY_H2D; + } else if (device_type == LLAISYS_DEVICE_CPU) { + kind = LLAISYS_MEMCPY_D2H; + } else { + kind = LLAISYS_MEMCPY_D2D; + } + + if (device_type != LLAISYS_DEVICE_CPU) { + core::context().setDevice(device_type, device); + } else { + core::context().setDevice(this->deviceType(), this->deviceId()); + } + + core::context().runtime().api()->memcpy_sync( + new_tensor->data(), + this->data(), + this->numel() * this->elementSize(), + kind); + + return new_tensor; } } // namespace llaisys diff --git a/test/benchmark_infer.py b/test/benchmark_infer.py new file mode 100644 index 000000000..6217a0cea --- /dev/null +++ b/test/benchmark_infer.py @@ -0,0 +1,74 @@ +import argparse +import contextlib +import time + +from test_utils import llaisys_device +import llaisys +from transformers import AutoTokenizer + + +def make_step_context(record, api, use_nvtx): + @contextlib.contextmanager + def step_context(step, is_prefill): + if use_nvtx: + import torch + + torch.cuda.nvtx.range_push(f"{'prefill' if is_prefill else 'decode'}_{step}") + start = time.perf_counter() + try: + yield + finally: + api.device_synchronize() + record.append((is_prefill, time.perf_counter() - start)) + if use_nvtx: + import torch + + torch.cuda.nvtx.range_pop() + + return step_context + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) + parser.add_argument("--prompt", default="Who are you?", type=str) + parser.add_argument("--max_steps", default=128, type=int) + parser.add_argument( + "--nvtx", + action="store_true", + help="wrap each prefill/decode step in an NVTX range for nsys (needs torch + an NVIDIA device)", + ) + args = parser.parse_args() + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + input_content = tokenizer.apply_chat_template( + conversation=[{"role": "user", "content": args.prompt}], + add_generation_prompt=True, + tokenize=False, + ) + inputs = tokenizer.encode(input_content) + + model = llaisys.models.Qwen2(args.model, llaisys_device(args.device)) + api = llaisys.RuntimeAPI(llaisys_device(args.device)) + + record = [] + step_context = make_step_context(record, api, args.nvtx) + + start = time.perf_counter() + outputs = model.generate(inputs, max_new_tokens=args.max_steps, step_context=step_context) + total_elapsed = time.perf_counter() - start + + prefill = [elapsed for is_prefill, elapsed in record if is_prefill] + decode = [elapsed for is_prefill, elapsed in record if not is_prefill] + + print(f"Prompt tokens: {len(inputs)}, generated: {len(outputs) - len(inputs)}") + print(f"Prefill: {prefill[0] * 1000:.2f} ms") + if decode: + avg_decode = sum(decode) / len(decode) + print(f"Decode: {len(decode)} steps, avg {avg_decode * 1000:.2f} ms/token, {1.0 / avg_decode:.2f} tokens/s") + print(f"Total: {total_elapsed:.2f}s") + + +if __name__ == "__main__": + main() diff --git a/test/ops/add.py b/test/ops/add.py index bb8bf8ca8..d5937bdf7 100644 --- a/test/ops/add.py +++ b/test/ops/add.py @@ -42,7 +42,7 @@ def test_op_add( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [(2, 3), (512, 4096)] diff --git a/test/ops/argmax.py b/test/ops/argmax.py index d0f7ee298..0ea040b05 100644 --- a/test/ops/argmax.py +++ b/test/ops/argmax.py @@ -43,7 +43,7 @@ def test_op_argmax( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [(4,), (4096,)] diff --git a/test/ops/embedding.py b/test/ops/embedding.py index 99cadc1b8..daa9c68b0 100644 --- a/test/ops/embedding.py +++ b/test/ops/embedding.py @@ -39,7 +39,7 @@ def test_op_embedding( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [ diff --git a/test/ops/linear.py b/test/ops/linear.py index 38897331f..e979124c9 100644 --- a/test/ops/linear.py +++ b/test/ops/linear.py @@ -49,7 +49,7 @@ def test_op_linear( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [ diff --git a/test/ops/rms_norm.py b/test/ops/rms_norm.py index 67b789e3f..d4bee23b4 100644 --- a/test/ops/rms_norm.py +++ b/test/ops/rms_norm.py @@ -48,7 +48,7 @@ def test_op_rms_norm( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [(1, 4), (512, 4096)] diff --git a/test/ops/rope.py b/test/ops/rope.py index fe59dd11c..90d326afd 100644 --- a/test/ops/rope.py +++ b/test/ops/rope.py @@ -63,7 +63,7 @@ def test_op_rope( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [ diff --git a/test/ops/self_attention.py b/test/ops/self_attention.py index a042b51be..0a753e8d0 100644 --- a/test/ops/self_attention.py +++ b/test/ops/self_attention.py @@ -15,7 +15,7 @@ def torch_self_attention(attn_val, query, key, value, scale): L, S = query.size(-2), key.size(-2) attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device) - temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=S-L) + temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=S-L) attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) attn_bias.to(query.dtype) @@ -65,7 +65,7 @@ def test_op_self_attention( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [ diff --git a/test/ops/swiglu.py b/test/ops/swiglu.py index 1fa08f739..f11f573e8 100644 --- a/test/ops/swiglu.py +++ b/test/ops/swiglu.py @@ -42,7 +42,7 @@ def test_op_swiglu( import argparse parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--profile", action="store_true") args = parser.parse_args() testShapes = [(2, 3), (512, 4096)] diff --git a/test/test_infer.py b/test/test_infer.py index 59d06b874..489cbde99 100644 --- a/test/test_infer.py +++ b/test/test_infer.py @@ -81,7 +81,7 @@ def llaisys_infer( if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) parser.add_argument("--model", default=None, type=str) parser.add_argument("--prompt", default="Who are you?", type=str) parser.add_argument("--max_steps", default=128, type=int) diff --git a/test/test_runtime.py b/test/test_runtime.py index e2ac218a1..4176fdee6 100644 --- a/test/test_runtime.py +++ b/test/test_runtime.py @@ -55,7 +55,7 @@ def test_memcpy(api, size_bytes: int): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia"], type=str) + parser.add_argument("--device", default="cpu", choices=["cpu", "nvidia", "iluvatar"], type=str) args = parser.parse_args() test_basic_runtime_api(args.device) diff --git a/test/test_tensor.py b/test/test_tensor.py index 9d2e9a075..3858762f6 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -48,6 +48,29 @@ def test_tensor(): assert llaisys_tensor.is_contiguous() == torch_tensor.is_contiguous() assert check_equal(llaisys_tensor_slice, torch_tensor_slice) + # Test reshape (contiguous fast path) + print("===Test reshape (contiguous)===") + torch_tensor_reshape = torch_tensor.reshape(6, 10) + llaisys_tensor_reshape = llaisys_tensor.reshape(6, 10) + llaisys_tensor_reshape.debug() + assert llaisys_tensor_reshape.shape() == torch_tensor_reshape.shape + assert llaisys_tensor_reshape.strides() == torch_tensor_reshape.stride() + assert llaisys_tensor_reshape.is_contiguous() == torch_tensor_reshape.is_contiguous() + assert check_equal(llaisys_tensor_reshape, torch_tensor_reshape) + + # Test reshape (non-contiguous, must go through contiguous()) + print("===Test reshape (non-contiguous)===") + torch_tensor_perm_reshape = torch_tensor_perm.reshape(5, 12) + llaisys_tensor_perm_reshape = llaisys_tensor_perm.reshape(5, 12) + llaisys_tensor_perm_reshape.debug() + # NOTE: our reshape() always materializes a contiguous copy when the + # source isn't contiguous, so strides may differ from torch's (which can + # sometimes reshape part of a non-contiguous tensor without copying). + # Shape, contiguity of the *result*, and the actual data must still match. + assert llaisys_tensor_perm_reshape.shape() == torch_tensor_perm_reshape.shape + assert llaisys_tensor_perm_reshape.is_contiguous() + assert check_equal(llaisys_tensor_perm_reshape, torch_tensor_perm_reshape) + if __name__ == "__main__": test_tensor() diff --git a/test/test_utils.py b/test/test_utils.py index 0f38f0c8e..4e2fcf206 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -188,6 +188,9 @@ def torch_device(device_name: str, device_id=0): return torch.device("cpu") elif device_name == "nvidia": return torch.device(f"cuda:{device_id}") + elif device_name == "iluvatar": + # Iluvatar's PyTorch build hooks into the standard torch.cuda namespace. + return torch.device(f"cuda:{device_id}") else: raise ValueError(f"Unsupported device name: {device_name}") @@ -197,6 +200,8 @@ def llaisys_device(device_name: str): return llaisys.DeviceType.CPU elif device_name == "nvidia": return llaisys.DeviceType.NVIDIA + elif device_name == "iluvatar": + return llaisys.DeviceType.ILUVATAR else: raise ValueError(f"Unsupported device name: {device_name}") @@ -206,6 +211,8 @@ def device_name(llaisys_device: llaisys.DeviceType): return "cpu" elif llaisys_device == llaisys.DeviceType.NVIDIA: return "nvidia" + elif llaisys_device == llaisys.DeviceType.ILUVATAR: + return "iluvatar" else: raise ValueError(f"Unsupported llaisys device: {llaisys_device}") diff --git a/xmake.lua b/xmake.lua index 1f65f7a95..16eb6877d 100644 --- a/xmake.lua +++ b/xmake.lua @@ -18,6 +18,18 @@ if has_config("nv-gpu") then includes("xmake/nvidia.lua") end + +option("iluvatar-gpu") + set_default(false) + set_showmenu(true) + set_description("Whether to compile implementations for Iluvatar GPU") +option_end() + +if has_config("iluvatar-gpu") then + add_defines("ENABLE_ILUVATAR_API") + includes("xmake/iluvatar.lua") +end + target("llaisys-utils") set_kind("static") @@ -37,6 +49,12 @@ target("llaisys-device") set_kind("static") add_deps("llaisys-utils") add_deps("llaisys-device-cpu") + if has_config("nv-gpu") then + add_deps("llaisys-device-nvidia") + end + if has_config("iluvatar-gpu") then + add_deps("llaisys-device-iluvatar") + end set_languages("cxx17") set_warnings("all", "error") @@ -83,6 +101,12 @@ target_end() target("llaisys-ops") set_kind("static") add_deps("llaisys-ops-cpu") + if has_config("nv-gpu") then + add_deps("llaisys-ops-nvidia") + end + if has_config("iluvatar-gpu") then + add_deps("llaisys-ops-iluvatar") + end set_languages("cxx17") set_warnings("all", "error") @@ -105,10 +129,14 @@ target("llaisys") set_languages("cxx17") set_warnings("all", "error") - add_files("src/llaisys/*.cc") + add_files("src/llaisys/*.cc", "src/llaisys/**/*.cc") set_installdir(".") - + if has_config("nv-gpu") then + add_links("cudart") + end + + after_install(function (target) -- copy shared library to python package print("Copying llaisys to python/llaisys/libllaisys/ ..") diff --git a/xmake/iluvatar.lua b/xmake/iluvatar.lua new file mode 100644 index 000000000..9b084d765 --- /dev/null +++ b/xmake/iluvatar.lua @@ -0,0 +1,64 @@ +toolchain("iluvatar") + set_kind("standalone") + + -- 编译 .cu 文件 + set_toolset("cu", "/usr/local/corex/bin/clang++") + + -- 编译普通 C 文件 + set_toolset("cc", "/usr/bin/gcc") + + -- 编译普通 C++ 文件 + set_toolset("cxx", "/usr/bin/g++") + + -- 链接 C++ 目标 + set_toolset("ld", "/usr/bin/g++") + set_toolset("sh", "/usr/bin/g++") + + on_check(function (toolchain) + local find_tool = import("lib.detect.find_tool") + + return find_tool( + "clang++", + {paths = "/usr/local/corex/bin"} + ) + end) +toolchain_end() + +target("llaisys-device-iluvatar") + set_kind("static") + set_toolchains("iluvatar") + set_languages("cxx17") + set_values("cuda.rdc", false) + add_cuflags("-x", "ivcore", {force = true}) + add_cuflags("-std=c++17", {force = true}) + add_cuflags("-fPIC", {force = true}) + add_links("cublas") + -- Iluvatar's SDK only ships libcudart.so, not libcudart_static.a. Linking + -- "cudart" explicitly here satisfies xmake's built-in cuda rule (which + -- auto-adds "cudart_static" unless a cudart link is already present) so + -- it picks up the real shared library instead of a static one that + -- doesn't exist on this platform. + add_links("cudart") + + add_files("../src/device/iluvatar/*.cu") + add_linkdirs("/usr/local/corex-4.4.0/lib64") + on_install(function (target) end) +target_end() + +target("llaisys-ops-iluvatar") + set_kind("static") + set_toolchains("iluvatar") + add_deps("llaisys-tensor") + set_languages("cxx17") + add_linkdirs("/usr/local/corex-4.4.0/lib64") + set_values("cuda.rdc", false) + add_cuflags("-x", "ivcore", {force = true}) + add_cuflags("-std=c++17", {force = true}) + add_cuflags("-fPIC", {force = true}) + add_links("cublas") + add_links("cudart") + + add_files("../src/ops/*/iluvatar/*.cu") + + on_install(function (target) end) +target_end() diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua new file mode 100644 index 000000000..1f2c825ca --- /dev/null +++ b/xmake/nvidia.lua @@ -0,0 +1,26 @@ +target("llaisys-device-nvidia") + set_kind("static") + set_languages("cxx17") + add_cugencodes("native", "sm_80") + set_values("cuda.rdc", false) + add_cuflags("-Xcompiler=-fPIC", {force = true}) + add_links("cublas") + + add_files("../src/device/nvidia/*.cu") + + on_install(function (target) end) +target_end() + +target("llaisys-ops-nvidia") + set_kind("static") + add_deps("llaisys-tensor") + set_languages("cxx17") + add_cugencodes("native", "sm_80") + set_values("cuda.rdc", false) + add_cuflags("-Xcompiler=-fPIC", {force = true}) + add_links("cublas") + + add_files("../src/ops/*/nvidia/*.cu") + + on_install(function (target) end) +target_end()